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

kiln_ai/adapters/vector_store/test_vector_store_registry.py

@@ -0,0 +1,199 @@
+from typing import Callable
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from kiln_ai.adapters.vector_store.vector_store_registry import (
+    vector_store_adapter_for_config,
+)
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+from kiln_ai.datamodel.embedding import EmbeddingConfig
+from kiln_ai.datamodel.rag import RagConfig
+from kiln_ai.datamodel.vector_store import VectorStoreConfig, VectorStoreType
+
+
+@pytest.fixture(autouse=True)
+def patch_settings_dir(tmp_path):
+    with patch("kiln_ai.utils.config.Config.settings_dir", return_value=tmp_path):
+        yield
+
+
+@pytest.fixture
+def create_rag_config_factory() -> Callable[
+    [VectorStoreConfig, EmbeddingConfig], RagConfig
+]:
+    def create_rag_config(
+        vector_store_config: VectorStoreConfig, embedding_config: EmbeddingConfig
+    ) -> RagConfig:
+        return RagConfig(
+            name="test_rag",
+            tool_name="test_rag_tool",
+            tool_description="A test RAG tool for registry testing",
+            extractor_config_id="test_extractor",
+            chunker_config_id="test_chunker",
+            embedding_config_id=embedding_config.id,
+            vector_store_config_id=vector_store_config.id,
+        )
+
+    return create_rag_config
+
+
+@pytest.fixture
+def embedding_config():
+    """Create an embedding config for testing."""
+    return EmbeddingConfig(
+        name="test_embedding",
+        model_provider_name=ModelProviderName.openai,
+        model_name="text-embedding-ada-002",
+        properties={},
+    )
+
+
+@pytest.fixture
+def lancedb_fts_vector_store_config():
+    """Create a vector store config for testing."""
+    config = VectorStoreConfig(
+        name="test_config",
+        store_type=VectorStoreType.LANCE_DB_FTS,
+        properties={
+            "similarity_top_k": 10,
+            "overfetch_factor": 20,
+            "vector_column_name": "vector",
+            "text_key": "text",
+            "doc_id_key": "doc_id",
+        },
+    )
+    # Set an ID for the config since build_lancedb_vector_store requires it
+    config.id = "test_config_id"
+    return config
+
+
+@pytest.fixture
+def lancedb_knn_vector_store_config():
+    """Create a vector store config for testing."""
+    config = VectorStoreConfig(
+        name="test_config",
+        store_type=VectorStoreType.LANCE_DB_VECTOR,
+        properties={
+            "similarity_top_k": 10,
+            "overfetch_factor": 20,
+            "vector_column_name": "vector",
+            "text_key": "text",
+            "doc_id_key": "doc_id",
+            "nprobes": 10,
+        },
+    )
+    # Set an ID for the config since build_lancedb_vector_store requires it
+    config.id = "test_config_id"
+    return config
+
+
+@pytest.fixture
+def lancedb_hybrid_vector_store_config():
+    """Create a vector store config for testing."""
+    config = VectorStoreConfig(
+        name="test_config",
+        store_type=VectorStoreType.LANCE_DB_HYBRID,
+        properties={
+            "similarity_top_k": 10,
+            "overfetch_factor": 20,
+            "vector_column_name": "vector",
+            "text_key": "text",
+            "doc_id_key": "doc_id",
+            "nprobes": 10,
+        },
+    )
+    # Set an ID for the config since build_lancedb_vector_store requires it
+    config.id = "test_config_id"
+    return config
+
+
+class TestVectorStoreAdapterForConfig:
+    """Test the vector_store_adapter_for_config function."""
+
+    @pytest.mark.asyncio
+    async def test_vector_store_adapter_for_config_unsupported_type(
+        self,
+        create_rag_config_factory,
+        lancedb_fts_vector_store_config,
+        embedding_config,
+    ):
+        """Test error handling for unsupported vector store types."""
+        # Create a mock config with an invalid store type
+        unsupported_config = MagicMock()
+        unsupported_config.store_type = "INVALID_TYPE"
+        unsupported_config.name = "unsupported"
+        unsupported_config.id = "test_config_id"
+
+        rag_config = create_rag_config_factory(
+            MagicMock(spec=VectorStoreConfig, id="test_config_id"), embedding_config
+        )
+        with pytest.raises(ValueError, match="Unhandled enum value"):
+            await vector_store_adapter_for_config(rag_config, unsupported_config)
+
+    async def test_lancedb_fts_vector_store_adapter_for_config(
+        self,
+        lancedb_fts_vector_store_config,
+        create_rag_config_factory,
+        embedding_config,
+    ):
+        rag_config = create_rag_config_factory(
+            lancedb_fts_vector_store_config, embedding_config
+        )
+        adapter = await vector_store_adapter_for_config(
+            rag_config, lancedb_fts_vector_store_config
+        )
+
+        assert adapter.vector_store_config == lancedb_fts_vector_store_config
+        assert adapter.vector_store_config.name == "test_config"
+        assert adapter.vector_store_config.store_type == VectorStoreType.LANCE_DB_FTS
+
+    async def test_lancedb_hybrid_vector_store_adapter_for_config(
+        self,
+        lancedb_hybrid_vector_store_config,
+        create_rag_config_factory,
+        embedding_config,
+    ):
+        rag_config = create_rag_config_factory(
+            lancedb_hybrid_vector_store_config, embedding_config
+        )
+        adapter = await vector_store_adapter_for_config(
+            rag_config, lancedb_hybrid_vector_store_config
+        )
+
+        assert adapter.vector_store_config == lancedb_hybrid_vector_store_config
+        assert adapter.vector_store_config.name == "test_config"
+        assert adapter.vector_store_config.store_type == VectorStoreType.LANCE_DB_HYBRID
+
+    async def test_lancedb_vector_vector_store_adapter_for_config(
+        self,
+        lancedb_knn_vector_store_config,
+        create_rag_config_factory,
+        embedding_config,
+    ):
+        rag_config = create_rag_config_factory(
+            lancedb_knn_vector_store_config, embedding_config
+        )
+        adapter = await vector_store_adapter_for_config(
+            rag_config, lancedb_knn_vector_store_config
+        )
+        assert adapter.vector_store_config == lancedb_knn_vector_store_config
+        assert adapter.vector_store_config.name == "test_config"
+        assert adapter.vector_store_config.store_type == VectorStoreType.LANCE_DB_VECTOR
+
+    async def test_vector_store_adapter_for_config_missing_id(
+        self,
+        create_rag_config_factory,
+        lancedb_fts_vector_store_config,
+        embedding_config,
+    ):
+        rag_config = create_rag_config_factory(
+            lancedb_fts_vector_store_config, embedding_config
+        )
+
+        lancedb_fts_vector_store_config.id = None
+
+        with pytest.raises(ValueError, match="Vector store config ID is required"):
+            await vector_store_adapter_for_config(
+                rag_config, lancedb_fts_vector_store_config
+            )

kiln_ai/adapters/vector_store/vector_store_registry.py

@@ -0,0 +1,33 @@
+import logging
+
+from kiln_ai.adapters.vector_store.base_vector_store_adapter import (
+    BaseVectorStoreAdapter,
+)
+from kiln_ai.adapters.vector_store.lancedb_adapter import LanceDBAdapter
+from kiln_ai.datamodel.rag import RagConfig
+from kiln_ai.datamodel.vector_store import VectorStoreConfig, VectorStoreType
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+logger = logging.getLogger(__name__)
+
+
+async def vector_store_adapter_for_config(
+    rag_config: RagConfig,
+    vector_store_config: VectorStoreConfig,
+) -> BaseVectorStoreAdapter:
+    vector_store_config_id = vector_store_config.id
+    if vector_store_config_id is None:
+        raise ValueError("Vector store config ID is required")
+
+    match vector_store_config.store_type:
+        case (
+            VectorStoreType.LANCE_DB_FTS
+            | VectorStoreType.LANCE_DB_HYBRID
+            | VectorStoreType.LANCE_DB_VECTOR
+        ):
+            return LanceDBAdapter(
+                rag_config,
+                vector_store_config,
+            )
+        case _:
+            raise_exhaustive_enum_error(vector_store_config.store_type)

kiln_ai/datamodel/__init__.py

@@ -11,21 +11,24 @@ User docs: https://docs.kiln.tech/developers/kiln-datamodel
 
 from __future__ import annotations
 
-from kiln_ai.datamodel import dataset_split, eval, strict_mode
+from kiln_ai.datamodel import (
+    chunk,
+    dataset_split,
+    embedding,
+    eval,
+    extraction,
+    rag,
+    strict_mode,
+)
 from kiln_ai.datamodel.datamodel_enums import (
     FineTuneStatusType,
     Priority,
     StructuredOutputMode,
     TaskOutputRatingType,
 )
-from kiln_ai.datamodel.dataset_split import (
-    DatasetSplit,
-    DatasetSplitDefinition,
-)
+from kiln_ai.datamodel.dataset_split import DatasetSplit, DatasetSplitDefinition
 from kiln_ai.datamodel.external_tool_server import ExternalToolServer
-from kiln_ai.datamodel.finetune import (
-    Finetune,
-)
+from kiln_ai.datamodel.finetune import Finetune
 from kiln_ai.datamodel.project import Project
 from kiln_ai.datamodel.prompt import BasePrompt, Prompt
 from kiln_ai.datamodel.prompt_id import (
@@ -42,10 +45,7 @@ from kiln_ai.datamodel.task_output import (
     TaskOutput,
     TaskOutputRating,
 )
-from kiln_ai.datamodel.task_run import (
-    TaskRun,
-    Usage,
-)
+from kiln_ai.datamodel.task_run import TaskRun, Usage
 
 __all__ = [
     "BasePrompt",
@@ -69,11 +69,14 @@ __all__ = [
     "TaskOutputRating",
     "TaskOutputRatingType",
     "TaskRequirement",
-    "TaskRequirement",
     "TaskRun",
     "Usage",
+    "chunk",
     "dataset_split",
+    "embedding",
     "eval",
+    "extraction",
     "prompt_generator_values",
+    "rag",
     "strict_mode",
 ]

kiln_ai/datamodel/basemodel.py

@@ -2,6 +2,7 @@ import json
 import os
 import re
 import shutil
+import tempfile
 import unicodedata
 import uuid
 from abc import ABCMeta
@@ -15,9 +16,11 @@ from pydantic import (
     BeforeValidator,
     ConfigDict,
     Field,
+    SerializationInfo,
     ValidationError,
     ValidationInfo,
     computed_field,
+    model_serializer,
     model_validator,
 )
 from pydantic_core import ErrorDetails
@@ -26,6 +29,7 @@ from typing_extensions import Annotated, Self
 from kiln_ai.datamodel.model_cache import ModelCache
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.formatting import snake_case
+from kiln_ai.utils.mime_type import guess_extension
 
 # ID is a 12 digit random integer string.
 # Should be unique per item, at least inside the context of a parent/child relationship.
@@ -85,6 +89,161 @@ def string_to_valid_name(name: str) -> str:
     return valid_name.strip("_").strip()
 
 
+class KilnAttachmentModel(BaseModel):
+    path: Path | None = Field(
+        default=None,
+        description="The path to the attachment relative to the parent model's path",
+    )
+
+    input_path: Path | None = Field(
+        default=None,
+        description="The input absolute path to the attachment. The file will be copied to its permanent location when the model is saved.",
+    )
+
+    @model_validator(mode="after")
+    def check_file_exists(self, info: ValidationInfo) -> Self:
+        context = info.context or {}
+
+        if self.path is None and self.input_path is None:
+            raise ValueError("Path or input path is not set")
+        if self.path is not None and self.input_path is not None:
+            raise ValueError("Path and input path cannot both be set")
+
+        # when loading from file, we only know the relative path so we cannot check if it exists
+        # without knowing the parent path
+        if context.get("loading_from_file", False):
+            if isinstance(self.path, str):
+                self.path = Path(self.path)
+            self.input_path = None
+            return self
+
+        # when creating a new attachment, the path is not set yet (it is set when the model is saved)
+        # so we only expect the absolute path to be set
+        if self.input_path is not None:
+            if isinstance(self.input_path, str):
+                self.input_path = Path(self.input_path)
+            if not self.input_path.is_absolute():
+                raise ValueError(f"Input path is not absolute: {self.input_path}")
+            if not os.path.exists(self.input_path):
+                raise ValueError(f"Input path does not exist: {self.input_path}")
+            if not os.path.isfile(self.input_path):
+                raise ValueError(f"Input path is not a file: {self.input_path}")
+
+            # this normalizes the path and resolves symlinks
+            self.input_path = self.input_path.resolve()
+
+            return self
+
+        if self.path is not None:
+            if isinstance(self.path, str):
+                self.path = Path(self.path)
+            if self.path.is_absolute():
+                raise ValueError(
+                    f"Path is absolute but should be relative: {self.path}"
+                )
+            if not os.path.exists(self.path):
+                raise ValueError(f"Path does not exist: {self.path}")
+            if not os.path.isfile(self.path):
+                raise ValueError(f"Path is not a file: {self.path}")
+
+        return self
+
+    @model_serializer
+    def serialize(self, info: SerializationInfo) -> dict[str, Path] | None:
+        # when the attachment is optional on the model, we get None here
+        if self is None:
+            return None
+
+        context = info.context or {}
+
+        # serialization may also be called by other parts of the system, the callers should
+        # explicitly set save_attachments to True if they want to save attachments
+        save_attachments: bool = context.get("save_attachments", False)
+        if not save_attachments:
+            path_val = self.path if self.path is not None else self.input_path
+            if path_val is None:
+                raise ValueError("Attachment has no path")
+            return {"path": path_val}
+
+        dest_path: Path | None = context.get("dest_path", None)
+        if not dest_path or not isinstance(dest_path, Path):
+            raise ValueError(
+                f"dest_path must be a valid Path object when saving attachments, got: {dest_path}"
+            )
+        if not dest_path.is_dir():
+            raise ValueError("dest_path must be a directory when saving attachments")
+
+        # the attachment is already in the parent folder, so we don't need to copy it
+        # if the path is already relative, we consider it has been copied already
+        if self.path is not None:
+            return {"path": self.path}
+
+        # copy file and update the path to be relative to the dest_path
+        new_path = self.copy_file_to(dest_path, context.get("filename_prefix", None))
+
+        self.path = new_path.relative_to(dest_path)
+        self.input_path = None
+
+        return {"path": self.path}
+
+    def copy_file_to(
+        self, dest_folder: Path, filename_prefix: str | None = None
+    ) -> Path:
+        if self.input_path is None:
+            raise ValueError("Attachment has no input path to copy")
+
+        filename = f"{str(uuid.uuid4().int)[:12]}{self.input_path.suffix}"
+        if filename_prefix:
+            filename = f"{filename_prefix}_{filename}"
+        target_path = dest_folder / filename
+        shutil.copy(self.input_path, target_path)
+        return target_path
+
+    @classmethod
+    def from_data(cls, data: str | bytes, mime_type: str) -> Self:
+        """Create an attachment from str or byte data, in a temp file. The attachment is persisted to
+        its permanent location when the model is saved.
+        """
+        extension = guess_extension(mime_type) or ".unknown"
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=extension)
+        if isinstance(data, str):
+            temp_file.write(data.encode("utf-8"))
+        else:
+            temp_file.write(data)
+        temp_file.close()
+        return cls(input_path=Path(temp_file.name))
+
+    @classmethod
+    def from_file(cls, path: Path | str) -> Self:
+        """Create an attachment from a file path. The attachment is persisted to
+        its permanent location when the model is saved.
+        """
+        if isinstance(path, str):
+            path = Path(path)
+        return cls(input_path=path)
+
+    def resolve_path(self, parent_path: Path | None = None) -> Path:
+        """
+        Resolve the path of the attachment relative to the parent path. The attachment does not know
+        its parent, so we need to call this to get the full path.
+        Args:
+            parent_path (Path): The absolute path to the parent folder. Must be provided if the model is saved to disk.
+        Returns:
+            Path: The resolved path of the attachment
+        """
+        if self.input_path is not None:
+            return self.input_path
+        if self.path is None:
+            raise ValueError("Attachment path is not set")
+        if parent_path is None:
+            raise ValueError("Parent path is not set")
+        if not parent_path.is_absolute():
+            raise ValueError(
+                f"Failed to resolve attachment path for {self.path} because parent path is not absolute: {parent_path}"
+            )
+        return (parent_path / self.path).resolve()
+
+
 # Usage:
 # class MyModel(KilnBaseModel):
 #     name: FilenameString = Field(description="The name of the model.")
@@ -223,7 +382,17 @@ class KilnBaseModel(BaseModel):
                 f"id: {getattr(self, 'id', None)}, path: {path}"
             )
         path.parent.mkdir(parents=True, exist_ok=True)
-        json_data = self.model_dump_json(indent=2, exclude={"path"})
+
+        json_data = self.model_dump_json(
+            indent=2,
+            exclude={"path"},
+            # dest_path is used by the attachment serializer to save attachments to the correct location
+            # and update the paths to be relative to path.parent
+            context={
+                "save_attachments": True,
+                "dest_path": path.parent,
+            },
+        )
         with open(path, "w", encoding="utf-8") as file:
             file.write(json_data)
         # save the path so even if something like name changes, the file doesn't move

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"