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

kiln_ai/adapters/model_adapters/test_structured_output.py

@@ -17,7 +17,7 @@ from kiln_ai.adapters.model_adapters.base_adapter import (
 from kiln_ai.adapters.ollama_tools import ollama_online
 from kiln_ai.adapters.test_prompt_adaptors import get_all_models_and_providers
 from kiln_ai.datamodel import PromptId
-from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.datamodel.task import RunConfig, RunConfigProperties
 from kiln_ai.datamodel.test_json_schema import json_joke_schema, json_triangle_schema
 
 
@@ -51,6 +51,7 @@ class MockAdapter(BaseAdapter):
                 model_name="phi_3_5",
                 model_provider_name="ollama",
                 prompt_id="simple_chain_of_thought_prompt_builder",
+                structured_output_mode="json_schema",
             ),
         )
         self.response = response
@@ -146,7 +147,15 @@ def build_structured_output_test_task(tmp_path: Path):
 
 async def run_structured_output_test(tmp_path: Path, model_name: str, provider: str):
     task = build_structured_output_test_task(tmp_path)
-    a = adapter_for_task(task, model_name=model_name, provider=provider)
+    a = adapter_for_task(
+        task,
+        run_config_properties=RunConfigProperties(
+            model_name=model_name,
+            model_provider_name=provider,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="unknown",
+        ),
+    )
     try:
         run = await a.invoke("Cows")  # a joke about cows
         parsed = json.loads(run.output.output)
@@ -197,10 +206,12 @@ def build_structured_input_test_task(tmp_path: Path):
     return task
 
 
-async def run_structured_input_test(tmp_path: Path, model_name: str, provider: str):
+async def run_structured_input_test(
+    tmp_path: Path, model_name: str, provider: str, prompt_id: PromptId
+):
     task = build_structured_input_test_task(tmp_path)
     try:
-        await run_structured_input_task(task, model_name, provider)
+        await run_structured_input_task(task, model_name, provider, prompt_id)
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
@@ -209,17 +220,20 @@ async def run_structured_input_test(tmp_path: Path, model_name: str, provider: s
         raise e
 
 
-async def run_structured_input_task(
+async def run_structured_input_task_no_validation(
     task: datamodel.Task,
     model_name: str,
     provider: str,
-    prompt_id: PromptId | None = None,
+    prompt_id: PromptId,
 ):
     a = adapter_for_task(
         task,
-        model_name=model_name,
-        provider=provider,
-        prompt_id=prompt_id,
+        run_config_properties=RunConfigProperties(
+            model_name=model_name,
+            model_provider_name=provider,
+            prompt_id=prompt_id,
+            structured_output_mode="unknown",
+        ),
     )
     with pytest.raises(ValueError):
         # not structured input in dictionary
@@ -231,18 +245,29 @@ async def run_structured_input_task(
     try:
         run = await a.invoke({"a": 2, "b": 2, "c": 2})
         response = run.output.output
+        return response, a, run
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
                 f"Skipping {model_name} {provider} because Ollama is not running"
             )
         raise e
+
+
+async def run_structured_input_task(
+    task: datamodel.Task,
+    model_name: str,
+    provider: str,
+    prompt_id: PromptId,
+):
+    response, a, run = await run_structured_input_task_no_validation(
+        task, model_name, provider, prompt_id
+    )
     assert response is not None
     if isinstance(response, str):
         assert "[[equilateral]]" in response
     else:
         assert response["is_equilateral"] is True
-
     expected_pb_name = "simple_prompt_builder"
     if prompt_id is not None:
         expected_pb_name = prompt_id
@@ -269,7 +294,9 @@ async def test_structured_input_gpt_4o_mini(tmp_path):
 async def test_all_built_in_models_structured_input(
     tmp_path, model_name, provider_name
 ):
-    await run_structured_input_test(tmp_path, model_name, provider_name)
+    await run_structured_input_test(
+        tmp_path, model_name, provider_name, "simple_prompt_builder"
+    )
 
 
 @pytest.mark.paid
@@ -323,6 +350,11 @@ When asked for a final result, this is the format (for an equilateral example):
 """
     task.output_json_schema = json.dumps(triangle_schema)
     task.save_to_file()
-    await run_structured_input_task(
+    response, adapter, _ = await run_structured_input_task_no_validation(
         task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
     )
+
+    formatted_response = json.loads(response)
+    assert formatted_response["is_equilateral"] is True
+    assert formatted_response["is_scalene"] is False
+    assert formatted_response["is_obtuse"] is False

kiln_ai/adapters/parsers/parser_registry.py

@@ -1,5 +1,3 @@
-from typing import Type
-
 from kiln_ai.adapters.ml_model_list import ModelParserID
 from kiln_ai.adapters.parsers.base_parser import BaseParser
 from kiln_ai.adapters.parsers.r1_parser import R1ThinkingParser

kiln_ai/adapters/parsers/r1_parser.py

@@ -1,5 +1,4 @@
 from kiln_ai.adapters.parsers.base_parser import BaseParser
-from kiln_ai.adapters.parsers.json_parser import parse_json_string
 from kiln_ai.adapters.run_output import RunOutput
 
 

kiln_ai/adapters/prompt_builders.py

@@ -1,6 +1,4 @@
-import json
 from abc import ABCMeta, abstractmethod
-from typing import Dict
 
 from kiln_ai.datamodel import PromptGenerators, PromptId, Task, TaskRun
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
@@ -53,20 +51,6 @@ class BasePromptBuilder(metaclass=ABCMeta):
         """
         pass
 
-    def build_user_message(self, input: Dict | str) -> str:
-        """Build a user message from the input.
-
-        Args:
-            input (Union[Dict, str]): The input to format into a message.
-
-        Returns:
-            str: The formatted user message.
-        """
-        if isinstance(input, Dict):
-            return f"The input is:\n{json.dumps(input, indent=2, ensure_ascii=False)}"
-
-        return f"The input is:\n{input}"
-
     def chain_of_thought_prompt(self) -> str | None:
         """Build and return the chain of thought prompt string.
 

kiln_ai/adapters/provider_tools.py

@@ -1,3 +1,4 @@
+import logging
 from dataclasses import dataclass
 from typing import Dict, List
 
@@ -16,11 +17,15 @@ from kiln_ai.adapters.model_adapters.litellm_config import (
 from kiln_ai.adapters.ollama_tools import (
     get_ollama_connection,
 )
-from kiln_ai.datamodel import Finetune, FinetuneDataStrategy, Task
+from kiln_ai.datamodel import Finetune, Task
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.registry import project_from_id
+from kiln_ai.datamodel.task import RunConfigProperties
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
+logger = logging.getLogger(__name__)
+
 
 async def provider_enabled(provider_name: ModelProviderName) -> bool:
     if provider_name == ModelProviderName.ollama:
@@ -163,6 +168,10 @@ def kiln_model_provider_from(
     # For custom registry, get the provider name and model name from the model id
     if provider_name == ModelProviderName.kiln_custom_registry:
         provider_name, name = parse_custom_model_id(name)
+    else:
+        logger.warning(
+            f"Unexpected model/provider pair. Will treat as custom model but check your model settings. Provider: {provider_name}/{name}"
+        )
 
     # Custom/untested model. Set untested, and build a ModelProvider at runtime
     if provider_name is None:
@@ -177,12 +186,15 @@ def kiln_model_provider_from(
         supports_data_gen=False,
         untested_model=True,
         model_id=name,
+        # We don't know the structured output mode for custom models, so we default to json_instructions which is the only one that works everywhere.
+        structured_output_mode=StructuredOutputMode.json_instructions,
     )
 
 
-def lite_llm_config(
-    model_id: str,
+def lite_llm_config_for_openai_compatible(
+    run_config_properties: RunConfigProperties,
 ) -> LiteLlmConfig:
+    model_id = run_config_properties.model_name
     try:
         openai_provider_name, model_id = model_id.split("::")
     except Exception:
@@ -206,10 +218,16 @@ def lite_llm_config(
             f"OpenAI compatible provider {openai_provider_name} has no base URL"
         )
 
+    # Update a copy of the run config properties to use the openai compatible provider
+    updated_run_config_properties = run_config_properties.model_copy(deep=True)
+    updated_run_config_properties.model_provider_name = (
+        ModelProviderName.openai_compatible
+    )
+    updated_run_config_properties.model_name = model_id
+
     return LiteLlmConfig(
         # OpenAI compatible, with a custom base URL
-        model_name=model_id,
-        provider_name=ModelProviderName.openai_compatible,
+        run_config_properties=updated_run_config_properties,
         base_url=base_url,
         additional_body_options={
             "api_key": api_key,
@@ -259,9 +277,9 @@ def finetune_from_id(model_id: str) -> Finetune:
 
 
 def parser_from_data_strategy(
-    data_strategy: FinetuneDataStrategy,
+    data_strategy: ChatStrategy,
 ) -> ModelParserID | None:
-    if data_strategy == FinetuneDataStrategy.final_and_intermediate_r1_compatible:
+    if data_strategy == ChatStrategy.single_turn_r1_thinking:
         return ModelParserID.r1_thinking
     return None
 
@@ -279,10 +297,10 @@ def finetune_provider_model(
         reasoning_capable=(
             fine_tune.data_strategy
             in [
-                FinetuneDataStrategy.final_and_intermediate,
-                FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+                ChatStrategy.single_turn_r1_thinking,
             ]
         ),
+        tuned_chat_strategy=fine_tune.data_strategy,
     )
 
     if provider == ModelProviderName.vertex and fine_tune.fine_tune_model_id:

kiln_ai/adapters/remote_config.py

@@ -0,0 +1,66 @@
+import argparse
+import json
+import logging
+import os
+import threading
+from pathlib import Path
+from typing import List
+
+import requests
+
+from .ml_model_list import KilnModel, built_in_models
+
+logger = logging.getLogger(__name__)
+
+
+def serialize_config(models: List[KilnModel], path: str | Path) -> None:
+    data = {"model_list": [m.model_dump(mode="json") for m in models]}
+    Path(path).write_text(json.dumps(data, indent=2, sort_keys=True))
+
+
+def deserialize_config(path: str | Path) -> List[KilnModel]:
+    raw = json.loads(Path(path).read_text())
+    model_data = raw.get("model_list", raw if isinstance(raw, list) else [])
+    return [KilnModel.model_validate(item) for item in model_data]
+
+
+def load_from_url(url: str) -> List[KilnModel]:
+    response = requests.get(url, timeout=10)
+    response.raise_for_status()
+    data = response.json()
+    if isinstance(data, list):
+        model_data = data
+    else:
+        model_data = data.get("model_list", [])
+    return [KilnModel.model_validate(item) for item in model_data]
+
+
+def dump_builtin_config(path: str | Path) -> None:
+    serialize_config(built_in_models, path)
+
+
+def load_remote_models(url: str) -> None:
+    if os.environ.get("KILN_SKIP_REMOTE_MODEL_LIST") == "true":
+        return
+
+    def fetch_and_replace() -> None:
+        try:
+            models = load_from_url(url)
+            built_in_models[:] = models
+        except Exception as exc:
+            # Do not crash startup, but surface the issue
+            logger.warning("Failed to fetch remote model list from %s: %s", url, exc)
+
+    thread = threading.Thread(target=fetch_and_replace, daemon=True)
+    thread.start()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("path", help="output path")
+    args = parser.parse_args()
+    dump_builtin_config(args.path)
+
+
+if __name__ == "__main__":
+    main()

kiln_ai/adapters/repair/repair_task.py

@@ -1,13 +1,8 @@
 import json
-from typing import Type
 
 from pydantic import BaseModel, Field
 
-from kiln_ai.adapters.prompt_builders import (
-    BasePromptBuilder,
-    SavedPromptBuilder,
-    prompt_builder_from_id,
-)
+from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_from_id
 from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
 
 

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,
     }
 
 
@@ -224,7 +236,13 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         )
 
         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())
@@ -240,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