kiln_ai.datamodel

113class Task(
114    KilnParentedModel,
115    KilnParentModel,
116    parent_of={
117        "runs": TaskRun,
118        "dataset_splits": DatasetSplit,
119        "finetunes": Finetune,
120        "prompts": Prompt,
121        "evals": Eval,
122        "run_configs": TaskRunConfig,
123    },
124):
125    """
126    Represents a specific task to be performed, with associated requirements and validation rules.
127
128    Contains the task definition, requirements, input/output schemas, and maintains
129    a collection of task runs.
130    """
131
132    name: str = NAME_FIELD
133    description: str | None = Field(
134        default=None,
135        description="A description of the task for you and your team. Will not be used in prompts/training/validation.",
136    )
137    instruction: str = Field(
138        min_length=1,
139        description="The instructions for the task. Will be used in prompts/training/validation.",
140    )
141    requirements: List[TaskRequirement] = Field(default=[])
142    output_json_schema: JsonObjectSchema | None = None
143    input_json_schema: JsonObjectSchema | None = None
144    thinking_instruction: str | None = Field(
145        default=None,
146        description="Instructions for the model 'thinking' about the requirement prior to answering. Used for chain of thought style prompting.",
147    )
148
149    def output_schema(self) -> Dict | None:
150        if self.output_json_schema is None:
151            return None
152        return schema_from_json_str(self.output_json_schema)
153
154    def input_schema(self) -> Dict | None:
155        if self.input_json_schema is None:
156            return None
157        return schema_from_json_str(self.input_json_schema)
158
159    # These wrappers help for typechecking. TODO P2: fix this in KilnParentModel
160    def runs(self, readonly: bool = False) -> list[TaskRun]:
161        return super().runs(readonly=readonly)  # type: ignore
162
163    def dataset_splits(self, readonly: bool = False) -> list[DatasetSplit]:
164        return super().dataset_splits(readonly=readonly)  # type: ignore
165
166    def finetunes(self, readonly: bool = False) -> list[Finetune]:
167        return super().finetunes(readonly=readonly)  # type: ignore
168
169    def prompts(self, readonly: bool = False) -> list[Prompt]:
170        return super().prompts(readonly=readonly)  # type: ignore
171
172    def evals(self, readonly: bool = False) -> list[Eval]:
173        return super().evals(readonly=readonly)  # type: ignore
174
175    def run_configs(self, readonly: bool = False) -> list[TaskRunConfig]:
176        return super().run_configs(readonly=readonly)  # type: ignore
177
178    # Workaround to return typed parent without importing Task
179    def parent_project(self) -> Union["Project", None]:
180        if self.parent is None or self.parent.__class__.__name__ != "Project":
181            return None
182        return self.parent  # type: ignore

 8class Project(KilnParentModel, parent_of={"tasks": Task}):
 9    """
10    A collection of related tasks.
11
12    Projects organize tasks into logical groups and provide high-level descriptions
13    of the overall goals.
14    """
15
16    name: str = NAME_FIELD
17    description: str | None = Field(
18        default=None,
19        description="A description of the project for you and your team. Will not be used in prompts/training/validation.",
20    )
21
22    # Needed for typechecking. TODO P2: fix this in KilnParentModel
23    def tasks(self) -> list[Task]:
24        return super().tasks()  # type: ignore

 19class TaskRun(KilnParentedModel):
 20    """
 21    Represents a single execution of a Task.
 22
 23    Contains the input used, its source, the output produced, and optional
 24    repair information if the output needed correction.
 25    """
 26
 27    input: str = Field(
 28        description="The inputs to the task. JSON formatted for structured input, plaintext for unstructured input."
 29    )
 30    input_source: DataSource | None = Field(
 31        default=None, description="The source of the input: human or synthetic."
 32    )
 33
 34    output: TaskOutput = Field(description="The output of the task run.")
 35    repair_instructions: str | None = Field(
 36        default=None,
 37        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.",
 38    )
 39    repaired_output: TaskOutput | None = Field(
 40        default=None,
 41        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.",
 42    )
 43    intermediate_outputs: Dict[str, str] | None = Field(
 44        default=None,
 45        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.",
 46    )
 47    tags: List[str] = Field(
 48        default=[],
 49        description="Tags for the task run. Tags are used to categorize task runs for filtering and reporting.",
 50    )
 51
 52    def has_thinking_training_data(self) -> bool:
 53        """
 54        Does this run have thinking data that we can use to train a thinking model?
 55        """
 56        if self.intermediate_outputs is None:
 57            return False
 58        return (
 59            "chain_of_thought" in self.intermediate_outputs
 60            or "reasoning" in self.intermediate_outputs
 61        )
 62
 63    # Workaround to return typed parent without importing Task
 64    def parent_task(self) -> Union["Task", None]:
 65        if self.parent is None or self.parent.__class__.__name__ != "Task":
 66            return None
 67        return self.parent  # type: ignore
 68
 69    @model_validator(mode="after")
 70    def validate_input_format(self, info: ValidationInfo) -> Self:
 71        # Don't validate if loading from file (not new). Too slow.
 72        # We don't allow changing task schema, so this is redundant validation.
 73        # Note: we still validate if editing a loaded model
 74        if self.loading_from_file(info):
 75            # Consider loading an existing model as validated.
 76            self._last_validated_input = self.input
 77            return self
 78
 79        # Don't validate if input has not changed. Too slow to run this every time.
 80        if (
 81            hasattr(self, "_last_validated_input")
 82            and self.input == self._last_validated_input
 83        ):
 84            return self
 85
 86        task = self.parent_task()
 87        if task is None:
 88            # don't validate this relationship until we have a path or parent. Give them time to build it (but will catch it before saving)
 89            return self
 90
 91        # validate output
 92        if task.input_json_schema is not None:
 93            try:
 94                validate_schema(json.loads(self.input), task.input_json_schema)
 95            except json.JSONDecodeError:
 96                raise ValueError("Input is not a valid JSON object")
 97            except jsonschema.exceptions.ValidationError as e:
 98                raise ValueError(f"Input does not match task input schema: {e}")
 99        self._last_validated_input = self.input
100        return self
101
102    @model_validator(mode="after")
103    def validate_output_format(self, info: ValidationInfo) -> Self:
104        # Don't validate if loading from file (not new). Too slow.
105        # Note: we still validate if editing a loaded model's output.
106        if self.loading_from_file(info):
107            # Consider loading an existing model as validated.
108            self._last_validated_output = self.output.output if self.output else None
109            return self
110
111        # Don't validate unless output has changed since last validation.
112        # The validator is slow and costly, don't want it running when setting other fields.
113        if (
114            hasattr(self, "_last_validated_output")
115            and self.output is not None
116            and self.output.output == self._last_validated_output
117        ):
118            return self
119
120        task = self.parent_task()
121        if task is None:
122            return self
123
124        self.output.validate_output_format(task)
125        self._last_validated_output = self.output.output if self.output else None
126        return self
127
128    @model_validator(mode="after")
129    def validate_repaired_output(self) -> Self:
130        if self.repaired_output is not None:
131            if self.repaired_output.rating is not None:
132                raise ValueError(
133                    "Repaired output rating must be None. Repaired outputs are assumed to have a perfect rating, as they have been fixed."
134                )
135        if self.repair_instructions is None and self.repaired_output is not None:
136            raise ValueError(
137                "Repair instructions are required if providing a repaired output."
138            )
139        if self.repair_instructions is not None and self.repaired_output is None:
140            raise ValueError(
141                "A repaired output is required if providing repair instructions."
142            )
143        return self
144
145    @model_validator(mode="after")
146    def validate_input_source(self, info: ValidationInfo) -> Self:
147        # On strict mode and not loaded from file, we validate input_source is not None.
148        # 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.
149        if not strict_mode():
150            return self
151        if self.loaded_from_file(info):
152            return self
153        if self.input_source is None:
154            raise ValueError("input_source is required when strict mode is enabled")
155        return self
156
157    @model_validator(mode="after")
158    def validate_tags(self) -> Self:
159        for tag in self.tags:
160            if not tag:
161                raise ValueError("Tags cannot be empty strings")
162            if " " in tag:
163                raise ValueError("Tags cannot contain spaces. Try underscores.")
164
165        return self

289class TaskOutput(KilnBaseModel):
290    """
291    An output for a specific task run.
292
293    Contains the actual output content, its source (human or synthetic),
294    and optional rating information.
295    """
296
297    output: str = Field(
298        description="The output of the task. JSON formatted for structured output, plaintext for unstructured output."
299    )
300    source: DataSource | None = Field(
301        description="The source of the output: human or synthetic.",
302        default=None,
303    )
304    rating: TaskOutputRating | None = Field(
305        default=None, description="The rating of the output"
306    )
307
308    def validate_output_format(self, task: "Task") -> Self:
309        # validate output
310        if task.output_json_schema is not None:
311            try:
312                validate_schema(json.loads(self.output), task.output_json_schema)
313            except json.JSONDecodeError:
314                raise ValueError("Output is not a valid JSON object")
315            except jsonschema.exceptions.ValidationError as e:
316                raise ValueError(f"Output does not match task output schema: {e}")
317        return self
318
319    @model_validator(mode="after")
320    def validate_output_source(self, info: ValidationInfo) -> Self:
321        # On strict mode and not loaded from file, we validate output_source is not None.
322        # 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.
323        if not strict_mode():
324            return self
325        if self.loaded_from_file(info):
326            return self
327        if self.source is None:
328            raise ValueError("Output source is required when strict mode is enabled")
329        return self

 5class Priority(IntEnum):
 6    """Defines priority levels for tasks and requirements, where P0 is highest priority."""
 7
 8    p0 = 0
 9    p1 = 1
10    p2 = 2
11    p3 = 3

192class DataSource(BaseModel):
193    """
194    Represents the origin of data, either human or synthetic, with associated properties.
195
196    Properties vary based on the source type - for synthetic sources this includes
197    model information, for human sources this includes creator information.
198    """
199
200    type: DataSourceType
201    properties: Dict[str, str | int | float] = Field(
202        default={},
203        description="Properties describing the data source. For synthetic things like model. For human, the human's name.",
204    )
205
206    _data_source_properties = [
207        DataSourceProperty(
208            name="created_by",
209            type=str,
210            required_for=[DataSourceType.human],
211            not_allowed_for=[DataSourceType.synthetic, DataSourceType.file_import],
212        ),
213        DataSourceProperty(
214            name="model_name",
215            type=str,
216            required_for=[DataSourceType.synthetic],
217            not_allowed_for=[DataSourceType.human, DataSourceType.file_import],
218        ),
219        DataSourceProperty(
220            name="model_provider",
221            type=str,
222            required_for=[DataSourceType.synthetic],
223            not_allowed_for=[DataSourceType.human, DataSourceType.file_import],
224        ),
225        DataSourceProperty(
226            name="adapter_name",
227            type=str,
228            required_for=[DataSourceType.synthetic],
229            not_allowed_for=[DataSourceType.human, DataSourceType.file_import],
230        ),
231        DataSourceProperty(
232            # Legacy field -- allow loading from old runs, but we shouldn't be setting it.
233            name="prompt_builder_name",
234            type=str,
235            not_allowed_for=[DataSourceType.human, DataSourceType.file_import],
236        ),
237        DataSourceProperty(
238            # The PromptId of the prompt. Can be a saved prompt, fine-tune, generator name, etc. See PromptId type for more details.
239            name="prompt_id",
240            type=str,
241            not_allowed_for=[DataSourceType.human, DataSourceType.file_import],
242        ),
243        DataSourceProperty(
244            name="file_name",
245            type=str,
246            required_for=[DataSourceType.file_import],
247            not_allowed_for=[DataSourceType.human, DataSourceType.synthetic],
248        ),
249    ]
250
251    @model_validator(mode="after")
252    def validate_type(self) -> "DataSource":
253        if self.type not in DataSourceType:
254            raise ValueError(f"Invalid data source type: {self.type}")
255        return self
256
257    @model_validator(mode="after")
258    def validate_properties(self) -> "DataSource":
259        for prop in self._data_source_properties:
260            # Check the property type is correct
261            if prop.name in self.properties:
262                if not isinstance(self.properties[prop.name], prop.type):
263                    raise ValueError(
264                        f"'{prop.name}' must be of type {prop.type.__name__} for {self.type} data source"
265                    )
266            # Check the property is required for the data source type
267            if self.type in prop.required_for:
268                if prop.name not in self.properties:
269                    raise ValueError(
270                        f"'{prop.name}' is required for {self.type} data source"
271                    )
272            # Check the property is not allowed for the data source type
273            elif self.type in prop.not_allowed_for and prop.name in self.properties:
274                raise ValueError(
275                    f"'{prop.name}' is not allowed for {self.type} data source"
276                )
277        return self
278
279    @model_validator(mode="after")
280    def validate_no_empty_properties(self) -> Self:
281        for prop, value in self.properties.items():
282            if isinstance(value, str) and value == "":
283                raise ValueError(
284                    f"Property '{prop}' must be a non-empty string for {self.type} data source"
285                )
286        return self

165class DataSourceType(str, Enum):
166    """
167    The source type of a piece of data.
168
169    Human: a human created the data
170    Synthetic: a model created the data
171    """
172
173    human = "human"
174    synthetic = "synthetic"
175    file_import = "file_import"

178class DataSourceProperty(BaseModel):
179    """
180    Defines a property that can be associated with a data source.
181
182    Includes validation rules for when properties are required or not allowed
183    based on the data source type.
184    """
185
186    name: str
187    type: Type[Union[str, int, float]]
188    required_for: List[DataSourceType] = []
189    not_allowed_for: List[DataSourceType] = []

 18class Finetune(KilnParentedModel):
 19    """
 20    The Kiln fine-tune datamodel.
 21
 22    Initially holds a reference to a training job, with needed identifiers to update the status. When complete, contains the new model ID.
 23    """
 24
 25    name: str = NAME_FIELD
 26    description: str | None = Field(
 27        default=None,
 28        description="A description of the fine-tune for you and your team. Not used in training.",
 29    )
 30    structured_output_mode: StructuredOutputMode | None = Field(
 31        default=None,
 32        description="The mode to use to train the model for structured output, if it was trained with structured output. Will determine how we call the tuned model, so we call with the matching mode.",
 33    )
 34    provider: str = Field(
 35        description="The provider to use for the fine-tune (e.g. 'openai')."
 36    )
 37    base_model_id: str = Field(
 38        description="The id of the base model to use for the fine-tune. This string relates to the provider's IDs for their own models, not Kiln IDs."
 39    )
 40    provider_id: str | None = Field(
 41        default=None,
 42        description="The ID of the fine-tune job on the provider's side. May not be the same as the fine_tune_model_id.",
 43    )
 44    fine_tune_model_id: str | None = Field(
 45        default=None,
 46        description="The ID of the fine-tuned model on the provider's side. May not be the same as the provider_id.",
 47    )
 48    dataset_split_id: str = Field(
 49        description="The ID of the dataset split to use for this fine-tune.",
 50    )
 51    train_split_name: str = Field(
 52        default="train",
 53        description="The name of the training split to use for this fine-tune.",
 54    )
 55    validation_split_name: str | None = Field(
 56        default=None,
 57        description="The name of the validation split to use for this fine-tune. Optional.",
 58    )
 59    parameters: dict[str, str | int | float | bool] = Field(
 60        default={},
 61        description="The parameters to use for this fine-tune. These are provider-specific.",
 62    )
 63    # These two fields are saved exactly used for training. Even if they map exactly to a custom prompt or generator, those can change, so we want to keep a record of the training prompt.
 64    system_message: str = Field(
 65        description="The system message to use for this fine-tune.",
 66    )
 67    thinking_instructions: str | None = Field(
 68        default=None,
 69        description="The thinking instructions to use for this fine-tune. Only used when data_strategy is final_and_intermediate.",
 70    )
 71    latest_status: FineTuneStatusType = Field(
 72        default=FineTuneStatusType.unknown,
 73        description="The latest known status of this fine-tune. Not updated in real time.",
 74    )
 75    properties: Dict[str, str | int | float] = Field(
 76        default={},
 77        description="Properties of the fine-tune. Different providers may use different properties.",
 78    )
 79    data_strategy: FinetuneDataStrategy = Field(
 80        default=FinetuneDataStrategy.final_only,
 81        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).",
 82    )
 83
 84    # Workaround to return typed parent without importing Task
 85    def parent_task(self) -> Union["Task", None]:
 86        if self.parent is None or self.parent.__class__.__name__ != "Task":
 87            return None
 88        return self.parent  # type: ignore
 89
 90    @model_validator(mode="after")
 91    def validate_thinking_instructions(self) -> Self:
 92        if (
 93            self.thinking_instructions is not None
 94            and self.data_strategy != FinetuneDataStrategy.final_and_intermediate
 95        ):
 96            raise ValueError(
 97                "Thinking instructions can only be used when data_strategy is final_and_intermediate"
 98            )
 99        if (
100            self.thinking_instructions is None
101            and self.data_strategy == FinetuneDataStrategy.final_and_intermediate
102        ):
103            raise ValueError(
104                "Thinking instructions are required when data_strategy is final_and_intermediate"
105            )
106        return self

47class FineTuneStatusType(str, Enum):
48    """
49    The status type of a fine-tune (running, completed, failed, etc).
50    """
51
52    unknown = "unknown"  # server error
53    pending = "pending"
54    running = "running"
55    completed = "completed"
56    failed = "failed"

15class TaskOutputRatingType(str, Enum):
16    """Defines the types of rating systems available for task outputs."""
17
18    five_star = "five_star"
19    pass_fail = "pass_fail"
20    pass_fail_critical = "pass_fail_critical"
21    custom = "custom"

27class TaskRequirement(BaseModel):
28    """
29    Defines a specific requirement that should be met by task outputs.
30
31    Includes an identifier, name, description, instruction for meeting the requirement,
32    priority level, and rating type (five_star, pass_fail, pass_fail_critical, custom).
33    """
34
35    id: ID_TYPE = ID_FIELD
36    name: str = SHORT_NAME_FIELD
37    description: str | None = Field(default=None)
38    instruction: str = Field(min_length=1)
39    priority: Priority = Field(default=Priority.p2)
40    type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)

23class DatasetSplitDefinition(BaseModel):
24    """
25    A definition of a split in a dataset.
26
27    Example: name="train", description="The training set", percentage=0.8 (80% of the dataset)
28    """
29
30    name: str = NAME_FIELD
31    description: str | None = Field(
32        default=None,
33        description="A description of the dataset for you and your team. Not used in training.",
34    )
35    percentage: float = Field(
36        ge=0.0,
37        le=1.0,
38        description="The percentage of the dataset that this split represents (between 0 and 1).",
39    )