kiln_ai.datamodel

 7class BasePrompt(BaseModel):
 8    """
 9    A prompt for a task. This is the basic data storage format which can be used throughout a project.
10
11    The "Prompt" model name is reserved for the custom prompts parented by a task.
12    """
13
14    name: FilenameString = Field(description="The name of the prompt.")
15    description: str | None = Field(
16        default=None,
17        description="A more detailed description of the prompt.",
18    )
19    generator_id: str | None = Field(
20        default=None,
21        description="The id of the generator that created this prompt.",
22    )
23    prompt: str = Field(
24        description="The prompt for the task.",
25        min_length=1,
26    )
27    chain_of_thought_instructions: str | None = Field(
28        default=None,
29        description="Instructions for the model 'thinking' about the requirement prior to answering. Used for chain of thought style prompting. COT will not be used unless this is provided.",
30    )

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

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] = []

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

 67class DatasetSplit(KilnParentedModel):
 68    """
 69    A collection of task runs, with optional splits (train, test, validation).
 70
 71    Used to freeze a dataset into train/test/validation splits for repeatable fine-tuning or other tasks.
 72
 73    Maintains a list of IDs for each split, to avoid data duplication.
 74    """
 75
 76    name: FilenameString = Field(description="The name of the dataset split.")
 77    description: str | None = Field(
 78        default=None,
 79        description="A description of the dataset for you and your team. Not used in training.",
 80    )
 81    splits: list[DatasetSplitDefinition] = Field(
 82        default_factory=list,
 83        description="The splits in the dataset.",
 84    )
 85    split_contents: dict[str, list[str]] = Field(
 86        description="The contents of each split in the dataset. The key is the split name, and the value is a list of task run IDs.",
 87    )
 88    filter: DatasetFilterId | None = Field(
 89        default=None,
 90        description="The filter used to build the dataset.",
 91    )
 92
 93    @model_validator(mode="after")
 94    def validate_split_percentages(self) -> "DatasetSplit":
 95        total = sum(split.percentage for split in self.splits)
 96        if not math.isclose(total, 1.0, rel_tol=1e-9):
 97            raise ValueError(f"The sum of split percentages must be 1.0 (got {total})")
 98        return self
 99
100    @classmethod
101    def from_task(
102        cls,
103        name: str,
104        task: "Task",
105        splits: list[DatasetSplitDefinition],
106        filter_id: DatasetFilterId = "all",
107        description: str | None = None,
108    ):
109        """
110        Build a dataset split from a task.
111        """
112        filter = dataset_filter_from_id(filter_id)
113        split_contents = cls.build_split_contents(task, splits, filter)
114        return cls(
115            parent=task,
116            name=name,
117            description=description,
118            splits=splits,
119            split_contents=split_contents,
120            filter=filter_id,
121        )
122
123    @classmethod
124    def build_split_contents(
125        cls,
126        task: "Task",
127        splits: list[DatasetSplitDefinition],
128        filter: DatasetFilter,
129    ) -> dict[str, list[str]]:
130        valid_ids = []
131        for task_run in task.runs():
132            if filter(task_run):
133                valid_ids.append(task_run.id)
134
135        # Shuffle and split by split percentage
136        random.shuffle(valid_ids)
137        split_contents = {}
138        start_idx = 0
139        remaining_items = len(valid_ids)
140
141        # Handle all splits except the last one
142        for split in splits[:-1]:
143            split_size = round(len(valid_ids) * split.percentage)
144            split_contents[split.name] = valid_ids[start_idx : start_idx + split_size]
145            start_idx += split_size
146            remaining_items -= split_size
147
148        # Last split gets all remaining items (for rounding)
149        if splits:
150            split_contents[splits[-1].name] = valid_ids[start_idx:]
151
152        return split_contents
153
154    def parent_task(self) -> "Task | None":
155        # inline import to avoid circular import
156        from kiln_ai.datamodel import Task
157
158        if not isinstance(self.parent, Task):
159            return None
160        return self.parent
161
162    def missing_count(self) -> int:
163        """
164        Returns:
165            int: the number of task runs that have an ID persisted in this dataset split, but no longer exist in the dataset
166        """
167        parent = self.parent_task()
168        if parent is None:
169            raise ValueError("DatasetSplit has no parent task")
170
171        runs = parent.runs(readonly=True)
172        all_ids = set(run.id for run in runs)
173        all_ids_in_splits = set()
174        for ids in self.split_contents.values():
175            all_ids_in_splits.update(ids)
176        missing = all_ids_in_splits - all_ids
177        return len(missing)

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: FilenameString = Field(
31        description="The name of the dataset split definition."
32    )
33    description: str | None = Field(
34        default=None,
35        description="A description of the dataset for you and your team. Not used in training.",
36    )
37    percentage: float = Field(
38        ge=0.0,
39        le=1.0,
40        description="The percentage of the dataset that this split represents (between 0 and 1).",
41    )

 52class ExternalToolServer(KilnParentedModel):
 53    """
 54    Configuration for communicating with a external MCP (Model Context Protocol) Server for LLM tool calls. External tool servers can be remote or local.
 55
 56    This model stores the necessary configuration to connect to and authenticate with
 57    external MCP servers that provide tools for LLM interactions.
 58    """
 59
 60    name: FilenameString = Field(description="The name of the external tool.")
 61    type: ToolServerType = Field(
 62        description="The type of external tool server. Remote tools are hosted on a remote server",
 63    )
 64    description: str | None = Field(
 65        default=None,
 66        description="A description of the external tool for you and your team. Will not be used in prompts/training/validation.",
 67    )
 68
 69    properties: (
 70        LocalServerProperties | RemoteServerProperties | KilnTaskServerProperties
 71    ) = Field(
 72        description="Configuration properties specific to the tool type.",
 73    )
 74
 75    # Private variable to store unsaved secrets
 76    _unsaved_secrets: dict[str, str] = PrivateAttr(default_factory=dict)
 77
 78    def model_post_init(self, __context: Any) -> None:
 79        # Process secrets after initialization (pydantic v2 hook)
 80        self._process_secrets_from_properties()
 81
 82    def _process_secrets_from_properties(self) -> None:
 83        """
 84        Extract secrets from properties and move them to _unsaved_secrets.
 85        This removes secrets from the properties dict so they aren't saved to file.
 86        Clears existing _unsaved_secrets first to handle property updates correctly.
 87        """
 88        # Clear existing unsaved secrets since we're reprocessing
 89        self._unsaved_secrets.clear()
 90
 91        secret_keys = self.get_secret_keys()
 92
 93        if not secret_keys:
 94            return
 95
 96        # Extract secret values from properties based on server type
 97        match self.type:
 98            case ToolServerType.remote_mcp:
 99                headers = self.properties.get("headers", {})
100                for key_name in secret_keys:
101                    if key_name in headers:
102                        self._unsaved_secrets[key_name] = headers[key_name]
103                        # Remove from headers immediately so they are not saved to file
104                        del headers[key_name]
105
106            case ToolServerType.local_mcp:
107                env_vars = self.properties.get("env_vars", {})
108                for key_name in secret_keys:
109                    if key_name in env_vars:
110                        self._unsaved_secrets[key_name] = env_vars[key_name]
111                        # Remove from env_vars immediately so they are not saved to file
112                        del env_vars[key_name]
113
114            case ToolServerType.kiln_task:
115                pass
116
117            case _:
118                raise_exhaustive_enum_error(self.type)
119
120    def __setattr__(self, name: str, value: Any) -> None:
121        """
122        Override __setattr__ to process secrets whenever properties are updated.
123        """
124        super().__setattr__(name, value)
125
126        # Process secrets whenever properties are updated
127        if name == "properties":
128            self._process_secrets_from_properties()
129
130    # Validation Helpers
131
132    @classmethod
133    def check_server_url(cls, server_url: str) -> None:
134        """Validate Server URL"""
135        if not isinstance(server_url, str):
136            raise ValueError("Server URL must be a string")
137
138        # Check for leading whitespace in URL
139        if server_url != server_url.lstrip():
140            raise ValueError("Server URL must not have leading whitespace")
141
142        parsed_url = urlparse(server_url)
143        if not parsed_url.netloc:
144            raise ValueError("Server URL is not a valid URL")
145        if parsed_url.scheme not in ["http", "https"]:
146            raise ValueError("Server URL must start with http:// or https://")
147
148    @classmethod
149    def check_headers(cls, headers: dict) -> None:
150        """Validate Headers"""
151        if not isinstance(headers, dict):
152            raise ValueError("headers must be a dictionary")
153
154        for key, value in headers.items():
155            if not key:
156                raise ValueError("Header name is required")
157            if not value:
158                raise ValueError("Header value is required")
159
160            # Reject invalid header names and CR/LF in names/values
161            token_re = re.compile(r"^[!#$%&'*+.^_`|~0-9A-Za-z-]+$")
162            if not token_re.match(key):
163                raise ValueError(f'Invalid header name: "{key}"')
164            if re.search(r"\r|\n", key) or re.search(r"\r|\n", value):
165                raise ValueError(
166                    "Header names/values must not contain invalid characters"
167                )
168
169    @classmethod
170    def check_secret_keys(
171        cls, secret_keys: list, key_type: str, tool_type: str
172    ) -> None:
173        """Validate Secret Keys (generic method for both header and env var keys)"""
174        if not isinstance(secret_keys, list):
175            raise ValueError(
176                f"{key_type} must be a list for external tools of type '{tool_type}'"
177            )
178        if not all(isinstance(k, str) for k in secret_keys):
179            raise ValueError(f"{key_type} must contain only strings")
180        if not all(key for key in secret_keys):
181            raise ValueError("Secret key is required")
182
183    @classmethod
184    def check_env_vars(cls, env_vars: dict) -> None:
185        """Validate Environment Variables"""
186        if not isinstance(env_vars, dict):
187            raise ValueError("environment variables must be a dictionary")
188
189        # Validate env_vars keys are in the correct format for Environment Variables
190        # According to POSIX specification, environment variable names must:
191        # - Start with a letter (a-z, A-Z) or underscore (_)
192        # - Contain only ASCII letters, digits, and underscores
193        for key, _ in env_vars.items():
194            if not key or not (
195                key[0].isascii() and (key[0].isalpha() or key[0] == "_")
196            ):
197                raise ValueError(
198                    f"Invalid environment variable key: {key}. Must start with a letter or underscore."
199                )
200
201            if not all(c.isascii() and (c.isalnum() or c == "_") for c in key):
202                raise ValueError(
203                    f"Invalid environment variable key: {key}. Can only contain letters, digits, and underscores."
204                )
205
206    @classmethod
207    def type_from_data(cls, data: dict) -> ToolServerType:
208        """Get the tool server type from the data for the the validators"""
209        raw_type = data.get("type")
210        if raw_type is None:
211            raise ValueError("type is required")
212        try:
213            return ToolServerType(raw_type)
214        except ValueError:
215            valid_types = ", ".join(type.value for type in ToolServerType)
216            raise ValueError(f"type must be one of: {valid_types}")
217
218    @model_validator(mode="before")
219    def upgrade_old_properties(cls, data: dict) -> dict:
220        """
221        Upgrade properties for backwards compatibility.
222        """
223        properties = data.get("properties")
224        if properties is not None and "is_archived" not in properties:
225            # Add is_archived field with default value back to data
226            properties["is_archived"] = False
227            data["properties"] = properties
228        return data
229
230    @model_validator(mode="before")
231    def validate_required_fields(cls, data: dict) -> dict:
232        """Validate that each tool type has the required configuration."""
233        server_type = ExternalToolServer.type_from_data(data)
234        properties = data.get("properties", {})
235
236        match server_type:
237            case ToolServerType.remote_mcp:
238                server_url = properties.get("server_url", None)
239                if server_url is None:
240                    raise ValueError(
241                        "Server URL is required to connect to a remote MCP server"
242                    )
243                ExternalToolServer.check_server_url(server_url)
244
245            case ToolServerType.local_mcp:
246                command = properties.get("command", None)
247                if command is None:
248                    raise ValueError("command is required to start a local MCP server")
249                if not isinstance(command, str):
250                    raise ValueError(
251                        "command must be a string to start a local MCP server"
252                    )
253                # Reject empty/whitespace-only command strings
254                if command.strip() == "":
255                    raise ValueError("command must be a non-empty string")
256
257                args = properties.get("args", None)
258                if args is not None:
259                    if not isinstance(args, list):
260                        raise ValueError(
261                            "arguments must be a list to start a local MCP server"
262                        )
263
264            case ToolServerType.kiln_task:
265                tool_name_validator(properties.get("name", ""))
266                err_msg_prefix = "Kiln task server properties:"
267                validate_return_dict_prop(
268                    properties, "description", str, err_msg_prefix
269                )
270                description = properties.get("description", "")
271                if len(description) > 128:
272                    raise ValueError("description must be 128 characters or less")
273                validate_return_dict_prop(
274                    properties, "is_archived", bool, err_msg_prefix
275                )
276                validate_return_dict_prop(properties, "task_id", str, err_msg_prefix)
277                validate_return_dict_prop(
278                    properties, "run_config_id", str, err_msg_prefix
279                )
280
281            case _:
282                # Type checking will catch missing cases
283                raise_exhaustive_enum_error(server_type)
284        return data
285
286    @model_validator(mode="before")
287    def validate_headers_and_env_vars(cls, data: dict) -> dict:
288        """
289        Validate secrets, these needs to be validated before model initlization because secrets will be processed and stripped
290        """
291        type = ExternalToolServer.type_from_data(data)
292
293        properties = data.get("properties", {})
294        if properties is None:
295            raise ValueError("properties is required")
296
297        match type:
298            case ToolServerType.remote_mcp:
299                # Validate headers
300                headers = properties.get("headers", None)
301                if headers is not None:
302                    ExternalToolServer.check_headers(headers)
303
304                # Secret header keys are optional, validate if they are set
305                secret_header_keys = properties.get("secret_header_keys", None)
306                if secret_header_keys is not None:
307                    ExternalToolServer.check_secret_keys(
308                        secret_header_keys, "secret_header_keys", "remote_mcp"
309                    )
310
311            case ToolServerType.local_mcp:
312                # Validate secret environment variable keys
313                env_vars = properties.get("env_vars", {})
314                if env_vars is not None:
315                    ExternalToolServer.check_env_vars(env_vars)
316
317                # Secret env var keys are optional, but if they are set, they must be a list of strings
318                secret_env_var_keys = properties.get("secret_env_var_keys", None)
319                if secret_env_var_keys is not None:
320                    ExternalToolServer.check_secret_keys(
321                        secret_env_var_keys, "secret_env_var_keys", "local_mcp"
322                    )
323
324            case ToolServerType.kiln_task:
325                pass
326
327            case _:
328                raise_exhaustive_enum_error(type)
329
330        return data
331
332    def get_secret_keys(self) -> list[str]:
333        """
334        Get the list of secret key names based on server type.
335
336        Returns:
337            List of secret key names (header names for remote, env var names for local)
338        """
339        match self.type:
340            case ToolServerType.remote_mcp:
341                return self.properties.get("secret_header_keys", [])
342            case ToolServerType.local_mcp:
343                return self.properties.get("secret_env_var_keys", [])
344            case ToolServerType.kiln_task:
345                return []
346            case _:
347                raise_exhaustive_enum_error(self.type)
348
349    def retrieve_secrets(self) -> tuple[dict[str, str], list[str]]:
350        """
351        Retrieve secrets from configuration system or in-memory storage.
352        Automatically determines which secret keys to retrieve based on the server type.
353        Config secrets take precedence over unsaved secrets.
354
355        Returns:
356            Tuple of (secrets_dict, missing_secrets_list) where:
357            - secrets_dict: Dictionary mapping key names to their secret values
358            - missing_secrets_list: List of secret key names that are missing values
359        """
360        secrets = {}
361        missing_secrets = []
362        secret_keys = self.get_secret_keys()
363
364        if secret_keys and len(secret_keys) > 0:
365            config = Config.shared()
366            mcp_secrets = config.get_value(MCP_SECRETS_KEY)
367
368            for key_name in secret_keys:
369                secret_value = None
370
371                # First check config secrets (persistent storage), key is mcp_server_id::key_name
372                secret_key = self._config_secret_key(key_name)
373                secret_value = mcp_secrets.get(secret_key) if mcp_secrets else None
374
375                # Fall back to unsaved secrets (in-memory storage)
376                if (
377                    not secret_value
378                    and hasattr(self, "_unsaved_secrets")
379                    and key_name in self._unsaved_secrets
380                ):
381                    secret_value = self._unsaved_secrets[key_name]
382
383                if secret_value:
384                    secrets[key_name] = secret_value
385                else:
386                    missing_secrets.append(key_name)
387
388        return secrets, missing_secrets
389
390    def _save_secrets(self) -> None:
391        """
392        Save unsaved secrets to the configuration system.
393        """
394        secret_keys = self.get_secret_keys()
395
396        # No secrets to save
397        if not secret_keys:
398            return
399
400        if self.id is None:
401            raise ValueError("Server ID cannot be None when saving secrets")
402
403        # Check if secrets are already saved
404        if not hasattr(self, "_unsaved_secrets") or not self._unsaved_secrets:
405            return
406
407        config = Config.shared()
408        mcp_secrets: dict[str, str] = config.get_value(MCP_SECRETS_KEY) or {}
409
410        # Store secrets with the pattern: mcp_server_id::key_name
411        for key_name, secret_value in self._unsaved_secrets.items():
412            secret_key = self._config_secret_key(key_name)
413            mcp_secrets[secret_key] = secret_value
414
415        config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
416
417        # Clear unsaved secrets after saving
418        self._unsaved_secrets.clear()
419
420    def delete_secrets(self) -> None:
421        """
422        Delete all secrets for this tool server from the configuration system.
423        """
424        secret_keys = self.get_secret_keys()
425
426        config = Config.shared()
427        mcp_secrets = config.get_value(MCP_SECRETS_KEY) or dict[str, str]()
428
429        # Remove secrets with the pattern: mcp_server_id::key_name
430        for key_name in secret_keys:
431            secret_key = self._config_secret_key(key_name)
432            if secret_key in mcp_secrets:
433                del mcp_secrets[secret_key]
434
435        # Always call update_settings to maintain consistency with the old behavior
436        config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
437
438    def save_to_file(self) -> None:
439        """
440        Override save_to_file to automatically save any unsaved secrets before saving to file.
441
442        This ensures that secrets are always saved when the object is saved,
443        preventing the issue where secrets could be lost if save_to_file is called
444        without explicitly saving secrets first.
445        """
446        # Save any unsaved secrets first
447        if hasattr(self, "_unsaved_secrets") and self._unsaved_secrets:
448            self._save_secrets()
449
450        # Call the parent save_to_file method
451        super().save_to_file()
452
453    #  Internal helpers
454
455    def _config_secret_key(self, key_name: str) -> str:
456        """
457        Generate the secret key pattern for storing/retrieving secrets.
458
459        Args:
460            key_name: The name of the secret key
461
462        Returns:
463            The formatted secret key: "{server_id}::{key_name}"
464        """
465        return f"{self.id}::{key_name}"