kiln-ai 0.20.1__py3-none-any.whl → 0.22.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,22 +2,25 @@ import json
 import os
 import re
 import shutil
+import tempfile
 import unicodedata
 import uuid
 from abc import ABCMeta
 from builtins import classmethod
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Type, TypeVar
+from typing import Any, Callable, Dict, List, Optional, Set, Type, TypeVar
 
 from pydantic import (
     BaseModel,
     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.
@@ -74,9 +78,9 @@ def string_to_valid_name(name: str) -> str:
     # https://docs.python.org/3/library/unicodedata.html#unicodedata.normalize
     valid_name = unicodedata.normalize("NFKD", name)
     # Replace any forbidden chars with an underscore
-    valid_name = re.sub(FORBIDDEN_CHARS_REGEX, "_", valid_name)
+    valid_name = re.sub(FORBIDDEN_CHARS_REGEX, " ", valid_name)
     # Replace control characters with an underscore
-    valid_name = re.sub(r"[\x00-\x1F]", "_", valid_name)
+    valid_name = re.sub(r"[\x00-\x1F]", " ", valid_name)
     # Replace consecutive whitespace with a single space
     valid_name = re.sub(r"\s+", " ", valid_name)
     # Replace consecutive underscores with a single underscore
@@ -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
@@ -425,6 +594,34 @@ class KilnParentedModel(KilnBaseModel, metaclass=ABCMeta):
                     return child
         return None
 
+    @classmethod
+    def from_ids_and_parent_path(
+        cls: Type[PT], ids: Set[str], parent_path: Path | None
+    ) -> Dict[str, PT]:
+        """
+        Bulk equivalent of from_id_and_parent_path, much faster for large collections.
+
+        It picks out the matching models from the directory only once. This avoids
+        doing individual costly lookups that scan the whole directory in scenarios
+        where we need to iterate over a large collection of models (e.g. bulk tagging).
+        """
+        if parent_path is None:
+            return {}
+
+        children = {}
+
+        # Note: we're using the in-file ID. We could make this faster using the path-ID if this becomes perf bottleneck, but it's better to have 1 source of truth.
+        for child_path in cls.iterate_children_paths_of_parent_path(parent_path):
+            child_id = ModelCache.shared().get_model_id(child_path, cls)
+            if child_id in ids:
+                children[child_id] = cls.load_from_file(child_path)
+            if child_id is None:
+                child = cls.load_from_file(child_path)
+                if child.id in ids:
+                    children[child.id] = child
+
+        return children
+
 
 # Parent create methods for all child relationships
 # You must pass in parent_of in the subclass definition, defining the child relationships