kiln-ai 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

kiln_ai/adapters/model_adapters/litellm_adapter.py

@@ -12,15 +12,13 @@ from kiln_ai.adapters.ml_model_list import (
     StructuredOutputMode,
 )
 from kiln_ai.adapters.model_adapters.base_adapter import (
-    COT_FINAL_ANSWER_PROMPT,
     AdapterConfig,
     BaseAdapter,
     RunOutput,
     Usage,
 )
 from kiln_ai.adapters.model_adapters.litellm_config import LiteLlmConfig
-from kiln_ai.datamodel import PromptGenerators, PromptId
-from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.datamodel.task import run_config_from_run_config_properties
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 logger = logging.getLogger(__name__)
@@ -31,7 +29,6 @@ class LiteLlmAdapter(BaseAdapter):
         self,
         config: LiteLlmConfig,
         kiln_task: datamodel.Task,
-        prompt_id: PromptId | None = None,
         base_adapter_config: AdapterConfig | None = None,
     ):
         self.config = config
@@ -40,11 +37,10 @@ class LiteLlmAdapter(BaseAdapter):
         self._headers = config.default_headers
         self._litellm_model_id: str | None = None
 
-        run_config = RunConfig(
+        # Create a RunConfig, adding the task to the RunConfigProperties
+        run_config = run_config_from_run_config_properties(
             task=kiln_task,
-            model_name=config.model_name,
-            model_provider_name=config.provider_name,
-            prompt_id=prompt_id or PromptGenerators.SIMPLE,
+            run_config_properties=config.run_config_properties,
         )
 
         super().__init__(
@@ -57,79 +53,69 @@ class LiteLlmAdapter(BaseAdapter):
         if not provider.model_id:
             raise ValueError("Model ID is required for OpenAI compatible models")
 
-        intermediate_outputs: dict[str, str] = {}
-        prompt = self.build_prompt()
-        user_msg = self.prompt_builder.build_user_message(input)
-        messages = [
-            {"role": "system", "content": prompt},
-            {"role": "user", "content": user_msg},
-        ]
-
-        run_strategy, cot_prompt = self.run_strategy()
-
-        if run_strategy == "cot_as_message":
-            # Used for reasoning-capable models that can output thinking and structured format
-            if not cot_prompt:
-                raise ValueError("cot_prompt is required for cot_as_message strategy")
-            messages.append({"role": "system", "content": cot_prompt})
-        elif run_strategy == "cot_two_call":
-            if not cot_prompt:
-                raise ValueError("cot_prompt is required for cot_two_call strategy")
-            messages.append({"role": "system", "content": cot_prompt})
-
-            # First call for chain of thought
-            # No response format as this request is for "thinking" in plain text
-            # No logprobs as only needed for final answer
+        chat_formatter = self.build_chat_formatter(input)
+
+        prior_output = None
+        prior_message = None
+        response = None
+        turns = 0
+        while True:
+            turns += 1
+            if turns > 10:
+                raise RuntimeError(
+                    "Too many turns. Stopping iteration to avoid using too many tokens."
+                )
+
+            turn = chat_formatter.next_turn(prior_output)
+            if turn is None:
+                break
+
+            skip_response_format = not turn.final_call
+            all_messages = chat_formatter.message_dicts()
             completion_kwargs = await self.build_completion_kwargs(
-                provider, messages, None, skip_response_format=True
+                provider,
+                all_messages,
+                self.base_adapter_config.top_logprobs if turn.final_call else None,
+                skip_response_format,
             )
-            cot_response = await litellm.acompletion(**completion_kwargs)
+            response = await litellm.acompletion(**completion_kwargs)
             if (
-                not isinstance(cot_response, ModelResponse)
-                or not cot_response.choices
-                or len(cot_response.choices) == 0
-                or not isinstance(cot_response.choices[0], Choices)
+                not isinstance(response, ModelResponse)
+                or not response.choices
+                or len(response.choices) == 0
+                or not isinstance(response.choices[0], Choices)
             ):
                 raise RuntimeError(
-                    f"Expected ModelResponse with Choices, got {type(cot_response)}."
+                    f"Expected ModelResponse with Choices, got {type(response)}."
                 )
-            cot_content = cot_response.choices[0].message.content
-            if cot_content is not None:
-                intermediate_outputs["chain_of_thought"] = cot_content
-
-            messages.extend(
-                [
-                    {"role": "assistant", "content": cot_content or ""},
-                    {"role": "user", "content": COT_FINAL_ANSWER_PROMPT},
-                ]
-            )
+            prior_message = response.choices[0].message
+            prior_output = prior_message.content
 
-        # 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)
+            # Fallback: Use args of first tool call to task_response if it exists
+            if (
+                not prior_output
+                and hasattr(prior_message, "tool_calls")
+                and prior_message.tool_calls
+            ):
+                tool_call = next(
+                    (
+                        tool_call
+                        for tool_call in prior_message.tool_calls
+                        if tool_call.function.name == "task_response"
+                    ),
+                    None,
+                )
+                if tool_call:
+                    prior_output = tool_call.function.arguments
 
-        if not isinstance(response, ModelResponse):
-            raise RuntimeError(f"Expected ModelResponse, got {type(response)}.")
+            if not prior_output:
+                raise RuntimeError("No output returned from model")
 
-        # Maybe remove this? There is no error attribute on the response object.
-        # # Keeping in typesafe way as we added it for a reason, but should investigate what that was and if it still applies.
-        if hasattr(response, "error") and response.__getattribute__("error"):
-            raise RuntimeError(
-                f"LLM API returned an error: {response.__getattribute__('error')}"
-            )
+        if response is None or prior_message is None:
+            raise RuntimeError("No response returned from model")
 
-        if (
-            not response.choices
-            or len(response.choices) == 0
-            or not isinstance(response.choices[0], Choices)
-        ):
-            raise RuntimeError(
-                "No message content returned in the response from LLM API"
-            )
+        intermediate_outputs = chat_formatter.intermediate_outputs()
 
-        message = response.choices[0].message
         logprobs = (
             response.choices[0].logprobs
             if hasattr(response.choices[0], "logprobs")
@@ -143,31 +129,15 @@ class LiteLlmAdapter(BaseAdapter):
 
         # Save reasoning if it exists and was parsed by LiteLLM (or openrouter, or anyone upstream)
         if (
-            hasattr(message, "reasoning_content")
-            and message.reasoning_content
-            and len(message.reasoning_content.strip()) > 0
+            prior_message is not None
+            and hasattr(prior_message, "reasoning_content")
+            and prior_message.reasoning_content
+            and len(prior_message.reasoning_content.strip()) > 0
         ):
-            intermediate_outputs["reasoning"] = message.reasoning_content.strip()
+            intermediate_outputs["reasoning"] = prior_message.reasoning_content.strip()
 
         # the string content of the response
-        response_content = message.content
-
-        # Fallback: Use args of first tool call to task_response if it exists
-        if (
-            not response_content
-            and hasattr(message, "tool_calls")
-            and message.tool_calls
-        ):
-            tool_call = next(
-                (
-                    tool_call
-                    for tool_call in message.tool_calls
-                    if tool_call.function.name == "task_response"
-                ),
-                None,
-            )
-            if tool_call:
-                response_content = tool_call.function.arguments
+        response_content = prior_output
 
         if not isinstance(response_content, str):
             raise RuntimeError(f"response is not a string: {response_content}")
@@ -186,8 +156,9 @@ class LiteLlmAdapter(BaseAdapter):
         if not self.has_structured_output():
             return {}
 
-        provider = self.model_provider()
-        match provider.structured_output_mode:
+        structured_output_mode = self.run_config.structured_output_mode
+
+        match structured_output_mode:
             case StructuredOutputMode.json_mode:
                 return {"response_format": {"type": "json_object"}}
             case StructuredOutputMode.json_schema:
@@ -206,16 +177,20 @@ class LiteLlmAdapter(BaseAdapter):
                 # We set response_format to json_object and also set json instructions in the prompt
                 return {"response_format": {"type": "json_object"}}
             case StructuredOutputMode.default:
-                if provider.name == ModelProviderName.ollama:
+                provider_name = self.run_config.model_provider_name
+                if provider_name == ModelProviderName.ollama:
                     # Ollama added json_schema to all models: https://ollama.com/blog/structured-outputs
                     return self.json_schema_response_format()
                 else:
                     # Default to function calling -- it's older than the other modes. Higher compatibility.
                     # Strict isn't widely supported yet, so we don't use it by default unless it's OpenAI.
-                    strict = provider.name == ModelProviderName.openai
+                    strict = provider_name == ModelProviderName.openai
                     return self.tool_call_params(strict=strict)
+            case StructuredOutputMode.unknown:
+                # See above, but this case should never happen.
+                raise ValueError("Structured output mode is unknown.")
             case _:
-                raise_exhaustive_enum_error(provider.structured_output_mode)
+                raise_exhaustive_enum_error(structured_output_mode)
 
     def json_schema_response_format(self) -> dict[str, Any]:
         output_schema = self.task().output_schema()
@@ -387,6 +362,13 @@ class LiteLlmAdapter(BaseAdapter):
             "messages": messages,
             "api_base": self._api_base,
             "headers": self._headers,
+            "temperature": self.run_config.temperature,
+            "top_p": self.run_config.top_p,
+            # This drops params that are not supported by the model. Only openai params like top_p, temperature -- not litellm params like model, etc.
+            # Not all models and providers support all openai params (for example, o3 doesn't support top_p)
+            # Better to ignore them than to fail the model call.
+            # https://docs.litellm.ai/docs/completion/input
+            "drop_params": True,
             **extra_body,
             **self._additional_body_options,
         }

kiln_ai/adapters/model_adapters/litellm_config.py

@@ -1,10 +1,11 @@
 from dataclasses import dataclass, field
 
+from kiln_ai.datamodel.task import RunConfigProperties
+
 
 @dataclass
 class LiteLlmConfig:
-    model_name: str
-    provider_name: str
+    run_config_properties: RunConfigProperties
     # If set, over rides the provider-name based URL from litellm
     base_url: str | None = None
     # Headers to send with every request

kiln_ai/adapters/model_adapters/test_base_adapter.py

@@ -6,7 +6,8 @@ from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMo
 from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput
 from kiln_ai.adapters.parsers.request_formatters import request_formatter_from_id
 from kiln_ai.datamodel import Task
-from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
+from kiln_ai.datamodel.task import RunConfig, RunConfigProperties
 
 
 class MockAdapter(BaseAdapter):
@@ -37,8 +38,9 @@ def adapter(base_task):
         run_config=RunConfig(
             task=base_task,
             model_name="test_model",
-            model_provider_name="test_provider",
+            model_provider_name="openai",
             prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
         ),
     )
 
@@ -88,7 +90,7 @@ async def test_model_provider_loads_and_caches(adapter, mock_provider):
         # First call should load and cache
         provider1 = adapter.model_provider()
         assert provider1 == mock_provider
-        mock_loader.assert_called_once_with("test_model", "test_provider")
+        mock_loader.assert_called_once_with("test_model", "openai")
 
         # Second call should use cache
         mock_loader.reset_mock()
@@ -97,29 +99,30 @@ async def test_model_provider_loads_and_caches(adapter, mock_provider):
         mock_loader.assert_not_called()
 
 
-async def test_model_provider_missing_names(base_task):
+async def test_model_provider_invalid_provider_model_name(base_task):
+    """Test error when model or provider name is missing"""
+    # Test with missing model name
+    with pytest.raises(ValueError, match="Input should be"):
+        adapter = MockAdapter(
+            run_config=RunConfig(
+                task=base_task,
+                model_name="test_model",
+                model_provider_name="invalid",
+                prompt_id="simple_prompt_builder",
+            ),
+        )
+
+
+async def test_model_provider_missing_model_names(base_task):
     """Test error when model or provider name is missing"""
     # Test with missing model name
     adapter = MockAdapter(
         run_config=RunConfig(
             task=base_task,
             model_name="",
-            model_provider_name="",
-            prompt_id="simple_prompt_builder",
-        ),
-    )
-    with pytest.raises(
-        ValueError, match="model_name and model_provider_name must be provided"
-    ):
-        await adapter.model_provider()
-
-    # Test with missing provider name
-    adapter = MockAdapter(
-        run_config=RunConfig(
-            task=base_task,
-            model_name="test_model",
-            model_provider_name="",
+            model_provider_name="openai",
             prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
         ),
     )
     with pytest.raises(
@@ -138,7 +141,7 @@ async def test_model_provider_not_found(adapter):
 
         with pytest.raises(
             ValueError,
-            match="model_provider_name test_provider not found for model test_model",
+            match="not found for model test_model",
         ):
             await adapter.model_provider()
 
@@ -168,11 +171,7 @@ async def test_prompt_builder_json_instructions(
     adapter.prompt_builder = mock_prompt_builder
     adapter.model_provider_name = "openai"
     adapter.has_structured_output = MagicMock(return_value=output_schema)
-
-    # provider mock
-    provider = MagicMock()
-    provider.structured_output_mode = structured_output_mode
-    adapter.model_provider = MagicMock(return_value=provider)
+    adapter.run_config.structured_output_mode = structured_output_mode
 
     # Test
     adapter.build_prompt()
@@ -181,41 +180,6 @@ async def test_prompt_builder_json_instructions(
     )
 
 
-@pytest.mark.parametrize(
-    "cot_prompt,has_structured_output,reasoning_capable,expected",
-    [
-        # COT and normal LLM
-        ("think carefully", False, False, ("cot_two_call", "think carefully")),
-        # Structured output with thinking-capable LLM
-        ("think carefully", True, True, ("cot_as_message", "think carefully")),
-        # Structured output with normal LLM
-        ("think carefully", True, False, ("cot_two_call", "think carefully")),
-        # Basic cases - no COT
-        (None, True, True, ("basic", None)),
-        (None, False, False, ("basic", None)),
-        (None, True, False, ("basic", None)),
-        (None, False, True, ("basic", None)),
-        # Edge case - COT prompt exists but structured output is False and reasoning_capable is True
-        ("think carefully", False, True, ("cot_as_message", "think carefully")),
-    ],
-)
-async def test_run_strategy(
-    adapter, cot_prompt, has_structured_output, reasoning_capable, expected
-):
-    """Test that run_strategy returns correct strategy based on conditions"""
-    # Mock dependencies
-    adapter.prompt_builder.chain_of_thought_prompt = MagicMock(return_value=cot_prompt)
-    adapter.has_structured_output = MagicMock(return_value=has_structured_output)
-
-    provider = MagicMock()
-    provider.reasoning_capable = reasoning_capable
-    adapter.model_provider = MagicMock(return_value=provider)
-
-    # Test
-    result = adapter.run_strategy()
-    assert result == expected
-
-
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "formatter_id,expected_input,expected_calls",
@@ -269,3 +233,214 @@ async def test_input_formatting(
         # Verify original input was preserved in the run
         if formatter_id:
             mock_formatter.format_input.assert_called_once_with(original_input)
+
+
+async def test_properties_for_task_output_includes_all_run_config_properties(adapter):
+    """Test that all properties from RunConfigProperties are saved in task output properties"""
+    # Get all field names from RunConfigProperties
+    run_config_properties_fields = set(RunConfigProperties.model_fields.keys())
+
+    # Get the properties saved by the adapter
+    saved_properties = adapter._properties_for_task_output()
+    saved_property_keys = set(saved_properties.keys())
+
+    # Check which RunConfigProperties fields are missing from saved properties
+    # Note: model_provider_name becomes model_provider in saved properties
+    expected_mappings = {
+        "model_name": "model_name",
+        "model_provider_name": "model_provider",
+        "prompt_id": "prompt_id",
+        "temperature": "temperature",
+        "top_p": "top_p",
+        "structured_output_mode": "structured_output_mode",
+    }
+
+    missing_properties = []
+    for field_name in run_config_properties_fields:
+        expected_key = expected_mappings.get(field_name, field_name)
+        if expected_key not in saved_property_keys:
+            missing_properties.append(
+                f"RunConfigProperties.{field_name} -> {expected_key}"
+            )
+
+    assert not missing_properties, (
+        f"The following RunConfigProperties fields are not saved by _properties_for_task_output: {missing_properties}. Please update the method to include them."
+    )
+
+
+async def test_properties_for_task_output_catches_missing_new_property(adapter):
+    """Test that demonstrates our test will catch when new properties are added to RunConfigProperties but not to _properties_for_task_output"""
+    # Simulate what happens if a new property was added to RunConfigProperties
+    # We'll mock the model_fields to include a fake new property
+    original_fields = RunConfigProperties.model_fields.copy()
+
+    # Create a mock field to simulate a new property being added
+    from pydantic.fields import FieldInfo
+
+    mock_field = FieldInfo(annotation=str, default="default_value")
+
+    try:
+        # Add a fake new field to simulate someone adding a property
+        RunConfigProperties.model_fields["new_fake_property"] = mock_field
+
+        # Get all field names from RunConfigProperties (now includes our fake property)
+        run_config_properties_fields = set(RunConfigProperties.model_fields.keys())
+
+        # Get the properties saved by the adapter (won't include our fake property)
+        saved_properties = adapter._properties_for_task_output()
+        saved_property_keys = set(saved_properties.keys())
+
+        # The mappings don't include our fake property
+        expected_mappings = {
+            "model_name": "model_name",
+            "model_provider_name": "model_provider",
+            "prompt_id": "prompt_id",
+            "temperature": "temperature",
+            "top_p": "top_p",
+            "structured_output_mode": "structured_output_mode",
+        }
+
+        missing_properties = []
+        for field_name in run_config_properties_fields:
+            expected_key = expected_mappings.get(field_name, field_name)
+            if expected_key not in saved_property_keys:
+                missing_properties.append(
+                    f"RunConfigProperties.{field_name} -> {expected_key}"
+                )
+
+        # This should find our missing fake property
+        assert missing_properties == [
+            "RunConfigProperties.new_fake_property -> new_fake_property"
+        ], f"Expected to find missing fake property, but got: {missing_properties}"
+
+    finally:
+        # Restore the original fields
+        RunConfigProperties.model_fields.clear()
+        RunConfigProperties.model_fields.update(original_fields)
+
+
+@pytest.mark.parametrize(
+    "cot_prompt,tuned_strategy,reasoning_capable,expected_formatter_class",
+    [
+        # No COT prompt -> always single turn
+        (None, None, False, "SingleTurnFormatter"),
+        (None, ChatStrategy.two_message_cot, False, "SingleTurnFormatter"),
+        (None, ChatStrategy.single_turn_r1_thinking, True, "SingleTurnFormatter"),
+        # With COT prompt:
+        # - Tuned strategy takes precedence (except single turn)
+        (
+            "think step by step",
+            ChatStrategy.two_message_cot,
+            False,
+            "TwoMessageCotFormatter",
+        ),
+        (
+            "think step by step",
+            ChatStrategy.single_turn_r1_thinking,
+            False,
+            "SingleTurnR1ThinkingFormatter",
+        ),
+        # - Tuned single turn is ignored when COT exists
+        (
+            "think step by step",
+            ChatStrategy.single_turn,
+            True,
+            "SingleTurnR1ThinkingFormatter",
+        ),
+        # - Reasoning capable -> single turn R1 thinking
+        ("think step by step", None, True, "SingleTurnR1ThinkingFormatter"),
+        # - Not reasoning capable -> two message COT
+        ("think step by step", None, False, "TwoMessageCotFormatter"),
+    ],
+)
+def test_build_chat_formatter(
+    adapter,
+    cot_prompt,
+    tuned_strategy,
+    reasoning_capable,
+    expected_formatter_class,
+):
+    """Test chat formatter strategy selection based on COT prompt, tuned strategy, and model capabilities"""
+    # Mock the prompt builder
+    mock_prompt_builder = MagicMock()
+    mock_prompt_builder.chain_of_thought_prompt.return_value = cot_prompt
+    mock_prompt_builder.build_prompt.return_value = "system message"
+    adapter.prompt_builder = mock_prompt_builder
+
+    # Mock the model provider
+    mock_provider = MagicMock()
+    mock_provider.tuned_chat_strategy = tuned_strategy
+    mock_provider.reasoning_capable = reasoning_capable
+    adapter.model_provider = MagicMock(return_value=mock_provider)
+
+    # Get the formatter
+    formatter = adapter.build_chat_formatter("test input")
+
+    # Verify the formatter type
+    assert formatter.__class__.__name__ == expected_formatter_class
+
+    # Verify the formatter was created with correct parameters
+    assert formatter.system_message == "system message"
+    assert formatter.user_input == "test input"
+    # Only check thinking_instructions for formatters that use it
+    if expected_formatter_class == "TwoMessageCotFormatter":
+        if cot_prompt:
+            assert formatter.thinking_instructions == cot_prompt
+        else:
+            assert formatter.thinking_instructions is None
+    # For other formatters, don't assert thinking_instructions
+
+    # Verify prompt builder was called correctly
+    mock_prompt_builder.build_prompt.assert_called_once()
+    mock_prompt_builder.chain_of_thought_prompt.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "initial_mode,expected_mode",
+    [
+        (
+            StructuredOutputMode.json_schema,
+            StructuredOutputMode.json_schema,
+        ),  # Should not change
+        (
+            StructuredOutputMode.unknown,
+            StructuredOutputMode.json_mode,
+        ),  # Should update to default
+    ],
+)
+async def test_update_run_config_unknown_structured_output_mode(
+    base_task, initial_mode, expected_mode
+):
+    """Test that unknown structured output mode is updated to the default for the model provider"""
+    # Create a run config with the initial mode
+    run_config = RunConfig(
+        task=base_task,
+        model_name="test_model",
+        model_provider_name="openai",
+        prompt_id="simple_prompt_builder",
+        structured_output_mode=initial_mode,
+        temperature=0.7,  # Add some other properties to verify they're preserved
+        top_p=0.9,
+    )
+
+    # Mock the default mode lookup
+    with patch(
+        "kiln_ai.adapters.model_adapters.base_adapter.default_structured_output_mode_for_model_provider"
+    ) as mock_default:
+        mock_default.return_value = StructuredOutputMode.json_mode
+
+        # Create the adapter
+        adapter = MockAdapter(run_config=run_config)
+
+        # Verify the mode was updated correctly
+        assert adapter.run_config.structured_output_mode == expected_mode
+
+        # Verify other properties were preserved
+        assert adapter.run_config.temperature == 0.7
+        assert adapter.run_config.top_p == 0.9
+
+        # Verify the default mode lookup was only called when needed
+        if initial_mode == StructuredOutputMode.unknown:
+            mock_default.assert_called_once_with("test_model", "openai")
+        else:
+            mock_default.assert_not_called()