feat(llm): improve environment variable handling for API keys in LiteLLMCaller

Author: garland3

backend/modules/llm/litellm_caller.py

@@ -27,7 +27,14 @@
 
 
 class LiteLLMCaller:
-    """Clean interface for all LLM calling patterns using LiteLLM."""
+    """Clean interface for all LLM calling patterns using LiteLLM.
+
+    Note: this class may set provider-specific LLM API key environment
+    variables (for example ``OPENAI_API_KEY``) to maintain compatibility
+    with LiteLLM's internal provider detection. These mutations are
+    best-effort only and are not intended to provide strong isolation
+    guarantees in multi-tenant or highly concurrent environments.
+    """
 
     def __init__(self, llm_config=None, debug_mode: bool = False):
         """Initialize with optional config dependency injection."""
@@ -91,20 +98,34 @@ def _get_model_kwargs(self, model_name: str, temperature: Optional[float] = None
             kwargs["api_key"] = api_key
 
             # Additionally set provider-specific env vars for LiteLLM's internal logic
+            def _set_env_var_if_needed(env_key: str, value: str) -> None:
+                existing = os.environ.get(env_key)
+                if existing is None:
+                    os.environ[env_key] = value
+                elif existing != value:
+                    logger.warning(
+                        "Overwriting existing environment variable %s for model %s", 
+                        env_key,
+                        model_name,
+                    )
+                    os.environ[env_key] = value
+
             if "openrouter" in model_config.model_url:
-                os.environ["OPENROUTER_API_KEY"] = api_key
+                _set_env_var_if_needed("OPENROUTER_API_KEY", api_key)
             elif "openai" in model_config.model_url:
-                os.environ["OPENAI_API_KEY"] = api_key
+                _set_env_var_if_needed("OPENAI_API_KEY", api_key)
             elif "anthropic" in model_config.model_url:
-                os.environ["ANTHROPIC_API_KEY"] = api_key
+                _set_env_var_if_needed("ANTHROPIC_API_KEY", api_key)
             elif "google" in model_config.model_url:
-                os.environ["GOOGLE_API_KEY"] = api_key
+                _set_env_var_if_needed("GOOGLE_API_KEY", api_key)
             elif "cerebras" in model_config.model_url:
-                os.environ["CEREBRAS_API_KEY"] = api_key
+                _set_env_var_if_needed("CEREBRAS_API_KEY", api_key)
             else:
-                # Custom endpoint - set OPENAI_API_KEY as fallback
-                # (most custom endpoints are OpenAI-compatible)
-                os.environ["OPENAI_API_KEY"] = api_key
+                # Custom endpoint - set OPENAI_API_KEY as fallback for
+                # OpenAI-compatible endpoints. This is a heuristic and
+                # only updates the env var if it is unset or already
+                # matches the same value.
+                _set_env_var_if_needed("OPENAI_API_KEY", api_key)
 
         # Set custom API base for non-standard endpoints
         if hasattr(model_config, 'model_url') and model_config.model_url:

backend/tests/test_llm_env_expansion.py

@@ -33,6 +33,9 @@ def test_litellm_caller_resolves_api_key_env_var(self, monkeypatch):
         import os
         assert os.environ.get("OPENAI_API_KEY") == "sk-test-12345"
 
+        # Cleanup to avoid leaking into other tests
+        monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+
     def test_litellm_caller_raises_on_missing_api_key_env_var(self):
         """LiteLLMCaller should raise ValueError when api_key env var is missing."""
         # Create LLM config with missing env var in api_key
@@ -53,7 +56,7 @@ def test_litellm_caller_raises_on_missing_api_key_env_var(self):
         with pytest.raises(ValueError, match="Environment variable 'MISSING_OPENAI_KEY' is not set"):
             caller._get_model_kwargs("test-model")
 
-    def test_litellm_caller_handles_literal_api_key(self):
+    def test_litellm_caller_handles_literal_api_key(self, monkeypatch):
         """LiteLLMCaller should handle literal api_key values."""
         # Create LLM config with literal api_key
         llm_config = LLMConfig(
@@ -71,11 +74,14 @@ def test_litellm_caller_handles_literal_api_key(self):
 
         # Get model kwargs - this should work without errors
         _ = caller._get_model_kwargs("test-model")
-        
+
         # Verify that the environment variable was set
         import os
         assert os.environ.get("OPENAI_API_KEY") == "sk-literal-key-12345"
 
+        # Cleanup to avoid leaking into other tests
+        monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+
     def test_litellm_caller_resolves_extra_headers_env_vars(self, monkeypatch):
         """LiteLLMCaller should resolve environment variables in extra_headers."""
         monkeypatch.setenv("TEST_REFERER", "https://myapp.com")
@@ -191,6 +197,9 @@ def test_custom_endpoint_with_env_var_api_key(self, monkeypatch):
         import os
         assert os.environ.get("OPENAI_API_KEY") == "sk-custom-12345"
 
+        # Cleanup to avoid leaking into other tests
+        monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+
     def test_custom_endpoint_with_literal_api_key(self):
         """Custom endpoint should pass api_key in kwargs when using literal value."""
         # Create LLM config for custom endpoint with literal api_key
@@ -218,6 +227,90 @@ def test_custom_endpoint_with_literal_api_key(self):
         assert "api_base" in model_kwargs
         assert model_kwargs["api_base"] == "https://custom-llm.example.com/v1"
 
+    def test_openai_env_not_overwritten_if_same_value(self, monkeypatch):
+        """OPENAI_API_KEY is left as-is when value matches."""
+        # Pre-set env to a specific value
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-same")
+
+        llm_config = LLMConfig(
+            models={
+                "openai-model": ModelConfig(
+                    model_name="gpt-4",
+                    model_url="https://api.openai.com/v1",
+                    api_key="sk-openai-same",
+                )
+            }
+        )
+
+        caller = LiteLLMCaller(llm_config, debug_mode=True)
+        model_kwargs = caller._get_model_kwargs("openai-model")
+
+        import os
+        # Still should have correct key in kwargs
+        assert model_kwargs["api_key"] == "sk-openai-same"
+        # Env var should remain the same
+        assert os.environ.get("OPENAI_API_KEY") == "sk-openai-same"
+
+    def test_openai_env_overwritten_with_warning(self, monkeypatch, caplog):
+        """OPENAI_API_KEY overwrite should occur with a warning when value differs."""
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-original")
+
+        llm_config = LLMConfig(
+            models={
+                "openai-model": ModelConfig(
+                    model_name="gpt-4",
+                    model_url="https://api.openai.com/v1",
+                    api_key="sk-openai-new",
+                )
+            }
+        )
+
+        caller = LiteLLMCaller(llm_config, debug_mode=True)
+
+        with caplog.at_level("WARNING"):
+            model_kwargs = caller._get_model_kwargs("openai-model")
+
+        import os
+        # kwargs should use the new key
+        assert model_kwargs["api_key"] == "sk-openai-new"
+        # Env var should be overwritten to the new value
+        assert os.environ.get("OPENAI_API_KEY") == "sk-openai-new"
+        # A warning about overwriting should be logged
+        assert any("Overwriting existing environment variable OPENAI_API_KEY" in rec.getMessage() for rec in caplog.records)
+
+    def test_openai_and_custom_models_resolved_in_succession(self, monkeypatch):
+        """Sequence of OpenAI then custom endpoint should keep last key in env while kwargs stay correct."""
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-preexisting")
+
+        llm_config = LLMConfig(
+            models={
+                "openai-model": ModelConfig(
+                    model_name="gpt-4",
+                    model_url="https://api.openai.com/v1",
+                    api_key="sk-openai-1",
+                ),
+                "custom-model": ModelConfig(
+                    model_name="custom-model-name",
+                    model_url="https://custom-llm.example.com/v1",
+                    api_key="sk-custom-2",
+                ),
+            }
+        )
+
+        caller = LiteLLMCaller(llm_config, debug_mode=True)
+
+        # First resolve OpenAI model
+        openai_kwargs = caller._get_model_kwargs("openai-model")
+        # Then resolve custom model
+        custom_kwargs = caller._get_model_kwargs("custom-model")
+
+        import os
+        # kwargs should always reflect model-specific keys
+        assert openai_kwargs["api_key"] == "sk-openai-1"
+        assert custom_kwargs["api_key"] == "sk-custom-2"
+        # Env var ends up with the last key used (custom model)
+        assert os.environ.get("OPENAI_API_KEY") == "sk-custom-2"
+
     def test_custom_endpoint_with_extra_headers(self, monkeypatch):
         """Custom endpoint should handle extra_headers correctly."""
         monkeypatch.setenv("CUSTOM_API_KEY", "sk-custom-auth")
@@ -256,7 +349,7 @@ def test_custom_endpoint_with_extra_headers(self, monkeypatch):
         # Verify api_base is set
         assert "api_base" in model_kwargs
 
-    def test_known_providers_still_get_api_key_in_kwargs(self):
+    def test_known_providers_still_get_api_key_in_kwargs(self, monkeypatch):
         """Verify that known providers also get api_key in kwargs (backward compatibility)."""
         # Test OpenAI
         llm_config = LLMConfig(
@@ -275,6 +368,9 @@ def test_known_providers_still_get_api_key_in_kwargs(self):
         assert "api_key" in model_kwargs
         assert model_kwargs["api_key"] == "sk-openai-test"
 
+        # cleanup any env var potentially set by implementation
+        monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+
         # Test OpenRouter
         llm_config = LLMConfig(
             models={
@@ -291,3 +387,6 @@ def test_known_providers_still_get_api_key_in_kwargs(self):
         # OpenRouter should get api_key in kwargs
         assert "api_key" in model_kwargs
         assert model_kwargs["api_key"] == "sk-or-test"
+
+        # cleanup any env var potentially set by implementation
+        monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)