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

kiln_ai/datamodel/extraction.py

@@ -0,0 +1,303 @@
+import logging
+from enum import Enum
+from typing import TYPE_CHECKING, Any, List, Union
+
+import anyio
+from pydantic import (
+    BaseModel,
+    Field,
+    SerializationInfo,
+    ValidationInfo,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
+from typing_extensions import Self
+
+from kiln_ai.datamodel.basemodel import (
+    ID_TYPE,
+    FilenameString,
+    KilnAttachmentModel,
+    KilnParentedModel,
+    KilnParentModel,
+)
+from kiln_ai.datamodel.chunk import ChunkedDocument
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.project import Project
+
+logger = logging.getLogger(__name__)
+
+
+class Kind(str, Enum):
+    DOCUMENT = "document"
+    IMAGE = "image"
+    VIDEO = "video"
+    AUDIO = "audio"
+
+
+class OutputFormat(str, Enum):
+    TEXT = "text/plain"
+    MARKDOWN = "text/markdown"
+
+
+class ExtractorType(str, Enum):
+    LITELLM = "litellm"
+
+
+SUPPORTED_MIME_TYPES = {
+    Kind.DOCUMENT: {
+        "application/pdf",
+        "text/plain",
+        "text/markdown",
+        "text/html",
+        "text/md",
+    },
+    Kind.IMAGE: {
+        "image/png",
+        "image/jpeg",
+    },
+    Kind.VIDEO: {
+        "video/mp4",
+        "video/quicktime",
+    },
+    Kind.AUDIO: {
+        "audio/wav",
+        "audio/mpeg",
+        "audio/ogg",
+    },
+}
+
+
+class ExtractionModel(BaseModel):
+    name: str
+    label: str
+
+
+def validate_prompt(prompt: Any, name: str):
+    if not isinstance(prompt, str):
+        raise ValueError(f"{name} must be a string.")
+    if prompt == "":
+        raise ValueError(f"{name} cannot be empty.")
+
+
+class ExtractionSource(str, Enum):
+    PROCESSED = "processed"
+    PASSTHROUGH = "passthrough"
+
+
+class Extraction(
+    KilnParentedModel, KilnParentModel, parent_of={"chunked_documents": ChunkedDocument}
+):
+    source: ExtractionSource = Field(
+        description="The source of the extraction.",
+    )
+    extractor_config_id: ID_TYPE = Field(
+        description="The ID of the extractor config used to extract the data.",
+    )
+    output: KilnAttachmentModel = Field(
+        description="The extraction output.",
+    )
+
+    def parent_document(self) -> Union["Document", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Document":
+            return None
+        return self.parent  # type: ignore
+
+    async def output_content(self) -> str | None:
+        if not self.path:
+            raise ValueError(
+                "Failed to resolve the path of extraction output attachment because the extraction does not have a path."
+            )
+
+        full_path = self.output.resolve_path(self.path.parent)
+
+        try:
+            return await anyio.Path(full_path).read_text(encoding="utf-8")
+        except Exception as e:
+            logger.error(
+                f"Failed to read extraction output for {full_path}: {e}", exc_info=True
+            )
+            raise ValueError(f"Failed to read extraction output: {e}")
+
+    def chunked_documents(self, readonly: bool = False) -> list[ChunkedDocument]:
+        return super().chunked_documents(readonly=readonly)  # type: ignore
+
+
+class ExtractorConfig(KilnParentedModel):
+    name: FilenameString = Field(
+        description="A name to identify the extractor config.",
+    )
+    is_archived: bool = Field(
+        default=False,
+        description="Whether the extractor config is archived. Archived extractor configs are not shown in the UI and are not available for use.",
+    )
+    description: str | None = Field(
+        default=None, description="The description of the extractor config"
+    )
+    model_provider_name: str = Field(
+        description="The name of the model provider to use for the extractor config.",
+    )
+    model_name: str = Field(
+        description="The name of the model to use for the extractor config.",
+    )
+    output_format: OutputFormat = Field(
+        default=OutputFormat.MARKDOWN,
+        description="The format to use for the output.",
+    )
+    passthrough_mimetypes: list[OutputFormat] = Field(
+        default_factory=list,
+        description="If the mimetype is in this list, the extractor will not be used and the text content of the file will be returned as is.",
+    )
+    extractor_type: ExtractorType = Field(
+        description="This is used to determine the type of extractor to use.",
+    )
+    properties: dict[str, str | int | float | bool | dict[str, str] | None] = Field(
+        default_factory=dict,
+        description="Properties to be used to execute the extractor config. This is extractor_type specific and should serialize to a json dict.",
+    )
+
+    @field_validator("properties")
+    @classmethod
+    def validate_properties(
+        cls, properties: dict[str, Any], info: ValidationInfo
+    ) -> dict[str, Any]:
+        def get_property(key: str) -> str:
+            value = properties.get(key)
+            if value is None or value == "" or not isinstance(value, str):
+                raise ValueError(f"Prompt for {key} must be a string")
+            return value
+
+        return {
+            "prompt_document": get_property(
+                "prompt_document",
+            ),
+            "prompt_image": get_property(
+                "prompt_image",
+            ),
+            "prompt_video": get_property(
+                "prompt_video",
+            ),
+            "prompt_audio": get_property(
+                "prompt_audio",
+            ),
+        }
+
+    def prompt_document(self) -> str | None:
+        prompt = self.properties.get("prompt_document")
+        if prompt is None:
+            return None
+        if not isinstance(prompt, str):
+            raise ValueError(
+                "Invalid prompt_document. prompt_document must be a string."
+            )
+        return prompt
+
+    def prompt_video(self) -> str | None:
+        prompt = self.properties.get("prompt_video")
+        if prompt is None:
+            return None
+        if not isinstance(prompt, str):
+            raise ValueError("Invalid prompt_video. prompt_video must be a string.")
+        return prompt
+
+    def prompt_audio(self) -> str | None:
+        prompt = self.properties.get("prompt_audio")
+        if prompt is None:
+            return None
+        if not isinstance(prompt, str):
+            raise ValueError("Invalid prompt_audio. prompt_audio must be a string.")
+        return prompt
+
+    def prompt_image(self) -> str | None:
+        prompt = self.properties.get("prompt_image")
+        if prompt is None:
+            return None
+        if not isinstance(prompt, str):
+            raise ValueError("Invalid prompt_image. prompt_image must be a string.")
+        return prompt
+
+    # 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
+
+
+class FileInfo(BaseModel):
+    filename: str = Field(description="The filename of the file")
+
+    size: int = Field(description="The size of the file in bytes")
+
+    mime_type: str = Field(description="The MIME type of the file")
+
+    attachment: KilnAttachmentModel = Field(
+        description="The attachment to the file",
+    )
+
+    @field_serializer("attachment")
+    def serialize_attachment(
+        self, attachment: KilnAttachmentModel, info: SerializationInfo
+    ) -> dict:
+        context = info.context or {}
+        context["filename_prefix"] = "attachment"
+        return attachment.model_dump(mode="json", context=context)
+
+    @field_validator("mime_type")
+    @classmethod
+    def validate_mime_type(cls, mime_type: str, info: ValidationInfo) -> str:
+        filename = info.data.get("filename") or ""
+
+        for mime_types in SUPPORTED_MIME_TYPES.values():
+            if mime_type in mime_types:
+                return mime_type
+        raise ValueError(f"MIME type is not supported: {mime_type} (for {filename})")
+
+
+class Document(
+    KilnParentedModel, KilnParentModel, parent_of={"extractions": Extraction}
+):
+    name: FilenameString = Field(
+        description="A name to identify the document.",
+    )
+
+    description: str = Field(description="A description for the file")
+
+    original_file: FileInfo = Field(description="The original file")
+
+    kind: Kind = Field(
+        description="The kind of document. The kind is a broad family of filetypes that can be handled in a similar way"
+    )
+
+    tags: List[str] = Field(
+        default_factory=list,
+        description="Tags for the document. Tags are used to categorize documents for filtering and reporting.",
+    )
+
+    @model_validator(mode="after")
+    def validate_tags(self) -> Self:
+        for tag in self.tags:
+            if not tag:
+                raise ValueError("Tags cannot be empty strings")
+            if " " in tag:
+                raise ValueError("Tags cannot contain spaces. Try underscores.")
+
+        return self
+
+    # 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
+
+    def extractions(self, readonly: bool = False) -> list[Extraction]:
+        return super().extractions(readonly=readonly)  # type: ignore
+
+
+def get_kind_from_mime_type(mime_type: str) -> Kind | None:
+    for kind, mime_types in SUPPORTED_MIME_TYPES.items():
+        if mime_type in mime_types:
+            return kind
+    return None

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,28 @@
 from pydantic import Field
 
 from kiln_ai.datamodel.basemodel import FilenameString, KilnParentModel
+from kiln_ai.datamodel.chunk import ChunkerConfig
+from kiln_ai.datamodel.embedding import EmbeddingConfig
+from kiln_ai.datamodel.external_tool_server import ExternalToolServer
+from kiln_ai.datamodel.extraction import Document, ExtractorConfig
+from kiln_ai.datamodel.rag import RagConfig
 from kiln_ai.datamodel.task import Task
+from kiln_ai.datamodel.vector_store import VectorStoreConfig
 
 
-class Project(KilnParentModel, parent_of={"tasks": Task}):
+class Project(
+    KilnParentModel,
+    parent_of={
+        "tasks": Task,
+        "documents": Document,
+        "extractor_configs": ExtractorConfig,
+        "chunker_configs": ChunkerConfig,
+        "embedding_configs": EmbeddingConfig,
+        "rag_configs": RagConfig,
+        "vector_store_configs": VectorStoreConfig,
+        "external_tool_servers": ExternalToolServer,
+    },
+):
     """
     A collection of related tasks.
 
@@ -21,3 +39,24 @@ 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 documents(self, readonly: bool = False) -> list[Document]:
+        return super().documents(readonly=readonly)  # type: ignore
+
+    def extractor_configs(self, readonly: bool = False) -> list[ExtractorConfig]:
+        return super().extractor_configs(readonly=readonly)  # type: ignore
+
+    def chunker_configs(self, readonly: bool = False) -> list[ChunkerConfig]:
+        return super().chunker_configs(readonly=readonly)  # type: ignore
+
+    def embedding_configs(self, readonly: bool = False) -> list[EmbeddingConfig]:
+        return super().embedding_configs(readonly=readonly)  # type: ignore
+
+    def vector_store_configs(self, readonly: bool = False) -> list[VectorStoreConfig]:
+        return super().vector_store_configs(readonly=readonly)  # type: ignore
+
+    def rag_configs(self, readonly: bool = False) -> list[RagConfig]:
+        return super().rag_configs(readonly=readonly)  # type: ignore
+
+    def external_tool_servers(self, readonly: bool = False) -> list[ExternalToolServer]:
+        return super().external_tool_servers(readonly=readonly)  # type: ignore

kiln_ai/datamodel/rag.py

@@ -0,0 +1,79 @@
+from typing import TYPE_CHECKING, Union
+
+from pydantic import Field, model_validator
+
+from kiln_ai.datamodel.basemodel import ID_TYPE, FilenameString, KilnParentedModel
+from kiln_ai.utils.validation import ToolNameString
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.project import Project
+
+
+class RagConfig(KilnParentedModel):
+    name: FilenameString = Field(
+        description="A name to identify this RAG configuration for your own reference.",
+    )
+
+    is_archived: bool = Field(
+        default=False,
+        description="Whether the RAG configuration is archived. Archived RAG configurations are not shown in the UI and are not available for use.",
+    )
+
+    description: str | None = Field(
+        default=None,
+        description="A description of the RAG configuration for you and your team. Will not be used in prompts/training/validation.",
+    )
+
+    tool_name: ToolNameString = Field(
+        description="A name for the model to identify the Search Tool in conversations.",
+    )
+
+    tool_description: str = Field(
+        description="A description of the purpose of the tool. The model will use this description to understand the tool's capabilities.",
+        max_length=128,
+    )
+
+    extractor_config_id: ID_TYPE = Field(
+        description="The ID of the extractor config used to extract the documents.",
+    )
+
+    chunker_config_id: ID_TYPE = Field(
+        description="The ID of the chunker config used to chunk the documents.",
+    )
+
+    embedding_config_id: ID_TYPE = Field(
+        description="The ID of the embedding config used to embed the documents.",
+    )
+
+    vector_store_config_id: ID_TYPE = Field(
+        description="The ID of the vector store config used to store the documents.",
+    )
+
+    tags: list[str] | None = Field(
+        default=None,
+        description="List of document tags to filter by. If None, all documents in the project are used.",
+    )
+
+    # 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_tags(self):
+        if self.tags is not None:
+            if len(self.tags) == 0:
+                raise ValueError("Tags cannot be an empty list.")
+            for tag in self.tags:
+                if not tag:
+                    raise ValueError("Tags cannot be empty.")
+                if " " in tag:
+                    raise ValueError("Tags cannot contain spaces. Try underscores.")
+
+        if self.tool_name.strip() == "":
+            raise ValueError("Tool name cannot be empty.")
+        if self.tool_description.strip() == "":
+            raise ValueError("Tool description cannot be empty.")
+
+        return self

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,