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

kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py

@@ -13,11 +13,11 @@ from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetF
 from kiln_ai.adapters.fine_tune.fireworks_finetune import FireworksFinetune
 from kiln_ai.datamodel import (
     DatasetSplit,
-    FinetuneDataStrategy,
     StructuredOutputMode,
     Task,
 )
 from kiln_ai.datamodel import Finetune as FinetuneModel
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config
 
@@ -232,8 +232,8 @@ def mock_task():
 @pytest.mark.parametrize(
     "data_strategy,thinking_instructions",
     [
-        (FinetuneDataStrategy.final_and_intermediate, "thinking instructions"),
-        (FinetuneDataStrategy.final_only, None),
+        (ChatStrategy.two_message_cot, "thinking instructions"),
+        (ChatStrategy.single_turn, None),
     ],
 )
 async def test_generate_and_upload_jsonl_success(

kiln_ai/adapters/fine_tune/test_openai_finetune.py

@@ -12,11 +12,11 @@ from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetF
 from kiln_ai.adapters.fine_tune.openai_finetune import OpenAIFinetune
 from kiln_ai.datamodel import (
     DatasetSplit,
-    FinetuneDataStrategy,
     StructuredOutputMode,
     Task,
 )
 from kiln_ai.datamodel import Finetune as FinetuneModel
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config
 
@@ -35,7 +35,7 @@ def openai_finetune(tmp_path):
             system_message="Test system message",
             fine_tune_model_id="ft-123",
             path=tmp_file,
-            data_strategy=FinetuneDataStrategy.final_only,
+            data_strategy=ChatStrategy.single_turn,
         ),
     )
     return finetune
@@ -247,7 +247,7 @@ async def test_generate_and_upload_jsonl_success(
         mock_formatter.dump_to_file.assert_called_once_with(
             "train",
             DatasetFormat.OPENAI_CHAT_JSONL,
-            FinetuneDataStrategy.final_only,
+            ChatStrategy.single_turn,
         )
 
         # Verify file was opened and uploaded
@@ -299,7 +299,7 @@ async def test_generate_and_upload_jsonl_schema_success(
         mock_formatter.dump_to_file.assert_called_once_with(
             "train",
             DatasetFormat.OPENAI_CHAT_JSON_SCHEMA_JSONL,
-            FinetuneDataStrategy.final_only,
+            ChatStrategy.single_turn,
         )
 
         # Verify file was opened and uploaded
@@ -555,8 +555,8 @@ async def test_status_updates_latest_status(openai_finetune, mock_response):
 @pytest.mark.parametrize(
     "data_strategy,thinking_instructions",
     [
-        (FinetuneDataStrategy.final_and_intermediate, "Custom thinking instructions"),
-        (FinetuneDataStrategy.final_only, None),
+        (ChatStrategy.two_message_cot, "Custom thinking instructions"),
+        (ChatStrategy.single_turn, None),
     ],
 )
 async def test_generate_and_upload_jsonl_with_data_strategy(

kiln_ai/adapters/fine_tune/test_together_finetune.py

@@ -183,6 +183,7 @@ async def test_status_job_states(
     # Mock the retrieve method of the fine_tuning object
     mock_job = MagicMock()
     mock_job.status = together_status
+    mock_job.output_name = None
     mock_together_client.fine_tuning.retrieve.return_value = mock_job
 
     status = await together_finetune.status()

kiln_ai/adapters/fine_tune/test_vertex_finetune.py

@@ -12,11 +12,11 @@ from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetF
 from kiln_ai.adapters.fine_tune.vertex_finetune import VertexFinetune
 from kiln_ai.datamodel import (
     DatasetSplit,
-    FinetuneDataStrategy,
     StructuredOutputMode,
     Task,
 )
 from kiln_ai.datamodel import Finetune as FinetuneModel
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config
 
@@ -35,7 +35,7 @@ def vertex_finetune(tmp_path):
             system_message="Test system message",
             fine_tune_model_id="ft-123",
             path=tmp_file,
-            data_strategy=FinetuneDataStrategy.final_only,
+            data_strategy=ChatStrategy.single_turn,
         ),
     )
     return finetune
@@ -252,8 +252,8 @@ async def test_status_model_id_update_exception(vertex_finetune, mock_response):
 @pytest.mark.parametrize(
     "data_strategy,thinking_instructions",
     [
-        (FinetuneDataStrategy.final_and_intermediate, "Custom thinking instructions"),
-        (FinetuneDataStrategy.final_only, None),
+        (ChatStrategy.two_message_cot, "Custom thinking instructions"),
+        (ChatStrategy.single_turn, None),
     ],
 )
 async def test_generate_and_upload_jsonl(

kiln_ai/adapters/fine_tune/together_finetune.py

@@ -66,6 +66,12 @@ class TogetherFinetune(BaseFinetuneAdapter):
             # retrieve the fine-tuning job
             together_finetune = self.client.fine_tuning.retrieve(id=fine_tuning_job_id)
 
+            # update the fine tune model ID if it has changed (sometimes it's not set at training time)
+            if self.datamodel.fine_tune_model_id != together_finetune.output_name:
+                self.datamodel.fine_tune_model_id = together_finetune.output_name
+                if self.datamodel.path:
+                    self.datamodel.save_to_file()
+
             status = together_finetune.status
             if status in _pending_statuses:
                 return FineTuneStatus(
@@ -135,8 +141,13 @@ class TogetherFinetune(BaseFinetuneAdapter):
             **self._build_finetune_parameters(),
         )
 
-        # 2 different IDs, output_name is the name of the model that results from the fine-tune job, the finetune_job_id is the ID of the fine-tune job
+        # 2 different IDs, output_name is the name of the model that results from the fine-tune job, while the id is the ID of the fine-tune job itself
+        if not together_finetune.id:
+            raise ValueError(
+                "Together failed to return a fine-tune job ID. While tuning job was dispatched, Kiln never received the ID so won't be able to reference it. Check for errors before dispatching more jobs."
+            )
         self.datamodel.provider_id = together_finetune.id
+        # Output name is sometimes returned here, and save it if it is. But it might be populated later by status call
         self.datamodel.fine_tune_model_id = together_finetune.output_name
 
         if self.datamodel.path:

kiln_ai/adapters/ml_model_list.py

@@ -3,7 +3,11 @@ from typing import Dict, List, Literal
 
 from pydantic import BaseModel
 
-from kiln_ai.datamodel import StructuredOutputMode
+from kiln_ai.datamodel.datamodel_enums import (
+    ChatStrategy,
+    ModelProviderName,
+    StructuredOutputMode,
+)
 
 """
 Provides model configuration and management for various LLM providers and models.
@@ -12,28 +16,6 @@ including configuration, validation, and instantiation of language models.
 """
 
 
-class ModelProviderName(str, Enum):
-    """
-    Enumeration of supported AI model providers.
-    """
-
-    openai = "openai"
-    groq = "groq"
-    amazon_bedrock = "amazon_bedrock"
-    ollama = "ollama"
-    openrouter = "openrouter"
-    fireworks_ai = "fireworks_ai"
-    kiln_fine_tune = "kiln_fine_tune"
-    kiln_custom_registry = "kiln_custom_registry"
-    openai_compatible = "openai_compatible"
-    anthropic = "anthropic"
-    gemini_api = "gemini_api"
-    azure_openai = "azure_openai"
-    huggingface = "huggingface"
-    vertex = "vertex"
-    together_ai = "together_ai"
-
-
 class ModelFamily(str, Enum):
     """
     Enumeration of supported model families/architectures.
@@ -102,6 +84,8 @@ class ModelName(str, Enum):
     claude_3_5_sonnet = "claude_3_5_sonnet"
     claude_3_7_sonnet = "claude_3_7_sonnet"
     claude_3_7_sonnet_thinking = "claude_3_7_sonnet_thinking"
+    claude_sonnet_4 = "claude_sonnet_4"
+    claude_opus_4 = "claude_opus_4"
     gemini_1_5_flash = "gemini_1_5_flash"
     gemini_1_5_flash_8b = "gemini_1_5_flash_8b"
     gemini_1_5_pro = "gemini_1_5_pro"
@@ -174,6 +158,7 @@ class KilnModelProvider(BaseModel):
         structured_output_mode: The mode we should use to call the model for structured output, if it was trained with structured output.
         parser: A parser to use for the model, if applicable
         reasoning_capable: Whether the model is designed to output thinking in a structured format (eg <think></think>). If so we don't use COT across 2 calls, and ask for thinking and final response in the same call.
+        tuned_chat_strategy: Used when a model is finetuned with a specific chat strategy, and it's best to use it at call time.
     """
 
     name: ModelProviderName
@@ -189,6 +174,7 @@ class KilnModelProvider(BaseModel):
     reasoning_capable: bool = False
     supports_logprobs: bool = False
     suggested_for_evals: bool = False
+    tuned_chat_strategy: ChatStrategy | None = None
 
     # TODO P1: Need a more generalized way to handle custom provider parameters.
     # Making them quite declarative here for now, isolating provider specific logic
@@ -666,14 +652,12 @@ built_in_models: List[KilnModel] = [
                 structured_output_mode=StructuredOutputMode.function_calling,
                 model_id="anthropic/claude-3.7-sonnet",
                 suggested_for_data_gen=True,
-                suggested_for_evals=True,
             ),
             KilnModelProvider(
                 name=ModelProviderName.anthropic,
                 model_id="claude-3-7-sonnet-20250219",
                 structured_output_mode=StructuredOutputMode.function_calling,
                 suggested_for_data_gen=True,
-                suggested_for_evals=True,
             ),
         ],
     ),
@@ -700,6 +684,46 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # Claude Sonnet 4
+    KilnModel(
+        family=ModelFamily.claude,
+        name=ModelName.claude_sonnet_4,
+        friendly_name="Claude Sonnet 4",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="anthropic/claude-sonnet-4",
+                structured_output_mode=StructuredOutputMode.function_calling,
+                suggested_for_data_gen=True,
+                suggested_for_evals=True,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.anthropic,
+                model_id="claude-sonnet-4-20250514",
+                structured_output_mode=StructuredOutputMode.function_calling,
+                suggested_for_data_gen=True,
+                suggested_for_evals=True,
+            ),
+        ],
+    ),
+    # Claude Opus 4
+    KilnModel(
+        family=ModelFamily.claude,
+        name=ModelName.claude_opus_4,
+        friendly_name="Claude Opus 4",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="anthropic/claude-opus-4",
+                structured_output_mode=StructuredOutputMode.function_calling,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.anthropic,
+                model_id="claude-opus-4-20250514",
+                structured_output_mode=StructuredOutputMode.function_calling,
+            ),
+        ],
+    ),
     # Gemini 2.5 Pro
     KilnModel(
         family=ModelFamily.gemini,
@@ -920,7 +944,7 @@ built_in_models: List[KilnModel] = [
                 model_id="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
                 supports_data_gen=False,
                 structured_output_mode=StructuredOutputMode.function_calling_weak,
-                provider_finetune_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
+                provider_finetune_id="meta-llama/Meta-Llama-3.1-8B-Instruct-Reference",
             ),
         ],
     ),
@@ -939,8 +963,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_data_gen=False,
-                # Need to not pass "strict=True" to the function call to get this to work with logprobs for some reason. Openrouter issue.
-                structured_output_mode=StructuredOutputMode.function_calling_weak,
+                structured_output_mode=StructuredOutputMode.json_schema,
                 model_id="meta-llama/llama-3.1-70b-instruct",
                 supports_logprobs=True,
                 logprobs_openrouter_options=True,
@@ -961,7 +984,7 @@ built_in_models: List[KilnModel] = [
                 model_id="meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
                 supports_data_gen=False,
                 structured_output_mode=StructuredOutputMode.function_calling_weak,
-                provider_finetune_id="meta-llama/Meta-Llama-3.1-70B-Instruct",
+                provider_finetune_id="meta-llama/Meta-Llama-3.1-70B-Instruct-Reference",
             ),
         ],
     ),
@@ -1068,10 +1091,6 @@ built_in_models: List[KilnModel] = [
                 supports_structured_output=False,
                 supports_data_gen=False,
             ),
-            KilnModelProvider(
-                name=ModelProviderName.together_ai,
-                provider_finetune_id="meta-llama/Llama-3.2-1B-Instruct",
-            ),
         ],
     ),
     # Llama 3.2 3B
@@ -1108,7 +1127,6 @@ built_in_models: List[KilnModel] = [
                 model_id="meta-llama/Llama-3.2-3B-Instruct-Turbo",
                 supports_structured_output=False,
                 supports_data_gen=False,
-                provider_finetune_id="meta-llama/Llama-3.2-3B-Instruct",
             ),
         ],
     ),
@@ -1158,10 +1176,6 @@ built_in_models: List[KilnModel] = [
         name=ModelName.llama_3_2_90b,
         friendly_name="Llama 3.2 90B",
         providers=[
-            KilnModelProvider(
-                name=ModelProviderName.groq,
-                model_id="llama-3.2-90b-vision-preview",
-            ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 structured_output_mode=StructuredOutputMode.json_instruction_and_object,
@@ -1434,11 +1448,6 @@ built_in_models: List[KilnModel] = [
                 structured_output_mode=StructuredOutputMode.json_instruction_and_object,
                 model_id="google/gemma-3-27b-it",
             ),
-            KilnModelProvider(
-                name=ModelProviderName.huggingface,
-                model_id="google/gemma-3-27b-it",
-                structured_output_mode=StructuredOutputMode.json_instructions,
-            ),
         ],
     ),
     # Mixtral 8x7B
@@ -2294,3 +2303,31 @@ def get_model_by_name(name: ModelName) -> KilnModel:
         if model.name == name:
             return model
     raise ValueError(f"Model {name} not found in the list of built-in models")
+
+
+def default_structured_output_mode_for_model_provider(
+    model_name: str,
+    provider: ModelProviderName,
+    default: StructuredOutputMode = StructuredOutputMode.default,
+    disallowed_modes: List[StructuredOutputMode] = [],
+) -> StructuredOutputMode:
+    """
+    We don't expose setting this manually in the UI, so pull a recommended mode from ml_model_list
+    """
+    try:
+        # Convert string to ModelName enum
+        model_name_enum = ModelName(model_name)
+        model = get_model_by_name(model_name_enum)
+    except (ValueError, KeyError):
+        # If model not found, return default
+        return default
+
+    # Find the provider within the model's providers
+    for model_provider in model.providers:
+        if model_provider.name == provider:
+            mode = model_provider.structured_output_mode
+            if mode not in disallowed_modes:
+                return mode
+
+    # If provider not found, return default
+    return default

kiln_ai/adapters/model_adapters/base_adapter.py

@@ -1,11 +1,17 @@
 import json
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
-from typing import Dict, Literal, Tuple
+from typing import Dict, Tuple
 
-import jsonschema
-
-from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMode
+from kiln_ai.adapters.chat.chat_formatter import (
+    ChatFormatter,
+    get_chat_formatter,
+)
+from kiln_ai.adapters.ml_model_list import (
+    KilnModelProvider,
+    StructuredOutputMode,
+    default_structured_output_mode_for_model_provider,
+)
 from kiln_ai.adapters.parsers.json_parser import parse_json_string
 from kiln_ai.adapters.parsers.parser_registry import model_parser_from_id
 from kiln_ai.adapters.parsers.request_formatters import request_formatter_from_id
@@ -20,6 +26,7 @@ from kiln_ai.datamodel import (
     TaskRun,
     Usage,
 )
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 from kiln_ai.datamodel.task import RunConfig
 from kiln_ai.utils.config import Config
@@ -38,9 +45,6 @@ class AdapterConfig:
     default_tags: list[str] | None = None
 
 
-COT_FINAL_ANSWER_PROMPT = "Considering the above, return a final result."
-
-
 class BaseAdapter(metaclass=ABCMeta):
     """Base class for AI model adapters that handle task execution.
 
@@ -61,6 +65,7 @@ class BaseAdapter(metaclass=ABCMeta):
         config: AdapterConfig | None = None,
     ):
         self.run_config = run_config
+        self.update_run_config_unknown_structured_output_mode()
         self.prompt_builder = prompt_builder_from_id(
             run_config.prompt_id, run_config.task
         )
@@ -188,10 +193,10 @@ class BaseAdapter(metaclass=ABCMeta):
 
     def build_prompt(self) -> str:
         # The prompt builder needs to know if we want to inject formatting instructions
-        provider = self.model_provider()
+        structured_output_mode = self.run_config.structured_output_mode
         add_json_instructions = self.has_structured_output() and (
-            provider.structured_output_mode == StructuredOutputMode.json_instructions
-            or provider.structured_output_mode
+            structured_output_mode == StructuredOutputMode.json_instructions
+            or structured_output_mode
             == StructuredOutputMode.json_instruction_and_object
         )
 
@@ -199,26 +204,51 @@ class BaseAdapter(metaclass=ABCMeta):
             include_json_instructions=add_json_instructions
         )
 
-    def run_strategy(
-        self,
-    ) -> Tuple[Literal["cot_as_message", "cot_two_call", "basic"], str | None]:
-        # Determine the run strategy for COT prompting. 3 options:
-        # 1. "Thinking" LLM designed to output thinking in a structured format plus a COT prompt: we make 1 call to the LLM, which outputs thinking in a structured format. We include the thinking instuctions as a message.
-        # 2. Normal LLM with COT prompt: we make 2 calls to the LLM - one for thinking and one for the final response. This helps us use the LLM's structured output modes (json_schema, tools, etc), which can't be used in a single call. It also separates the thinking from the final response.
-        # 3. Non chain of thought: we make 1 call to the LLM, with no COT prompt.
+    def build_chat_formatter(self, input: Dict | str) -> ChatFormatter:
+        # Determine the chat strategy to use based on the prompt the user selected, the model's capabilities, and if the model was finetuned with a specific chat strategy.
+
         cot_prompt = self.prompt_builder.chain_of_thought_prompt()
-        reasoning_capable = self.model_provider().reasoning_capable
+        system_message = self.build_prompt()
+
+        # If no COT prompt, use the single turn strategy. Even when a tuned strategy is set, as the tuned strategy is either already single turn, or won't work without a COT prompt.
+        if not cot_prompt:
+            return get_chat_formatter(
+                strategy=ChatStrategy.single_turn,
+                system_message=system_message,
+                user_input=input,
+            )
 
-        if cot_prompt and reasoning_capable:
-            # 1: "Thinking" LLM designed to output thinking in a structured format
+        # Some models like finetunes are trained with a specific chat strategy. Use that.
+        # However, don't use that if it is single turn. The user selected a COT prompt, and we give explicit prompt selection priority over the tuned strategy.
+        tuned_chat_strategy = self.model_provider().tuned_chat_strategy
+        if tuned_chat_strategy and tuned_chat_strategy != ChatStrategy.single_turn:
+            return get_chat_formatter(
+                strategy=tuned_chat_strategy,
+                system_message=system_message,
+                user_input=input,
+                thinking_instructions=cot_prompt,
+            )
+
+        # Pick the best chat strategy for the model given it has a cot prompt.
+        reasoning_capable = self.model_provider().reasoning_capable
+        if reasoning_capable:
+            # "Thinking" LLM designed to output thinking in a structured format. We'll use it's native format.
             # A simple message with the COT prompt appended to the message list is sufficient
-            return "cot_as_message", cot_prompt
-        elif cot_prompt:
-            # 2: Unstructured output with COT
-            # Two calls to separate the thinking from the final response
-            return "cot_two_call", cot_prompt
+            return get_chat_formatter(
+                strategy=ChatStrategy.single_turn_r1_thinking,
+                system_message=system_message,
+                user_input=input,
+                thinking_instructions=cot_prompt,
+            )
         else:
-            return "basic", None
+            # Unstructured output with COT
+            # Two calls to separate the thinking from the final response
+            return get_chat_formatter(
+                strategy=ChatStrategy.two_message_cot,
+                system_message=system_message,
+                user_input=input,
+                thinking_instructions=cot_prompt,
+            )
 
     # create a run and task output
     def generate_run(
@@ -272,5 +302,22 @@ class BaseAdapter(metaclass=ABCMeta):
         props["model_name"] = self.run_config.model_name
         props["model_provider"] = self.run_config.model_provider_name
         props["prompt_id"] = self.run_config.prompt_id
+        props["structured_output_mode"] = self.run_config.structured_output_mode
+        props["temperature"] = self.run_config.temperature
+        props["top_p"] = self.run_config.top_p
 
         return props
+
+    def update_run_config_unknown_structured_output_mode(self) -> None:
+        structured_output_mode = self.run_config.structured_output_mode
+
+        # Old datamodels didn't save the structured output mode. Some clients (tests, end users) might not set it.
+        # Look up our recommended mode from ml_model_list if we have one
+        if structured_output_mode == StructuredOutputMode.unknown:
+            new_run_config = self.run_config.model_copy(deep=True)
+            structured_output_mode = default_structured_output_mode_for_model_provider(
+                self.run_config.model_name,
+                self.run_config.model_provider_name,
+            )
+            new_run_config.structured_output_mode = structured_output_mode
+            self.run_config = new_run_config