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.

Author: zeke

src/replicate/_module_client.py

@@ -91,10 +91,11 @@ def _use(ref, *, hint=None, streaming=False, use_async=False, **kwargs):
         if use_async:
             # For async, we need to use AsyncReplicate instead
             from ._client import AsyncReplicate
+            from .lib._predictions_use import use
 
-            client = AsyncReplicate()
-            return client.use(ref, hint=hint, streaming=streaming, **kwargs)
-        return _load_client().use(ref, hint=hint, streaming=streaming, **kwargs)
+            return use(lambda: AsyncReplicate(), ref, hint=hint, streaming=streaming, **kwargs)
+        from .lib._predictions_use import use
+        return use(_load_client, ref, hint=hint, streaming=streaming, **kwargs)
 
     run = _run
     use = _use

src/replicate/lib/_predictions_use.py

@@ -436,15 +436,20 @@ class Function(Generic[Input, Output]):
     A wrapper for a Replicate model that can be called as a function.
     """
 
-    _client: Client
     _ref: str
     _streaming: bool
 
-    def __init__(self, client: Client, ref: str, *, streaming: bool) -> None:
-        self._client = client
+    def __init__(self, client: Union[Client, Callable[[], Client]], ref: str, *, streaming: bool) -> None:
+        self._client_or_factory = client
         self._ref = ref
         self._streaming = streaming
 
+    @property
+    def _client(self) -> Client:
+        if callable(self._client_or_factory):
+            return self._client_or_factory()
+        return self._client_or_factory
+
     def __call__(self, *args: Input.args, **inputs: Input.kwargs) -> Output:
         return self.create(*args, **inputs).output()
 
@@ -666,16 +671,21 @@ class AsyncFunction(Generic[Input, Output]):
     An async wrapper for a Replicate model that can be called as a function.
     """
 
-    _client: AsyncClient
     _ref: str
     _streaming: bool
     _openapi_schema: Optional[Dict[str, Any]] = None
 
-    def __init__(self, client: AsyncClient, ref: str, *, streaming: bool) -> None:
-        self._client = client
+    def __init__(self, client: Union[AsyncClient, Callable[[], AsyncClient]], ref: str, *, streaming: bool) -> None:
+        self._client_or_factory = client
         self._ref = ref
         self._streaming = streaming
 
+    @property
+    def _client(self) -> AsyncClient:
+        if callable(self._client_or_factory):
+            return self._client_or_factory()
+        return self._client_or_factory
+
     @cached_property
     def _parsed_ref(self) -> Tuple[str, str, Optional[str]]:
         return ModelVersionIdentifier.parse(self._ref)
@@ -804,7 +814,7 @@ async def openapi_schema(self) -> Dict[str, Any]:
 
 @overload
 def use(
-    client: Client,
+    client: Union[Client, Callable[[], Client]],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -814,7 +824,7 @@ def use(
 
 @overload
 def use(
-    client: Client,
+    client: Union[Client, Callable[[], Client]],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -824,7 +834,7 @@ def use(
 
 @overload
 def use(
-    client: AsyncClient,
+    client: Union[AsyncClient, Callable[[], AsyncClient]],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -834,7 +844,7 @@ def use(
 
 @overload
 def use(
-    client: AsyncClient,
+    client: Union[AsyncClient, Callable[[], AsyncClient]],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -843,7 +853,7 @@ def use(
 
 
 def use(
-    client: Union[Client, AsyncClient],
+    client: Union[Client, AsyncClient, Callable[[], Client], Callable[[], AsyncClient]],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,  # pylint: disable=unused-argument # noqa: ARG001 # required for type inference
@@ -868,9 +878,14 @@ def use(
     except AttributeError:
         pass
 
-    if isinstance(client, AsyncClient):
+    # Determine if this is async by checking the type
+    is_async = isinstance(client, AsyncClient) or (
+        callable(client) and isinstance(client(), AsyncClient)
+    )
+    
+    if is_async:
         # TODO: Fix type inference for AsyncFunction return type
         return AsyncFunction(client, str(ref), streaming=streaming)  # type: ignore[return-value]
-
-    # TODO: Fix type inference for Function return type
-    return Function(client, str(ref), streaming=streaming)  # type: ignore[return-value]
+    else:
+        # TODO: Fix type inference for Function return type
+        return Function(client, str(ref), streaming=streaming)  # type: ignore[return-value]

tests/test_simple_lazy.py

@@ -0,0 +1,79 @@
+"""Simple test showing the lazy client fix works."""
+
+import os
+from unittest.mock import MagicMock, patch
+import sys
+
+
+def test_use_does_not_create_client_immediately():
+    """Test that replicate.use() does not create a client until the model is called."""
+    sys.path.insert(0, 'src')
+    
+    # Clear any existing token to simulate the original error condition
+    with patch.dict(os.environ, {}, clear=True):
+        with patch.dict(sys.modules, {"cog": None}):
+            try:
+                import replicate
+                # This should work now - no client is created yet
+                model = replicate.use("test/model")
+                
+                # Verify we got a Function object back
+                from replicate.lib._predictions_use import Function
+                assert isinstance(model, Function)
+                print("✓ replicate.use() works without immediate client creation")
+                
+                # Verify the client is stored as a callable (factory function)
+                assert callable(model._client)
+                print("✓ Client is stored as factory function")
+                
+            except Exception as e:
+                print(f"✗ Test failed: {e}")
+                raise
+
+
+def test_client_created_when_model_called():
+    """Test that the client is created when the model is called."""
+    sys.path.insert(0, 'src')
+    
+    # Mock the client creation to track when it happens
+    created_clients = []
+    
+    def track_client_creation(*args, **kwargs):
+        client = MagicMock()
+        client.bearer_token = kwargs.get('bearer_token', 'no-token')
+        created_clients.append(client)
+        return client
+    
+    # Mock cog to provide a token
+    mock_scope = MagicMock()
+    mock_scope.context.items.return_value = [("REPLICATE_API_TOKEN", "cog-token")]
+    mock_cog = MagicMock()
+    mock_cog.current_scope.return_value = mock_scope
+    
+    with patch.dict(os.environ, {}, clear=True):
+        with patch.dict(sys.modules, {"cog": mock_cog}):
+            with patch('replicate._module_client._ModuleClient', side_effect=track_client_creation):
+                import replicate
+                
+                # Create model function - should not create client yet
+                model = replicate.use("test/model")
+                assert len(created_clients) == 0
+                print("✓ No client created when use() is called")
+                
+                # Try to call the model - this should create a client
+                try:
+                    model(prompt="test")
+                except Exception:
+                    # Expected to fail due to mocking, but client should be created
+                    pass
+                
+                # Verify client was created with the cog token
+                assert len(created_clients) == 1
+                assert created_clients[0].bearer_token == "cog-token"
+                print("✓ Client created with correct token when model is called")
+
+
+if __name__ == "__main__":
+    test_use_does_not_create_client_immediately()
+    test_client_created_when_model_called()
+    print("\n✓ All tests passed! The lazy client fix works correctly.")