kiln-ai 0.8.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

kiln_ai/datamodel/task.py

@@ -0,0 +1,181 @@
+from typing import TYPE_CHECKING, Dict, List, Union
+
+from pydantic import BaseModel, Field
+
+from kiln_ai.datamodel import Finetune
+from kiln_ai.datamodel.basemodel import (
+    ID_FIELD,
+    ID_TYPE,
+    NAME_FIELD,
+    SHORT_NAME_FIELD,
+    KilnParentedModel,
+    KilnParentModel,
+)
+from kiln_ai.datamodel.datamodel_enums import Priority, 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
+from kiln_ai.datamodel.prompt import BasePrompt, Prompt
+from kiln_ai.datamodel.prompt_id import PromptId
+from kiln_ai.datamodel.task_run import TaskRun
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.project import Project
+
+
+class TaskRequirement(BaseModel):
+    """
+    Defines a specific requirement that should be met by task outputs.
+
+    Includes an identifier, name, description, instruction for meeting the requirement,
+    priority level, and rating type (five_star, pass_fail, pass_fail_critical, custom).
+    """
+
+    id: ID_TYPE = ID_FIELD
+    name: str = SHORT_NAME_FIELD
+    description: str | None = Field(default=None)
+    instruction: str = Field(min_length=1)
+    priority: Priority = Field(default=Priority.p2)
+    type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
+
+
+class RunConfigProperties(BaseModel):
+    """
+    A configuration for running a task.
+
+    This includes everything needed to run a task, except the input and task ID. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
+    """
+
+    model_name: str = Field(description="The model to use for this run config.")
+    model_provider_name: str = 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.",
+    )
+
+
+class RunConfig(RunConfigProperties):
+    """
+    A configuration for running a task.
+
+    This includes everything needed to run a task, except the input. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
+
+    For example: task, model, provider, prompt, etc.
+    """
+
+    task: "Task" = Field(description="The task to run.")
+
+
+class TaskRunConfig(KilnParentedModel):
+    """
+    A Kiln model for persisting a run config in a Kiln Project, nested under a task.
+
+    Typically used to save a method of running a task for evaluation.
+
+    A run config includes everything needed to run a task, except the input. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
+    """
+
+    name: str = NAME_FIELD
+    description: str | None = Field(
+        default=None, description="The description of the task run config."
+    )
+    run_config_properties: RunConfigProperties = Field(
+        description="The run config properties to use for this task run."
+    )
+    # The prompt_id in the run_config_properties is the prompt ID to use for this task run.
+    # However, we want the prompt to be perfectly consistent, and some prompt_ids are dynamic.
+    # If we need to "freeze" a prompt, we can do so here (then point the prompt_id to this frozen prompt).
+    prompt: BasePrompt | None = Field(
+        default=None,
+        description="A prompt to use for run config.",
+    )
+
+    # Workaround to return typed parent without importing Task
+    def parent_task(self) -> Union["Task", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Task":
+            return None
+        return self.parent  # type: ignore
+
+    def run_config(self) -> RunConfig:
+        parent_task = self.parent_task()
+        if parent_task is None:
+            raise ValueError("Run config must be parented to a task")
+        return RunConfig(
+            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,
+        )
+
+
+class Task(
+    KilnParentedModel,
+    KilnParentModel,
+    parent_of={
+        "runs": TaskRun,
+        "dataset_splits": DatasetSplit,
+        "finetunes": Finetune,
+        "prompts": Prompt,
+        "evals": Eval,
+        "run_configs": TaskRunConfig,
+    },
+):
+    """
+    Represents a specific task to be performed, with associated requirements and validation rules.
+
+    Contains the task definition, requirements, input/output schemas, and maintains
+    a collection of task runs.
+    """
+
+    name: str = NAME_FIELD
+    description: str | None = Field(
+        default=None,
+        description="A description of the task for you and your team. Will not be used in prompts/training/validation.",
+    )
+    instruction: str = Field(
+        min_length=1,
+        description="The instructions for the task. Will be used in prompts/training/validation.",
+    )
+    requirements: List[TaskRequirement] = Field(default=[])
+    output_json_schema: JsonObjectSchema | None = None
+    input_json_schema: JsonObjectSchema | None = None
+    thinking_instruction: str | None = Field(
+        default=None,
+        description="Instructions for the model 'thinking' about the requirement prior to answering. Used for chain of thought style prompting.",
+    )
+
+    def output_schema(self) -> Dict | None:
+        if self.output_json_schema is None:
+            return None
+        return schema_from_json_str(self.output_json_schema)
+
+    def input_schema(self) -> Dict | None:
+        if self.input_json_schema is None:
+            return None
+        return schema_from_json_str(self.input_json_schema)
+
+    # These wrappers help for typechecking. TODO P2: fix this in KilnParentModel
+    def runs(self, readonly: bool = False) -> list[TaskRun]:
+        return super().runs(readonly=readonly)  # type: ignore
+
+    def dataset_splits(self, readonly: bool = False) -> list[DatasetSplit]:
+        return super().dataset_splits(readonly=readonly)  # type: ignore
+
+    def finetunes(self, readonly: bool = False) -> list[Finetune]:
+        return super().finetunes(readonly=readonly)  # type: ignore
+
+    def prompts(self, readonly: bool = False) -> list[Prompt]:
+        return super().prompts(readonly=readonly)  # type: ignore
+
+    def evals(self, readonly: bool = False) -> list[Eval]:
+        return super().evals(readonly=readonly)  # type: ignore
+
+    def run_configs(self, readonly: bool = False) -> list[TaskRunConfig]:
+        return super().run_configs(readonly=readonly)  # type: ignore
+
+    # Workaround to return typed parent without importing Task
+    def parent_project(self) -> Union["Project", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Project":
+            return None
+        return self.parent  # type: ignore

kiln_ai/datamodel/task_output.py

@@ -0,0 +1,321 @@
+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
+
+from kiln_ai.datamodel.basemodel import ID_TYPE, KilnBaseModel
+from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
+from kiln_ai.datamodel.json_schema import validate_schema
+from kiln_ai.datamodel.strict_mode import strict_mode
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.task import Task
+
+
+class RequirementRating(BaseModel):
+    """Rating for a specific requirement within a task output."""
+
+    value: float = Field(
+        description="The rating value. Interpretation depends on rating type"
+    )
+    type: TaskOutputRatingType = Field(description="The type of rating")
+
+
+def normalize_rating(rating: float, rating_type: TaskOutputRatingType) -> float:
+    """Normalize a rating to a 0-1 scale. Simple normalization, not z-score."""
+    match rating_type:
+        case TaskOutputRatingType.five_star:
+            if rating < 1 or rating > 5:
+                raise ValueError("Five star rating must be between 1 and 5")
+            return (rating - 1) / 4
+        case TaskOutputRatingType.pass_fail:
+            if rating < 0 or rating > 1:
+                raise ValueError("Pass fail rating must 0 to 1")
+            return rating
+        case TaskOutputRatingType.pass_fail_critical:
+            if rating < -1 or rating > 1:
+                raise ValueError("Pass fail critical rating must -1 to 1")
+            return (rating + 1) / 2  # -1 to 1
+        case TaskOutputRatingType.custom:
+            raise ValueError("Custom rating type can not be normalized")
+        case _:
+            raise_exhaustive_enum_error(rating_type)
+
+
+class TaskOutputRating(KilnBaseModel):
+    """
+    A rating for a task output, including an overall rating and ratings for each requirement.
+
+    Supports:
+    - five_star: 1-5 star ratings
+    - pass_fail: boolean pass/fail (1.0 = pass, 0.0 = fail)
+    - pass_fail_critical: tri-state (1.0 = pass, 0.0 = fail, -1.0 = critical fail)
+    """
+
+    type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
+    value: float | None = Field(
+        description="The rating value. Interpretation depends on rating type:\n- five_star: 1-5 stars\n- pass_fail: 1.0 (pass) or 0.0 (fail)\n- pass_fail_critical: 1.0 (pass), 0.0 (fail), or -1.0 (critical fail)",
+        default=None,
+    )
+    requirement_ratings: Dict[ID_TYPE, RequirementRating] = Field(
+        default={},
+        description="The ratings of the requirements of the task.",
+    )
+
+    # Previously we stored rating values as a dict of floats, but now we store them as RequirementRating objects.
+    @model_validator(mode="before")
+    def upgrade_old_format(cls, data: dict) -> dict:
+        if not isinstance(data, dict):
+            return data
+
+        # Check if we have the old format (dict of floats)
+        req_ratings = data.get("requirement_ratings", {})
+        if req_ratings and all(
+            isinstance(v, (int, float)) for v in req_ratings.values()
+        ):
+            # Convert each float to a RequirementRating object
+            # all ratings are five star at the point we used this format
+            data["requirement_ratings"] = {
+                k: {"value": v, "type": TaskOutputRatingType.five_star}
+                for k, v in req_ratings.items()
+            }
+
+        return data
+
+    # Used to select high quality outputs for example selection (MultiShotPromptBuilder, etc)
+    def is_high_quality(self) -> bool:
+        if self.value is None:
+            return False
+
+        if self.type == TaskOutputRatingType.five_star:
+            return self.value >= 4
+        elif self.type == TaskOutputRatingType.pass_fail:
+            return self.value == 1.0
+        elif self.type == TaskOutputRatingType.pass_fail_critical:
+            return self.value == 1.0
+        return False
+
+    @model_validator(mode="after")
+    def validate_rating(self) -> Self:
+        if self.type not in TaskOutputRatingType:
+            raise ValueError(f"Invalid rating type: {self.type}")
+
+        # Overall rating is optional
+        if self.value is not None:
+            self._validate_rating(self.type, self.value, "overall rating")
+
+        for req_id, req_rating in self.requirement_ratings.items():
+            self._validate_rating(
+                req_rating.type,
+                req_rating.value,
+                f"requirement rating for req ID: {req_id}",
+            )
+
+        return self
+
+    def _validate_rating(
+        self, type: TaskOutputRatingType, rating: float | None, rating_name: str
+    ) -> None:
+        if type == TaskOutputRatingType.five_star:
+            self._validate_five_star(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail:
+            self._validate_pass_fail(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail_critical:
+            self._validate_pass_fail_critical(rating, rating_name)
+
+    def _validate_five_star(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type five_star must be an integer value (1-5)"
+            )
+        if rating < 1 or rating > 5:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type five_star must be between 1 and 5 stars"
+            )
+
+    def _validate_pass_fail(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be an integer value (0 or 1)"
+            )
+        if rating not in [0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be 0 (fail) or 1 (pass)"
+            )
+
+    def _validate_pass_fail_critical(
+        self, rating: float | None, rating_name: str
+    ) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be an integer value (-1, 0, or 1)"
+            )
+        if rating not in [-1, 0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be -1 (critical fail), 0 (fail), or 1 (pass)"
+            )
+
+
+class DataSourceType(str, Enum):
+    """
+    The source type of a piece of data.
+
+    Human: a human created the data
+    Synthetic: a model created the data
+    """
+
+    human = "human"
+    synthetic = "synthetic"
+
+
+class DataSourceProperty(BaseModel):
+    """
+    Defines a property that can be associated with a data source.
+
+    Includes validation rules for when properties are required or not allowed
+    based on the data source type.
+    """
+
+    name: str
+    type: Type[Union[str, int, float]]
+    required_for: List[DataSourceType] = []
+    not_allowed_for: List[DataSourceType] = []
+
+
+class DataSource(BaseModel):
+    """
+    Represents the origin of data, either human or synthetic, with associated properties.
+
+    Properties vary based on the source type - for synthetic sources this includes
+    model information, for human sources this includes creator information.
+    """
+
+    type: DataSourceType
+    properties: Dict[str, str | int | float] = Field(
+        default={},
+        description="Properties describing the data source. For synthetic things like model. For human, the human's name.",
+    )
+
+    _data_source_properties = [
+        DataSourceProperty(
+            name="created_by",
+            type=str,
+            required_for=[DataSourceType.human],
+            not_allowed_for=[DataSourceType.synthetic],
+        ),
+        DataSourceProperty(
+            name="model_name",
+            type=str,
+            required_for=[DataSourceType.synthetic],
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            name="model_provider",
+            type=str,
+            required_for=[DataSourceType.synthetic],
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            name="adapter_name",
+            type=str,
+            required_for=[DataSourceType.synthetic],
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            # Legacy field -- allow loading from old runs, but we shouldn't be setting it.
+            name="prompt_builder_name",
+            type=str,
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            # The PromptId of the prompt. Can be a saved prompt, fine-tune, generator name, etc. See PromptId type for more details.
+            name="prompt_id",
+            type=str,
+            not_allowed_for=[DataSourceType.human],
+        ),
+    ]
+
+    @model_validator(mode="after")
+    def validate_type(self) -> "DataSource":
+        if self.type not in DataSourceType:
+            raise ValueError(f"Invalid data source type: {self.type}")
+        return self
+
+    @model_validator(mode="after")
+    def validate_properties(self) -> "DataSource":
+        for prop in self._data_source_properties:
+            # Check the property type is correct
+            if prop.name in self.properties:
+                if not isinstance(self.properties[prop.name], prop.type):
+                    raise ValueError(
+                        f"'{prop.name}' must be of type {prop.type.__name__} for {self.type} data source"
+                    )
+            # Check the property is required for the data source type
+            if self.type in prop.required_for:
+                if prop.name not in self.properties:
+                    raise ValueError(
+                        f"'{prop.name}' is required for {self.type} data source"
+                    )
+            # Check the property is not allowed for the data source type
+            elif self.type in prop.not_allowed_for and prop.name in self.properties:
+                raise ValueError(
+                    f"'{prop.name}' is not allowed for {self.type} data source"
+                )
+        return self
+
+    @model_validator(mode="after")
+    def validate_no_empty_properties(self) -> Self:
+        for prop, value in self.properties.items():
+            if isinstance(value, str) and value == "":
+                raise ValueError(
+                    f"Property '{prop}' must be a non-empty string for {self.type} data source"
+                )
+        return self
+
+
+class TaskOutput(KilnBaseModel):
+    """
+    An output for a specific task run.
+
+    Contains the actual output content, its source (human or synthetic),
+    and optional rating information.
+    """
+
+    output: str = Field(
+        description="The output of the task. JSON formatted for structured output, plaintext for unstructured output."
+    )
+    source: DataSource | None = Field(
+        description="The source of the output: human or synthetic.",
+        default=None,
+    )
+    rating: TaskOutputRating | None = Field(
+        default=None, description="The rating of the output"
+    )
+
+    def validate_output_format(self, task: "Task") -> Self:
+        # validate output
+        if task.output_json_schema is not None:
+            try:
+                validate_schema(json.loads(self.output), task.output_json_schema)
+            except json.JSONDecodeError:
+                raise ValueError("Output is not a valid JSON object")
+            except jsonschema.exceptions.ValidationError as e:
+                raise ValueError(f"Output does not match task output schema: {e}")
+        return self
+
+    @model_validator(mode="after")
+    def validate_output_source(self, info: ValidationInfo) -> Self:
+        # On strict mode and not loaded from file, we validate output_source is not None.
+        # We want to be able to load any data, even if it's not perfect. But we want to create perfect data when adding new data.
+        if not strict_mode():
+            return self
+        if self.loaded_from_file(info):
+            return self
+        if self.source is None:
+            raise ValueError("Output source is required when strict mode is enabled")
+        return self

kiln_ai/datamodel/task_run.py

@@ -0,0 +1,164 @@
+import json
+from typing import TYPE_CHECKING, Dict, List, Union
+
+import jsonschema
+import jsonschema.exceptions
+from pydantic import Field, ValidationInfo, model_validator
+from typing_extensions import Self
+
+from kiln_ai.datamodel.basemodel import KilnParentedModel
+from kiln_ai.datamodel.json_schema import validate_schema
+from kiln_ai.datamodel.strict_mode import strict_mode
+from kiln_ai.datamodel.task_output import DataSource, TaskOutput
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.task import Task
+
+
+class TaskRun(KilnParentedModel):
+    """
+    Represents a single execution of a Task.
+
+    Contains the input used, its source, the output produced, and optional
+    repair information if the output needed correction.
+    """
+
+    input: str = Field(
+        description="The inputs to the task. JSON formatted for structured input, plaintext for unstructured input."
+    )
+    input_source: DataSource | None = Field(
+        default=None, description="The source of the input: human or synthetic."
+    )
+
+    output: TaskOutput = Field(description="The output of the task run.")
+    repair_instructions: str | None = Field(
+        default=None,
+        description="Instructions for fixing the output. Should define what is wrong, and how to fix it. Will be used by models for both generating a fixed output, and evaluating future models.",
+    )
+    repaired_output: TaskOutput | None = Field(
+        default=None,
+        description="An version of the output with issues fixed. This must be a 'fixed' version of the existing output, and not an entirely new output. If you wish to generate an ideal curatorial output for this task unrelated to this output, generate a new TaskOutput with type 'human' instead of using this field.",
+    )
+    intermediate_outputs: Dict[str, str] | None = Field(
+        default=None,
+        description="Intermediate outputs from the task run. Keys are the names of the intermediate output steps (cot=chain of thought, etc), values are the output data.",
+    )
+    tags: List[str] = Field(
+        default=[],
+        description="Tags for the task run. Tags are used to categorize task runs for filtering and reporting.",
+    )
+
+    def has_thinking_training_data(self) -> bool:
+        """
+        Does this run have thinking data that we can use to train a thinking model?
+        """
+        if self.intermediate_outputs is None:
+            return False
+        return (
+            "chain_of_thought" in self.intermediate_outputs
+            or "reasoning" in self.intermediate_outputs
+        )
+
+    # Workaround to return typed parent without importing Task
+    def parent_task(self) -> Union["Task", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Task":
+            return None
+        return self.parent  # type: ignore
+
+    @model_validator(mode="after")
+    def validate_input_format(self, info: ValidationInfo) -> Self:
+        # Don't validate if loading from file (not new). Too slow.
+        # We don't allow changing task schema, so this is redundant validation.
+        # Note: we still validate if editing a loaded model
+        if self.loading_from_file(info):
+            # Consider loading an existing model as validated.
+            self._last_validated_input = self.input
+            return self
+
+        # Don't validate if input has not changed. Too slow to run this every time.
+        if (
+            hasattr(self, "_last_validated_input")
+            and self.input == self._last_validated_input
+        ):
+            return self
+
+        task = self.parent_task()
+        if task is None:
+            # don't validate this relationship until we have a path or parent. Give them time to build it (but will catch it before saving)
+            return self
+
+        # validate output
+        if task.input_json_schema is not None:
+            try:
+                validate_schema(json.loads(self.input), task.input_json_schema)
+            except json.JSONDecodeError:
+                raise ValueError("Input is not a valid JSON object")
+            except jsonschema.exceptions.ValidationError as e:
+                raise ValueError(f"Input does not match task input schema: {e}")
+        self._last_validated_input = self.input
+        return self
+
+    @model_validator(mode="after")
+    def validate_output_format(self, info: ValidationInfo) -> Self:
+        # Don't validate if loading from file (not new). Too slow.
+        # Note: we still validate if editing a loaded model's output.
+        if self.loading_from_file(info):
+            # Consider loading an existing model as validated.
+            self._last_validated_output = self.output.output if self.output else None
+            return self
+
+        # Don't validate unless output has changed since last validation.
+        # The validator is slow and costly, don't want it running when setting other fields.
+        if (
+            hasattr(self, "_last_validated_output")
+            and self.output is not None
+            and self.output.output == self._last_validated_output
+        ):
+            return self
+
+        task = self.parent_task()
+        if task is None:
+            return self
+
+        self.output.validate_output_format(task)
+        self._last_validated_output = self.output.output if self.output else None
+        return self
+
+    @model_validator(mode="after")
+    def validate_repaired_output(self) -> Self:
+        if self.repaired_output is not None:
+            if self.repaired_output.rating is not None:
+                raise ValueError(
+                    "Repaired output rating must be None. Repaired outputs are assumed to have a perfect rating, as they have been fixed."
+                )
+        if self.repair_instructions is None and self.repaired_output is not None:
+            raise ValueError(
+                "Repair instructions are required if providing a repaired output."
+            )
+        if self.repair_instructions is not None and self.repaired_output is None:
+            raise ValueError(
+                "A repaired output is required if providing repair instructions."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def validate_input_source(self, info: ValidationInfo) -> Self:
+        # On strict mode and not loaded from file, we validate input_source is not None.
+        # We want to be able to load any data, even if it's not perfect. But we want to create perfect data when adding new data.
+        if not strict_mode():
+            return self
+        if self.loaded_from_file(info):
+            return self
+        if self.input_source is None:
+            raise ValueError("input_source is required when strict mode is enabled")
+        return self
+
+    @model_validator(mode="after")
+    def validate_tags(self) -> Self:
+        for tag in self.tags:
+            if not tag:
+                raise ValueError("Tags cannot be empty strings")
+            if " " in tag:
+                raise ValueError("Tags cannot contain spaces. Try underscores.")
+
+        return self