kiln_ai.datamodel

174class Task(
175    KilnParentedModel,
176    KilnParentModel,
177    parent_of={
178        "runs": TaskRun,
179        "dataset_splits": DatasetSplit,
180        "finetunes": Finetune,
181        "prompts": Prompt,
182        "evals": Eval,
183        "run_configs": TaskRunConfig,
184    },
185):
186    """
187    Represents a specific task to be performed, with associated requirements and validation rules.
188
189    Contains the task definition, requirements, input/output schemas, and maintains
190    a collection of task runs.
191    """
192
193    name: FilenameString = Field(description="The name of the task.")
194    description: str | None = Field(
195        default=None,
196        description="A description of the task for you and your team. Will not be used in prompts/training/validation.",
197    )
198    instruction: str = Field(
199        min_length=1,
200        description="The instructions for the task. Will be used in prompts/training/validation.",
201    )
202    requirements: List[TaskRequirement] = Field(default=[])
203    output_json_schema: JsonObjectSchema | None = None
204    input_json_schema: JsonObjectSchema | None = None
205    thinking_instruction: str | None = Field(
206        default=None,
207        description="Instructions for the model 'thinking' about the requirement prior to answering. Used for chain of thought style prompting.",
208    )
209
210    def output_schema(self) -> Dict | None:
211        if self.output_json_schema is None:
212            return None
213        return schema_from_json_str(self.output_json_schema)
214
215    def input_schema(self) -> Dict | None:
216        if self.input_json_schema is None:
217            return None
218        return schema_from_json_str(self.input_json_schema)
219
220    # These wrappers help for typechecking. We should fix this in KilnParentModel
221    def runs(self, readonly: bool = False) -> list[TaskRun]:
222        return super().runs(readonly=readonly)  # type: ignore
223
224    def dataset_splits(self, readonly: bool = False) -> list[DatasetSplit]:
225        return super().dataset_splits(readonly=readonly)  # type: ignore
226
227    def finetunes(self, readonly: bool = False) -> list[Finetune]:
228        return super().finetunes(readonly=readonly)  # type: ignore
229
230    def prompts(self, readonly: bool = False) -> list[Prompt]:
231        return super().prompts(readonly=readonly)  # type: ignore
232
233    def evals(self, readonly: bool = False) -> list[Eval]:
234        return super().evals(readonly=readonly)  # type: ignore
235
236    def run_configs(self, readonly: bool = False) -> list[TaskRunConfig]:
237        return super().run_configs(readonly=readonly)  # type: ignore
238
239    # Workaround to return typed parent without importing Task
240    def parent_project(self) -> Union["Project", None]:
241        if self.parent is None or self.parent.__class__.__name__ != "Project":
242            return None
243        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: FilenameString = Field(description="The name of the project.")
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. We should fix this in KilnParentModel
23    def tasks(self) -> list[Task]:
24        return super().tasks()  # type: ignore

 40class TaskRun(KilnParentedModel):
 41    """
 42    Represents a single execution of a Task.
 43
 44    Contains the input used, its source, the output produced, and optional
 45    repair information if the output needed correction.
 46    """
 47
 48    input: str = Field(
 49        description="The inputs to the task. JSON formatted for structured input, plaintext for unstructured input."
 50    )
 51    input_source: DataSource | None = Field(
 52        default=None, description="The source of the input: human or synthetic."
 53    )
 54
 55    output: TaskOutput = Field(description="The output of the task run.")
 56    repair_instructions: str | None = Field(
 57        default=None,
 58        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.",
 59    )
 60    repaired_output: TaskOutput | None = Field(
 61        default=None,
 62        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.",
 63    )
 64    intermediate_outputs: Dict[str, str] | None = Field(
 65        default=None,
 66        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.",
 67    )
 68    tags: List[str] = Field(
 69        default=[],
 70        description="Tags for the task run. Tags are used to categorize task runs for filtering and reporting.",
 71    )
 72    usage: Usage | None = Field(
 73        default=None,
 74        description="Usage information for the task run. This includes the number of input tokens, output tokens, and total tokens used.",
 75    )
 76
 77    def thinking_training_data(self) -> str | None:
 78        """
 79        Get the thinking training data from the task run.
 80        """
 81        if self.intermediate_outputs is None:
 82            return None
 83        return self.intermediate_outputs.get(
 84            "reasoning"
 85        ) or self.intermediate_outputs.get("chain_of_thought")
 86
 87    def has_thinking_training_data(self) -> bool:
 88        """
 89        Does this run have thinking data that we can use to train a thinking model?
 90        """
 91        return self.thinking_training_data() is not None
 92
 93    # Workaround to return typed parent without importing Task
 94    def parent_task(self) -> Union["Task", None]:
 95        if self.parent is None or self.parent.__class__.__name__ != "Task":
 96            return None
 97        return self.parent  # type: ignore
 98
 99    @model_validator(mode="after")
100    def validate_input_format(self, info: ValidationInfo) -> Self:
101        # Don't validate if loading from file (not new). Too slow.
102        # We don't allow changing task schema, so this is redundant validation.
103        # Note: we still validate if editing a loaded model
104        if self.loading_from_file(info):
105            # Consider loading an existing model as validated.
106            self._last_validated_input = self.input
107            return self
108
109        # Don't validate if input has not changed. Too slow to run this every time.
110        if (
111            hasattr(self, "_last_validated_input")
112            and self.input == self._last_validated_input
113        ):
114            return self
115
116        task = self.parent_task()
117        if task is None:
118            # don't validate this relationship until we have a path or parent. Give them time to build it (but will catch it before saving)
119            return self
120
121        # validate input
122        if task.input_json_schema is not None:
123            try:
124                input_parsed = json.loads(self.input)
125            except json.JSONDecodeError:
126                raise ValueError("Input is not a valid JSON object")
127
128            validate_schema_with_value_error(
129                input_parsed,
130                task.input_json_schema,
131                "Input does not match task input schema.",
132            )
133
134        self._last_validated_input = self.input
135        return self
136
137    @model_validator(mode="after")
138    def validate_output_format(self, info: ValidationInfo) -> Self:
139        # Don't validate if loading from file (not new). Too slow.
140        # Note: we still validate if editing a loaded model's output.
141        if self.loading_from_file(info):
142            # Consider loading an existing model as validated.
143            self._last_validated_output = self.output.output if self.output else None
144            return self
145
146        # Don't validate unless output has changed since last validation.
147        # The validator is slow and costly, don't want it running when setting other fields.
148        if (
149            hasattr(self, "_last_validated_output")
150            and self.output is not None
151            and self.output.output == self._last_validated_output
152        ):
153            return self
154
155        task = self.parent_task()
156        if task is None:
157            return self
158
159        self.output.validate_output_format(task)
160        self._last_validated_output = self.output.output if self.output else None
161        return self
162
163    @model_validator(mode="after")
164    def validate_repaired_output(self) -> Self:
165        if self.repaired_output is not None:
166            if self.repaired_output.rating is not None:
167                raise ValueError(
168                    "Repaired output rating must be None. Repaired outputs are assumed to have a perfect rating, as they have been fixed."
169                )
170
171            task = self.parent_task()
172            if (
173                task is not None
174                and self.repaired_output.output is not None
175                and task.output_json_schema is not None
176            ):
177                try:
178                    output_parsed = json.loads(self.repaired_output.output)
179                except json.JSONDecodeError:
180                    raise ValueError("Repaired output is not a valid JSON object")
181
182                validate_schema_with_value_error(
183                    output_parsed,
184                    task.output_json_schema,
185                    "Repaired output does not match task output schema.",
186                )
187
188        if self.repair_instructions is None and self.repaired_output is not None:
189            raise ValueError(
190                "Repair instructions are required if providing a repaired output."
191            )
192        if self.repair_instructions is not None and self.repaired_output is None:
193            raise ValueError(
194                "A repaired output is required if providing repair instructions."
195            )
196
197        return self
198
199    @model_validator(mode="after")
200    def validate_input_source(self, info: ValidationInfo) -> Self:
201        # On strict mode and not loaded from file, we validate input_source is not None.
202        # 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.
203        if not strict_mode():
204            return self
205        if self.loaded_from_file(info):
206            return self
207        if self.input_source is None:
208            raise ValueError("input_source is required when strict mode is enabled")
209        return self
210
211    @model_validator(mode="after")
212    def validate_tags(self) -> Self:
213        for tag in self.tags:
214            if not tag:
215                raise ValueError("Tags cannot be empty strings")
216            if " " in tag:
217                raise ValueError("Tags cannot contain spaces. Try underscores.")
218
219        return self

287class TaskOutput(KilnBaseModel):
288    """
289    An output for a specific task run.
290
291    Contains the actual output content, its source (human or synthetic),
292    and optional rating information.
293    """
294
295    output: str = Field(
296        description="The output of the task. JSON formatted for structured output, plaintext for unstructured output."
297    )
298    source: DataSource | None = Field(
299        description="The source of the output: human or synthetic.",
300        default=None,
301    )
302    rating: TaskOutputRating | None = Field(
303        default=None, description="The rating of the output"
304    )
305
306    def validate_output_format(self, task: "Task") -> Self:
307        # validate output
308        if task.output_json_schema is not None:
309            try:
310                output_parsed = json.loads(self.output)
311            except json.JSONDecodeError:
312                raise ValueError("Output is not a valid JSON object")
313
314            validate_schema_with_value_error(
315                output_parsed,
316                task.output_json_schema,
317                "This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema. Search 'Troubleshooting Structured Data Issues' in our docs for more information.",
318            )
319        return self
320
321    @model_validator(mode="after")
322    def validate_output_source(self, info: ValidationInfo) -> Self:
323        # On strict mode and not loaded from file, we validate output_source is not None.
324        # 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.
325        if not strict_mode():
326            return self
327        if self.loaded_from_file(info):
328            return self
329        if self.source is None:
330            raise ValueError("Output source is required when strict mode is enabled")
331        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

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

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

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

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