kiln-ai 0.18.0__py3-none-any.whl → 0.20.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of kiln-ai might be problematic. Click here for more details.

Files changed (89) hide show
  1. kiln_ai/adapters/__init__.py +2 -2
  2. kiln_ai/adapters/adapter_registry.py +46 -0
  3. kiln_ai/adapters/chat/chat_formatter.py +8 -12
  4. kiln_ai/adapters/chat/test_chat_formatter.py +6 -2
  5. kiln_ai/adapters/data_gen/data_gen_task.py +2 -2
  6. kiln_ai/adapters/data_gen/test_data_gen_task.py +7 -3
  7. kiln_ai/adapters/docker_model_runner_tools.py +119 -0
  8. kiln_ai/adapters/eval/base_eval.py +2 -2
  9. kiln_ai/adapters/eval/eval_runner.py +3 -1
  10. kiln_ai/adapters/eval/g_eval.py +2 -2
  11. kiln_ai/adapters/eval/test_base_eval.py +1 -1
  12. kiln_ai/adapters/eval/test_eval_runner.py +6 -12
  13. kiln_ai/adapters/eval/test_g_eval.py +3 -4
  14. kiln_ai/adapters/eval/test_g_eval_data.py +1 -1
  15. kiln_ai/adapters/fine_tune/__init__.py +1 -1
  16. kiln_ai/adapters/fine_tune/base_finetune.py +1 -0
  17. kiln_ai/adapters/fine_tune/fireworks_finetune.py +32 -20
  18. kiln_ai/adapters/fine_tune/openai_finetune.py +14 -4
  19. kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py +30 -21
  20. kiln_ai/adapters/fine_tune/test_openai_finetune.py +108 -111
  21. kiln_ai/adapters/ml_model_list.py +1009 -111
  22. kiln_ai/adapters/model_adapters/base_adapter.py +62 -28
  23. kiln_ai/adapters/model_adapters/litellm_adapter.py +397 -80
  24. kiln_ai/adapters/model_adapters/test_base_adapter.py +194 -18
  25. kiln_ai/adapters/model_adapters/test_litellm_adapter.py +428 -4
  26. kiln_ai/adapters/model_adapters/test_litellm_adapter_tools.py +1103 -0
  27. kiln_ai/adapters/model_adapters/test_saving_adapter_results.py +5 -5
  28. kiln_ai/adapters/model_adapters/test_structured_output.py +120 -14
  29. kiln_ai/adapters/parsers/__init__.py +1 -1
  30. kiln_ai/adapters/parsers/test_r1_parser.py +1 -1
  31. kiln_ai/adapters/provider_tools.py +35 -20
  32. kiln_ai/adapters/remote_config.py +57 -10
  33. kiln_ai/adapters/repair/repair_task.py +1 -1
  34. kiln_ai/adapters/repair/test_repair_task.py +12 -9
  35. kiln_ai/adapters/run_output.py +3 -0
  36. kiln_ai/adapters/test_adapter_registry.py +109 -2
  37. kiln_ai/adapters/test_docker_model_runner_tools.py +305 -0
  38. kiln_ai/adapters/test_ml_model_list.py +51 -1
  39. kiln_ai/adapters/test_prompt_adaptors.py +13 -6
  40. kiln_ai/adapters/test_provider_tools.py +73 -12
  41. kiln_ai/adapters/test_remote_config.py +470 -16
  42. kiln_ai/datamodel/__init__.py +23 -21
  43. kiln_ai/datamodel/basemodel.py +54 -28
  44. kiln_ai/datamodel/datamodel_enums.py +3 -0
  45. kiln_ai/datamodel/dataset_split.py +5 -3
  46. kiln_ai/datamodel/eval.py +4 -4
  47. kiln_ai/datamodel/external_tool_server.py +298 -0
  48. kiln_ai/datamodel/finetune.py +2 -2
  49. kiln_ai/datamodel/json_schema.py +25 -10
  50. kiln_ai/datamodel/project.py +11 -4
  51. kiln_ai/datamodel/prompt.py +2 -2
  52. kiln_ai/datamodel/prompt_id.py +4 -4
  53. kiln_ai/datamodel/registry.py +0 -15
  54. kiln_ai/datamodel/run_config.py +62 -0
  55. kiln_ai/datamodel/task.py +8 -83
  56. kiln_ai/datamodel/task_output.py +7 -2
  57. kiln_ai/datamodel/task_run.py +41 -0
  58. kiln_ai/datamodel/test_basemodel.py +213 -21
  59. kiln_ai/datamodel/test_eval_model.py +6 -6
  60. kiln_ai/datamodel/test_example_models.py +175 -0
  61. kiln_ai/datamodel/test_external_tool_server.py +691 -0
  62. kiln_ai/datamodel/test_model_perf.py +1 -1
  63. kiln_ai/datamodel/test_prompt_id.py +5 -1
  64. kiln_ai/datamodel/test_registry.py +8 -3
  65. kiln_ai/datamodel/test_task.py +20 -47
  66. kiln_ai/datamodel/test_tool_id.py +239 -0
  67. kiln_ai/datamodel/tool_id.py +83 -0
  68. kiln_ai/tools/__init__.py +8 -0
  69. kiln_ai/tools/base_tool.py +82 -0
  70. kiln_ai/tools/built_in_tools/__init__.py +13 -0
  71. kiln_ai/tools/built_in_tools/math_tools.py +124 -0
  72. kiln_ai/tools/built_in_tools/test_math_tools.py +204 -0
  73. kiln_ai/tools/mcp_server_tool.py +95 -0
  74. kiln_ai/tools/mcp_session_manager.py +243 -0
  75. kiln_ai/tools/test_base_tools.py +199 -0
  76. kiln_ai/tools/test_mcp_server_tool.py +457 -0
  77. kiln_ai/tools/test_mcp_session_manager.py +1585 -0
  78. kiln_ai/tools/test_tool_registry.py +473 -0
  79. kiln_ai/tools/tool_registry.py +64 -0
  80. kiln_ai/utils/config.py +32 -0
  81. kiln_ai/utils/open_ai_types.py +94 -0
  82. kiln_ai/utils/project_utils.py +17 -0
  83. kiln_ai/utils/test_config.py +138 -1
  84. kiln_ai/utils/test_open_ai_types.py +131 -0
  85. {kiln_ai-0.18.0.dist-info → kiln_ai-0.20.1.dist-info}/METADATA +37 -6
  86. kiln_ai-0.20.1.dist-info/RECORD +138 -0
  87. kiln_ai-0.18.0.dist-info/RECORD +0 -115
  88. {kiln_ai-0.18.0.dist-info → kiln_ai-0.20.1.dist-info}/WHEEL +0 -0
  89. {kiln_ai-0.18.0.dist-info → kiln_ai-0.20.1.dist-info}/licenses/LICENSE.txt +0 -0
@@ -2,22 +2,17 @@ import json
2
2
  import os
3
3
  import re
4
4
  import shutil
5
+ import unicodedata
5
6
  import uuid
6
7
  from abc import ABCMeta
7
8
  from builtins import classmethod
8
9
  from datetime import datetime
9
10
  from pathlib import Path
10
- from typing import (
11
- Any,
12
- Dict,
13
- List,
14
- Optional,
15
- Type,
16
- TypeVar,
17
- )
11
+ from typing import Any, Callable, Dict, List, Optional, Type, TypeVar
18
12
 
19
13
  from pydantic import (
20
14
  BaseModel,
15
+ BeforeValidator,
21
16
  ConfigDict,
22
17
  Field,
23
18
  ValidationError,
@@ -26,7 +21,7 @@ from pydantic import (
26
21
  model_validator,
27
22
  )
28
23
  from pydantic_core import ErrorDetails
29
- from typing_extensions import Self
24
+ from typing_extensions import Annotated, Self
30
25
 
31
26
  from kiln_ai.datamodel.model_cache import ModelCache
32
27
  from kiln_ai.utils.config import Config
@@ -44,33 +39,64 @@ PT = TypeVar("PT", bound="KilnParentedModel")
44
39
 
45
40
  # Naming conventions:
46
41
  # 1) Names are filename safe as they may be used as file names. They are informational and not to be used in prompts/training/validation.
47
- # 2) Descrptions are for Kiln users to describe/understanding the purpose of this object. They must never be used in prompts/training/validation. Use "instruction/requirements" instead.
48
-
49
- # Filename compatible names
50
- NAME_REGEX = r"^[A-Za-z0-9 _-]+$"
51
- NAME_FIELD = Field(
52
- min_length=1,
53
- max_length=120,
54
- pattern=NAME_REGEX,
55
- description="A name for this entity.",
56
- )
57
- SHORT_NAME_FIELD = Field(
58
- min_length=1,
59
- max_length=32,
60
- pattern=NAME_REGEX,
61
- description="A name for this entity",
62
- )
42
+ # 2) Descriptions are for Kiln users to describe/understanding the purpose of this object. They must never be used in prompts/training/validation. Use "instruction/requirements" instead.
43
+
44
+ # Forbidden chars are not allowed in filenames on one or more platforms.
45
+ # ref: https://en.wikipedia.org/wiki/Filename#Problematic_characters
46
+ FORBIDDEN_CHARS_REGEX = r"[/\\?%*:|\"<>.,;=\n]"
47
+ FORBIDDEN_CHARS = "/ \\ ? % * : | < > . , ; = \\n"
48
+
49
+
50
+ def name_validator(*, min_length: int, max_length: int) -> Callable[[Any], str]:
51
+ def fn(name: Any) -> str:
52
+ if name is None:
53
+ raise ValueError("Name is required")
54
+ if not isinstance(name, str):
55
+ raise ValueError(f"Input should be a valid string, got {type(name)}")
56
+ if len(name) < min_length:
57
+ raise ValueError(
58
+ f"Name is too short. Min length is {min_length} characters, got {len(name)}"
59
+ )
60
+ if len(name) > max_length:
61
+ raise ValueError(
62
+ f"Name is too long. Max length is {max_length} characters, got {len(name)}"
63
+ )
64
+ if string_to_valid_name(name) != name:
65
+ raise ValueError(
66
+ f"Name is invalid. The name cannot contain any of the following characters: {FORBIDDEN_CHARS}, consecutive whitespace/underscores, or leading/trailing whitespace/underscores"
67
+ )
68
+ return name
69
+
70
+ return fn
63
71
 
64
72
 
65
73
  def string_to_valid_name(name: str) -> str:
66
- # Replace any character not allowed by NAME_REGEX with an underscore
67
- valid_name = re.sub(r"[^A-Za-z0-9 _-]", "_", name)
74
+ # https://docs.python.org/3/library/unicodedata.html#unicodedata.normalize
75
+ valid_name = unicodedata.normalize("NFKD", name)
76
+ # Replace any forbidden chars with an underscore
77
+ valid_name = re.sub(FORBIDDEN_CHARS_REGEX, "_", valid_name)
78
+ # Replace control characters with an underscore
79
+ valid_name = re.sub(r"[\x00-\x1F]", "_", valid_name)
80
+ # Replace consecutive whitespace with a single space
81
+ valid_name = re.sub(r"\s+", " ", valid_name)
68
82
  # Replace consecutive underscores with a single underscore
69
83
  valid_name = re.sub(r"_+", "_", valid_name)
70
84
  # Remove leading and trailing underscores or whitespace
71
85
  return valid_name.strip("_").strip()
72
86
 
73
87
 
88
+ # Usage:
89
+ # class MyModel(KilnBaseModel):
90
+ # name: FilenameString = Field(description="The name of the model.")
91
+ # name_short: FilenameStringShort = Field(description="The short name of the model.")
92
+ FilenameString = Annotated[
93
+ str, BeforeValidator(name_validator(min_length=1, max_length=120))
94
+ ]
95
+ FilenameStringShort = Annotated[
96
+ str, BeforeValidator(name_validator(min_length=1, max_length=32))
97
+ ]
98
+
99
+
74
100
  class KilnBaseModel(BaseModel):
75
101
  """Base model for all Kiln data models with common functionality for persistence and versioning.
76
102
 
@@ -470,7 +496,7 @@ class KilnParentModel(KilnBaseModel, metaclass=ABCMeta):
470
496
  ValidationError: If validation fails for the model or any of its children
471
497
  """
472
498
  # Validate first, then save. Don't want error half way through, and partly persisted
473
- # TODO P2: save to tmp dir, then move atomically. But need to merge directories so later.
499
+ # We should save to a tmp dir and move atomically, but need to merge directories later.
474
500
  cls._validate_nested(data, save=False, path=path, parent=parent)
475
501
  instance = cls._validate_nested(data, save=True, path=path, parent=parent)
476
502
  return instance
@@ -97,3 +97,6 @@ class ModelProviderName(str, Enum):
97
97
  huggingface = "huggingface"
98
98
  vertex = "vertex"
99
99
  together_ai = "together_ai"
100
+ siliconflow_cn = "siliconflow_cn"
101
+ cerebras = "cerebras"
102
+ docker_model_runner = "docker_model_runner"
@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING
8
8
 
9
9
  from pydantic import BaseModel, Field, model_validator
10
10
 
11
- from kiln_ai.datamodel.basemodel import NAME_FIELD, KilnParentedModel
11
+ from kiln_ai.datamodel.basemodel import FilenameString, KilnParentedModel
12
12
  from kiln_ai.datamodel.dataset_filters import (
13
13
  DatasetFilter,
14
14
  DatasetFilterId,
@@ -26,7 +26,9 @@ class DatasetSplitDefinition(BaseModel):
26
26
  Example: name="train", description="The training set", percentage=0.8 (80% of the dataset)
27
27
  """
28
28
 
29
- name: str = NAME_FIELD
29
+ name: FilenameString = Field(
30
+ description="The name of the dataset split definition."
31
+ )
30
32
  description: str | None = Field(
31
33
  default=None,
32
34
  description="A description of the dataset for you and your team. Not used in training.",
@@ -70,7 +72,7 @@ class DatasetSplit(KilnParentedModel):
70
72
  Maintains a list of IDs for each split, to avoid data duplication.
71
73
  """
72
74
 
73
- name: str = NAME_FIELD
75
+ name: FilenameString = Field(description="The name of the dataset split.")
74
76
  description: str | None = Field(
75
77
  default=None,
76
78
  description="A description of the dataset for you and your team. Not used in training.",
kiln_ai/datamodel/eval.py CHANGED
@@ -7,7 +7,7 @@ from typing_extensions import Self
7
7
 
8
8
  from kiln_ai.datamodel.basemodel import (
9
9
  ID_TYPE,
10
- NAME_FIELD,
10
+ FilenameString,
11
11
  KilnParentedModel,
12
12
  KilnParentModel,
13
13
  )
@@ -202,7 +202,7 @@ class EvalConfig(KilnParentedModel, KilnParentModel, parent_of={"runs": EvalRun}
202
202
  A eval might have many configs, example running the same eval with 2 different models. Comparing eval results is only valid within the scope of the same config.
203
203
  """
204
204
 
205
- name: str = NAME_FIELD
205
+ name: FilenameString = Field(description="The name of the eval config.")
206
206
  model_name: str = Field(
207
207
  description="The name of the model to use for this eval config. ",
208
208
  )
@@ -252,12 +252,12 @@ class EvalConfig(KilnParentedModel, KilnParentModel, parent_of={"runs": EvalRun}
252
252
  # This will raise a TypeError if the dict contains non-JSON-serializable objects
253
253
  json.dumps(self.properties)
254
254
  except TypeError as e:
255
- raise ValueError(f"Properties must be JSON serializable: {str(e)}")
255
+ raise ValueError(f"Properties must be JSON serializable: {e!s}")
256
256
  return self
257
257
 
258
258
 
259
259
  class Eval(KilnParentedModel, KilnParentModel, parent_of={"configs": EvalConfig}):
260
- name: str = NAME_FIELD
260
+ name: FilenameString = Field(description="The name of the eval.")
261
261
  description: str | None = Field(
262
262
  default=None, description="The description of the eval"
263
263
  )
@@ -0,0 +1,298 @@
1
+ from enum import Enum
2
+ from typing import Any, Dict
3
+
4
+ from pydantic import Field, PrivateAttr, model_validator
5
+
6
+ from kiln_ai.datamodel.basemodel import (
7
+ FilenameString,
8
+ KilnParentedModel,
9
+ )
10
+ from kiln_ai.utils.config import MCP_SECRETS_KEY, Config
11
+ from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
12
+
13
+
14
+ class ToolServerType(str, Enum):
15
+ """
16
+ Enumeration of supported external tool server types.
17
+ """
18
+
19
+ remote_mcp = "remote_mcp"
20
+ local_mcp = "local_mcp"
21
+
22
+
23
+ class ExternalToolServer(KilnParentedModel):
24
+ """
25
+ Configuration for communicating with a external MCP (Model Context Protocol) Server for LLM tool calls. External tool servers can be remote or local.
26
+
27
+ This model stores the necessary configuration to connect to and authenticate with
28
+ external MCP servers that provide tools for LLM interactions.
29
+ """
30
+
31
+ name: FilenameString = Field(description="The name of the external tool.")
32
+ type: ToolServerType = Field(
33
+ description="The type of external tool server. Remote tools are hosted on a remote server",
34
+ )
35
+ description: str | None = Field(
36
+ default=None,
37
+ description="A description of the external tool for you and your team. Will not be used in prompts/training/validation.",
38
+ )
39
+ properties: Dict[str, Any] = Field(
40
+ default={},
41
+ description="Configuration properties specific to the tool type.",
42
+ )
43
+
44
+ # Private variable to store unsaved secrets
45
+ _unsaved_secrets: dict[str, str] = PrivateAttr(default_factory=dict)
46
+
47
+ def model_post_init(self, __context: Any) -> None:
48
+ # Process secrets after initialization (pydantic v2 hook)
49
+ self._process_secrets_from_properties()
50
+
51
+ def _process_secrets_from_properties(self) -> None:
52
+ """
53
+ Extract secrets from properties and move them to _unsaved_secrets.
54
+ This removes secrets from the properties dict so they aren't saved to file.
55
+ Clears existing _unsaved_secrets first to handle property updates correctly.
56
+ """
57
+ # Clear existing unsaved secrets since we're reprocessing
58
+ self._unsaved_secrets.clear()
59
+
60
+ secret_keys = self.get_secret_keys()
61
+
62
+ if not secret_keys:
63
+ return
64
+
65
+ # Extract secret values from properties based on server type
66
+ match self.type:
67
+ case ToolServerType.remote_mcp:
68
+ headers = self.properties.get("headers", {})
69
+ for key_name in secret_keys:
70
+ if key_name in headers:
71
+ self._unsaved_secrets[key_name] = headers[key_name]
72
+ # Remove from headers immediately so they are not saved to file
73
+ del headers[key_name]
74
+
75
+ case ToolServerType.local_mcp:
76
+ env_vars = self.properties.get("env_vars", {})
77
+ for key_name in secret_keys:
78
+ if key_name in env_vars:
79
+ self._unsaved_secrets[key_name] = env_vars[key_name]
80
+ # Remove from env_vars immediately so they are not saved to file
81
+ del env_vars[key_name]
82
+
83
+ case _:
84
+ raise_exhaustive_enum_error(self.type)
85
+
86
+ def __setattr__(self, name: str, value: Any) -> None:
87
+ """
88
+ Override __setattr__ to process secrets whenever properties are updated.
89
+ """
90
+ super().__setattr__(name, value)
91
+
92
+ # Process secrets whenever properties are updated
93
+ if name == "properties":
94
+ self._process_secrets_from_properties()
95
+
96
+ @model_validator(mode="after")
97
+ def validate_required_fields(self) -> "ExternalToolServer":
98
+ """Validate that each tool type has the required configuration."""
99
+ match self.type:
100
+ case ToolServerType.remote_mcp:
101
+ server_url = self.properties.get("server_url", None)
102
+ if not isinstance(server_url, str):
103
+ raise ValueError(
104
+ "server_url must be a string for external tools of type 'remote_mcp'"
105
+ )
106
+ if not server_url:
107
+ raise ValueError(
108
+ "server_url is required to connect to a remote MCP server"
109
+ )
110
+
111
+ headers = self.properties.get("headers", None)
112
+ if headers is None:
113
+ raise ValueError("headers must be set when type is 'remote_mcp'")
114
+ if not isinstance(headers, dict):
115
+ raise ValueError(
116
+ "headers must be a dictionary for external tools of type 'remote_mcp'"
117
+ )
118
+
119
+ secret_header_keys = self.properties.get("secret_header_keys", None)
120
+ # Secret header keys are optional, but if they are set, they must be a list of strings
121
+ if secret_header_keys is not None:
122
+ if not isinstance(secret_header_keys, list):
123
+ raise ValueError(
124
+ "secret_header_keys must be a list for external tools of type 'remote_mcp'"
125
+ )
126
+ if not all(isinstance(k, str) for k in secret_header_keys):
127
+ raise ValueError("secret_header_keys must contain only strings")
128
+
129
+ case ToolServerType.local_mcp:
130
+ command = self.properties.get("command", None)
131
+ if not isinstance(command, str):
132
+ raise ValueError(
133
+ "command must be a string to start a local MCP server"
134
+ )
135
+ if not command.strip():
136
+ raise ValueError("command is required to start a local MCP server")
137
+
138
+ args = self.properties.get("args", None)
139
+ if not isinstance(args, list):
140
+ raise ValueError(
141
+ "arguments must be a list to start a local MCP server"
142
+ )
143
+
144
+ env_vars = self.properties.get("env_vars", {})
145
+ if not isinstance(env_vars, dict):
146
+ raise ValueError(
147
+ "environment variables must be a dictionary for external tools of type 'local_mcp'"
148
+ )
149
+
150
+ secret_env_var_keys = self.properties.get("secret_env_var_keys", None)
151
+ # Secret env var keys are optional, but if they are set, they must be a list of strings
152
+ if secret_env_var_keys is not None:
153
+ if not isinstance(secret_env_var_keys, list):
154
+ raise ValueError(
155
+ "secret_env_var_keys must be a list for external tools of type 'local_mcp'"
156
+ )
157
+ if not all(isinstance(k, str) for k in secret_env_var_keys):
158
+ raise ValueError(
159
+ "secret_env_var_keys must contain only strings"
160
+ )
161
+
162
+ case _:
163
+ # Type checking will catch missing cases
164
+ raise_exhaustive_enum_error(self.type)
165
+ return self
166
+
167
+ def get_secret_keys(self) -> list[str]:
168
+ """
169
+ Get the list of secret key names based on server type.
170
+
171
+ Returns:
172
+ List of secret key names (header names for remote, env var names for local)
173
+ """
174
+ match self.type:
175
+ case ToolServerType.remote_mcp:
176
+ return self.properties.get("secret_header_keys", [])
177
+ case ToolServerType.local_mcp:
178
+ return self.properties.get("secret_env_var_keys", [])
179
+ case _:
180
+ raise_exhaustive_enum_error(self.type)
181
+
182
+ def retrieve_secrets(self) -> tuple[dict[str, str], list[str]]:
183
+ """
184
+ Retrieve secrets from configuration system or in-memory storage.
185
+ Automatically determines which secret keys to retrieve based on the server type.
186
+ Config secrets take precedence over unsaved secrets.
187
+
188
+ Returns:
189
+ Tuple of (secrets_dict, missing_secrets_list) where:
190
+ - secrets_dict: Dictionary mapping key names to their secret values
191
+ - missing_secrets_list: List of secret key names that are missing values
192
+ """
193
+ secrets = {}
194
+ missing_secrets = []
195
+ secret_keys = self.get_secret_keys()
196
+
197
+ if secret_keys and len(secret_keys) > 0:
198
+ config = Config.shared()
199
+ mcp_secrets = config.get_value(MCP_SECRETS_KEY)
200
+
201
+ for key_name in secret_keys:
202
+ secret_value = None
203
+
204
+ # First check config secrets (persistent storage), key is mcp_server_id::key_name
205
+ secret_key = self._config_secret_key(key_name)
206
+ secret_value = mcp_secrets.get(secret_key) if mcp_secrets else None
207
+
208
+ # Fall back to unsaved secrets (in-memory storage)
209
+ if (
210
+ not secret_value
211
+ and hasattr(self, "_unsaved_secrets")
212
+ and key_name in self._unsaved_secrets
213
+ ):
214
+ secret_value = self._unsaved_secrets[key_name]
215
+
216
+ if secret_value:
217
+ secrets[key_name] = secret_value
218
+ else:
219
+ missing_secrets.append(key_name)
220
+
221
+ return secrets, missing_secrets
222
+
223
+ def _save_secrets(self) -> None:
224
+ """
225
+ Save unsaved secrets to the configuration system.
226
+ """
227
+ secret_keys = self.get_secret_keys()
228
+
229
+ # No secrets to save
230
+ if not secret_keys:
231
+ return
232
+
233
+ if self.id is None:
234
+ raise ValueError("Server ID cannot be None when saving secrets")
235
+
236
+ # Check if secrets are already saved
237
+ if not hasattr(self, "_unsaved_secrets") or not self._unsaved_secrets:
238
+ return
239
+
240
+ config = Config.shared()
241
+ mcp_secrets: dict[str, str] = config.get_value(MCP_SECRETS_KEY) or {}
242
+
243
+ # Store secrets with the pattern: mcp_server_id::key_name
244
+ for key_name, secret_value in self._unsaved_secrets.items():
245
+ secret_key = self._config_secret_key(key_name)
246
+ mcp_secrets[secret_key] = secret_value
247
+
248
+ config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
249
+
250
+ # Clear unsaved secrets after saving
251
+ self._unsaved_secrets.clear()
252
+
253
+ def delete_secrets(self) -> None:
254
+ """
255
+ Delete all secrets for this tool server from the configuration system.
256
+ """
257
+ secret_keys = self.get_secret_keys()
258
+
259
+ config = Config.shared()
260
+ mcp_secrets = config.get_value(MCP_SECRETS_KEY) or dict[str, str]()
261
+
262
+ # Remove secrets with the pattern: mcp_server_id::key_name
263
+ for key_name in secret_keys:
264
+ secret_key = self._config_secret_key(key_name)
265
+ if secret_key in mcp_secrets:
266
+ del mcp_secrets[secret_key]
267
+
268
+ # Always call update_settings to maintain consistency with the old behavior
269
+ config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
270
+
271
+ def save_to_file(self) -> None:
272
+ """
273
+ Override save_to_file to automatically save any unsaved secrets before saving to file.
274
+
275
+ This ensures that secrets are always saved when the object is saved,
276
+ preventing the issue where secrets could be lost if save_to_file is called
277
+ without explicitly saving secrets first.
278
+ """
279
+ # Save any unsaved secrets first
280
+ if hasattr(self, "_unsaved_secrets") and self._unsaved_secrets:
281
+ self._save_secrets()
282
+
283
+ # Call the parent save_to_file method
284
+ super().save_to_file()
285
+
286
+ # Internal helpers
287
+
288
+ def _config_secret_key(self, key_name: str) -> str:
289
+ """
290
+ Generate the secret key pattern for storing/retrieving secrets.
291
+
292
+ Args:
293
+ key_name: The name of the secret key
294
+
295
+ Returns:
296
+ The formatted secret key: "{server_id}::{key_name}"
297
+ """
298
+ return f"{self.id}::{key_name}"
@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Dict, Union
3
3
  from pydantic import Field, model_validator
4
4
  from typing_extensions import Self
5
5
 
6
- from kiln_ai.datamodel.basemodel import NAME_FIELD, KilnParentedModel
6
+ from kiln_ai.datamodel.basemodel import FilenameString, KilnParentedModel
7
7
  from kiln_ai.datamodel.datamodel_enums import (
8
8
  ChatStrategy,
9
9
  FineTuneStatusType,
@@ -26,7 +26,7 @@ class Finetune(KilnParentedModel):
26
26
  Initially holds a reference to a training job, with needed identifiers to update the status. When complete, contains the new model ID.
27
27
  """
28
28
 
29
- name: str = NAME_FIELD
29
+ name: FilenameString = Field(description="The name of the fine-tune.")
30
30
  description: str | None = Field(
31
31
  default=None,
32
32
  description="A description of the fine-tune for you and your team. Not used in training.",
@@ -84,25 +84,40 @@ def schema_from_json_str(v: str) -> Dict:
84
84
  """
85
85
  try:
86
86
  parsed = json.loads(v)
87
- jsonschema.Draft202012Validator.check_schema(parsed)
88
87
  if not isinstance(parsed, dict):
89
88
  raise ValueError(f"JSON schema must be a dict, not {type(parsed)}")
90
- # Top level arrays are valid JSON schemas, but we don't want to allow them here as they often cause issues
91
- if (
92
- "type" not in parsed
93
- or parsed["type"] != "object"
94
- or "properties" not in parsed
95
- ):
96
- raise ValueError(f"JSON schema must be an object with properties: {v}")
89
+
90
+ validate_schema_dict(parsed)
97
91
  return parsed
98
- except jsonschema.exceptions.SchemaError as e:
99
- raise ValueError(f"Invalid JSON schema: {v} \n{e}")
100
92
  except json.JSONDecodeError as e:
101
93
  raise ValueError(f"Invalid JSON: {v}\n {e}")
102
94
  except Exception as e:
103
95
  raise ValueError(f"Unexpected error parsing JSON schema: {v}\n {e}")
104
96
 
105
97
 
98
+ def validate_schema_dict(v: Dict):
99
+ """Parse and validate a JSON schema dictionary.
100
+
101
+ Args:
102
+ v: Dictionary containing a JSON schema definition
103
+
104
+ Returns:
105
+ Dict containing the parsed JSON schema
106
+
107
+ Raises:
108
+ ValueError: If the input is not a valid JSON schema object with required properties
109
+ """
110
+ try:
111
+ jsonschema.Draft202012Validator.check_schema(v)
112
+ # Top level arrays are valid JSON schemas, but we don't want to allow them here as they often cause issues
113
+ if "type" not in v or v["type"] != "object" or "properties" not in v:
114
+ raise ValueError(f"JSON schema must be an object with properties: {v}")
115
+ except jsonschema.exceptions.SchemaError as e:
116
+ raise ValueError(f"Invalid JSON schema: {v} \n{e}")
117
+ except Exception as e:
118
+ raise ValueError(f"Unexpected error validating dict JSON schema: {v}\n {e}")
119
+
120
+
106
121
  def string_to_json_key(s: str) -> str:
107
122
  """Convert a string to a valid JSON key."""
108
123
  return re.sub(r"[^a-z0-9_]", "", s.strip().lower().replace(" ", "_"))
@@ -1,10 +1,14 @@
1
1
  from pydantic import Field
2
2
 
3
- from kiln_ai.datamodel.basemodel import NAME_FIELD, KilnParentModel
3
+ from kiln_ai.datamodel.basemodel import FilenameString, KilnParentModel
4
+ from kiln_ai.datamodel.external_tool_server import ExternalToolServer
4
5
  from kiln_ai.datamodel.task import Task
5
6
 
6
7
 
7
- class Project(KilnParentModel, parent_of={"tasks": Task}):
8
+ class Project(
9
+ KilnParentModel,
10
+ parent_of={"tasks": Task, "external_tool_servers": ExternalToolServer},
11
+ ):
8
12
  """
9
13
  A collection of related tasks.
10
14
 
@@ -12,12 +16,15 @@ class Project(KilnParentModel, parent_of={"tasks": Task}):
12
16
  of the overall goals.
13
17
  """
14
18
 
15
- name: str = NAME_FIELD
19
+ name: FilenameString = Field(description="The name of the project.")
16
20
  description: str | None = Field(
17
21
  default=None,
18
22
  description="A description of the project for you and your team. Will not be used in prompts/training/validation.",
19
23
  )
20
24
 
21
- # Needed for typechecking. TODO P2: fix this in KilnParentModel
25
+ # Needed for typechecking. We should fix this in KilnParentModel
22
26
  def tasks(self) -> list[Task]:
23
27
  return super().tasks() # type: ignore
28
+
29
+ def external_tool_servers(self, readonly: bool = False) -> list[ExternalToolServer]:
30
+ return super().external_tool_servers(readonly=readonly) # type: ignore
@@ -1,6 +1,6 @@
1
1
  from pydantic import BaseModel, Field
2
2
 
3
- from kiln_ai.datamodel.basemodel import NAME_FIELD, KilnParentedModel
3
+ from kiln_ai.datamodel.basemodel import FilenameString, KilnParentedModel
4
4
 
5
5
 
6
6
  class BasePrompt(BaseModel):
@@ -10,7 +10,7 @@ class BasePrompt(BaseModel):
10
10
  The "Prompt" model name is reserved for the custom prompts parented by a task.
11
11
  """
12
12
 
13
- name: str = NAME_FIELD
13
+ name: FilenameString = Field(description="The name of the prompt.")
14
14
  description: str | None = Field(
15
15
  default=None,
16
16
  description="A more detailed description of the prompt.",
@@ -60,11 +60,11 @@ def _check_prompt_id(id: str) -> str:
60
60
  return id
61
61
 
62
62
  if id.startswith("fine_tune_prompt::"):
63
- # check it had a fine_tune_id after the :: -- 'fine_tune_prompt::fine_tune_id'
64
- fine_tune_id = id[18:]
65
- if len(fine_tune_id) == 0:
63
+ # check it had a fine_tune_id after the :: -- 'fine_tune_prompt::[project_id]::[task_id]::fine_tune_id'
64
+ parts = id.split("::")
65
+ if len(parts) != 4 or len(parts[3]) == 0:
66
66
  raise ValueError(
67
- f"Invalid fine-tune prompt ID: {id}. Expected format: 'fine_tune_prompt::[fine_tune_id]'."
67
+ f"Invalid fine-tune prompt ID: {id}. Expected format: 'fine_tune_prompt::[project_id]::[task_id]::[fine_tune_id]'."
68
68
  )
69
69
  return id
70
70
 
@@ -14,18 +14,3 @@ def all_projects() -> list[Project]:
14
14
  # deleted files are possible continue with the rest
15
15
  continue
16
16
  return projects
17
-
18
-
19
- def project_from_id(project_id: str) -> Project | None:
20
- project_paths = Config.shared().projects
21
- if project_paths is not None:
22
- for project_path in project_paths:
23
- try:
24
- project = Project.load_from_file(project_path)
25
- if project.id == project_id:
26
- return project
27
- except Exception:
28
- # deleted files are possible continue with the rest
29
- continue
30
-
31
- return None