refactor(llm): move provider transport helpers into LLMProvider

Co-authored-by: openhands <openhands@all-hands.dev>

Author: openhands-agent

openhands-sdk/openhands/sdk/llm/llm.py

@@ -869,15 +869,15 @@ def _one_attempt(**retry_kwargs) -> ResponsesAPIResponse:
                     typed_input: ResponseInputParam | str = (
                         cast(ResponseInputParam, input_items) if input_items else ""
                     )
-                    api_key_value = self._get_litellm_api_key_value()
                     provider_info = self._get_litellm_provider_info()
 
                     ret = litellm_responses(
-                        **provider_info.as_litellm_call_kwargs(),
+                        **provider_info.as_litellm_call_kwargs(
+                            api_key=self._get_api_key_value()
+                        ),
                         input=typed_input,
                         instructions=instructions,
                         tools=resp_tools,
-                        api_key=api_key_value,
                         api_base=self.base_url,
                         api_version=self.api_version,
                         timeout=self.timeout,
@@ -985,28 +985,21 @@ def _one_attempt(**retry_kwargs) -> ResponsesAPIResponse:
     # =========================================================================
 
     def _get_litellm_provider_info(self) -> LLMProvider:
-        cache_key = (self.model, self.base_url)
-        if self._provider_info is None or self._provider_info_cache_key != cache_key:
-            self._provider_info = LLMProvider.from_model(
-                model=self.model, api_base=self.base_url
-            )
-            self._provider_info_cache_key = cache_key
+        self._provider_info, self._provider_info_cache_key = LLMProvider.resolve_cached(
+            model=self.model,
+            api_base=self.base_url,
+            cached_provider=self._provider_info,
+            cached_key=self._provider_info_cache_key,
+        )
+        assert self._provider_info is not None
         return self._provider_info
 
-    def _get_litellm_api_key_value(self) -> str | None:
-        api_key_value: str | None = None
-        if self.api_key:
-            assert isinstance(self.api_key, SecretStr)
-            api_key_value = self.api_key.get_secret_value()
-
-        # LiteLLM treats api_key for Bedrock as an AWS bearer token.
-        # Passing a non-Bedrock key (e.g. OpenAI/Anthropic) can cause Bedrock
-        # to reject the request with an "Invalid API Key format" error.
-        # For IAM/SigV4 auth (the default Bedrock path), do not forward api_key.
-        if api_key_value is not None and self._get_litellm_provider_info().is_bedrock:
+    def _get_api_key_value(self) -> str | None:
+        if self.api_key is None:
             return None
 
-        return api_key_value
+        assert isinstance(self.api_key, SecretStr)
+        return self.api_key.get_secret_value()
 
     def _transport_call(
         self,
@@ -1041,13 +1034,13 @@ def _transport_call(
                     category=DeprecationWarning,
                     message="Accessing the 'model_fields' attribute.*",
                 )
-                api_key_value = self._get_litellm_api_key_value()
                 provider_info = self._get_litellm_provider_info()
 
                 # Some providers need renames handled in _normalize_call_kwargs.
                 ret = litellm_completion(
-                    **provider_info.as_litellm_call_kwargs(),
-                    api_key=api_key_value,
+                    **provider_info.as_litellm_call_kwargs(
+                        api_key=self._get_api_key_value()
+                    ),
                     api_base=self.base_url,
                     api_version=self.api_version,
                     timeout=self.timeout,

openhands-sdk/openhands/sdk/llm/utils/litellm_provider.py

@@ -25,6 +25,24 @@ class LLMProvider:
     name: str | None
     resolved_api_base: str | None
 
+    @staticmethod
+    def cache_key(*, model: str, api_base: str | None) -> tuple[str, str | None]:
+        return (model, api_base)
+
+    @classmethod
+    def resolve_cached(
+        cls,
+        *,
+        model: str,
+        api_base: str | None,
+        cached_provider: LLMProvider | None,
+        cached_key: tuple[str, str | None] | None,
+    ) -> tuple[LLMProvider, tuple[str, str | None]]:
+        cache_key = cls.cache_key(model=model, api_base=api_base)
+        if cached_provider is not None and cached_key == cache_key:
+            return cached_provider, cache_key
+        return cls.from_model(model=model, api_base=api_base), cache_key
+
     @classmethod
     def from_model(cls, *, model: str, api_base: str | None) -> LLMProvider:
         """Parse a model string using LiteLLM's provider inference logic."""
@@ -81,10 +99,22 @@ def canonical_name(self) -> str:
     def is_bedrock(self) -> bool:
         return self.name == "bedrock"
 
-    def as_litellm_call_kwargs(self) -> dict[str, str]:
+    def api_key_for_litellm(self, api_key: str | None) -> str | None:
+        # LiteLLM treats api_key for Bedrock as an AWS bearer token.
+        # Passing a non-Bedrock key (e.g. OpenAI/Anthropic) can cause Bedrock
+        # to reject the request with an "Invalid API Key format" error.
+        # For IAM/SigV4 auth (the default Bedrock path), do not forward api_key.
+        if api_key is not None and self.is_bedrock:
+            return None
+        return api_key
+
+    def as_litellm_call_kwargs(self, *, api_key: str | None = None) -> dict[str, str]:
         kwargs = {"model": self.model}
         if self.name is not None:
             kwargs["custom_llm_provider"] = self.name
+        normalized_api_key = self.api_key_for_litellm(api_key)
+        if normalized_api_key is not None:
+            kwargs["api_key"] = normalized_api_key
         return kwargs
 
     def infer_api_base(self) -> str | None:

tests/sdk/llm/test_api_key_validation.py

@@ -85,7 +85,9 @@ def test_bedrock_model_with_api_key_not_forwarded_to_litellm():
         api_key=SecretStr("sk-ant-not-a-bedrock-key"),
     )
     assert llm.api_key is not None
-    assert llm._get_litellm_api_key_value() is None
+    assert isinstance(llm.api_key, SecretStr)
+    provider = llm._get_litellm_provider_info()
+    assert provider.api_key_for_litellm(llm.api_key.get_secret_value()) is None
 
 
 def test_non_bedrock_model_with_valid_key():

tests/sdk/llm/test_litellm_provider.py

@@ -30,6 +30,47 @@ def test_llm_provider_parses_bedrock_model():
     )
 
 
+def test_llm_provider_strips_api_key_for_bedrock_calls():
+    provider = LLMProvider.from_model(
+        model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
+        api_base=None,
+    )
+
+    assert provider.api_key_for_litellm("sk-ant-not-a-bedrock-key") is None
+    assert provider.as_litellm_call_kwargs(api_key="sk-ant-not-a-bedrock-key") == {
+        "model": "anthropic.claude-3-5-sonnet-20241022-v2:0",
+        "custom_llm_provider": "bedrock",
+    }
+
+
+def test_llm_provider_reuses_cached_instance_for_same_request():
+    provider = LLMProvider.from_model(model="gpt-4o", api_base=None)
+
+    cached_provider, cached_key = LLMProvider.resolve_cached(
+        model="gpt-4o",
+        api_base=None,
+        cached_provider=provider,
+        cached_key=LLMProvider.cache_key(model="gpt-4o", api_base=None),
+    )
+
+    assert cached_provider is provider
+    assert cached_key == ("gpt-4o", None)
+
+
+def test_llm_provider_refreshes_cached_instance_when_request_changes():
+    provider = LLMProvider.from_model(model="gpt-4o", api_base=None)
+
+    refreshed_provider, refreshed_key = LLMProvider.resolve_cached(
+        model="proxy/test-renamed-model",
+        api_base="http://localhost:8000",
+        cached_provider=provider,
+        cached_key=LLMProvider.cache_key(model="gpt-4o", api_base=None),
+    )
+
+    assert refreshed_provider is not provider
+    assert refreshed_key == ("proxy/test-renamed-model", "http://localhost:8000")
+
+
 def test_llm_provider_handles_unknown_model_without_provider():
     provider = LLMProvider.from_model(model="unknown-model", api_base=None)
 

tests/sdk/llm/test_llm.py

@@ -390,6 +390,31 @@ def test_llm_responses_forwards_extra_headers_to_litellm(mock_responses):
     assert kwargs.get("extra_headers") == headers
 
 
+@patch("openhands.sdk.llm.llm.litellm_completion")
+def test_llm_completion_does_not_forward_bedrock_api_key(mock_completion):
+    mock_response = create_mock_litellm_response("ok")
+    mock_completion.return_value = mock_response
+
+    llm = LLM(
+        usage_id="test-llm",
+        model="us.anthropic.claude-3-sonnet-20240229-v1:0",
+        api_key=SecretStr("sk-ant-not-a-bedrock-key"),
+        num_retries=0,
+    )
+
+    provider_info = llm._get_litellm_provider_info()
+
+    messages = [Message(role="user", content=[TextContent(text="Hi")])]
+    _ = llm.completion(messages=messages)
+
+    assert mock_completion.call_count == 1
+    _, kwargs = mock_completion.call_args
+    assert kwargs["model"] == provider_info.model
+    if provider_info.name is not None:
+        assert kwargs["custom_llm_provider"] == provider_info.name
+    assert "api_key" not in kwargs
+
+
 @patch("openhands.sdk.llm.llm.litellm_completion")
 def test_llm_model_copy_recomputes_transport_provider_for_proxy_alias(mock_completion):
     mock_response = create_mock_litellm_response("ok")