feat: improve cog integration with error logging and better organization

- Add debug-level logging for cog.current_scope() failures
- Move cog logic to dedicated lib/cog.py module for better organization
- Update imports and tests to use new module structure

Addresses @dgellow's review suggestions:
1. Added error logging at debug level for better troubleshooting
2. Moved cog-specific code to lib/cog.py to make it obvious it's custom code
   and reduce risk of git conflicts

Co-authored-by: Zeke Sikelianos <[email protected]>

Author: claude[bot]

src/replicate/_client.py

@@ -23,6 +23,7 @@
 
 from replicate.lib._files import FileEncodingStrategy
 from replicate.lib._predictions_run import Model, Version, ModelVersionIdentifier
+from replicate.lib.cog import get_api_token_from_environment
 from replicate.types.prediction_create_params import PredictionCreateParamsWithoutVersion
 
 from . import _exceptions
@@ -77,37 +78,6 @@
 ]
 
 
-def _get_api_token_from_environment() -> str | None:
-    """Get API token from cog current scope if available, otherwise from environment."""
-    try:
-        import cog  # type: ignore[import-untyped, import-not-found]
-        
-        # Get the current scope - this might return None or raise an exception
-        scope = getattr(cog, "current_scope", lambda: None)()
-        if scope is None:
-            return os.environ.get("REPLICATE_API_TOKEN")
-            
-        # Get the context from the scope
-        context = getattr(scope, "context", None)
-        if context is None:
-            return os.environ.get("REPLICATE_API_TOKEN")
-            
-        # Get the items method and call it
-        items_method = getattr(context, "items", None)
-        if not callable(items_method):
-            return os.environ.get("REPLICATE_API_TOKEN")
-            
-        # Iterate through context items looking for the API token
-        items = cast(Iterator[tuple[Any, Any]], items_method())
-        for key, value in items:
-            if str(key).upper() == "REPLICATE_API_TOKEN":
-                return str(value) if value is not None else value
-                
-    except Exception:  # Catch all exceptions to ensure robust fallback
-        pass
-
-    return os.environ.get("REPLICATE_API_TOKEN")
-
 
 class Replicate(SyncAPIClient):
     # client options
@@ -141,7 +111,7 @@ def __init__(
         This automatically infers the `bearer_token` argument from the `REPLICATE_API_TOKEN` environment variable if it is not provided.
         """
         if bearer_token is None:
-            bearer_token = _get_api_token_from_environment()
+            bearer_token = get_api_token_from_environment()
         if bearer_token is None:
             raise ReplicateError(
                 "The bearer_token client option must be set either by passing bearer_token to the client or by setting the REPLICATE_API_TOKEN environment variable"
@@ -452,7 +422,7 @@ def __init__(
         This automatically infers the `bearer_token` argument from the `REPLICATE_API_TOKEN` environment variable if it is not provided.
         """
         if bearer_token is None:
-            bearer_token = _get_api_token_from_environment()
+            bearer_token = get_api_token_from_environment()
         if bearer_token is None:
             raise ReplicateError(
                 "The bearer_token client option must be set either by passing bearer_token to the client or by setting the REPLICATE_API_TOKEN environment variable"

src/replicate/lib/cog.py

@@ -0,0 +1,58 @@
+"""Integration with cog's current_scope functionality."""
+
+import os
+import logging
+from typing import Iterator, Any, cast
+
+logger = logging.getLogger("replicate")
+
+
+def get_api_token_from_current_scope() -> str | None:
+    """Get API token from cog's current_scope if available, otherwise return None.
+    
+    This function attempts to retrieve the API token from cog's current_scope context.
+    It gracefully handles all errors and returns None if cog is not available or if
+    any part of the retrieval process fails.
+    
+    Returns:
+        str | None: The API token from cog's current_scope, or None if not available.
+    """
+    try:
+        import cog  # type: ignore[import-untyped, import-not-found]
+        
+        # Get the current scope - this might return None or raise an exception
+        scope = getattr(cog, "current_scope", lambda: None)()
+        if scope is None:
+            return None
+            
+        # Get the context from the scope
+        context = getattr(scope, "context", None)
+        if context is None:
+            return None
+            
+        # Get the items method and call it
+        items_method = getattr(context, "items", None)
+        if not callable(items_method):
+            return None
+            
+        # Iterate through context items looking for the API token
+        items = cast(Iterator[tuple[Any, Any]], items_method())
+        for key, value in items:
+            if str(key).upper() == "REPLICATE_API_TOKEN":
+                return str(value) if value is not None else value
+                
+    except Exception as e:  # Catch all exceptions to ensure robust fallback
+        logger.debug("Failed to retrieve API token from cog.current_scope(): %s", e)
+
+    return None
+
+
+def get_api_token_from_environment() -> str | None:
+    """Get API token from cog current scope if available, otherwise from environment."""
+    # Try to get token from cog's current_scope first
+    token = get_api_token_from_current_scope()
+    if token is not None:
+        return token
+    
+    # Fall back to environment variable
+    return os.environ.get("REPLICATE_API_TOKEN")

tests/test_current_scope.py

@@ -7,20 +7,20 @@
 import pytest
 
 from replicate import Replicate, AsyncReplicate
-from replicate._client import _get_api_token_from_environment
+from replicate.lib.cog import get_api_token_from_environment
 from replicate._exceptions import ReplicateError
 
 
 class TestGetApiTokenFromEnvironment:
-    """Test the _get_api_token_from_environment function."""
+    """Test the get_api_token_from_environment function."""
 
     def test_cog_no_current_scope_method_falls_back_to_env(self):
         """Test fallback when cog exists but has no current_scope method."""
         mock_cog = mock.MagicMock()
         del mock_cog.current_scope  # Remove the method
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "env-token"
 
     def test_cog_current_scope_returns_none_falls_back_to_env(self):
@@ -29,7 +29,7 @@ def test_cog_current_scope_returns_none_falls_back_to_env(self):
         mock_cog.current_scope.return_value = None
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "env-token"
 
     def test_cog_scope_no_context_attr_falls_back_to_env(self):
@@ -40,7 +40,7 @@ def test_cog_scope_no_context_attr_falls_back_to_env(self):
         mock_cog.current_scope.return_value = mock_scope
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "env-token"
 
     def test_cog_scope_context_not_dict_falls_back_to_env(self):
@@ -51,7 +51,7 @@ def test_cog_scope_context_not_dict_falls_back_to_env(self):
         mock_cog.current_scope.return_value = mock_scope
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "env-token"
 
     def test_cog_scope_no_replicate_api_token_key_falls_back_to_env(self):
@@ -62,7 +62,7 @@ def test_cog_scope_no_replicate_api_token_key_falls_back_to_env(self):
         mock_cog.current_scope.return_value = mock_scope
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "env-token"
 
     def test_cog_scope_replicate_api_token_valid_string(self):
@@ -73,7 +73,7 @@ def test_cog_scope_replicate_api_token_valid_string(self):
         mock_cog.current_scope.return_value = mock_scope
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "cog-token"
 
     def test_cog_scope_replicate_api_token_case_insensitive(self):
@@ -84,7 +84,7 @@ def test_cog_scope_replicate_api_token_case_insensitive(self):
         mock_cog.current_scope.return_value = mock_scope
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "cog-token"
 
     def test_cog_scope_replicate_api_token_empty_string(self):
@@ -95,7 +95,7 @@ def test_cog_scope_replicate_api_token_empty_string(self):
         mock_cog.current_scope.return_value = mock_scope
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == ""  # Should return empty string, not env token
 
     def test_cog_scope_replicate_api_token_none(self):
@@ -106,7 +106,7 @@ def test_cog_scope_replicate_api_token_none(self):
         mock_cog.current_scope.return_value = mock_scope
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token is None  # Should return None, not env token
 
     def test_cog_current_scope_raises_exception_falls_back_to_env(self):
@@ -115,28 +115,28 @@ def test_cog_current_scope_raises_exception_falls_back_to_env(self):
         mock_cog.current_scope.side_effect = RuntimeError("Scope error")
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": mock_cog}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "env-token"
 
     def test_no_env_token_returns_none(self):
         """Test that None is returned when no environment token is set and cog unavailable."""
         with mock.patch.dict(os.environ, {}, clear=True):  # Clear all env vars
             with mock.patch.dict(sys.modules, {"cog": None}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token is None
 
     def test_env_token_empty_string(self):
         """Test that empty string from environment is returned."""
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": ""}):
             with mock.patch.dict(sys.modules, {"cog": None}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == ""
 
     def test_env_token_valid_string(self):
         """Test that valid token from environment is returned."""
         with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
             with mock.patch.dict(sys.modules, {"cog": None}):
-                token = _get_api_token_from_environment()
+                token = get_api_token_from_environment()
                 assert token == "env-token"