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

kiln_ai/datamodel/external_tool_server.py

@@ -0,0 +1,298 @@
+from enum import Enum
+from typing import Any, Dict
+
+from pydantic import Field, PrivateAttr, model_validator
+
+from kiln_ai.datamodel.basemodel import (
+    FilenameString,
+    KilnParentedModel,
+)
+from kiln_ai.utils.config import MCP_SECRETS_KEY, Config
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+
+class ToolServerType(str, Enum):
+    """
+    Enumeration of supported external tool server types.
+    """
+
+    remote_mcp = "remote_mcp"
+    local_mcp = "local_mcp"
+
+
+class ExternalToolServer(KilnParentedModel):
+    """
+    Configuration for communicating with a external MCP (Model Context Protocol) Server for LLM tool calls. External tool servers can be remote or local.
+
+    This model stores the necessary configuration to connect to and authenticate with
+    external MCP servers that provide tools for LLM interactions.
+    """
+
+    name: FilenameString = Field(description="The name of the external tool.")
+    type: ToolServerType = Field(
+        description="The type of external tool server. Remote tools are hosted on a remote server",
+    )
+    description: str | None = Field(
+        default=None,
+        description="A description of the external tool for you and your team. Will not be used in prompts/training/validation.",
+    )
+    properties: Dict[str, Any] = Field(
+        default={},
+        description="Configuration properties specific to the tool type.",
+    )
+
+    # Private variable to store unsaved secrets
+    _unsaved_secrets: dict[str, str] = PrivateAttr(default_factory=dict)
+
+    def model_post_init(self, __context: Any) -> None:
+        # Process secrets after initialization (pydantic v2 hook)
+        self._process_secrets_from_properties()
+
+    def _process_secrets_from_properties(self) -> None:
+        """
+        Extract secrets from properties and move them to _unsaved_secrets.
+        This removes secrets from the properties dict so they aren't saved to file.
+        Clears existing _unsaved_secrets first to handle property updates correctly.
+        """
+        # Clear existing unsaved secrets since we're reprocessing
+        self._unsaved_secrets.clear()
+
+        secret_keys = self.get_secret_keys()
+
+        if not secret_keys:
+            return
+
+        # Extract secret values from properties based on server type
+        match self.type:
+            case ToolServerType.remote_mcp:
+                headers = self.properties.get("headers", {})
+                for key_name in secret_keys:
+                    if key_name in headers:
+                        self._unsaved_secrets[key_name] = headers[key_name]
+                        # Remove from headers immediately so they are not saved to file
+                        del headers[key_name]
+
+            case ToolServerType.local_mcp:
+                env_vars = self.properties.get("env_vars", {})
+                for key_name in secret_keys:
+                    if key_name in env_vars:
+                        self._unsaved_secrets[key_name] = env_vars[key_name]
+                        # Remove from env_vars immediately so they are not saved to file
+                        del env_vars[key_name]
+
+            case _:
+                raise_exhaustive_enum_error(self.type)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """
+        Override __setattr__ to process secrets whenever properties are updated.
+        """
+        super().__setattr__(name, value)
+
+        # Process secrets whenever properties are updated
+        if name == "properties":
+            self._process_secrets_from_properties()
+
+    @model_validator(mode="after")
+    def validate_required_fields(self) -> "ExternalToolServer":
+        """Validate that each tool type has the required configuration."""
+        match self.type:
+            case ToolServerType.remote_mcp:
+                server_url = self.properties.get("server_url", None)
+                if not isinstance(server_url, str):
+                    raise ValueError(
+                        "server_url must be a string for external tools of type 'remote_mcp'"
+                    )
+                if not server_url:
+                    raise ValueError(
+                        "server_url is required to connect to a remote MCP server"
+                    )
+
+                headers = self.properties.get("headers", None)
+                if headers is None:
+                    raise ValueError("headers must be set when type is 'remote_mcp'")
+                if not isinstance(headers, dict):
+                    raise ValueError(
+                        "headers must be a dictionary for external tools of type 'remote_mcp'"
+                    )
+
+                secret_header_keys = self.properties.get("secret_header_keys", None)
+                # Secret header keys are optional, but if they are set, they must be a list of strings
+                if secret_header_keys is not None:
+                    if not isinstance(secret_header_keys, list):
+                        raise ValueError(
+                            "secret_header_keys must be a list for external tools of type 'remote_mcp'"
+                        )
+                    if not all(isinstance(k, str) for k in secret_header_keys):
+                        raise ValueError("secret_header_keys must contain only strings")
+
+            case ToolServerType.local_mcp:
+                command = self.properties.get("command", None)
+                if not isinstance(command, str):
+                    raise ValueError(
+                        "command must be a string to start a local MCP server"
+                    )
+                if not command.strip():
+                    raise ValueError("command is required to start a local MCP server")
+
+                args = self.properties.get("args", None)
+                if not isinstance(args, list):
+                    raise ValueError(
+                        "arguments must be a list to start a local MCP server"
+                    )
+
+                env_vars = self.properties.get("env_vars", {})
+                if not isinstance(env_vars, dict):
+                    raise ValueError(
+                        "environment variables must be a dictionary for external tools of type 'local_mcp'"
+                    )
+
+                secret_env_var_keys = self.properties.get("secret_env_var_keys", None)
+                # Secret env var keys are optional, but if they are set, they must be a list of strings
+                if secret_env_var_keys is not None:
+                    if not isinstance(secret_env_var_keys, list):
+                        raise ValueError(
+                            "secret_env_var_keys must be a list for external tools of type 'local_mcp'"
+                        )
+                    if not all(isinstance(k, str) for k in secret_env_var_keys):
+                        raise ValueError(
+                            "secret_env_var_keys must contain only strings"
+                        )
+
+            case _:
+                # Type checking will catch missing cases
+                raise_exhaustive_enum_error(self.type)
+        return self
+
+    def get_secret_keys(self) -> list[str]:
+        """
+        Get the list of secret key names based on server type.
+
+        Returns:
+            List of secret key names (header names for remote, env var names for local)
+        """
+        match self.type:
+            case ToolServerType.remote_mcp:
+                return self.properties.get("secret_header_keys", [])
+            case ToolServerType.local_mcp:
+                return self.properties.get("secret_env_var_keys", [])
+            case _:
+                raise_exhaustive_enum_error(self.type)
+
+    def retrieve_secrets(self) -> tuple[dict[str, str], list[str]]:
+        """
+        Retrieve secrets from configuration system or in-memory storage.
+        Automatically determines which secret keys to retrieve based on the server type.
+        Config secrets take precedence over unsaved secrets.
+
+        Returns:
+            Tuple of (secrets_dict, missing_secrets_list) where:
+            - secrets_dict: Dictionary mapping key names to their secret values
+            - missing_secrets_list: List of secret key names that are missing values
+        """
+        secrets = {}
+        missing_secrets = []
+        secret_keys = self.get_secret_keys()
+
+        if secret_keys and len(secret_keys) > 0:
+            config = Config.shared()
+            mcp_secrets = config.get_value(MCP_SECRETS_KEY)
+
+            for key_name in secret_keys:
+                secret_value = None
+
+                # First check config secrets (persistent storage), key is mcp_server_id::key_name
+                secret_key = self._config_secret_key(key_name)
+                secret_value = mcp_secrets.get(secret_key) if mcp_secrets else None
+
+                # Fall back to unsaved secrets (in-memory storage)
+                if (
+                    not secret_value
+                    and hasattr(self, "_unsaved_secrets")
+                    and key_name in self._unsaved_secrets
+                ):
+                    secret_value = self._unsaved_secrets[key_name]
+
+                if secret_value:
+                    secrets[key_name] = secret_value
+                else:
+                    missing_secrets.append(key_name)
+
+        return secrets, missing_secrets
+
+    def _save_secrets(self) -> None:
+        """
+        Save unsaved secrets to the configuration system.
+        """
+        secret_keys = self.get_secret_keys()
+
+        # No secrets to save
+        if not secret_keys:
+            return
+
+        if self.id is None:
+            raise ValueError("Server ID cannot be None when saving secrets")
+
+        # Check if secrets are already saved
+        if not hasattr(self, "_unsaved_secrets") or not self._unsaved_secrets:
+            return
+
+        config = Config.shared()
+        mcp_secrets: dict[str, str] = config.get_value(MCP_SECRETS_KEY) or {}
+
+        # Store secrets with the pattern: mcp_server_id::key_name
+        for key_name, secret_value in self._unsaved_secrets.items():
+            secret_key = self._config_secret_key(key_name)
+            mcp_secrets[secret_key] = secret_value
+
+        config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
+
+        # Clear unsaved secrets after saving
+        self._unsaved_secrets.clear()
+
+    def delete_secrets(self) -> None:
+        """
+        Delete all secrets for this tool server from the configuration system.
+        """
+        secret_keys = self.get_secret_keys()
+
+        config = Config.shared()
+        mcp_secrets = config.get_value(MCP_SECRETS_KEY) or dict[str, str]()
+
+        # Remove secrets with the pattern: mcp_server_id::key_name
+        for key_name in secret_keys:
+            secret_key = self._config_secret_key(key_name)
+            if secret_key in mcp_secrets:
+                del mcp_secrets[secret_key]
+
+        # Always call update_settings to maintain consistency with the old behavior
+        config.update_settings({MCP_SECRETS_KEY: mcp_secrets})
+
+    def save_to_file(self) -> None:
+        """
+        Override save_to_file to automatically save any unsaved secrets before saving to file.
+
+        This ensures that secrets are always saved when the object is saved,
+        preventing the issue where secrets could be lost if save_to_file is called
+        without explicitly saving secrets first.
+        """
+        # Save any unsaved secrets first
+        if hasattr(self, "_unsaved_secrets") and self._unsaved_secrets:
+            self._save_secrets()
+
+        # Call the parent save_to_file method
+        super().save_to_file()
+
+    #  Internal helpers
+
+    def _config_secret_key(self, key_name: str) -> str:
+        """
+        Generate the secret key pattern for storing/retrieving secrets.
+
+        Args:
+            key_name: The name of the secret key
+
+        Returns:
+            The formatted secret key: "{server_id}::{key_name}"
+        """
+        return f"{self.id}::{key_name}"

kiln_ai/datamodel/json_schema.py

@@ -84,25 +84,40 @@ def schema_from_json_str(v: str) -> Dict:
     """
     try:
         parsed = json.loads(v)
-        jsonschema.Draft202012Validator.check_schema(parsed)
         if not isinstance(parsed, dict):
             raise ValueError(f"JSON schema must be a dict, not {type(parsed)}")
-        # Top level arrays are valid JSON schemas, but we don't want to allow them here as they often cause issues
-        if (
-            "type" not in parsed
-            or parsed["type"] != "object"
-            or "properties" not in parsed
-        ):
-            raise ValueError(f"JSON schema must be an object with properties: {v}")
+
+        validate_schema_dict(parsed)
         return parsed
-    except jsonschema.exceptions.SchemaError as e:
-        raise ValueError(f"Invalid JSON schema: {v} \n{e}")
     except json.JSONDecodeError as e:
         raise ValueError(f"Invalid JSON: {v}\n {e}")
     except Exception as e:
         raise ValueError(f"Unexpected error parsing JSON schema: {v}\n {e}")
 
 
+def validate_schema_dict(v: Dict):
+    """Parse and validate a JSON schema dictionary.
+
+    Args:
+        v: Dictionary containing a JSON schema definition
+
+    Returns:
+        Dict containing the parsed JSON schema
+
+    Raises:
+        ValueError: If the input is not a valid JSON schema object with required properties
+    """
+    try:
+        jsonschema.Draft202012Validator.check_schema(v)
+        # Top level arrays are valid JSON schemas, but we don't want to allow them here as they often cause issues
+        if "type" not in v or v["type"] != "object" or "properties" not in v:
+            raise ValueError(f"JSON schema must be an object with properties: {v}")
+    except jsonschema.exceptions.SchemaError as e:
+        raise ValueError(f"Invalid JSON schema: {v} \n{e}")
+    except Exception as e:
+        raise ValueError(f"Unexpected error validating dict JSON schema: {v}\n {e}")
+
+
 def string_to_json_key(s: str) -> str:
     """Convert a string to a valid JSON key."""
     return re.sub(r"[^a-z0-9_]", "", s.strip().lower().replace(" ", "_"))

kiln_ai/datamodel/project.py

@@ -1,10 +1,14 @@
 from pydantic import Field
 
 from kiln_ai.datamodel.basemodel import FilenameString, KilnParentModel
+from kiln_ai.datamodel.external_tool_server import ExternalToolServer
 from kiln_ai.datamodel.task import Task
 
 
-class Project(KilnParentModel, parent_of={"tasks": Task}):
+class Project(
+    KilnParentModel,
+    parent_of={"tasks": Task, "external_tool_servers": ExternalToolServer},
+):
     """
     A collection of related tasks.
 
@@ -21,3 +25,6 @@ class Project(KilnParentModel, parent_of={"tasks": Task}):
     # Needed for typechecking. We should fix this in KilnParentModel
     def tasks(self) -> list[Task]:
         return super().tasks()  # type: ignore
+
+    def external_tool_servers(self, readonly: bool = False) -> list[ExternalToolServer]:
+        return super().external_tool_servers(readonly=readonly)  # type: ignore

kiln_ai/datamodel/registry.py

@@ -14,18 +14,3 @@ def all_projects() -> list[Project]:
             # deleted files are possible continue with the rest
             continue
     return projects
-
-
-def project_from_id(project_id: str) -> Project | None:
-    project_paths = Config.shared().projects
-    if project_paths is not None:
-        for project_path in project_paths:
-            try:
-                project = Project.load_from_file(project_path)
-                if project.id == project_id:
-                    return project
-            except Exception:
-                # deleted files are possible continue with the rest
-                continue
-
-    return None

kiln_ai/datamodel/run_config.py

@@ -0,0 +1,62 @@
+from typing import List
+
+from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
+
+from kiln_ai.datamodel.datamodel_enums import (
+    ModelProviderName,
+    StructuredOutputMode,
+)
+from kiln_ai.datamodel.prompt_id import PromptId
+from kiln_ai.datamodel.tool_id import ToolId
+
+
+class ToolsRunConfig(BaseModel):
+    """
+    A config describing which tools are available to a task.
+    """
+
+    tools: List[ToolId] = Field(
+        description="The IDs of the tools available to the task."
+    )
+
+
+class RunConfigProperties(BaseModel):
+    """
+    A configuration for running a task.
+
+    This includes everything needed to run a task, except the input and task ID. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
+    """
+
+    model_name: str = Field(description="The model to use for this run config.")
+    model_provider_name: ModelProviderName = Field(
+        description="The provider to use for this run config."
+    )
+    prompt_id: PromptId = Field(
+        description="The prompt to use for this run config. Defaults to building a simple prompt from the task if not provided.",
+    )
+    top_p: float = Field(
+        default=1.0,
+        description="The top-p value to use for this run config. Defaults to 1.0.",
+    )
+    temperature: float = Field(
+        default=1.0,
+        description="The temperature to use for this run config. Defaults to 1.0.",
+    )
+    structured_output_mode: StructuredOutputMode = Field(
+        description="The structured output mode to use for this run config.",
+    )
+    tools_config: ToolsRunConfig | None = Field(
+        default=None,
+        description="The tools config to use for this run config, defining which tools are available to the model.",
+    )
+
+    @model_validator(mode="after")
+    def validate_required_fields(self) -> Self:
+        if not (0 <= self.top_p <= 1):
+            raise ValueError("top_p must be between 0 and 1")
+
+        elif self.temperature < 0 or self.temperature > 2:
+            raise ValueError("temperature must be between 0 and 2")
+
+        return self

kiln_ai/datamodel/task.py

@@ -1,9 +1,7 @@
 from typing import TYPE_CHECKING, Dict, List, Union
 
 from pydantic import BaseModel, Field, ValidationInfo, model_validator
-from typing_extensions import Self
 
-from kiln_ai.datamodel import Finetune
 from kiln_ai.datamodel.basemodel import (
     ID_FIELD,
     ID_TYPE,
@@ -13,16 +11,16 @@ from kiln_ai.datamodel.basemodel import (
     KilnParentModel,
 )
 from kiln_ai.datamodel.datamodel_enums import (
-    ModelProviderName,
     Priority,
     StructuredOutputMode,
     TaskOutputRatingType,
 )
 from kiln_ai.datamodel.dataset_split import DatasetSplit
 from kiln_ai.datamodel.eval import Eval
+from kiln_ai.datamodel.finetune import Finetune
 from kiln_ai.datamodel.json_schema import JsonObjectSchema, schema_from_json_str
 from kiln_ai.datamodel.prompt import BasePrompt, Prompt
-from kiln_ai.datamodel.prompt_id import PromptId
+from kiln_ai.datamodel.run_config import RunConfigProperties
 from kiln_ai.datamodel.task_run import TaskRun
 
 if TYPE_CHECKING:
@@ -45,55 +43,6 @@ class TaskRequirement(BaseModel):
     type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
 
 
-class RunConfigProperties(BaseModel):
-    """
-    A configuration for running a task.
-
-    This includes everything needed to run a task, except the input and task ID. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
-    """
-
-    model_name: str = Field(description="The model to use for this run config.")
-    model_provider_name: ModelProviderName = Field(
-        description="The provider to use for this run config."
-    )
-    prompt_id: PromptId = Field(
-        description="The prompt to use for this run config. Defaults to building a simple prompt from the task if not provided.",
-    )
-    top_p: float = Field(
-        default=1.0,
-        description="The top-p value to use for this run config. Defaults to 1.0.",
-    )
-    temperature: float = Field(
-        default=1.0,
-        description="The temperature to use for this run config. Defaults to 1.0.",
-    )
-    structured_output_mode: StructuredOutputMode = Field(
-        description="The structured output mode to use for this run config.",
-    )
-
-    @model_validator(mode="after")
-    def validate_required_fields(self) -> Self:
-        if not (0 <= self.top_p <= 1):
-            raise ValueError("top_p must be between 0 and 1")
-
-        elif self.temperature < 0 or self.temperature > 2:
-            raise ValueError("temperature must be between 0 and 2")
-
-        return self
-
-
-class RunConfig(RunConfigProperties):
-    """
-    A configuration for running a task.
-
-    This includes everything needed to run a task, except the input. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
-
-    For example: task, model, provider, prompt, etc.
-    """
-
-    task: "Task" = Field(description="The task to run.")
-
-
 class TaskRunConfig(KilnParentedModel):
     """
     A Kiln model for persisting a run config in a Kiln Project, nested under a task.
@@ -124,15 +73,6 @@ class TaskRunConfig(KilnParentedModel):
             return None
         return self.parent  # type: ignore
 
-    def run_config(self) -> RunConfig:
-        parent_task = self.parent_task()
-        if parent_task is None:
-            raise ValueError("Run config must be parented to a task")
-        return run_config_from_run_config_properties(
-            task=parent_task,
-            run_config_properties=self.run_config_properties,
-        )
-
     # Previously we didn't store structured_output_mode in the run_config_properties. Updgrade old models when loading from file.
     @model_validator(mode="before")
     def upgrade_old_entries(cls, data: dict, info: ValidationInfo) -> dict:
@@ -155,21 +95,6 @@ class TaskRunConfig(KilnParentedModel):
         return data
 
 
-def run_config_from_run_config_properties(
-    task: "Task",
-    run_config_properties: RunConfigProperties,
-) -> RunConfig:
-    return RunConfig(
-        task=task,
-        model_name=run_config_properties.model_name,
-        model_provider_name=run_config_properties.model_provider_name,
-        prompt_id=run_config_properties.prompt_id,
-        top_p=run_config_properties.top_p,
-        temperature=run_config_properties.temperature,
-        structured_output_mode=run_config_properties.structured_output_mode,
-    )
-
-
 class Task(
     KilnParentedModel,
     KilnParentModel,

kiln_ai/datamodel/task_output.py

@@ -1,6 +1,6 @@
 import json
 from enum import Enum
-from typing import TYPE_CHECKING, Dict, List, Type, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel, Field, ValidationInfo, model_validator
 from typing_extensions import Self
@@ -8,6 +8,7 @@ from typing_extensions import Self
 from kiln_ai.datamodel.basemodel import ID_TYPE, KilnBaseModel
 from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
 from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
+from kiln_ai.datamodel.run_config import RunConfigProperties
 from kiln_ai.datamodel.strict_mode import strict_mode
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
@@ -199,6 +200,10 @@ class DataSource(BaseModel):
         default={},
         description="Properties describing the data source. For synthetic things like model. For human, the human's name.",
     )
+    run_config: Optional[RunConfigProperties] = Field(
+        default=None,
+        description="The run config used to generate the data, if generated by a running a model in Kiln (only true for type=synthetic).",
+    )
 
     _data_source_properties = [
         DataSourceProperty(

kiln_ai/datamodel/task_run.py

@@ -8,6 +8,7 @@ from kiln_ai.datamodel.basemodel import KilnParentedModel
 from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 from kiln_ai.datamodel.strict_mode import strict_mode
 from kiln_ai.datamodel.task_output import DataSource, TaskOutput
+from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam
 
 if TYPE_CHECKING:
     from kiln_ai.datamodel.task import Task
@@ -35,6 +36,42 @@ class Usage(BaseModel):
         ge=0,
     )
 
+    def __add__(self, other: "Usage") -> "Usage":
+        """Add two Usage objects together, handling None values gracefully.
+
+        None + None = None
+        None + value = value
+        value + None = value
+        value1 + value2 = value1 + value2
+        """
+        if not isinstance(other, Usage):
+            raise TypeError(f"Cannot add Usage with {type(other).__name__}")
+
+        def _add_optional_int(a: int | None, b: int | None) -> int | None:
+            if a is None and b is None:
+                return None
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return a + b
+
+        def _add_optional_float(a: float | None, b: float | None) -> float | None:
+            if a is None and b is None:
+                return None
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return a + b
+
+        return Usage(
+            input_tokens=_add_optional_int(self.input_tokens, other.input_tokens),
+            output_tokens=_add_optional_int(self.output_tokens, other.output_tokens),
+            total_tokens=_add_optional_int(self.total_tokens, other.total_tokens),
+            cost=_add_optional_float(self.cost, other.cost),
+        )
+
 
 class TaskRun(KilnParentedModel):
     """
@@ -72,6 +109,10 @@ class TaskRun(KilnParentedModel):
         default=None,
         description="Usage information for the task run. This includes the number of input tokens, output tokens, and total tokens used.",
     )
+    trace: list[ChatCompletionMessageParam] | None = Field(
+        default=None,
+        description="The trace of the task run in OpenAI format. This is the list of messages that were sent to/from the model.",
+    )
 
     def thinking_training_data(self) -> str | None:
         """

kiln_ai/datamodel/test_basemodel.py

@@ -17,7 +17,7 @@ from kiln_ai.datamodel.basemodel import (
     string_to_valid_name,
 )
 from kiln_ai.datamodel.model_cache import ModelCache
-from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 @pytest.fixture
@@ -552,8 +552,8 @@ def base_task():
 @pytest.fixture
 def adapter(base_task):
     return MockAdapter(
-        run_config=RunConfig(
-            task=base_task,
+        task=base_task,
+        run_config=RunConfigProperties(
             model_name="test_model",
             model_provider_name="openai",
             prompt_id="simple_prompt_builder",