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

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/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/project.py

@@ -1,13 +1,27 @@
 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, "external_tool_servers": ExternalToolServer},
+    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.
@@ -26,5 +40,23 @@ class Project(
     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