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

kiln_ai/datamodel/datamodel_enums.py

@@ -24,13 +24,14 @@ class StructuredOutputMode(str, Enum):
     """
     Enumeration of supported structured output modes.
 
-    - default: let the adapter decide
     - json_schema: request json using API capabilities for json_schema
     - function_calling: request json using API capabilities for function calling
     - json_mode: request json using API's JSON mode, which should return valid JSON, but isn't checking/passing the schema
     - json_instructions: append instructions to the prompt to request json matching the schema. No API capabilities are used. You should have a custom parser on these models as they will be returning strings.
     - json_instruction_and_object: append instructions to the prompt to request json matching the schema. Also request the response as json_mode via API capabilities (returning dictionaries).
     - json_custom_instructions: The model should output JSON, but custom instructions are already included in the system prompt. Don't append additional JSON instructions.
+    - default: let the adapter decide (legacy, do not use for new use cases)
+    - unknown: used for cases where the structured output mode is not known (on old models where it wasn't saved). Should lookup best option at runtime.
     """
 
     default = "default"
@@ -41,6 +42,7 @@ class StructuredOutputMode(str, Enum):
     json_instructions = "json_instructions"
     json_instruction_and_object = "json_instruction_and_object"
     json_custom_instructions = "json_custom_instructions"
+    unknown = "unknown"
 
 
 class FineTuneStatusType(str, Enum):
@@ -55,20 +57,43 @@ class FineTuneStatusType(str, Enum):
     failed = "failed"
 
 
-class FinetuneDataStrategy(str, Enum):
-    """Strategy for what data to include when fine-tuning a model."""
+class ChatStrategy(str, Enum):
+    """Strategy for how a chat is structured."""
 
-    # Only train on the final response, ignoring any intermediate steps or chain of thought
-    final_only = "final_only"
+    # Single turn, immediately return the answer
+    single_turn = "final_only"
+    # Two turn, first turn is the thinking, second turn is the answer. Legacy format - used for old fine tunes but not new trains.
+    two_message_cot_legacy = "final_and_intermediate"
+    # Two turn, first turn is the thinking, second turn is the answer. New format - used for new trains.
+    two_message_cot = "two_message_cot"
+    # Single turn, with both the thinking and the answer in the same message, using R1-style thinking format in <think> tags
+    single_turn_r1_thinking = "final_and_intermediate_r1_compatible"
 
-    # Train on both the final response and any intermediate steps/chain of thought
-    final_and_intermediate = "final_and_intermediate"
 
-    # Train using R1-style thinking format, which includes the reasoning in <think> tags in the message
-    final_and_intermediate_r1_compatible = "final_and_intermediate_r1_compatible"
+THINKING_DATA_STRATEGIES: list[ChatStrategy] = [
+    ChatStrategy.two_message_cot_legacy,
+    ChatStrategy.single_turn_r1_thinking,
+    ChatStrategy.two_message_cot,
+]
 
 
-THINKING_DATA_STRATEGIES: list[FinetuneDataStrategy] = [
-    FinetuneDataStrategy.final_and_intermediate,
-    FinetuneDataStrategy.final_and_intermediate_r1_compatible,
-]
+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"

kiln_ai/datamodel/eval.py

@@ -14,6 +14,7 @@ from kiln_ai.datamodel.basemodel import (
 from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
 from kiln_ai.datamodel.dataset_filters import DatasetFilterId
 from kiln_ai.datamodel.json_schema import string_to_json_key
+from kiln_ai.datamodel.task_run import Usage
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 if TYPE_CHECKING:
@@ -28,6 +29,7 @@ class EvalTemplateId(str, Enum):
     """
 
     kiln_requirements = "kiln_requirements"
+    issue = "kiln_issue"
     toxicity = "toxicity"
     bias = "bias"
     maliciousness = "maliciousness"
@@ -110,6 +112,10 @@ class EvalRun(KilnParentedModel):
     scores: EvalScores = Field(
         description="The output scores of the evaluator (aligning to those required by the grand-parent Eval this object is a child of)."
     )
+    task_run_usage: Usage | None = Field(
+        default=None,
+        description="The usage of the task run that produced this eval run output (not the usage by the evaluation model).",
+    )
 
     def parent_eval_config(self) -> Union["EvalConfig", None]:
         if self.parent is not None and self.parent.__class__.__name__ != "EvalConfig":
@@ -280,6 +286,10 @@ class Eval(KilnParentedModel, KilnParentModel, parent_of={"configs": EvalConfig}
         default=False,
         description="Whether this eval is a favourite of the user. Rendered as a star icon in the UI.",
     )
+    template_properties: dict[str, str | int | bool | float] = Field(
+        default={},
+        description="Properties to be used to execute the eval. This is template_type specific and should serialize to a json dict.",
+    )
 
     # Workaround to return typed parent without importing Task
     def parent_task(self) -> Union["Task", None]:
@@ -304,3 +314,25 @@ class Eval(KilnParentedModel, KilnParentModel, parent_of={"configs": EvalConfig}
                 f"output_scores must have unique names (once transformed to JSON keys). Got: [{', '.join(output_score_keys)}]"
             )
         return self
+
+    @model_validator(mode="after")
+    def validate_template_properties(self) -> Self:
+        # Check for properties that are required for the issue template
+        if self.template == EvalTemplateId.issue:
+            if "issue_prompt" not in self.template_properties or not isinstance(
+                self.template_properties["issue_prompt"], str
+            ):
+                raise ValueError("issue_prompt is required for issue template")
+            if "failure_example" in self.template_properties and not isinstance(
+                self.template_properties["failure_example"], str
+            ):
+                raise ValueError(
+                    "failure_example is optional for issue template, but if provided must be a string"
+                )
+            if "pass_example" in self.template_properties and not isinstance(
+                self.template_properties["pass_example"], str
+            ):
+                raise ValueError(
+                    "pass_example is optional for issue template, but if provided must be a string"
+                )
+        return self

kiln_ai/datamodel/finetune.py

@@ -5,8 +5,7 @@ from typing_extensions import Self
 
 from kiln_ai.datamodel.basemodel import NAME_FIELD, KilnParentedModel
 from kiln_ai.datamodel.datamodel_enums import (
-    THINKING_DATA_STRATEGIES,
-    FinetuneDataStrategy,
+    ChatStrategy,
     FineTuneStatusType,
     StructuredOutputMode,
 )
@@ -14,6 +13,11 @@ from kiln_ai.datamodel.datamodel_enums import (
 if TYPE_CHECKING:
     from kiln_ai.datamodel.task import Task
 
+DATA_STRATIGIES_REQUIRED_THINKING_INSTRUCTIONS = [
+    ChatStrategy.two_message_cot_legacy,
+    ChatStrategy.two_message_cot,
+]
+
 
 class Finetune(KilnParentedModel):
     """
@@ -76,8 +80,8 @@ class Finetune(KilnParentedModel):
         default={},
         description="Properties of the fine-tune. Different providers may use different properties.",
     )
-    data_strategy: FinetuneDataStrategy = Field(
-        default=FinetuneDataStrategy.final_only,
+    data_strategy: ChatStrategy = Field(
+        default=ChatStrategy.single_turn,
         description="The strategy to use for training the model. 'final_only' will only train on the final response. 'final_and_intermediate' will train on the final response and intermediate outputs (chain of thought or reasoning).",
     )
 
@@ -91,16 +95,16 @@ class Finetune(KilnParentedModel):
     def validate_thinking_instructions(self) -> Self:
         if (
             self.thinking_instructions is not None
-            and self.data_strategy != FinetuneDataStrategy.final_and_intermediate
+            and self.data_strategy not in DATA_STRATIGIES_REQUIRED_THINKING_INSTRUCTIONS
         ):
             raise ValueError(
-                "Thinking instructions can only be used when data_strategy is final_and_intermediate"
+                f"Thinking instructions can only be used when data_strategy is one of the following: {DATA_STRATIGIES_REQUIRED_THINKING_INSTRUCTIONS}"
             )
         if (
             self.thinking_instructions is None
-            and self.data_strategy == FinetuneDataStrategy.final_and_intermediate
+            and self.data_strategy in DATA_STRATIGIES_REQUIRED_THINKING_INSTRUCTIONS
         ):
             raise ValueError(
-                "Thinking instructions are required when data_strategy is final_and_intermediate"
+                f"Thinking instructions are required when data_strategy is one of the following: {DATA_STRATIGIES_REQUIRED_THINKING_INSTRUCTIONS}"
             )
         return self

kiln_ai/datamodel/task.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Dict, List, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationInfo, model_validator
+from typing_extensions import Self
 
 from kiln_ai.datamodel import Finetune
 from kiln_ai.datamodel.basemodel import (
@@ -11,7 +12,12 @@ from kiln_ai.datamodel.basemodel import (
     KilnParentedModel,
     KilnParentModel,
 )
-from kiln_ai.datamodel.datamodel_enums import Priority, TaskOutputRatingType
+from kiln_ai.datamodel.datamodel_enums import (
+    ModelProviderName,
+    Priority,
+    StructuredOutputMode,
+    TaskOutputRatingType,
+)
 from kiln_ai.datamodel.dataset_split import DatasetSplit
 from kiln_ai.datamodel.eval import Eval
 from kiln_ai.datamodel.json_schema import JsonObjectSchema, schema_from_json_str
@@ -47,12 +53,33 @@ class RunConfigProperties(BaseModel):
     """
 
     model_name: str = Field(description="The model to use for this run config.")
-    model_provider_name: str = Field(
+    model_provider_name: ModelProviderName = Field(
         description="The provider to use for this run config."
     )
     prompt_id: PromptId = Field(
         description="The prompt to use for this run config. Defaults to building a simple prompt from the task if not provided.",
     )
+    top_p: float = Field(
+        default=1.0,
+        description="The top-p value to use for this run config. Defaults to 1.0.",
+    )
+    temperature: float = Field(
+        default=1.0,
+        description="The temperature to use for this run config. Defaults to 1.0.",
+    )
+    structured_output_mode: StructuredOutputMode = Field(
+        description="The structured output mode to use for this run config.",
+    )
+
+    @model_validator(mode="after")
+    def validate_required_fields(self) -> Self:
+        if not (0 <= self.top_p <= 1):
+            raise ValueError("top_p must be between 0 and 1")
+
+        elif self.temperature < 0 or self.temperature > 2:
+            raise ValueError("temperature must be between 0 and 2")
+
+        return self
 
 
 class RunConfig(RunConfigProperties):
@@ -101,12 +128,46 @@ class TaskRunConfig(KilnParentedModel):
         parent_task = self.parent_task()
         if parent_task is None:
             raise ValueError("Run config must be parented to a task")
-        return RunConfig(
+        return run_config_from_run_config_properties(
             task=parent_task,
-            model_name=self.run_config_properties.model_name,
-            model_provider_name=self.run_config_properties.model_provider_name,
-            prompt_id=self.run_config_properties.prompt_id,
+            run_config_properties=self.run_config_properties,
+        )
+
+    # Previously we didn't store structured_output_mode in the run_config_properties. Updgrade old models when loading from file.
+    @model_validator(mode="before")
+    def upgrade_old_entries(cls, data: dict, info: ValidationInfo) -> dict:
+        if not info.context or not info.context.get("loading_from_file", False):
+            # Not loading from file, so no need to upgrade
+            return data
+
+        if not isinstance(data, dict):
+            return data
+
+        structured_output_mode = data.get("run_config_properties", {}).get(
+            "structured_output_mode", None
         )
+        if structured_output_mode is None and "run_config_properties" in data:
+            # Default to unknown. Adapter will have to guess at runtime.
+            data["run_config_properties"]["structured_output_mode"] = (
+                StructuredOutputMode.unknown
+            )
+
+        return data
+
+
+def run_config_from_run_config_properties(
+    task: "Task",
+    run_config_properties: RunConfigProperties,
+) -> RunConfig:
+    return RunConfig(
+        task=task,
+        model_name=run_config_properties.model_name,
+        model_provider_name=run_config_properties.model_provider_name,
+        prompt_id=run_config_properties.prompt_id,
+        top_p=run_config_properties.top_p,
+        temperature=run_config_properties.temperature,
+        structured_output_mode=run_config_properties.structured_output_mode,
+    )
 
 
 class Task(

kiln_ai/datamodel/task_output.py

@@ -2,8 +2,6 @@ import json
 from enum import Enum
 from typing import TYPE_CHECKING, Dict, List, Type, Union
 
-import jsonschema
-import jsonschema.exceptions
 from pydantic import BaseModel, Field, ValidationInfo, model_validator
 from typing_extensions import Self
 

kiln_ai/datamodel/task_run.py

@@ -1,8 +1,6 @@
 import json
 from typing import TYPE_CHECKING, Dict, List, Union
 
-import jsonschema
-import jsonschema.exceptions
 from pydantic import BaseModel, Field, ValidationInfo, model_validator
 from typing_extensions import Self
 

kiln_ai/datamodel/test_basemodel.py

@@ -500,8 +500,9 @@ def adapter(base_task):
         run_config=RunConfig(
             task=base_task,
             model_name="test_model",
-            model_provider_name="test_provider",
+            model_provider_name="openai",
             prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
         ),
     )
 

kiln_ai/datamodel/test_dataset_split.py

@@ -72,14 +72,6 @@ def sample_task_runs(sample_task):
     return task_runs
 
 
-@pytest.fixture
-def standard_splitstandard_splitss():
-    return [
-        DatasetSplitDefinition(name="train", percentage=0.8),
-        DatasetSplitDefinition(name="test", percentage=0.2),
-    ]
-
-
 @pytest.fixture
 def task_run():
     return TaskRun(

kiln_ai/datamodel/test_eval_model.py

@@ -1,7 +1,6 @@
 import pytest
 from pydantic import ValidationError
 
-from kiln_ai.datamodel import BasePrompt
 from kiln_ai.datamodel.basemodel import KilnParentModel
 from kiln_ai.datamodel.eval import (
     Eval,
@@ -9,11 +8,10 @@ from kiln_ai.datamodel.eval import (
     EvalConfigType,
     EvalOutputScore,
     EvalRun,
+    EvalTemplateId,
 )
 from kiln_ai.datamodel.task import Task
-from kiln_ai.datamodel.task_output import (
-    TaskOutputRatingType,
-)
+from kiln_ai.datamodel.task_output import TaskOutputRatingType
 
 
 @pytest.fixture
@@ -633,3 +631,147 @@ def test_eval_run_eval_config_eval_validation():
             output="test output",
             scores={"score": 1.0},
         )
+
+
+@pytest.mark.parametrize(
+    "template_properties,should_raise,expected_error",
+    [
+        # Valid cases
+        (
+            {"issue_prompt": "Test issue prompt"},
+            False,
+            None,
+        ),
+        (
+            {
+                "issue_prompt": "Test issue prompt",
+                "failure_example": "Test failure example",
+            },
+            False,
+            None,
+        ),
+        (
+            {
+                "issue_prompt": "Test issue prompt",
+                "failure_example": "Test failure example",
+                "pass_example": "Test pass example",
+            },
+            False,
+            None,
+        ),
+        (
+            {
+                "issue_prompt": "",
+                "failure_example": "",
+                "pass_example": "",
+            },
+            False,
+            None,
+        ),
+        # Invalid cases
+        (
+            {},
+            True,
+            "issue_prompt is required for issue template",
+        ),
+        (
+            {"failure_example": "Test failure example"},
+            True,
+            "issue_prompt is required for issue template",
+        ),
+        (
+            {"issue_prompt": 123},
+            True,
+            "issue_prompt is required for issue template",
+        ),
+        (
+            {
+                "issue_prompt": "Test issue prompt",
+                "failure_example": 456,
+            },
+            True,
+            "failure_example is optional for issue template, but if provided must be a string",
+        ),
+        (
+            {
+                "issue_prompt": "Test issue prompt",
+                "failure_example": "Test failure example",
+                "pass_example": 789,
+            },
+            True,
+            "pass_example is optional for issue template, but if provided must be a string",
+        ),
+    ],
+)
+def test_eval_template_properties_issue_template_validation(
+    template_properties, should_raise, expected_error
+):
+    """Test issue template validation with various property combinations"""
+    if should_raise:
+        with pytest.raises(ValueError, match=expected_error):
+            Eval(
+                name="Test Eval",
+                template=EvalTemplateId.issue,
+                eval_set_filter_id="tag::tag1",
+                eval_configs_filter_id="tag::tag2",
+                output_scores=[
+                    EvalOutputScore(
+                        name="score",
+                        type=TaskOutputRatingType.pass_fail,
+                    )
+                ],
+                template_properties=template_properties,
+            )
+    else:
+        eval = Eval(
+            name="Test Eval",
+            template=EvalTemplateId.issue,
+            eval_set_filter_id="tag::tag1",
+            eval_configs_filter_id="tag::tag2",
+            output_scores=[
+                EvalOutputScore(
+                    name="score",
+                    type=TaskOutputRatingType.pass_fail,
+                )
+            ],
+            template_properties=template_properties,
+        )
+        assert eval.template == EvalTemplateId.issue
+        for key, value in template_properties.items():
+            assert eval.template_properties[key] == value
+
+
+@pytest.mark.parametrize(
+    "template,template_properties",
+    [
+        (EvalTemplateId.kiln_requirements, {"random_property": "random_value"}),
+        (EvalTemplateId.toxicity, {}),
+        (EvalTemplateId.bias, {"some_property": 123}),
+        (EvalTemplateId.maliciousness, {"test": True}),
+        (EvalTemplateId.factual_correctness, {"score": 4.5}),
+        (EvalTemplateId.jailbreak, {"prompt": "test"}),
+        (
+            None,
+            {"issue_prompt": "This should not be validated", "failure_example": 123},
+        ),
+    ],
+)
+def test_eval_template_properties_non_issue_templates(template, template_properties):
+    """Test that non-issue templates pass validation regardless of template_properties"""
+    eval = Eval(
+        name="Test Eval",
+        template=template,
+        eval_set_filter_id="tag::tag1",
+        eval_configs_filter_id="tag::tag2",
+        output_scores=[
+            EvalOutputScore(
+                name="score",
+                type=TaskOutputRatingType.pass_fail,
+            )
+        ],
+        template_properties=template_properties,
+    )
+
+    assert eval.template == template
+    for key, value in template_properties.items():
+        assert eval.template_properties[key] == value

kiln_ai/datamodel/test_models.py

@@ -9,13 +9,13 @@ from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
     Finetune,
-    FinetuneDataStrategy,
     Project,
     Prompt,
     Task,
     TaskOutput,
     TaskRun,
 )
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.test_json_schema import json_joke_schema
 
 
@@ -536,44 +536,58 @@ def test_prompt_parent_task():
         # Test 1: Valid case - no thinking instructions with final_only
         (
             None,
-            FinetuneDataStrategy.final_only,
+            ChatStrategy.single_turn,
             False,
             None,
         ),
         # Test 2: Valid case - thinking instructions with final_and_intermediate
         (
             "Think step by step",
-            FinetuneDataStrategy.final_and_intermediate,
+            ChatStrategy.two_message_cot_legacy,
             False,
             None,
         ),
         # Test 3: Valid case - no thinking instructions with final_and_intermediate_r1_compatible
         (
             None,
-            FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+            ChatStrategy.single_turn_r1_thinking,
             False,
             None,
         ),
         # Test 4: Invalid case - thinking instructions with final_only
         (
             "Think step by step",
-            FinetuneDataStrategy.final_only,
+            ChatStrategy.single_turn,
             True,
-            "Thinking instructions can only be used when data_strategy is final_and_intermediate",
+            "Thinking instructions can only be used when data_strategy is",
         ),
         # Test 5: Invalid case - no thinking instructions with final_and_intermediate
         (
             None,
-            FinetuneDataStrategy.final_and_intermediate,
+            ChatStrategy.two_message_cot_legacy,
             True,
-            "Thinking instructions are required when data_strategy is final_and_intermediate",
+            "Thinking instructions are required when data_strategy is",
         ),
         # Test 6: Invalid case - thinking instructions with final_and_intermediate_r1_compatible
         (
             "Think step by step",
-            FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+            ChatStrategy.single_turn_r1_thinking,
             True,
-            "Thinking instructions can only be used when data_strategy is final_and_intermediate",
+            "Thinking instructions can only be used when data_strategy is",
+        ),
+        # Test 7: new COT format
+        (
+            "Think step by step",
+            ChatStrategy.two_message_cot,
+            False,
+            None,
+        ),
+        # Test 8: new COT format
+        (
+            None,
+            ChatStrategy.two_message_cot,
+            True,
+            "Thinking instructions are required when data_strategy is",
         ),
     ],
 )
@@ -665,3 +679,12 @@ def test_task_run_thinking_training_data(intermediate_outputs, expected):
         intermediate_outputs=intermediate_outputs,
     )
     assert task_run.thinking_training_data() == expected
+
+
+def test_chat_strategy_enum():
+    # This has to align to the old FinetuneDataStrategy enum
+    assert ChatStrategy.single_turn == "final_only"
+    assert ChatStrategy.two_message_cot_legacy == "final_and_intermediate"
+    assert (
+        ChatStrategy.single_turn_r1_thinking == "final_and_intermediate_r1_compatible"
+    )