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

kiln_ai/adapters/model_adapters/test_base_adapter.py

@@ -4,9 +4,9 @@ import pytest
 
 from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMode
 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 +37,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 +89,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 +98,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 +140,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 +170,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 +179,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 +232,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()

kiln_ai/adapters/model_adapters/test_litellm_adapter.py

@@ -11,6 +11,7 @@ from kiln_ai.adapters.model_adapters.litellm_config import (
     LiteLlmConfig,
 )
 from kiln_ai.datamodel import Project, Task, Usage
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 @pytest.fixture
@@ -41,8 +42,12 @@ def mock_task(tmp_path):
 def config():
     return LiteLlmConfig(
         base_url="https://api.test.com",
-        model_name="test-model",
-        provider_name="openrouter",
+        run_config_properties=RunConfigProperties(
+            model_name="test-model",
+            model_provider_name="openrouter",
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
         default_headers={"X-Test": "test"},
         additional_body_options={"api_key": "test_key"},
     )
@@ -52,7 +57,6 @@ def test_initialization(config, mock_task):
     adapter = LiteLlmAdapter(
         config=config,
         kiln_task=mock_task,
-        prompt_id="simple_prompt_builder",
         base_adapter_config=AdapterConfig(default_tags=["test-tag"]),
     )
 
@@ -60,8 +64,11 @@ def test_initialization(config, mock_task):
     assert adapter.run_config.task == mock_task
     assert adapter.run_config.prompt_id == "simple_prompt_builder"
     assert adapter.base_adapter_config.default_tags == ["test-tag"]
-    assert adapter.run_config.model_name == config.model_name
-    assert adapter.run_config.model_provider_name == config.provider_name
+    assert adapter.run_config.model_name == config.run_config_properties.model_name
+    assert (
+        adapter.run_config.model_provider_name
+        == config.run_config_properties.model_provider_name
+    )
     assert adapter.config.additional_body_options["api_key"] == "test_key"
     assert adapter._api_base == config.base_url
     assert adapter._headers == config.default_headers
@@ -72,8 +79,11 @@ def test_adapter_info(config, mock_task):
 
     assert adapter.adapter_name() == "kiln_openai_compatible_adapter"
 
-    assert adapter.run_config.model_name == config.model_name
-    assert adapter.run_config.model_provider_name == config.provider_name
+    assert adapter.run_config.model_name == config.run_config_properties.model_name
+    assert (
+        adapter.run_config.model_provider_name
+        == config.run_config_properties.model_provider_name
+    )
     assert adapter.run_config.prompt_id == "simple_prompt_builder"
 
 
@@ -96,14 +106,12 @@ async def test_response_format_options_unstructured(config, mock_task):
 )
 @pytest.mark.asyncio
 async def test_response_format_options_json_mode(config, mock_task, mode):
+    config.run_config_properties.structured_output_mode = mode
     adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
 
     with (
         patch.object(adapter, "has_structured_output", return_value=True),
-        patch.object(adapter, "model_provider") as mock_provider,
     ):
-        mock_provider.return_value.structured_output_mode = mode
-
         options = await adapter.response_format_options()
         assert options == {"response_format": {"type": "json_object"}}
 
@@ -117,14 +125,12 @@ async def test_response_format_options_json_mode(config, mock_task, mode):
 )
 @pytest.mark.asyncio
 async def test_response_format_options_function_calling(config, mock_task, mode):
+    config.run_config_properties.structured_output_mode = mode
     adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
 
     with (
         patch.object(adapter, "has_structured_output", return_value=True),
-        patch.object(adapter, "model_provider") as mock_provider,
     ):
-        mock_provider.return_value.structured_output_mode = mode
-
         options = await adapter.response_format_options()
         assert "tools" in options
         # full tool structure validated below
@@ -139,30 +145,26 @@ async def test_response_format_options_function_calling(config, mock_task, mode)
 )
 @pytest.mark.asyncio
 async def test_response_format_options_json_instructions(config, mock_task, mode):
+    config.run_config_properties.structured_output_mode = mode
     adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
 
     with (
         patch.object(adapter, "has_structured_output", return_value=True),
-        patch.object(adapter, "model_provider") as mock_provider,
     ):
-        mock_provider.return_value.structured_output_mode = (
-            StructuredOutputMode.json_instructions
-        )
         options = await adapter.response_format_options()
         assert options == {}
 
 
 @pytest.mark.asyncio
 async def test_response_format_options_json_schema(config, mock_task):
+    config.run_config_properties.structured_output_mode = (
+        StructuredOutputMode.json_schema
+    )
     adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
 
     with (
         patch.object(adapter, "has_structured_output", return_value=True),
-        patch.object(adapter, "model_provider") as mock_provider,
     ):
-        mock_provider.return_value.structured_output_mode = (
-            StructuredOutputMode.json_schema
-        )
         options = await adapter.response_format_options()
         assert options == {
             "response_format": {
@@ -350,6 +352,69 @@ def test_litellm_model_id_unknown_provider(config, mock_task):
                 adapter.litellm_model_id()
 
 
+@pytest.mark.parametrize(
+    "provider_name,expected_usage_param",
+    [
+        (ModelProviderName.openrouter, {"usage": {"include": True}}),
+        (ModelProviderName.openai, {}),
+        (ModelProviderName.anthropic, {}),
+        (ModelProviderName.groq, {}),
+    ],
+)
+def test_build_extra_body_openrouter_usage(
+    config, mock_task, provider_name, expected_usage_param
+):
+    """Test build_extra_body includes usage parameter for OpenRouter providers"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Create a mock provider with the specified name and minimal required attributes
+    mock_provider = Mock()
+    mock_provider.name = provider_name
+    mock_provider.thinking_level = None
+    mock_provider.require_openrouter_reasoning = False
+    mock_provider.anthropic_extended_thinking = False
+    mock_provider.r1_openrouter_options = False
+    mock_provider.logprobs_openrouter_options = False
+    mock_provider.openrouter_skip_required_parameters = False
+
+    # Call build_extra_body
+    extra_body = adapter.build_extra_body(mock_provider)
+
+    # Verify the usage parameter is included only for OpenRouter
+    for key, value in expected_usage_param.items():
+        assert extra_body.get(key) == value
+
+    # Verify non-OpenRouter providers don't have the usage parameter
+    if provider_name != ModelProviderName.openrouter:
+        assert "usage" not in extra_body
+
+
+@pytest.mark.asyncio
+async def test_build_completion_kwargs_custom_temperature_top_p(config, mock_task):
+    """Test build_completion_kwargs with custom temperature and top_p values"""
+    # Create config with custom temperature and top_p
+    config.run_config_properties.temperature = 0.7
+    config.run_config_properties.top_p = 0.9
+
+    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={}),
+        patch.object(adapter, "response_format_options", return_value={}),
+    ):
+        kwargs = await adapter.build_completion_kwargs(mock_provider, messages, None)
+
+    # Verify custom temperature and top_p are passed through
+    assert kwargs["temperature"] == 0.7
+    assert kwargs["top_p"] == 0.9
+    # Verify drop_params is set correctly
+    assert kwargs["drop_params"] is True
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "top_logprobs,response_format,extra_body",
@@ -391,6 +456,13 @@ async def test_build_completion_kwargs(
     assert kwargs["messages"] == messages
     assert kwargs["api_base"] == config.base_url
 
+    # Verify temperature and top_p are included with default values
+    assert kwargs["temperature"] == 1.0  # Default from RunConfigProperties
+    assert kwargs["top_p"] == 1.0  # Default from RunConfigProperties
+
+    # Verify drop_params is set correctly
+    assert kwargs["drop_params"] is True
+
     # Verify optional parameters
     if top_logprobs is not None:
         assert kwargs["logprobs"] is True
@@ -439,6 +511,17 @@ async def test_build_completion_kwargs(
         ({"prompt_tokens": 10}, None, None),
         # Invalid cost type (should be ignored)
         (None, "0.5", None),
+        # Cost in OpenRouter format
+        (
+            litellm.types.utils.Usage(
+                prompt_tokens=10,
+                completion_tokens=20,
+                total_tokens=30,
+                cost=0.5,
+            ),
+            None,
+            Usage(input_tokens=10, output_tokens=20, total_tokens=30, cost=0.5),
+        ),
     ],
 )
 def test_usage_from_response(config, mock_task, litellm_usage, cost, expected_usage):

kiln_ai/adapters/model_adapters/test_saving_adapter_results.py

@@ -46,6 +46,7 @@ def adapter(test_task):
             model_name="phi_3_5",
             model_provider_name="ollama",
             prompt_id="simple_chain_of_thought_prompt_builder",
+            structured_output_mode="json_schema",
         ),
     )
 
@@ -102,6 +103,9 @@ def test_save_run_isolation(test_task, adapter):
         reloaded_output.source.properties["prompt_id"]
         == "simple_chain_of_thought_prompt_builder"
     )
+    assert reloaded_output.source.properties["structured_output_mode"] == "json_schema"
+    assert reloaded_output.source.properties["temperature"] == 1.0
+    assert reloaded_output.source.properties["top_p"] == 1.0
     # Run again, with same input and different output. Should create a new TaskRun.
     different_run_output = RunOutput(
         output="Different output", intermediate_outputs=None
@@ -228,3 +232,40 @@ async def test_autosave_true(test_task, adapter):
             output.source.properties["prompt_id"]
             == "simple_chain_of_thought_prompt_builder"
         )
+        assert output.source.properties["structured_output_mode"] == "json_schema"
+        assert output.source.properties["temperature"] == 1.0
+        assert output.source.properties["top_p"] == 1.0
+
+
+def test_properties_for_task_output_custom_values(test_task):
+    """Test that _properties_for_task_output includes custom temperature, top_p, and structured_output_mode"""
+    adapter = MockAdapter(
+        run_config=RunConfig(
+            task=test_task,
+            model_name="gpt-4",
+            model_provider_name="openai",
+            prompt_id="simple_prompt_builder",
+            temperature=0.7,
+            top_p=0.9,
+            structured_output_mode="json_schema",
+        ),
+    )
+
+    input_data = "Test input"
+    output_data = "Test output"
+    run_output = RunOutput(output=output_data, intermediate_outputs=None)
+
+    task_run = adapter.generate_run(
+        input=input_data, input_source=None, run_output=run_output
+    )
+    task_run.save_to_file()
+
+    # Verify custom values are preserved in properties
+    output = task_run.output
+    assert output.source.properties["adapter_name"] == "mock_adapter"
+    assert output.source.properties["model_name"] == "gpt-4"
+    assert output.source.properties["model_provider"] == "openai"
+    assert output.source.properties["prompt_id"] == "simple_prompt_builder"
+    assert output.source.properties["structured_output_mode"] == "json_schema"
+    assert output.source.properties["temperature"] == 0.7
+    assert output.source.properties["top_p"] == 0.9