fix: resolve linting and type errors in LLMWithGateway

- Fix line length (E501) errors by splitting long strings
- Fix pyright type errors with proper type assertions and overrides
- Fix method overrides to use completion/responses instead of _completion/_responses
- Update tests to use proper Message objects instead of dicts
- Fix test assertions to match actual error messages
- All 25 tests now passing

Author: ak684

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

@@ -16,6 +16,7 @@
 from openhands.sdk.llm.llm import LLM
 from openhands.sdk.logger import get_logger
 
+
 logger = get_logger(__name__)
 
 __all__ = ["LLMWithGateway"]
@@ -63,7 +64,10 @@ class LLMWithGateway(LLM):
     )
     gateway_auth_token_path: str = Field(
         default="access_token",
-        description="Dot-notation path to the token in the OAuth response (e.g., 'access_token' or 'data.token').",
+        description=(
+            "Dot-notation path to the token in the OAuth response "
+            "(e.g., 'access_token' or 'data.token')."
+        ),
     )
     gateway_auth_token_ttl: int | None = Field(
         default=None,
@@ -98,19 +102,15 @@ def model_post_init(self, __context: Any) -> None:
         self._gateway_token = None
         self._gateway_token_expiry = None
 
-    def _completion(
-        self, messages: list[dict], **kwargs
-    ) -> Any:  # Returns ModelResponse
+    def completion(self, *args, **kwargs):
         """Override to inject gateway authentication before calling LiteLLM."""
         self._prepare_gateway_call(kwargs)
-        return super()._completion(messages, **kwargs)
+        return super().completion(*args, **kwargs)
 
-    def _responses(
-        self, messages: list[dict], **kwargs
-    ) -> Any:  # Returns ResponsesAPIResponse
+    def responses(self, *args, **kwargs):
         """Override to inject gateway authentication before calling LiteLLM."""
         self._prepare_gateway_call(kwargs)
-        return super()._responses(messages, **kwargs)
+        return super().responses(*args, **kwargs)
 
     def _prepare_gateway_call(self, call_kwargs: dict[str, Any]) -> None:
         """Augment LiteLLM kwargs with gateway headers and token.
@@ -199,13 +199,13 @@ def _refresh_gateway_token(self) -> str:
         Raises:
             Exception: If token fetch fails.
         """
+        assert self.gateway_auth_url is not None, "gateway_auth_url must be set"
         method = self.gateway_auth_method.upper()
         headers = self._render_templates(self.gateway_auth_headers or {})
         body = self._render_templates(self.gateway_auth_body or {})
 
         logger.debug(
-            f"Fetching gateway token from {self.gateway_auth_url} "
-            f"(method={method})"
+            f"Fetching gateway token from {self.gateway_auth_url} (method={method})"
         )
 
         try:
@@ -243,9 +243,7 @@ def _refresh_gateway_token(self) -> str:
         self._gateway_token = token_value.strip()
         self._gateway_token_expiry = time.time() + max(ttl_seconds, 1.0)
 
-        logger.info(
-            f"Gateway token refreshed successfully (expires in {ttl_seconds}s)"
-        )
+        logger.info(f"Gateway token refreshed successfully (expires in {ttl_seconds}s)")
         return self._gateway_token
 
     def _render_templates(self, value: Any) -> Any:
@@ -310,24 +308,28 @@ def _extract_from_path(payload: Any, path: str) -> Any:
                 current = current.get(part)
                 if current is None:
                     raise ValueError(
-                        f'Key "{part}" not found in response while traversing path "{path}".'
+                        f'Key "{part}" not found in response '
+                        f'while traversing path "{path}".'
                     )
             elif isinstance(current, list):
                 try:
                     index = int(part)
                 except (ValueError, TypeError):
                     raise ValueError(
-                        f'Invalid list index "{part}" while traversing response path "{path}".'
+                        f'Invalid list index "{part}" '
+                        f'while traversing response path "{path}".'
                     ) from None
                 try:
                     current = current[index]
                 except (IndexError, TypeError):
                     raise ValueError(
-                        f'Index {index} out of range while traversing response path "{path}".'
+                        f"Index {index} out of range "
+                        f'while traversing response path "{path}".'
                     ) from None
             else:
                 raise ValueError(
-                    f'Cannot traverse path "{path}"; segment "{part}" not found or not accessible.'
+                    f'Cannot traverse path "{path}"; '
+                    f'segment "{part}" not found or not accessible.'
                 )
 
         return current

tests/sdk/llm/test_llm_with_gateway.py

@@ -8,7 +8,7 @@
 import pytest
 from pydantic import SecretStr
 
-from openhands.sdk.llm import LLMWithGateway
+from openhands.sdk.llm import LLMWithGateway, Message, TextContent
 from tests.conftest import create_mock_litellm_response
 
 
@@ -178,7 +178,7 @@ def test_token_fetch_missing_token(self, mock_request, gateway_llm):
         mock_response.raise_for_status = Mock()
         mock_request.return_value = mock_response
 
-        with pytest.raises(ValueError, match="did not contain token"):
+        with pytest.raises(ValueError, match="not found in response"):
             gateway_llm._ensure_gateway_token()
 
 
@@ -203,7 +203,7 @@ def test_custom_headers_injected(
         )
 
         # Make completion request
-        messages = [{"role": "user", "content": "test"}]
+        messages = [Message(role="user", content=[TextContent(text="test")])]
         gateway_llm.completion(messages)
 
         # Verify custom headers were passed
@@ -231,7 +231,7 @@ def test_gateway_token_header_injected(
         )
 
         # Make completion request
-        messages = [{"role": "user", "content": "test"}]
+        messages = [Message(role="user", content=[TextContent(text="test")])]
         gateway_llm.completion(messages)
 
         # Verify gateway token was passed
@@ -370,7 +370,7 @@ def test_full_gateway_flow(
         )
 
         # Make completion request
-        messages = [{"role": "user", "content": "Hello"}]
+        messages = [Message(role="user", content=[TextContent(text="Hello")])]
         response = llm.completion(messages)
 
         # Verify OAuth was called
@@ -387,7 +387,8 @@ def test_full_gateway_flow(
         assert headers["X-Gateway-Key"] == "gateway789"
 
         # Verify response
-        assert response.content[0].text == "Hello from gateway!"
+        assert isinstance(response.message.content[0], TextContent)
+        assert response.message.content[0].text == "Hello from gateway!"
 
     def test_gateway_disabled_when_no_config(self):
         """Test that gateway logic is skipped when not configured."""