- 
                Notifications
    
You must be signed in to change notification settings  - Fork 0
 
fix: implement lazy client creation in replicate.use() #57
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Fixes issue where replicate.use() would fail if no API token was available at call time, even when token becomes available later (e.g., from cog.current_scope). Changes: - Modified Function/AsyncFunction classes to accept client factories - Added _client property that creates client on demand - Updated module client to pass factory functions instead of instances - Token is now retrieved from current scope when model is called This maintains full backward compatibility while enabling use in Cog pipelines where tokens are provided through the execution context.
- Remove unused *args parameter in test function - Fix formatting issues from linter
- Fix async detection to not call client factory prematurely - Add use_async parameter to explicitly indicate async mode - Update test to avoid creating client during verification - Fix test mocking to use correct module path
Replace complex mocking test with simpler verification that: - use() works without token initially - Lazy client factory is properly configured - Client can be created when needed This avoids complex mocking while still verifying the core functionality.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks fine, but I think we can simplify it a little, I've left some suggestions inline.
Address PR feedback by removing Union types and using a single consistent approach: - Change Function/AsyncFunction constructors to accept Type[Client] only - Remove Union[Client, Type[Client]] in favor of just Type[Client] - Simplify _client property logic by removing isinstance checks - Update all use() overloads to accept class types only - Use issubclass() for async client detection instead of complex logic - Update tests to check for _client_class attribute This maintains the same lazy client creation behavior while being much simpler and more consistent.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks much cleaner, thanks. Left some comments about the unit tests, consider them suggestions.
Co-authored-by: Aron Carroll <aron@aroncarroll.com>
- Remove verbose comments and print statements - Focus on observable behavior rather than internal implementation - Use proper mocking that matches actual cog integration - Test that cog.current_scope() is called on client creation - Address code review feedback from PR discussion
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice, looks good.
| # Access the client property - this should trigger client creation and cog.current_scope call | ||
| _ = model._client | ||
| 
               | 
          ||
| assert mock_cog.current_scope.call_count == 1 | ||
| 
               | 
          ||
| # Change the token and access client again - should trigger another call | ||
| mock_context.items.return_value = [("REPLICATE_API_TOKEN", "test-token-2")] | ||
| 
               | 
          ||
| # Create a new model to trigger another client creation | ||
| model2 = replicate.use("test/model2") # type: ignore[misc] | ||
| _ = model2._client | 
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Minor tweak, we want to assert that the same model refreshes the client between calls. If we're looking at the client we can look at the bearer_token to verify it has actually picked up the token correctly.
| # Access the client property - this should trigger client creation and cog.current_scope call | |
| _ = model._client | |
| assert mock_cog.current_scope.call_count == 1 | |
| # Change the token and access client again - should trigger another call | |
| mock_context.items.return_value = [("REPLICATE_API_TOKEN", "test-token-2")] | |
| # Create a new model to trigger another client creation | |
| model2 = replicate.use("test/model2") # type: ignore[misc] | |
| _ = model2._client | |
| # Access the client property - this should trigger client creation and cog.current_scope call | |
| assert model._client.bearer_token == "test-token-1" | |
| assert mock_cog.current_scope.call_count == 1 | |
| # Change the token and access client again - should trigger another call | |
| mock_context.items.return_value = [("REPLICATE_API_TOKEN", "test-token-2")] | |
| # Assert that the second time picks up the new token | |
| assert model2._client.bearer_token == "test-token-2" | 
* fix: implement lazy client creation in replicate.use() (#57) * fix: implement lazy client creation in replicate.use() Fixes issue where replicate.use() would fail if no API token was available at call time, even when token becomes available later (e.g., from cog.current_scope). Changes: - Modified Function/AsyncFunction classes to accept client factories - Added _client property that creates client on demand - Updated module client to pass factory functions instead of instances - Token is now retrieved from current scope when model is called This maintains full backward compatibility while enabling use in Cog pipelines where tokens are provided through the execution context. * style: fix linter issues - Remove unused *args parameter in test function - Fix formatting issues from linter * fix: resolve async detection and test issues - Fix async detection to not call client factory prematurely - Add use_async parameter to explicitly indicate async mode - Update test to avoid creating client during verification - Fix test mocking to use correct module path * test: simplify lazy client test Replace complex mocking test with simpler verification that: - use() works without token initially - Lazy client factory is properly configured - Client can be created when needed This avoids complex mocking while still verifying the core functionality. * lint * fix: add type ignore for final linter warning * fix: add arg-type ignore for type checker warnings * refactor: simplify lazy client creation to use Type[Client] only Address PR feedback by removing Union types and using a single consistent approach: - Change Function/AsyncFunction constructors to accept Type[Client] only - Remove Union[Client, Type[Client]] in favor of just Type[Client] - Simplify _client property logic by removing isinstance checks - Update all use() overloads to accept class types only - Use issubclass() for async client detection instead of complex logic - Update tests to check for _client_class attribute This maintains the same lazy client creation behavior while being much simpler and more consistent. * Update tests/test_simple_lazy.py Co-authored-by: Aron Carroll <aron@aroncarroll.com> * test: improve lazy client test to follow project conventions - Remove verbose comments and print statements - Focus on observable behavior rather than internal implementation - Use proper mocking that matches actual cog integration - Test that cog.current_scope() is called on client creation - Address code review feedback from PR discussion * lint * lint --------- Co-authored-by: Aron Carroll <aron@aroncarroll.com> * release: 2.0.0-alpha.22 --------- Co-authored-by: Zeke Sikelianos <zeke@sikelianos.com> Co-authored-by: Aron Carroll <aron@aroncarroll.com> Co-authored-by: stainless-app[bot] <142633134+stainless-app[bot]@users.noreply.github.com>
Problem
replicate.use()was creating a client immediately, causing errors when no API token was available at call time. This prevented usage in Cog pipelines where tokens are provided throughcog.current_scope()at runtime.Solution
Modified
FunctionandAsyncFunctionclasses to support lazy client creation by always accepting client classes and instantiating them on demand.Changes
_module_client.py: PassReplicateandAsyncReplicateclasses directly instead of factory functionslib/_predictions_use.py: Simplified to always useType[Client]approach:Union[Client, Type[Client]]in favor of justType[Client]_client_classinstead of mixed types_clientproperty to always instantiate from classuse()overloads to accept class types onlyissubclass()for async client detectionResult
✅
replicate.use("model")works even when no token is initially available✅ Client created lazily when model is actually called
✅ Token retrieved from current cog scope at call time
✅ Simplified, consistent approach using only class types
✅ Clean type inference without complex Union handling
Fixes the original error: