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.
- kiln_ai/adapters/__init__.py +2 -2
- kiln_ai/adapters/adapter_registry.py +46 -0
- kiln_ai/adapters/chat/chat_formatter.py +8 -12
- kiln_ai/adapters/chat/test_chat_formatter.py +6 -2
- kiln_ai/adapters/data_gen/data_gen_task.py +2 -2
- kiln_ai/adapters/data_gen/test_data_gen_task.py +7 -3
- kiln_ai/adapters/docker_model_runner_tools.py +119 -0
- kiln_ai/adapters/eval/base_eval.py +2 -2
- kiln_ai/adapters/eval/eval_runner.py +3 -1
- kiln_ai/adapters/eval/g_eval.py +2 -2
- kiln_ai/adapters/eval/test_base_eval.py +1 -1
- kiln_ai/adapters/eval/test_eval_runner.py +6 -12
- kiln_ai/adapters/eval/test_g_eval.py +3 -4
- kiln_ai/adapters/eval/test_g_eval_data.py +1 -1
- kiln_ai/adapters/fine_tune/__init__.py +1 -1
- kiln_ai/adapters/fine_tune/base_finetune.py +1 -0
- kiln_ai/adapters/fine_tune/fireworks_finetune.py +32 -20
- kiln_ai/adapters/fine_tune/openai_finetune.py +14 -4
- kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py +30 -21
- kiln_ai/adapters/fine_tune/test_openai_finetune.py +108 -111
- kiln_ai/adapters/ml_model_list.py +1009 -111
- kiln_ai/adapters/model_adapters/base_adapter.py +62 -28
- kiln_ai/adapters/model_adapters/litellm_adapter.py +397 -80
- kiln_ai/adapters/model_adapters/test_base_adapter.py +194 -18
- kiln_ai/adapters/model_adapters/test_litellm_adapter.py +428 -4
- kiln_ai/adapters/model_adapters/test_litellm_adapter_tools.py +1103 -0
- kiln_ai/adapters/model_adapters/test_saving_adapter_results.py +5 -5
- kiln_ai/adapters/model_adapters/test_structured_output.py +120 -14
- kiln_ai/adapters/parsers/__init__.py +1 -1
- kiln_ai/adapters/parsers/test_r1_parser.py +1 -1
- kiln_ai/adapters/provider_tools.py +35 -20
- kiln_ai/adapters/remote_config.py +57 -10
- kiln_ai/adapters/repair/repair_task.py +1 -1
- kiln_ai/adapters/repair/test_repair_task.py +12 -9
- kiln_ai/adapters/run_output.py +3 -0
- kiln_ai/adapters/test_adapter_registry.py +109 -2
- kiln_ai/adapters/test_docker_model_runner_tools.py +305 -0
- kiln_ai/adapters/test_ml_model_list.py +51 -1
- kiln_ai/adapters/test_prompt_adaptors.py +13 -6
- kiln_ai/adapters/test_provider_tools.py +73 -12
- kiln_ai/adapters/test_remote_config.py +470 -16
- kiln_ai/datamodel/__init__.py +23 -21
- kiln_ai/datamodel/basemodel.py +54 -28
- kiln_ai/datamodel/datamodel_enums.py +3 -0
- kiln_ai/datamodel/dataset_split.py +5 -3
- kiln_ai/datamodel/eval.py +4 -4
- kiln_ai/datamodel/external_tool_server.py +298 -0
- kiln_ai/datamodel/finetune.py +2 -2
- kiln_ai/datamodel/json_schema.py +25 -10
- kiln_ai/datamodel/project.py +11 -4
- kiln_ai/datamodel/prompt.py +2 -2
- kiln_ai/datamodel/prompt_id.py +4 -4
- kiln_ai/datamodel/registry.py +0 -15
- kiln_ai/datamodel/run_config.py +62 -0
- kiln_ai/datamodel/task.py +8 -83
- kiln_ai/datamodel/task_output.py +7 -2
- kiln_ai/datamodel/task_run.py +41 -0
- kiln_ai/datamodel/test_basemodel.py +213 -21
- kiln_ai/datamodel/test_eval_model.py +6 -6
- kiln_ai/datamodel/test_example_models.py +175 -0
- kiln_ai/datamodel/test_external_tool_server.py +691 -0
- kiln_ai/datamodel/test_model_perf.py +1 -1
- kiln_ai/datamodel/test_prompt_id.py +5 -1
- kiln_ai/datamodel/test_registry.py +8 -3
- kiln_ai/datamodel/test_task.py +20 -47
- kiln_ai/datamodel/test_tool_id.py +239 -0
- kiln_ai/datamodel/tool_id.py +83 -0
- kiln_ai/tools/__init__.py +8 -0
- kiln_ai/tools/base_tool.py +82 -0
- kiln_ai/tools/built_in_tools/__init__.py +13 -0
- kiln_ai/tools/built_in_tools/math_tools.py +124 -0
- kiln_ai/tools/built_in_tools/test_math_tools.py +204 -0
- kiln_ai/tools/mcp_server_tool.py +95 -0
- kiln_ai/tools/mcp_session_manager.py +243 -0
- kiln_ai/tools/test_base_tools.py +199 -0
- kiln_ai/tools/test_mcp_server_tool.py +457 -0
- kiln_ai/tools/test_mcp_session_manager.py +1585 -0
- kiln_ai/tools/test_tool_registry.py +473 -0
- kiln_ai/tools/tool_registry.py +64 -0
- kiln_ai/utils/config.py +32 -0
- kiln_ai/utils/open_ai_types.py +94 -0
- kiln_ai/utils/project_utils.py +17 -0
- kiln_ai/utils/test_config.py +138 -1
- kiln_ai/utils/test_open_ai_types.py +131 -0
- {kiln_ai-0.18.0.dist-info → kiln_ai-0.20.1.dist-info}/METADATA +37 -6
- kiln_ai-0.20.1.dist-info/RECORD +138 -0
- kiln_ai-0.18.0.dist-info/RECORD +0 -115
- {kiln_ai-0.18.0.dist-info → kiln_ai-0.20.1.dist-info}/WHEEL +0 -0
- {kiln_ai-0.18.0.dist-info → kiln_ai-0.20.1.dist-info}/licenses/LICENSE.txt +0 -0
kiln_ai/datamodel/basemodel.py
CHANGED
|
@@ -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)
|
|
48
|
-
|
|
49
|
-
#
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
)
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
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
|
-
#
|
|
67
|
-
valid_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
|
-
#
|
|
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
|
|
@@ -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
|
|
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:
|
|
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:
|
|
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
|
-
|
|
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:
|
|
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: {
|
|
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:
|
|
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}"
|
kiln_ai/datamodel/finetune.py
CHANGED
|
@@ -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
|
|
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:
|
|
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.",
|
kiln_ai/datamodel/json_schema.py
CHANGED
|
@@ -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
|
-
|
|
91
|
-
|
|
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(" ", "_"))
|
kiln_ai/datamodel/project.py
CHANGED
|
@@ -1,10 +1,14 @@
|
|
|
1
1
|
from pydantic import Field
|
|
2
2
|
|
|
3
|
-
from kiln_ai.datamodel.basemodel import
|
|
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(
|
|
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:
|
|
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.
|
|
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
|
kiln_ai/datamodel/prompt.py
CHANGED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
from pydantic import BaseModel, Field
|
|
2
2
|
|
|
3
|
-
from kiln_ai.datamodel.basemodel import
|
|
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:
|
|
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.",
|
kiln_ai/datamodel/prompt_id.py
CHANGED
|
@@ -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
|
-
|
|
65
|
-
if len(
|
|
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
|
|
kiln_ai/datamodel/registry.py
CHANGED
|
@@ -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
|