Better code reuse and tests for LiteLLM params

Author: scosman

.cursorrules

@@ -1,4 +1,9 @@
+ - call me "boss"
  - Always assume pydantic 2 (not pydantic 1)
- - Always use pytest for tests
  - The project supports Python 3.10 and above
+ - When writing tests:
+   1) Always use pytest for tests in python code
+   2) assume an appriopirate test file already exists, find it, and suggest tests get appended to that file. If no such file exists, ask me before assuming a new test file is the correct route.
+   3) Test brevity is important. Use approaches for re-use and brevity including using fixtures for repeated code, and using pytest parameterize for similar tests
+
 

libs/core/kiln_ai/adapters/model_adapters/litellm_adapter.py

@@ -74,15 +74,11 @@ async def _run(self, input: Dict | str) -> RunOutput:
                 raise ValueError("cot_prompt is required for cot_two_call strategy")
             messages.append({"role": "system", "content": cot_prompt})
 
-            # First call for chain of thought
-            cot_response = await litellm.acompletion(
-                model=self.litellm_model_id(),
-                messages=messages,
-                api_base=self._api_base,
-                headers=self._headers,
-                # TODO P1 - remove type ignore
-                **self._additional_body_options,  # type: ignore
+            # First call for chain of thought - No logprobs as only needed for final answer
+            completion_kwargs = await self.build_completion_kwargs(
+                provider, messages, None
             )
+            cot_response = await litellm.acompletion(**completion_kwargs)
             if (
                 not isinstance(cot_response, ModelResponse)
                 or not cot_response.choices
@@ -103,32 +99,10 @@ async def _run(self, input: Dict | str) -> RunOutput:
                 ]
             )
 
-        # Build custom request params based on model provider
-        extra_body = self.build_extra_body(provider)
-
-        # Main completion call
-        response_format_options = await self.response_format_options()
-
-        # Merge all parameters into a single kwargs dict for litellm
-        # TODO P0 - make this shared
-        completion_kwargs = {
-            "model": self.litellm_model_id(),
-            "messages": messages,
-            "api_base": self._api_base,
-            "headers": self._headers,
-            **extra_body,
-            **self._additional_body_options,
-        }
-
-        # Add logprobs if requested
-        if self.base_adapter_config.top_logprobs is not None:
-            completion_kwargs["logprobs"] = True
-            completion_kwargs["top_logprobs"] = self.base_adapter_config.top_logprobs
-
-        # Add response format options
-        completion_kwargs.update(response_format_options)
-
         # Make the API call using litellm
+        completion_kwargs = await self.build_completion_kwargs(
+            provider, messages, self.base_adapter_config.top_logprobs
+        )
         response = await litellm.acompletion(**completion_kwargs)
 
         if not isinstance(response, ModelResponse):
@@ -379,3 +353,31 @@ def litellm_model_id(self) -> str:
 
         self._litellm_model_id = litellm_provider_name + "/" + provider.model_id
         return self._litellm_model_id
+
+    async def build_completion_kwargs(
+        self,
+        provider: KilnModelProvider,
+        messages: list[dict[str, Any]],
+        top_logprobs: int | None,
+    ) -> dict[str, Any]:
+        extra_body = self.build_extra_body(provider)
+
+        # Merge all parameters into a single kwargs dict for litellm
+        completion_kwargs = {
+            "model": self.litellm_model_id(),
+            "messages": messages,
+            "api_base": self._api_base,
+            "headers": self._headers,
+            **extra_body,
+            **self._additional_body_options,
+        }
+
+        # Response format: json_schema, json_instructions, json_mode, function_calling, etc
+        response_format_options = await self.response_format_options()
+        completion_kwargs.update(response_format_options)
+
+        if top_logprobs is not None:
+            completion_kwargs["logprobs"] = True
+            completion_kwargs["top_logprobs"] = top_logprobs
+
+        return completion_kwargs

libs/core/kiln_ai/adapters/model_adapters/test_litellm_adapter.py

@@ -337,3 +337,61 @@ def test_litellm_model_id_unknown_provider(config, mock_task):
 
             with pytest.raises(Exception, match="Test error"):
                 adapter.litellm_model_id()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "top_logprobs,response_format,extra_body",
+    [
+        (None, {}, {}),  # Basic case
+        (5, {}, {}),  # With logprobs
+        (
+            None,
+            {"response_format": {"type": "json_object"}},
+            {},
+        ),  # With response format
+        (
+            3,
+            {"tools": [{"type": "function"}]},
+            {"reasoning_effort": 0.8},
+        ),  # Combined options
+    ],
+)
+async def test_build_completion_kwargs(
+    config, mock_task, top_logprobs, response_format, extra_body
+):
+    """Test build_completion_kwargs with various configurations"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+    mock_provider = Mock()
+    messages = [{"role": "user", "content": "Hello"}]
+
+    with (
+        patch.object(adapter, "model_provider", return_value=mock_provider),
+        patch.object(adapter, "litellm_model_id", return_value="openai/test-model"),
+        patch.object(adapter, "build_extra_body", return_value=extra_body),
+        patch.object(adapter, "response_format_options", return_value=response_format),
+    ):
+        kwargs = await adapter.build_completion_kwargs(
+            mock_provider, messages, top_logprobs
+        )
+
+    # Verify core functionality
+    assert kwargs["model"] == "openai/test-model"
+    assert kwargs["messages"] == messages
+    assert kwargs["api_base"] == config.base_url
+
+    # Verify optional parameters
+    if top_logprobs is not None:
+        assert kwargs["logprobs"] is True
+        assert kwargs["top_logprobs"] == top_logprobs
+    else:
+        assert "logprobs" not in kwargs
+        assert "top_logprobs" not in kwargs
+
+    # Verify response format is included
+    for key, value in response_format.items():
+        assert kwargs[key] == value
+
+    # Verify extra body is included
+    for key, value in extra_body.items():
+        assert kwargs[key] == value