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

kiln_ai/datamodel/chunk.py

@@ -0,0 +1,158 @@
+import logging
+from enum import Enum
+from typing import TYPE_CHECKING, List, Union
+
+import anyio
+from pydantic import (
+    BaseModel,
+    Field,
+    SerializationInfo,
+    ValidationInfo,
+    field_serializer,
+    field_validator,
+)
+
+from kiln_ai.datamodel.basemodel import (
+    ID_TYPE,
+    FilenameString,
+    KilnAttachmentModel,
+    KilnParentedModel,
+    KilnParentModel,
+)
+from kiln_ai.datamodel.embedding import ChunkEmbeddings
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.extraction import Extraction
+    from kiln_ai.datamodel.project import Project
+
+
+def validate_fixed_window_chunker_properties(
+    properties: dict[str, str | int | float | bool],
+) -> dict[str, str | int | float | bool]:
+    """Validate the properties for the fixed window chunker and set defaults if needed."""
+    chunk_overlap = properties.get("chunk_overlap")
+    if chunk_overlap is None:
+        raise ValueError("Chunk overlap is required.")
+
+    chunk_size = properties.get("chunk_size")
+    if chunk_size is None:
+        raise ValueError("Chunk size is required.")
+
+    if not isinstance(chunk_overlap, int):
+        raise ValueError("Chunk overlap must be an integer.")
+    if chunk_overlap < 0:
+        raise ValueError("Chunk overlap must be greater than or equal to 0.")
+
+    if not isinstance(chunk_size, int):
+        raise ValueError("Chunk size must be an integer.")
+    if chunk_size <= 0:
+        raise ValueError("Chunk size must be greater than 0.")
+
+    if chunk_overlap >= chunk_size:
+        raise ValueError("Chunk overlap must be less than chunk size.")
+
+    return properties
+
+
+class ChunkerType(str, Enum):
+    FIXED_WINDOW = "fixed_window"
+
+
+class ChunkerConfig(KilnParentedModel):
+    name: FilenameString = Field(
+        description="A name to identify the chunker config.",
+    )
+    description: str | None = Field(
+        default=None, description="The description of the chunker config"
+    )
+    chunker_type: ChunkerType = Field(
+        description="This is used to determine the type of chunker to use.",
+    )
+    properties: dict[str, str | int | float | bool] = Field(
+        description="Properties to be used to execute the chunker config. This is chunker_type specific and should serialize to a json dict.",
+    )
+
+    # Workaround to return typed parent without importing Project
+    def parent_project(self) -> Union["Project", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Project":
+            return None
+        return self.parent  # type: ignore
+
+    @field_validator("properties")
+    @classmethod
+    def validate_properties(
+        cls, properties: dict[str, str | int | float | bool], info: ValidationInfo
+    ) -> dict[str, str | int | float | bool]:
+        if info.data.get("chunker_type") == ChunkerType.FIXED_WINDOW:
+            # do not trigger revalidation of properties
+            return validate_fixed_window_chunker_properties(properties)
+        return properties
+
+    def chunk_size(self) -> int | None:
+        if self.properties.get("chunk_size") is None:
+            return None
+        if not isinstance(self.properties["chunk_size"], int):
+            raise ValueError("Chunk size must be an integer.")
+        return self.properties["chunk_size"]
+
+    def chunk_overlap(self) -> int | None:
+        if self.properties.get("chunk_overlap") is None:
+            return None
+        if not isinstance(self.properties["chunk_overlap"], int):
+            raise ValueError("Chunk overlap must be an integer.")
+        return self.properties["chunk_overlap"]
+
+
+class Chunk(BaseModel):
+    content: KilnAttachmentModel = Field(
+        description="The content of the chunk, stored as an attachment."
+    )
+
+    @field_serializer("content")
+    def serialize_content(
+        self, content: KilnAttachmentModel, info: SerializationInfo
+    ) -> dict:
+        context = info.context or {}
+        context["filename_prefix"] = "content"
+        return content.model_dump(mode="json", context=context)
+
+
+class ChunkedDocument(
+    KilnParentedModel, KilnParentModel, parent_of={"chunk_embeddings": ChunkEmbeddings}
+):
+    chunker_config_id: ID_TYPE = Field(
+        description="The ID of the chunker config used to chunk the document.",
+    )
+    chunks: List[Chunk] = Field(description="The chunks of the document.")
+
+    def parent_extraction(self) -> Union["Extraction", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Extraction":
+            return None
+        return self.parent  # type: ignore
+
+    def chunk_embeddings(self, readonly: bool = False) -> list[ChunkEmbeddings]:
+        return super().chunk_embeddings(readonly=readonly)  # type: ignore
+
+    async def load_chunks_text(self) -> list[str]:
+        """Utility to return a list of text for each chunk, loaded from each chunk's content attachment."""
+        if not self.path:
+            raise ValueError(
+                "Failed to resolve the path of chunk content attachment because the chunk does not have a path."
+            )
+
+        chunks_text: list[str] = []
+        for chunk in self.chunks:
+            full_path = chunk.content.resolve_path(self.path.parent)
+
+            try:
+                chunks_text.append(
+                    await anyio.Path(full_path).read_text(encoding="utf-8")
+                )
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to read chunk content for {full_path}: {e}"
+                ) from e
+
+        return chunks_text

kiln_ai/datamodel/datamodel_enums.py

@@ -100,3 +100,30 @@ class ModelProviderName(str, Enum):
     siliconflow_cn = "siliconflow_cn"
     cerebras = "cerebras"
     docker_model_runner = "docker_model_runner"
+
+
+class KilnMimeType(str, Enum):
+    """
+    Enumeration of supported mime types.
+    """
+
+    # documents
+    PDF = "application/pdf"
+    CSV = "text/csv"
+    TXT = "text/plain"
+    HTML = "text/html"
+    MD = "text/markdown"
+
+    # images
+    PNG = "image/png"
+    JPG = "image/jpeg"
+    JPEG = "image/jpeg"
+
+    # audio
+    MP3 = "audio/mpeg"
+    WAV = "audio/wav"
+    OGG = "audio/ogg"
+
+    # video
+    MP4 = "video/mp4"
+    MOV = "video/quicktime"

kiln_ai/datamodel/embedding.py

@@ -0,0 +1,64 @@
+from typing import TYPE_CHECKING, List, Union
+
+from pydantic import BaseModel, Field, model_validator
+
+from kiln_ai.datamodel.basemodel import ID_TYPE, FilenameString, KilnParentedModel
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.chunk import ChunkedDocument
+    from kiln_ai.datamodel.project import Project
+
+
+class EmbeddingConfig(KilnParentedModel):
+    name: FilenameString = Field(
+        description="A name to identify the embedding config.",
+    )
+    description: str | None = Field(
+        default=None,
+        description="A description for your reference, not shared with embedding models.",
+    )
+    model_provider_name: ModelProviderName = Field(
+        description="The provider to use to generate embeddings.",
+    )
+    model_name: str = Field(
+        description="The model to use to generate embeddings.",
+    )
+    properties: dict[str, str | int | float | bool] = Field(
+        description="Properties to be used to execute the embedding config.",
+    )
+
+    # Workaround to return typed parent without importing Project
+    def parent_project(self) -> Union["Project", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Project":
+            return None
+        return self.parent  # type: ignore
+
+    @model_validator(mode="after")
+    def validate_properties(self):
+        if "dimensions" in self.properties:
+            if (
+                not isinstance(self.properties["dimensions"], int)
+                or self.properties["dimensions"] <= 0
+            ):
+                raise ValueError("Dimensions must be a positive integer")
+
+        return self
+
+
+class Embedding(BaseModel):
+    vector: List[float] = Field(description="The vector of the embedding.")
+
+
+class ChunkEmbeddings(KilnParentedModel):
+    embedding_config_id: ID_TYPE = Field(
+        description="The ID of the embedding config used to generate the embeddings.",
+    )
+    embeddings: List[Embedding] = Field(
+        description="The embeddings of the chunks. The embedding at index i corresponds to the chunk at index i in the parent chunked document."
+    )
+
+    def parent_chunked_document(self) -> Union["ChunkedDocument", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "ChunkedDocument":
+            return None
+        return self.parent  # type: ignore

kiln_ai/datamodel/external_tool_server.py

@@ -1,7 +1,10 @@
+import re
 from enum import Enum
-from typing import Any, Dict
+from typing import Any
+from urllib.parse import urlparse
 
 from pydantic import Field, PrivateAttr, model_validator
+from typing_extensions import NotRequired, TypedDict
 
 from kiln_ai.datamodel.basemodel import (
     FilenameString,
@@ -9,6 +12,7 @@ from kiln_ai.datamodel.basemodel import (
 )
 from kiln_ai.utils.config import MCP_SECRETS_KEY, Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+from kiln_ai.utils.validation import tool_name_validator, validate_return_dict_prop
 
 
 class ToolServerType(str, Enum):
@@ -18,6 +22,28 @@ class ToolServerType(str, Enum):
 
     remote_mcp = "remote_mcp"
     local_mcp = "local_mcp"
+    kiln_task = "kiln_task"
+
+
+class LocalServerProperties(TypedDict, total=True):
+    command: str
+    args: NotRequired[list[str]]
+    env_vars: NotRequired[dict[str, str]]
+    secret_env_var_keys: NotRequired[list[str]]
+
+
+class RemoteServerProperties(TypedDict, total=True):
+    server_url: str
+    headers: NotRequired[dict[str, str]]
+    secret_header_keys: NotRequired[list[str]]
+
+
+class KilnTaskServerProperties(TypedDict, total=True):
+    task_id: str
+    run_config_id: str
+    name: str
+    description: str
+    is_archived: bool
 
 
 class ExternalToolServer(KilnParentedModel):
@@ -36,8 +62,10 @@ class ExternalToolServer(KilnParentedModel):
         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={},
+
+    properties: (
+        LocalServerProperties | RemoteServerProperties | KilnTaskServerProperties
+    ) = Field(
         description="Configuration properties specific to the tool type.",
     )
 
@@ -80,6 +108,9 @@ class ExternalToolServer(KilnParentedModel):
                         # Remove from env_vars immediately so they are not saved to file
                         del env_vars[key_name]
 
+            case ToolServerType.kiln_task:
+                pass
+
             case _:
                 raise_exhaustive_enum_error(self.type)
 
@@ -93,76 +124,195 @@ class ExternalToolServer(KilnParentedModel):
         if name == "properties":
             self._process_secrets_from_properties()
 
-    @model_validator(mode="after")
-    def validate_required_fields(self) -> "ExternalToolServer":
+    # Validation Helpers
+
+    @classmethod
+    def check_server_url(cls, server_url: str) -> None:
+        """Validate Server URL"""
+        if not isinstance(server_url, str):
+            raise ValueError("Server URL must be a string")
+
+        # Check for leading whitespace in URL
+        if server_url != server_url.lstrip():
+            raise ValueError("Server URL must not have leading whitespace")
+
+        parsed_url = urlparse(server_url)
+        if not parsed_url.netloc:
+            raise ValueError("Server URL is not a valid URL")
+        if parsed_url.scheme not in ["http", "https"]:
+            raise ValueError("Server URL must start with http:// or https://")
+
+    @classmethod
+    def check_headers(cls, headers: dict) -> None:
+        """Validate Headers"""
+        if not isinstance(headers, dict):
+            raise ValueError("headers must be a dictionary")
+
+        for key, value in headers.items():
+            if not key:
+                raise ValueError("Header name is required")
+            if not value:
+                raise ValueError("Header value is required")
+
+            # Reject invalid header names and CR/LF in names/values
+            token_re = re.compile(r"^[!#$%&'*+.^_`|~0-9A-Za-z-]+$")
+            if not token_re.match(key):
+                raise ValueError(f'Invalid header name: "{key}"')
+            if re.search(r"\r|\n", key) or re.search(r"\r|\n", value):
+                raise ValueError(
+                    "Header names/values must not contain invalid characters"
+                )
+
+    @classmethod
+    def check_secret_keys(
+        cls, secret_keys: list, key_type: str, tool_type: str
+    ) -> None:
+        """Validate Secret Keys (generic method for both header and env var keys)"""
+        if not isinstance(secret_keys, list):
+            raise ValueError(
+                f"{key_type} must be a list for external tools of type '{tool_type}'"
+            )
+        if not all(isinstance(k, str) for k in secret_keys):
+            raise ValueError(f"{key_type} must contain only strings")
+        if not all(key for key in secret_keys):
+            raise ValueError("Secret key is required")
+
+    @classmethod
+    def check_env_vars(cls, env_vars: dict) -> None:
+        """Validate Environment Variables"""
+        if not isinstance(env_vars, dict):
+            raise ValueError("environment variables must be a dictionary")
+
+        # Validate env_vars keys are in the correct format for Environment Variables
+        # According to POSIX specification, environment variable names must:
+        # - Start with a letter (a-z, A-Z) or underscore (_)
+        # - Contain only ASCII letters, digits, and underscores
+        for key, _ in env_vars.items():
+            if not key or not (
+                key[0].isascii() and (key[0].isalpha() or key[0] == "_")
+            ):
+                raise ValueError(
+                    f"Invalid environment variable key: {key}. Must start with a letter or underscore."
+                )
+
+            if not all(c.isascii() and (c.isalnum() or c == "_") for c in key):
+                raise ValueError(
+                    f"Invalid environment variable key: {key}. Can only contain letters, digits, and underscores."
+                )
+
+    @classmethod
+    def type_from_data(cls, data: dict) -> ToolServerType:
+        """Get the tool server type from the data for the the validators"""
+        raw_type = data.get("type")
+        if raw_type is None:
+            raise ValueError("type is required")
+        try:
+            return ToolServerType(raw_type)
+        except ValueError:
+            valid_types = ", ".join(type.value for type in ToolServerType)
+            raise ValueError(f"type must be one of: {valid_types}")
+
+    @model_validator(mode="before")
+    def validate_required_fields(cls, data: dict) -> dict:
         """Validate that each tool type has the required configuration."""
-        match self.type:
+        server_type = ExternalToolServer.type_from_data(data)
+        properties = data.get("properties", {})
+
+        match server_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:
+                server_url = properties.get("server_url", None)
+                if server_url is None:
                     raise ValueError(
-                        "server_url is required to connect to a remote MCP server"
+                        "Server URL is required to connect to a remote MCP server"
                     )
+                ExternalToolServer.check_server_url(server_url)
 
-                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):
+            case ToolServerType.local_mcp:
+                command = properties.get("command", None)
+                if command is None:
+                    raise ValueError("command is required to start a local MCP server")
+                if not isinstance(command, str):
                     raise ValueError(
-                        "headers must be a dictionary for external tools of type 'remote_mcp'"
+                        "command must be a string to start a local MCP server"
                     )
+                # Reject empty/whitespace-only command strings
+                if command.strip() == "":
+                    raise ValueError("command must be a non-empty string")
 
-                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):
+                args = properties.get("args", None)
+                if args is not None:
+                    if not isinstance(args, list):
                         raise ValueError(
-                            "secret_header_keys must be a list for external tools of type 'remote_mcp'"
+                            "arguments must be a list to start a local MCP server"
                         )
-                    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")
+            case ToolServerType.kiln_task:
+                tool_name_validator(properties.get("name", ""))
+                err_msg_prefix = "Kiln task server properties:"
+                validate_return_dict_prop(
+                    properties, "description", str, err_msg_prefix
+                )
+                description = properties.get("description", "")
+                if len(description) > 128:
+                    raise ValueError("description must be 128 characters or less")
+                validate_return_dict_prop(
+                    properties, "is_archived", bool, err_msg_prefix
+                )
+                validate_return_dict_prop(properties, "task_id", str, err_msg_prefix)
+                validate_return_dict_prop(
+                    properties, "run_config_id", str, err_msg_prefix
+                )
 
-                args = self.properties.get("args", None)
-                if not isinstance(args, list):
-                    raise ValueError(
-                        "arguments must be a list to start a local MCP server"
-                    )
+            case _:
+                # Type checking will catch missing cases
+                raise_exhaustive_enum_error(server_type)
+        return data
 
-                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'"
+    @model_validator(mode="before")
+    def validate_headers_and_env_vars(cls, data: dict) -> dict:
+        """
+        Validate secrets, these needs to be validated before model initlization because secrets will be processed and stripped
+        """
+        type = ExternalToolServer.type_from_data(data)
+
+        properties = data.get("properties", {})
+        if properties is None:
+            raise ValueError("properties is required")
+
+        match type:
+            case ToolServerType.remote_mcp:
+                # Validate headers
+                headers = properties.get("headers", None)
+                if headers is not None:
+                    ExternalToolServer.check_headers(headers)
+
+                # Secret header keys are optional, validate if they are set
+                secret_header_keys = properties.get("secret_header_keys", None)
+                if secret_header_keys is not None:
+                    ExternalToolServer.check_secret_keys(
+                        secret_header_keys, "secret_header_keys", "remote_mcp"
                     )
 
-                secret_env_var_keys = self.properties.get("secret_env_var_keys", None)
+            case ToolServerType.local_mcp:
+                # Validate secret environment variable keys
+                env_vars = properties.get("env_vars", {})
+                if env_vars is not None:
+                    ExternalToolServer.check_env_vars(env_vars)
+
                 # Secret env var keys are optional, but if they are set, they must be a list of strings
+                secret_env_var_keys = properties.get("secret_env_var_keys", None)
                 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"
-                        )
+                    ExternalToolServer.check_secret_keys(
+                        secret_env_var_keys, "secret_env_var_keys", "local_mcp"
+                    )
+
+            case ToolServerType.kiln_task:
+                pass
 
             case _:
-                # Type checking will catch missing cases
-                raise_exhaustive_enum_error(self.type)
-        return self
+                raise_exhaustive_enum_error(type)
+
+        return data
 
     def get_secret_keys(self) -> list[str]:
         """
@@ -176,6 +326,8 @@ class ExternalToolServer(KilnParentedModel):
                 return self.properties.get("secret_header_keys", [])
             case ToolServerType.local_mcp:
                 return self.properties.get("secret_env_var_keys", [])
+            case ToolServerType.kiln_task:
+                return []
             case _:
                 raise_exhaustive_enum_error(self.type)