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    )

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