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

kiln_ai/adapters/test_ml_model_list.py

@@ -0,0 +1,176 @@
+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
+
+
+def test_uncensored():
+    """Test that uncensored is set correctly"""
+    model = get_model_by_name(ModelName.grok_3_mini)
+    for provider in model.providers:
+        assert provider.uncensored
+        assert not provider.suggested_for_uncensored_data_gen
+
+    model = get_model_by_name(ModelName.gpt_4_1_nano)
+    for provider in model.providers:
+        assert not provider.uncensored
+        assert not provider.suggested_for_uncensored_data_gen
+
+    model = get_model_by_name(ModelName.grok_4)
+    for provider in model.providers:
+        assert provider.uncensored
+        assert provider.suggested_for_uncensored_data_gen

kiln_ai/adapters/test_prompt_adaptors.py

@@ -13,11 +13,8 @@ from kiln_ai.adapters.model_adapters.litellm_adapter import (
     LiteLlmConfig,
 )
 from kiln_ai.adapters.ollama_tools import ollama_online
-from kiln_ai.adapters.prompt_builders import (
-    BasePromptBuilder,
-    SimpleChainOfThoughtPromptBuilder,
-)
 from kiln_ai.datamodel import PromptId
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 def get_all_models_and_providers():
@@ -124,8 +121,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 +146,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 +216,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

@@ -27,7 +27,6 @@ from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
     Finetune,
-    FinetuneDataStrategy,
     Project,
     Prompt,
     Task,
@@ -36,6 +35,7 @@ from kiln_ai.datamodel import (
     TaskRun,
     Usage,
 )
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.task import RunConfigProperties, TaskRunConfig
 
 logger = logging.getLogger(__name__)
@@ -54,9 +54,6 @@ 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
 
 
@@ -93,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
@@ -404,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 = (
@@ -619,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",

kiln_ai/adapters/test_provider_tools.py

@@ -18,7 +18,7 @@ from kiln_ai.adapters.provider_tools import (
     finetune_provider_model,
     get_model_and_provider,
     kiln_model_provider_from,
-    lite_llm_config,
+    lite_llm_config_for_openai_compatible,
     lite_llm_provider_model,
     parse_custom_model_id,
     provider_enabled,
@@ -27,10 +27,11 @@ from kiln_ai.adapters.provider_tools import (
 )
 from kiln_ai.datamodel import (
     Finetune,
-    FinetuneDataStrategy,
     StructuredOutputMode,
     Task,
 )
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 @pytest.fixture(autouse=True)
@@ -71,7 +72,7 @@ def mock_finetune():
         finetune.provider = ModelProviderName.openai
         finetune.fine_tune_model_id = "ft:gpt-3.5-turbo:custom:model-123"
         finetune.structured_output_mode = StructuredOutputMode.json_schema
-        finetune.data_strategy = FinetuneDataStrategy.final_only
+        finetune.data_strategy = ChatStrategy.single_turn
         mock.return_value = finetune
         yield mock
 
@@ -83,7 +84,7 @@ def mock_finetune_final_and_intermediate():
         finetune.provider = ModelProviderName.openai
         finetune.fine_tune_model_id = "ft:gpt-3.5-turbo:custom:model-123"
         finetune.structured_output_mode = StructuredOutputMode.json_schema
-        finetune.data_strategy = FinetuneDataStrategy.final_and_intermediate
+        finetune.data_strategy = ChatStrategy.two_message_cot
         mock.return_value = finetune
         yield mock
 
@@ -95,9 +96,7 @@ def mock_finetune_r1_compatible():
         finetune.provider = ModelProviderName.ollama
         finetune.fine_tune_model_id = "ft:deepseek-r1:671b:custom:model-123"
         finetune.structured_output_mode = StructuredOutputMode.json_schema
-        finetune.data_strategy = (
-            FinetuneDataStrategy.final_and_intermediate_r1_compatible
-        )
+        finetune.data_strategy = ChatStrategy.single_turn_r1_thinking
         mock.return_value = finetune
         yield mock
 
@@ -357,6 +356,7 @@ async def test_kiln_model_provider_from_custom_model_valid(mock_config):
     assert provider.supports_data_gen is False
     assert provider.untested_model is True
     assert provider.model_id == "custom_model"
+    assert provider.structured_output_mode == StructuredOutputMode.json_instructions
 
 
 @pytest.mark.asyncio
@@ -374,6 +374,7 @@ async def test_kiln_model_provider_from_custom_registry(mock_config):
     assert provider.supports_data_gen is False
     assert provider.untested_model is True
     assert provider.model_id == "gpt-4-turbo"
+    assert provider.structured_output_mode == StructuredOutputMode.json_instructions
 
 
 @pytest.mark.asyncio
@@ -474,7 +475,7 @@ def test_finetune_provider_model_success_final_and_intermediate(
     assert provider.name == ModelProviderName.openai
     assert provider.model_id == "ft:gpt-3.5-turbo:custom:model-123"
     assert provider.structured_output_mode == StructuredOutputMode.json_schema
-    assert provider.reasoning_capable is True
+    assert provider.reasoning_capable is False
     assert provider.parser == None
 
 
@@ -580,7 +581,7 @@ def test_finetune_provider_model_structured_mode(
     finetune.provider = provider_name
     finetune.fine_tune_model_id = "fireworks-model-123"
     finetune.structured_output_mode = structured_output_mode
-    finetune.data_strategy = FinetuneDataStrategy.final_only
+    finetune.data_strategy = ChatStrategy.single_turn
     mock_finetune.return_value = finetune
 
     provider = finetune_provider_model("project-123::task-456::finetune-789")
@@ -596,10 +597,20 @@ def test_openai_compatible_provider_config(mock_shared_config):
     """Test successful creation of an OpenAI compatible provider"""
     model_id = "test_provider::gpt-4"
 
-    config = lite_llm_config(model_id)
+    config = lite_llm_config_for_openai_compatible(
+        RunConfigProperties(
+            model_name=model_id,
+            model_provider_name=ModelProviderName.openai_compatible,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        )
+    )
 
-    assert config.provider_name == ModelProviderName.openai_compatible
-    assert config.model_name == "gpt-4"
+    assert (
+        config.run_config_properties.model_provider_name
+        == ModelProviderName.openai_compatible
+    )
+    assert config.run_config_properties.model_name == "gpt-4"
     assert config.additional_body_options == {"api_key": "test-key"}
     assert config.base_url == "https://api.test.com"
 
@@ -621,10 +632,20 @@ def test_lite_llm_config_no_api_key(mock_shared_config):
     """Test provider creation without API key (should work as some providers don't require it, but should pass NA to LiteLLM as it requires one)"""
     model_id = "no_key_provider::gpt-4"
 
-    config = lite_llm_config(model_id)
+    config = lite_llm_config_for_openai_compatible(
+        RunConfigProperties(
+            model_name=model_id,
+            model_provider_name=ModelProviderName.openai,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        )
+    )
 
-    assert config.provider_name == ModelProviderName.openai_compatible
-    assert config.model_name == "gpt-4"
+    assert (
+        config.run_config_properties.model_provider_name
+        == ModelProviderName.openai_compatible
+    )
+    assert config.run_config_properties.model_name == "gpt-4"
     assert config.additional_body_options == {"api_key": "NA"}
     assert config.base_url == "https://api.nokey.com"
 
@@ -632,7 +653,14 @@ def test_lite_llm_config_no_api_key(mock_shared_config):
 def test_lite_llm_config_invalid_id():
     """Test handling of invalid model ID format"""
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config("invalid-id-format")
+        lite_llm_config_for_openai_compatible(
+            RunConfigProperties(
+                model_name="invalid-id-format",
+                model_provider_name=ModelProviderName.openai_compatible,
+                prompt_id="simple_prompt_builder",
+                structured_output_mode="json_schema",
+            )
+        )
     assert (
         str(exc_info.value) == "Invalid openai compatible model ID: invalid-id-format"
     )
@@ -643,14 +671,28 @@ def test_lite_llm_config_no_providers(mock_shared_config):
     mock_shared_config.return_value.openai_compatible_providers = None
 
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config("test_provider::gpt-4")
+        lite_llm_config_for_openai_compatible(
+            RunConfigProperties(
+                model_name="test_provider::gpt-4",
+                model_provider_name=ModelProviderName.openai_compatible,
+                prompt_id="simple_prompt_builder",
+                structured_output_mode="json_schema",
+            )
+        )
     assert str(exc_info.value) == "OpenAI compatible provider test_provider not found"
 
 
 def test_lite_llm_config_provider_not_found(mock_shared_config):
     """Test handling of non-existent provider"""
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config("unknown_provider::gpt-4")
+        lite_llm_config_for_openai_compatible(
+            RunConfigProperties(
+                model_name="unknown_provider::gpt-4",
+                model_provider_name=ModelProviderName.openai_compatible,
+                prompt_id="simple_prompt_builder",
+                structured_output_mode="json_schema",
+            )
+        )
     assert (
         str(exc_info.value) == "OpenAI compatible provider unknown_provider not found"
     )
@@ -666,7 +708,14 @@ def test_lite_llm_config_no_base_url(mock_shared_config):
     ]
 
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config("test_provider::gpt-4")
+        lite_llm_config_for_openai_compatible(
+            RunConfigProperties(
+                model_name="test_provider::gpt-4",
+                model_provider_name=ModelProviderName.openai_compatible,
+                prompt_id="simple_prompt_builder",
+                structured_output_mode="json_schema",
+            )
+        )
     assert (
         str(exc_info.value)
         == "OpenAI compatible provider test_provider has no base URL"
@@ -867,7 +916,7 @@ def test_finetune_provider_model_vertex_ai(mock_project, mock_task, mock_finetun
     finetune.provider = ModelProviderName.vertex
     finetune.fine_tune_model_id = "projects/123/locations/us-central1/endpoints/456"
     finetune.structured_output_mode = StructuredOutputMode.json_mode
-    finetune.data_strategy = FinetuneDataStrategy.final_only
+    finetune.data_strategy = ChatStrategy.single_turn
     mock_finetune.return_value = finetune
 
     provider = finetune_provider_model("project-123::task-456::finetune-789")

kiln_ai/adapters/test_remote_config.py

@@ -0,0 +1,100 @@
+import asyncio
+import os
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai.adapters.ml_model_list import built_in_models
+from kiln_ai.adapters.remote_config import (
+    deserialize_config,
+    dump_builtin_config,
+    load_from_url,
+    load_remote_models,
+    serialize_config,
+)
+
+
+def test_round_trip(tmp_path):
+    path = tmp_path / "models.json"
+    serialize_config(built_in_models, path)
+    loaded = deserialize_config(path)
+    assert [m.model_dump(mode="json") for m in loaded] == [
+        m.model_dump(mode="json") for m in built_in_models
+    ]
+
+
+def test_load_from_url():
+    sample = [built_in_models[0].model_dump(mode="json")]
+
+    class FakeResponse:
+        def raise_for_status(self):
+            pass
+
+        def json(self):
+            return {"model_list": sample}
+
+    with patch(
+        "kiln_ai.adapters.remote_config.requests.get", return_value=FakeResponse()
+    ):
+        models = load_from_url("http://example.com/models.json")
+    assert [m.model_dump(mode="json") for m in models] == sample
+
+
+def test_dump_builtin_config(tmp_path):
+    path = tmp_path / "out.json"
+    dump_builtin_config(path)
+    loaded = deserialize_config(path)
+    assert [m.model_dump(mode="json") for m in loaded] == [
+        m.model_dump(mode="json") for m in built_in_models
+    ]
+
+
+@pytest.mark.asyncio
+async def test_load_remote_models_success(monkeypatch):
+    del os.environ["KILN_SKIP_REMOTE_MODEL_LIST"]
+    original = built_in_models.copy()
+    sample_models = [built_in_models[0]]
+
+    def fake_fetch(url):
+        return sample_models
+
+    monkeypatch.setattr("kiln_ai.adapters.remote_config.load_from_url", fake_fetch)
+
+    load_remote_models("http://example.com/models.json")
+    await asyncio.sleep(0.01)
+    assert built_in_models == sample_models
+    built_in_models[:] = original
+
+
+@pytest.mark.asyncio
+async def test_load_remote_models_failure(monkeypatch):
+    original = built_in_models.copy()
+
+    def fake_fetch(url):
+        raise RuntimeError("fail")
+
+    monkeypatch.setattr("kiln_ai.adapters.remote_config.load_from_url", fake_fetch)
+
+    load_remote_models("http://example.com/models.json")
+    await asyncio.sleep(0.01)
+    assert built_in_models == original
+
+
+def test_deserialize_config_with_extra_keys(tmp_path):
+    # Take a valid model and add an extra key, ensure it is ignored and still loads
+    import json
+
+    from kiln_ai.adapters.ml_model_list import built_in_models
+
+    model_dict = built_in_models[0].model_dump(mode="json")
+    model_dict["extra_key"] = "should be ignored or error"
+    model_dict["providers"][0]["extra_key"] = "should be ignored or error"
+    data = {"model_list": [model_dict]}
+    path = tmp_path / "extra.json"
+    path.write_text(json.dumps(data))
+    # Should NOT raise, and extra key should be ignored
+    models = deserialize_config(path)
+    assert hasattr(models[0], "family")
+    assert not hasattr(models[0], "extra_key")
+    assert hasattr(models[0], "providers")
+    assert not hasattr(models[0].providers[0], "extra_key")

kiln_ai/datamodel/__init__.py

@@ -13,7 +13,6 @@ from __future__ import annotations
 
 from kiln_ai.datamodel import dataset_split, eval, strict_mode
 from kiln_ai.datamodel.datamodel_enums import (
-    FinetuneDataStrategy,
     FineTuneStatusType,
     Priority,
     StructuredOutputMode,
@@ -71,7 +70,6 @@ __all__ = [
     "Prompt",
     "TaskOutputRating",
     "StructuredOutputMode",
-    "FinetuneDataStrategy",
     "PromptId",
     "PromptGenerators",
     "prompt_generator_values",