refactor openai chat kwargs (#933)

Author: burtenshaw

src/openenv/core/llm_client.py

@@ -178,6 +178,18 @@ def __init__(
             api_key=api_key if api_key is not None else "not-needed",
         )
 
+    def _chat_completion_kwargs(
+        self, messages: list[dict[str, Any]], **kwargs: Any
+    ) -> dict[str, Any]:
+        create_kwargs: dict[str, Any] = {
+            "model": self.model,
+            "messages": messages,
+            self._tokens_param: kwargs.get("max_tokens", self.max_tokens),
+        }
+        if not self._omit_temperature:
+            create_kwargs["temperature"] = kwargs.get("temperature", self.temperature)
+        return create_kwargs
+
     async def complete(self, prompt: str, **kwargs) -> str:
         """Send a chat completion request.
 
@@ -195,13 +207,7 @@ async def complete(self, prompt: str, **kwargs) -> str:
             messages.append({"role": "system", "content": self.system_prompt})
         messages.append({"role": "user", "content": prompt})
 
-        call_kwargs: dict[str, Any] = {
-            "model": self.model,
-            "messages": messages,
-            self._tokens_param: kwargs.get("max_tokens", self.max_tokens),
-        }
-        if not self._omit_temperature:
-            call_kwargs["temperature"] = kwargs.get("temperature", self.temperature)
+        call_kwargs = self._chat_completion_kwargs(messages, **kwargs)
         response = await self._client.chat.completions.create(**call_kwargs)
         return response.choices[0].message.content or ""
 
@@ -211,13 +217,7 @@ async def complete_with_tools(
         tools: list[dict[str, Any]],
         **kwargs: Any,
     ) -> LLMResponse:
-        create_kwargs: dict[str, Any] = {
-            "model": self.model,
-            "messages": messages,
-            self._tokens_param: kwargs.get("max_tokens", self.max_tokens),
-        }
-        if not self._omit_temperature:
-            create_kwargs["temperature"] = kwargs.get("temperature", self.temperature)
+        create_kwargs = self._chat_completion_kwargs(messages, **kwargs)
         openai_tools = _mcp_tools_to_openai(tools)
         if openai_tools:
             create_kwargs["tools"] = openai_tools

tests/core/test_llm_client.py

@@ -236,6 +236,34 @@ async def test_no_tool_calls(self):
         assert result.content == "Hello there"
         assert result.tool_calls == []
 
+    @pytest.mark.asyncio
+    async def test_kwargs_override(self):
+        """Keyword arguments override default temperature and max_tokens."""
+        mock_openai = MagicMock()
+        mock_msg = MagicMock()
+        mock_msg.content = "Hello there"
+        mock_msg.tool_calls = None
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock()]
+        mock_response.choices[0].message = mock_msg
+        mock_openai.chat.completions.create = AsyncMock(return_value=mock_response)
+
+        with patch("openenv.core.llm_client.AsyncOpenAI", return_value=mock_openai):
+            client = OpenAIClient("http://localhost", 8000, model="gpt-4")
+            await client.complete_with_tools(
+                [{"role": "user", "content": "hi"}],
+                [],
+                temperature=0.9,
+                max_tokens=100,
+            )
+
+        mock_openai.chat.completions.create.assert_called_once_with(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "hi"}],
+            temperature=0.9,
+            max_tokens=100,
+        )
+
     @pytest.mark.asyncio
     async def test_with_tool_calls(self):
         """Response with tool calls are parsed into ToolCall objects."""