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

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,32 @@ def test_litellm_model_id_unknown_provider(config, mock_task):
                 adapter.litellm_model_id()
 
 
+@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 +419,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

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

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
     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 = 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/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/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