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

kiln_ai/adapters/repair/test_repair_task.py

@@ -21,6 +21,7 @@ from kiln_ai.datamodel import (
     TaskRequirement,
     TaskRun,
 )
+from kiln_ai.datamodel.task import RunConfigProperties
 
 json_joke_schema = """{
   "type": "object",
@@ -189,7 +190,15 @@ async def test_live_run(sample_task, sample_task_run, sample_repair_data):
     repair_task_input = RepairTaskRun.build_repair_task_input(**sample_repair_data)
     assert isinstance(repair_task_input, RepairTaskInput)
 
-    adapter = adapter_for_task(repair_task, model_name="llama_3_1_8b", provider="groq")
+    adapter = adapter_for_task(
+        repair_task,
+        RunConfigProperties(
+            model_name="llama_3_1_8b",
+            model_provider_name="groq",
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="default",
+        ),
+    )
 
     run = await adapter.invoke(repair_task_input.model_dump())
     assert run is not None
@@ -198,10 +207,13 @@ async def test_live_run(sample_task, sample_task_run, sample_repair_data):
     assert "setup" in parsed_output
     assert "punchline" in parsed_output
     assert run.output.source.properties == {
-        "adapter_name": "kiln_langchain_adapter",
+        "adapter_name": "kiln_openai_compatible_adapter",
         "model_name": "llama_3_1_8b",
         "model_provider": "groq",
         "prompt_id": "simple_prompt_builder",
+        "structured_output_mode": "default",
+        "temperature": 1.0,
+        "top_p": 1.0,
     }
 
 
@@ -218,12 +230,19 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     }
 
     with patch.object(LiteLlmAdapter, "_run", new_callable=AsyncMock) as mock_run:
-        mock_run.return_value = RunOutput(
-            output=mocked_output, intermediate_outputs=None
+        mock_run.return_value = (
+            RunOutput(output=mocked_output, intermediate_outputs=None),
+            None,
         )
 
         adapter = adapter_for_task(
-            repair_task, model_name="llama_3_1_8b", provider="ollama"
+            repair_task,
+            RunConfigProperties(
+                model_name="llama_3_1_8b",
+                model_provider_name="ollama",
+                prompt_id="simple_prompt_builder",
+                structured_output_mode="json_schema",
+            ),
         )
 
         run = await adapter.invoke(repair_task_input.model_dump())
@@ -239,6 +258,9 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         "model_name": "llama_3_1_8b",
         "model_provider": "ollama",
         "prompt_id": "simple_prompt_builder",
+        "structured_output_mode": "json_schema",
+        "temperature": 1.0,
+        "top_p": 1.0,
     }
     assert run.input_source.type == DataSourceType.human
     assert "created_by" in run.input_source.properties

kiln_ai/adapters/test_adapter_registry.py

@@ -7,8 +7,8 @@ from kiln_ai.adapters.adapter_registry import adapter_for_task
 from kiln_ai.adapters.ml_model_list import ModelProviderName
 from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
 from kiln_ai.adapters.model_adapters.litellm_adapter import LiteLlmAdapter
-from kiln_ai.adapters.prompt_builders import BasePromptBuilder
 from kiln_ai.adapters.provider_tools import kiln_model_provider_from
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 @pytest.fixture
@@ -35,18 +35,28 @@ def mock_finetune_from_id():
     with patch("kiln_ai.adapters.provider_tools.finetune_from_id") as mock:
         mock.return_value.provider = ModelProviderName.openai
         mock.return_value.fine_tune_model_id = "test-model"
+        mock.return_value.data_strategy = "final_only"
         yield mock
 
 
 def test_openai_adapter_creation(mock_config, basic_task):
     adapter = adapter_for_task(
-        kiln_task=basic_task, model_name="gpt-4", provider=ModelProviderName.openai
+        kiln_task=basic_task,
+        run_config_properties=RunConfigProperties(
+            model_name="gpt-4",
+            model_provider_name=ModelProviderName.openai,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
     )
 
     assert isinstance(adapter, LiteLlmAdapter)
-    assert adapter.config.model_name == "gpt-4"
+    assert adapter.config.run_config_properties.model_name == "gpt-4"
     assert adapter.config.additional_body_options == {"api_key": "test-openai-key"}
-    assert adapter.config.provider_name == ModelProviderName.openai
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.openai
+    )
     assert adapter.config.base_url is None  # OpenAI url is default
     assert adapter.config.default_headers is None
 
@@ -54,14 +64,21 @@ def test_openai_adapter_creation(mock_config, basic_task):
 def test_openrouter_adapter_creation(mock_config, basic_task):
     adapter = adapter_for_task(
         kiln_task=basic_task,
-        model_name="anthropic/claude-3-opus",
-        provider=ModelProviderName.openrouter,
+        run_config_properties=RunConfigProperties(
+            model_name="anthropic/claude-3-opus",
+            model_provider_name=ModelProviderName.openrouter,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
     )
 
     assert isinstance(adapter, LiteLlmAdapter)
-    assert adapter.config.model_name == "anthropic/claude-3-opus"
+    assert adapter.config.run_config_properties.model_name == "anthropic/claude-3-opus"
     assert adapter.config.additional_body_options == {"api_key": "test-openrouter-key"}
-    assert adapter.config.provider_name == ModelProviderName.openrouter
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.openrouter
+    )
     assert adapter.config.default_headers == {
         "HTTP-Referer": "https://getkiln.ai/openrouter",
         "X-Title": "KilnAI",
@@ -79,7 +96,13 @@ def test_openrouter_adapter_creation(mock_config, basic_task):
 )
 def test_openai_compatible_adapter_creation(mock_config, basic_task, provider):
     adapter = adapter_for_task(
-        kiln_task=basic_task, model_name="test-model", provider=provider
+        kiln_task=basic_task,
+        run_config_properties=RunConfigProperties(
+            model_name="test-model",
+            model_provider_name=provider,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
     )
 
     assert isinstance(adapter, LiteLlmAdapter)
@@ -90,9 +113,12 @@ def test_openai_compatible_adapter_creation(mock_config, basic_task, provider):
 def test_custom_prompt_builder(mock_config, basic_task):
     adapter = adapter_for_task(
         kiln_task=basic_task,
-        model_name="gpt-4",
-        provider=ModelProviderName.openai,
-        prompt_id="simple_chain_of_thought_prompt_builder",
+        run_config_properties=RunConfigProperties(
+            model_name="gpt-4",
+            model_provider_name=ModelProviderName.openai,
+            prompt_id="simple_chain_of_thought_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
     )
 
     assert adapter.run_config.prompt_id == "simple_chain_of_thought_prompt_builder"
@@ -103,8 +129,12 @@ def test_tags_passed_through(mock_config, basic_task):
     tags = ["test-tag-1", "test-tag-2"]
     adapter = adapter_for_task(
         kiln_task=basic_task,
-        model_name="gpt-4",
-        provider=ModelProviderName.openai,
+        run_config_properties=RunConfigProperties(
+            model_name="gpt-4",
+            model_provider_name=ModelProviderName.openai,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
         base_adapter_config=AdapterConfig(
             default_tags=tags,
         ),
@@ -114,13 +144,19 @@ def test_tags_passed_through(mock_config, basic_task):
 
 
 def test_invalid_provider(mock_config, basic_task):
-    with pytest.raises(ValueError, match="Unhandled enum value"):
+    with pytest.raises(ValueError, match="Input should be"):
         adapter_for_task(
-            kiln_task=basic_task, model_name="test-model", provider="invalid"
+            kiln_task=basic_task,
+            run_config_properties=RunConfigProperties(
+                model_name="test-model",
+                model_provider_name="invalid",
+                prompt_id="simple_prompt_builder",
+                structured_output_mode="json_schema",
+            ),
         )
 
 
-@patch("kiln_ai.adapters.adapter_registry.lite_llm_config")
+@patch("kiln_ai.adapters.adapter_registry.lite_llm_config_for_openai_compatible")
 def test_openai_compatible_adapter(mock_compatible_config, mock_config, basic_task):
     mock_compatible_config.return_value.model_name = "test-model"
     mock_compatible_config.return_value.additional_body_options = {
@@ -128,44 +164,68 @@ def test_openai_compatible_adapter(mock_compatible_config, mock_config, basic_ta
     }
     mock_compatible_config.return_value.base_url = "https://test.com/v1"
     mock_compatible_config.return_value.provider_name = "CustomProvider99"
+    mock_compatible_config.return_value.run_config_properties = RunConfigProperties(
+        model_name="provider::test-model",
+        model_provider_name=ModelProviderName.openai_compatible,
+        prompt_id="simple_prompt_builder",
+        structured_output_mode="json_schema",
+    )
 
     adapter = adapter_for_task(
         kiln_task=basic_task,
-        model_name="provider::test-model",
-        provider=ModelProviderName.openai_compatible,
+        run_config_properties=RunConfigProperties(
+            model_name="provider::test-model",
+            model_provider_name=ModelProviderName.openai_compatible,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
     )
 
     assert isinstance(adapter, LiteLlmAdapter)
-    mock_compatible_config.assert_called_once_with("provider::test-model")
+    mock_compatible_config.assert_called_once()
     assert adapter.config == mock_compatible_config.return_value
 
 
 def test_custom_openai_compatible_provider(mock_config, basic_task):
     adapter = adapter_for_task(
         kiln_task=basic_task,
-        model_name="openai::test-model",
-        provider=ModelProviderName.kiln_custom_registry,
+        run_config_properties=RunConfigProperties(
+            model_name="openai::test-model",
+            model_provider_name=ModelProviderName.kiln_custom_registry,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
     )
 
     assert isinstance(adapter, LiteLlmAdapter)
-    assert adapter.config.model_name == "openai::test-model"
+    assert adapter.config.run_config_properties.model_name == "openai::test-model"
     assert adapter.config.additional_body_options == {"api_key": "test-openai-key"}
     assert adapter.config.base_url is None  # openai is none
-    assert adapter.config.provider_name == ModelProviderName.kiln_custom_registry
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.kiln_custom_registry
+    )
 
 
 async def test_fine_tune_provider(mock_config, basic_task, mock_finetune_from_id):
     adapter = adapter_for_task(
         kiln_task=basic_task,
-        model_name="proj::task::tune",
-        provider=ModelProviderName.kiln_fine_tune,
+        run_config_properties=RunConfigProperties(
+            model_name="proj::task::tune",
+            model_provider_name=ModelProviderName.kiln_fine_tune,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
     )
 
     mock_finetune_from_id.assert_called_once_with("proj::task::tune")
     assert isinstance(adapter, LiteLlmAdapter)
-    assert adapter.config.provider_name == ModelProviderName.kiln_fine_tune
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.kiln_fine_tune
+    )
     # Kiln model name here, but the underlying openai model id below
-    assert adapter.config.model_name == "proj::task::tune"
+    assert adapter.config.run_config_properties.model_name == "proj::task::tune"
 
     provider = kiln_model_provider_from(
         "proj::task::tune", provider_name=ModelProviderName.kiln_fine_tune

kiln_ai/adapters/test_ml_model_list.py

@@ -0,0 +1,158 @@
+import pytest
+
+from kiln_ai.adapters.ml_model_list import (
+    ModelName,
+    default_structured_output_mode_for_model_provider,
+    get_model_by_name,
+)
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName, StructuredOutputMode
+
+
+class TestDefaultStructuredOutputModeForModelProvider:
+    """Test cases for default_structured_output_mode_for_model_provider function"""
+
+    def test_valid_model_and_provider_returns_provider_mode(self):
+        """Test that valid model and provider returns the provider's structured output mode"""
+        # GPT 4.1 has OpenAI provider with json_schema mode
+        result = default_structured_output_mode_for_model_provider(
+            model_name="gpt_4_1",
+            provider=ModelProviderName.openai,
+        )
+        assert result == StructuredOutputMode.json_schema
+
+    def test_valid_model_different_provider_modes(self):
+        """Test that different providers for the same model return different modes"""
+        # Claude 3.5 Sonnet has different modes for different providers
+        # Anthropic provider uses function_calling
+        result_anthropic = default_structured_output_mode_for_model_provider(
+            model_name="claude_3_5_sonnet",
+            provider=ModelProviderName.anthropic,
+        )
+        assert result_anthropic == StructuredOutputMode.function_calling
+
+        # Vertex provider uses function_calling_weak
+        result_vertex = default_structured_output_mode_for_model_provider(
+            model_name="claude_3_5_sonnet",
+            provider=ModelProviderName.vertex,
+        )
+        assert result_vertex == StructuredOutputMode.function_calling_weak
+
+    def test_invalid_model_name_returns_default(self):
+        """Test that invalid model name returns the default value"""
+        result = default_structured_output_mode_for_model_provider(
+            model_name="invalid_model_name",
+            provider=ModelProviderName.openai,
+        )
+        assert result == StructuredOutputMode.default
+
+    def test_invalid_model_name_returns_custom_default(self):
+        """Test that invalid model name returns custom default when specified"""
+        custom_default = StructuredOutputMode.json_instructions
+        result = default_structured_output_mode_for_model_provider(
+            model_name="invalid_model_name",
+            provider=ModelProviderName.openai,
+            default=custom_default,
+        )
+        assert result == custom_default
+
+    def test_valid_model_invalid_provider_returns_default(self):
+        """Test that valid model but invalid provider returns default"""
+        result = default_structured_output_mode_for_model_provider(
+            model_name="gpt_4_1",
+            provider=ModelProviderName.gemini_api,  # GPT 4.1 doesn't have gemini_api provider
+        )
+        assert result == StructuredOutputMode.default
+
+    def test_disallowed_modes_returns_default(self):
+        """Test that when provider's mode is in disallowed_modes, returns default"""
+        # GPT 4.1 OpenAI provider uses json_schema, but we disallow it
+        result = default_structured_output_mode_for_model_provider(
+            model_name="gpt_4_1",
+            provider=ModelProviderName.openai,
+            disallowed_modes=[StructuredOutputMode.json_schema],
+        )
+        assert result == StructuredOutputMode.default
+
+    def test_disallowed_modes_with_custom_default(self):
+        """Test disallowed modes with custom default value"""
+        custom_default = StructuredOutputMode.json_instructions
+        result = default_structured_output_mode_for_model_provider(
+            model_name="gpt_4_1",
+            provider=ModelProviderName.openai,
+            default=custom_default,
+            disallowed_modes=[StructuredOutputMode.json_schema],
+        )
+        assert result == custom_default
+
+    def test_empty_disallowed_modes_list(self):
+        """Test that empty disallowed_modes list works correctly"""
+        result = default_structured_output_mode_for_model_provider(
+            model_name="gpt_4_1",
+            provider=ModelProviderName.openai,
+            disallowed_modes=[],
+        )
+        assert result == StructuredOutputMode.json_schema
+
+    def test_multiple_disallowed_modes(self):
+        """Test with multiple disallowed modes"""
+        result = default_structured_output_mode_for_model_provider(
+            model_name="gpt_4_1",
+            provider=ModelProviderName.openai,
+            disallowed_modes=[
+                StructuredOutputMode.json_schema,
+                StructuredOutputMode.function_calling,
+            ],
+        )
+        assert result == StructuredOutputMode.default
+
+    def test_reasoning_model_with_different_providers(self):
+        """Test reasoning models that have different structured output modes"""
+        # DeepSeek R1 uses json_instructions for reasoning
+        result = default_structured_output_mode_for_model_provider(
+            model_name="deepseek_r1",
+            provider=ModelProviderName.openrouter,
+        )
+        assert result == StructuredOutputMode.json_instructions
+
+    @pytest.mark.parametrize(
+        "model_name,provider,expected_mode",
+        [
+            ("gpt_4o", ModelProviderName.openai, StructuredOutputMode.json_schema),
+            (
+                "claude_3_5_haiku",
+                ModelProviderName.anthropic,
+                StructuredOutputMode.function_calling,
+            ),
+            (
+                "gemini_2_5_pro",
+                ModelProviderName.gemini_api,
+                StructuredOutputMode.json_schema,
+            ),
+            ("llama_3_1_8b", ModelProviderName.groq, StructuredOutputMode.default),
+            (
+                "qwq_32b",
+                ModelProviderName.fireworks_ai,
+                StructuredOutputMode.json_instructions,
+            ),
+        ],
+    )
+    def test_parametrized_valid_combinations(self, model_name, provider, expected_mode):
+        """Test multiple valid model/provider combinations"""
+        result = default_structured_output_mode_for_model_provider(
+            model_name=model_name,
+            provider=provider,
+        )
+        assert result == expected_mode
+
+    def test_model_with_single_provider(self):
+        """Test model that only has one provider"""
+        # Find a model with only one provider for this test
+        model = get_model_by_name(ModelName.gpt_4_1_nano)
+        assert len(model.providers) >= 1  # Verify it has providers
+
+        first_provider = model.providers[0]
+        result = default_structured_output_mode_for_model_provider(
+            model_name="gpt_4_1_nano",
+            provider=first_provider.name,
+        )
+        assert result == first_provider.structured_output_mode

kiln_ai/adapters/test_prompt_adaptors.py

@@ -18,6 +18,7 @@ from kiln_ai.adapters.prompt_builders import (
     SimpleChainOfThoughtPromptBuilder,
 )
 from kiln_ai.datamodel import PromptId
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 def get_all_models_and_providers():
@@ -124,8 +125,12 @@ async def test_mock_returning_run(tmp_path):
 
         adapter = LiteLlmAdapter(
             config=LiteLlmConfig(
-                model_name="custom_model",
-                provider_name="ollama",
+                run_config_properties=RunConfigProperties(
+                    model_name="custom_model",
+                    model_provider_name="ollama",
+                    prompt_id="simple_prompt_builder",
+                    structured_output_mode="json_schema",
+                ),
                 base_url="http://localhost:11434",
                 additional_body_options={"api_key": "test_key"},
             ),
@@ -145,6 +150,9 @@ async def test_mock_returning_run(tmp_path):
         "model_name": "custom_model",
         "model_provider": "ollama",
         "prompt_id": "simple_prompt_builder",
+        "structured_output_mode": "json_schema",
+        "temperature": 1.0,
+        "top_p": 1.0,
     }
 
 
@@ -212,7 +220,13 @@ async def run_simple_task(
     prompt_id: PromptId | None = None,
 ) -> datamodel.TaskRun:
     adapter = adapter_for_task(
-        task, model_name=model_name, provider=provider, prompt_id=prompt_id
+        task,
+        RunConfigProperties(
+            structured_output_mode="json_schema",
+            model_name=model_name,
+            model_provider_name=provider,
+            prompt_id=prompt_id or "simple_prompt_builder",
+        ),
     )
 
     run = await adapter.invoke(

kiln_ai/adapters/test_prompt_builders.py

@@ -3,7 +3,7 @@ import logging
 
 import pytest
 
-from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter
+from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput
 from kiln_ai.adapters.model_adapters.test_structured_output import (
     build_structured_output_test_task,
 )
@@ -15,6 +15,7 @@ from kiln_ai.adapters.prompt_builders import (
     MultiShotPromptBuilder,
     RepairsPromptBuilder,
     SavedPromptBuilder,
+    ShortPromptBuilder,
     SimpleChainOfThoughtPromptBuilder,
     SimplePromptBuilder,
     TaskRunConfigPromptBuilder,
@@ -26,14 +27,15 @@ from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
     Finetune,
-    FinetuneDataStrategy,
     Project,
     Prompt,
     Task,
     TaskOutput,
     TaskOutputRating,
     TaskRun,
+    Usage,
 )
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.task import RunConfigProperties, TaskRunConfig
 
 logger = logging.getLogger(__name__)
@@ -52,15 +54,31 @@ def test_simple_prompt_builder(tmp_path):
     assert "1) " + task.requirements[0].instruction in prompt
     assert "2) " + task.requirements[1].instruction in prompt
     assert "3) " + task.requirements[2].instruction in prompt
-
-    user_msg = builder.build_user_message(input)
-    assert input in user_msg
     assert input not in prompt
 
 
+def test_short_prompt_builder(tmp_path):
+    task = build_test_task(tmp_path)
+    builder = ShortPromptBuilder(task=task)
+    prompt = builder.build_prompt(include_json_instructions=False)
+
+    # Should only include the instruction, not requirements
+    assert task.instruction == prompt
+    assert task.requirements[0].instruction not in prompt
+    assert task.requirements[1].instruction not in prompt
+    assert task.requirements[2].instruction not in prompt
+
+    # Should handle JSON instructions correctly
+    prompt_with_json = builder.build_prompt(include_json_instructions=True)
+    assert task.instruction in prompt_with_json
+    if task.output_schema():
+        assert "# Format Instructions" in prompt_with_json
+        assert task.output_schema() in prompt_with_json
+
+
 class MockAdapter(BaseAdapter):
-    def _run(self, input: str) -> str:
-        return "mock response"
+    async def _run(self, input: str) -> tuple[RunOutput, Usage | None]:
+        return RunOutput(output="mock response", intermediate_outputs=None), None
 
     def adapter_name(self) -> str:
         return "mock_adapter"
@@ -72,20 +90,9 @@ def test_simple_prompt_builder_structured_output(tmp_path):
     input = "Cows"
     prompt = builder.build_prompt(include_json_instructions=False)
     assert "You are an assistant which tells a joke, given a subject." in prompt
-
-    user_msg = builder.build_user_message(input)
-    assert input in user_msg
     assert input not in prompt
 
 
-def test_simple_prompt_builder_structured_input_non_ascii(tmp_path):
-    task = build_structured_output_test_task(tmp_path)
-    builder = SimplePromptBuilder(task=task)
-    input = {"key": "你好👋"}
-    user_msg = builder.build_user_message(input)
-    assert "你好👋" in user_msg
-
-
 @pytest.fixture
 def task_with_examples(tmp_path):
     # Create a project and task hierarchy
@@ -383,7 +390,7 @@ def test_prompt_builder_from_id(task_with_examples):
         base_model_id="test_base_model_id",
         dataset_split_id="asdf",
         provider="test_provider",
-        data_strategy=FinetuneDataStrategy.final_and_intermediate,
+        data_strategy=ChatStrategy.two_message_cot,
     )
     finetune.save_to_file()
     nested_fine_tune_id = (
@@ -598,6 +605,7 @@ def test_task_run_config_prompt_builder(tmp_path):
             model_name="gpt-4",
             model_provider_name="openai",
             prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
         ),
         prompt=Prompt(
             name="test prompt name",