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

kiln_ai/adapters/vector_store/lancedb_helpers.py

@@ -0,0 +1,101 @@
+from typing import Any, Dict, List, Literal
+
+from llama_index.core.schema import NodeRelationship, RelatedNodeInfo, TextNode
+from llama_index.vector_stores.lancedb import LanceDBVectorStore
+
+from kiln_ai.datamodel.vector_store import (
+    VectorStoreConfig,
+    VectorStoreType,
+    raise_exhaustive_enum_error,
+)
+from kiln_ai.utils.uuid import string_to_uuid
+
+
+def store_type_to_lancedb_query_type(
+    store_type: VectorStoreType,
+) -> Literal["fts", "hybrid", "vector"]:
+    match store_type:
+        case VectorStoreType.LANCE_DB_FTS:
+            return "fts"
+        case VectorStoreType.LANCE_DB_HYBRID:
+            return "hybrid"
+        case VectorStoreType.LANCE_DB_VECTOR:
+            return "vector"
+        case _:
+            raise_exhaustive_enum_error(store_type)
+
+
+def lancedb_construct_from_config(
+    vector_store_config: VectorStoreConfig,
+    uri: str,
+    **extra_params: Any,
+) -> LanceDBVectorStore:
+    """Construct a LanceDBVectorStore from a VectorStoreConfig."""
+    kwargs: Dict[str, Any] = {**extra_params}
+    if (
+        vector_store_config.lancedb_properties.nprobes is not None
+        and "nprobes" not in kwargs
+    ):
+        kwargs["nprobes"] = vector_store_config.lancedb_properties.nprobes
+
+    return LanceDBVectorStore(
+        mode="create",
+        query_type=store_type_to_lancedb_query_type(vector_store_config.store_type),
+        overfetch_factor=vector_store_config.lancedb_properties.overfetch_factor,
+        vector_column_name=vector_store_config.lancedb_properties.vector_column_name,
+        text_key=vector_store_config.lancedb_properties.text_key,
+        doc_id_key=vector_store_config.lancedb_properties.doc_id_key,
+        uri=uri,
+        **kwargs,
+    )
+
+
+def convert_to_llama_index_node(
+    document_id: str,
+    chunk_idx: int,
+    node_id: str,
+    text: str,
+    vector: List[float],
+) -> TextNode:
+    return TextNode(
+        id_=node_id,
+        text=text,
+        embedding=vector,
+        metadata={
+            # metadata is populated by some internal llama_index logic
+            # that uses for example the source_node relationship
+            "kiln_doc_id": document_id,
+            "kiln_chunk_idx": chunk_idx,
+            #
+            # llama_index lancedb vector store automatically sets these metadata:
+            # "doc_id": "UUID node_id of the Source Node relationship",
+            # "document_id": "UUID node_id of the Source Node relationship",
+            # "ref_doc_id": "UUID node_id of the Source Node relationship"
+            #
+            # llama_index file loaders set these metadata, which would be useful to also support:
+            # "creation_date": "2025-09-03",
+            # "file_name": "file.pdf",
+            # "file_path": "/absolute/path/to/the/file.pdf",
+            # "file_size": 395154,
+            # "file_type": "application\/pdf",
+            # "last_modified_date": "2025-09-03",
+            # "page_label": "1",
+        },
+        relationships={
+            # when using the llama_index loaders, llama_index groups Nodes under Documents
+            # and relationships point to the Document (which is also a Node), which confusingly
+            # enough does not map to an actual file (for a PDF, a Document is a page of the PDF)
+            # the Document structure is not something that is persisted, so it is fine here
+            # if we have a relationship to a node_id that does not exist in the db
+            NodeRelationship.SOURCE: RelatedNodeInfo(
+                node_id=document_id,
+                node_type="1",
+                metadata={},
+            ),
+        },
+    )
+
+
+def deterministic_chunk_id(document_id: str, chunk_idx: int) -> str:
+    # the id_ of the Node must be a UUID string, otherwise llama_index / LanceDB fails downstream
+    return str(string_to_uuid(f"{document_id}::{chunk_idx}"))

kiln_ai/adapters/vector_store/test_lancedb_adapter.py

@@ -17,6 +17,7 @@ from kiln_ai.adapters.vector_store.base_vector_store_adapter import (
     VectorStoreQuery,
 )
 from kiln_ai.adapters.vector_store.lancedb_adapter import LanceDBAdapter
+from kiln_ai.adapters.vector_store.lancedb_helpers import deterministic_chunk_id
 from kiln_ai.adapters.vector_store.vector_store_registry import (
     vector_store_adapter_for_config,
 )
@@ -925,9 +926,7 @@ async def test_get_nodes_by_ids_functionality(
     await adapter.add_chunks_with_embeddings([mock_chunked_documents[0]])  # doc_001
 
     # Test getting nodes by IDs - compute expected IDs
-    expected_ids = [
-        adapter.compute_deterministic_chunk_id("doc_001", i) for i in range(4)
-    ]
+    expected_ids = [deterministic_chunk_id("doc_001", i) for i in range(4)]
 
     # Get nodes by IDs
     retrieved_nodes = await adapter.get_nodes_by_ids(expected_ids)
@@ -943,7 +942,7 @@ async def test_get_nodes_by_ids_functionality(
         assert len(node.get_content()) > 0
 
     # Test with non-existent IDs
-    fake_ids = [adapter.compute_deterministic_chunk_id("fake_doc", i) for i in range(2)]
+    fake_ids = [deterministic_chunk_id("fake_doc", i) for i in range(2)]
     retrieved_fake = await adapter.get_nodes_by_ids(fake_ids)
     assert len(retrieved_fake) == 0
 
@@ -1019,7 +1018,7 @@ async def test_uuid_scheme_retrieval_and_node_properties(
     # Test the UUID scheme: document_id::chunk_idx
     for chunk_idx in range(4):
         # Compute expected ID using the same scheme as the adapter
-        expected_id = adapter.compute_deterministic_chunk_id("doc_001", chunk_idx)
+        expected_id = deterministic_chunk_id("doc_001", chunk_idx)
 
         # Retrieve the specific node by ID
         retrieved_nodes = await adapter.get_nodes_by_ids([expected_id])
@@ -1053,7 +1052,7 @@ async def test_uuid_scheme_retrieval_and_node_properties(
 
     # Test retrieval of doc_002 chunks
     for chunk_idx in range(4):
-        expected_id = adapter.compute_deterministic_chunk_id("doc_002", chunk_idx)
+        expected_id = deterministic_chunk_id("doc_002", chunk_idx)
         retrieved_nodes = await adapter.get_nodes_by_ids([expected_id])
         assert len(retrieved_nodes) == 1
 
@@ -1080,25 +1079,19 @@ async def test_deterministic_chunk_id_consistency(
     create_rag_config_factory,
 ):
     """Test that the deterministic chunk ID generation is consistent."""
-    rag_config = create_rag_config_factory(fts_vector_store_config, embedding_config)
-
-    adapter = LanceDBAdapter(
-        rag_config,
-        fts_vector_store_config,
-    )
 
     # Test that the same document_id and chunk_idx always produce the same UUID
     doc_id = "test_doc_123"
     chunk_idx = 5
 
-    id1 = adapter.compute_deterministic_chunk_id(doc_id, chunk_idx)
-    id2 = adapter.compute_deterministic_chunk_id(doc_id, chunk_idx)
+    id1 = deterministic_chunk_id(doc_id, chunk_idx)
+    id2 = deterministic_chunk_id(doc_id, chunk_idx)
 
     assert id1 == id2
 
     # Test that different inputs produce different UUIDs
-    id3 = adapter.compute_deterministic_chunk_id(doc_id, chunk_idx + 1)
-    id4 = adapter.compute_deterministic_chunk_id(doc_id + "_different", chunk_idx)
+    id3 = deterministic_chunk_id(doc_id, chunk_idx + 1)
+    id4 = deterministic_chunk_id(doc_id + "_different", chunk_idx)
 
     assert id1 != id3
     assert id1 != id4

kiln_ai/adapters/vector_store/test_lancedb_helpers.py

@@ -0,0 +1,142 @@
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai.adapters.vector_store.lancedb_helpers import (
+    convert_to_llama_index_node,
+    deterministic_chunk_id,
+    lancedb_construct_from_config,
+    store_type_to_lancedb_query_type,
+)
+from kiln_ai.datamodel.vector_store import VectorStoreConfig, VectorStoreType
+from kiln_ai.utils.uuid import string_to_uuid
+
+
+class _FakeLanceDBVectorStore:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+
+def _base_properties(nprobes: int | None = None) -> dict[str, str | int | float | None]:
+    props: dict[str, str | int | float | None] = {
+        "similarity_top_k": 5,
+        "overfetch_factor": 2,
+        "vector_column_name": "vec",
+        "text_key": "text",
+        "doc_id_key": "doc_id",
+    }
+    if nprobes is not None:
+        props["nprobes"] = nprobes
+    return props
+
+
+def _make_config(
+    store_type: VectorStoreType, nprobes: int | None = None
+) -> VectorStoreConfig:
+    return VectorStoreConfig(
+        name="test_store",
+        description=None,
+        store_type=store_type,
+        properties=_base_properties(nprobes),
+    )
+
+
+def test_store_type_to_lancedb_query_type_mapping():
+    assert store_type_to_lancedb_query_type(VectorStoreType.LANCE_DB_FTS) == "fts"
+    assert store_type_to_lancedb_query_type(VectorStoreType.LANCE_DB_HYBRID) == "hybrid"
+    assert store_type_to_lancedb_query_type(VectorStoreType.LANCE_DB_VECTOR) == "vector"
+
+
+def test_store_type_to_lancedb_query_type_unsupported_raises():
+    with pytest.raises(Exception):
+        store_type_to_lancedb_query_type("unsupported")  # type: ignore[arg-type]
+
+
+def test_lancedb_construct_from_config_includes_nprobes():
+    with patch(
+        "kiln_ai.adapters.vector_store.lancedb_helpers.LanceDBVectorStore",
+        new=_FakeLanceDBVectorStore,
+    ):
+        cfg = _make_config(VectorStoreType.LANCE_DB_VECTOR, nprobes=7)
+
+        result = lancedb_construct_from_config(
+            vector_store_config=cfg,
+            uri="memory://",
+            api_key="k",
+            region="r",
+            table_name="t",
+        )
+
+    assert isinstance(result, _FakeLanceDBVectorStore)
+    kwargs = result.kwargs
+
+    assert kwargs["mode"] == "create"
+    assert kwargs["uri"] == "memory://"
+    assert kwargs["query_type"] == "vector"
+    assert kwargs["overfetch_factor"] == 2
+    assert kwargs["vector_column_name"] == "vec"
+    assert kwargs["text_key"] == "text"
+    assert kwargs["doc_id_key"] == "doc_id"
+    assert kwargs["api_key"] == "k"
+    assert kwargs["region"] == "r"
+    assert kwargs["table_name"] == "t"
+    # extra optional kwarg present when provided
+    assert kwargs["nprobes"] == 7
+
+
+def test_lancedb_construct_from_config_omits_nprobes_when_none():
+    with patch(
+        "kiln_ai.adapters.vector_store.lancedb_helpers.LanceDBVectorStore",
+        new=_FakeLanceDBVectorStore,
+    ):
+        cfg = _make_config(VectorStoreType.LANCE_DB_FTS, nprobes=None)
+
+        result = lancedb_construct_from_config(
+            vector_store_config=cfg,
+            uri="memory://",
+            api_key=None,
+            region=None,
+            table_name=None,
+        )
+
+    assert isinstance(result, _FakeLanceDBVectorStore)
+    kwargs = result.kwargs
+
+    assert kwargs["query_type"] == "fts"
+    assert "nprobes" not in kwargs
+
+
+def test_convert_to_llama_index_node_builds_expected_structure():
+    node = convert_to_llama_index_node(
+        document_id="doc-123",
+        chunk_idx=0,
+        node_id="11111111-1111-5111-8111-111111111111",
+        text="hello",
+        vector=[0.1, 0.2],
+    )
+
+    assert node.id_ == "11111111-1111-5111-8111-111111111111"
+    assert node.text == "hello"
+    assert node.embedding == [0.1, 0.2]
+    assert node.metadata["kiln_doc_id"] == "doc-123"
+    assert node.metadata["kiln_chunk_idx"] == 0
+
+    # relationship exists and points to the source document id
+    from llama_index.core.schema import NodeRelationship, RelatedNodeInfo
+
+    assert NodeRelationship.SOURCE in node.relationships
+    related = node.relationships[NodeRelationship.SOURCE]
+    assert isinstance(related, RelatedNodeInfo)
+    assert related.node_id == "doc-123"
+    assert related.node_type == "1"
+    assert isinstance(related.metadata, dict)
+
+
+def test_deterministic_chunk_id_uses_uuid_v5_namespace():
+    doc_id = "doc-abc"
+    idx = 3
+    expected = str(string_to_uuid(f"{doc_id}::{idx}"))
+    assert deterministic_chunk_id(doc_id, idx) == expected
+
+    # call again to ensure the same value is returned
+    assert deterministic_chunk_id(doc_id, idx) == expected

kiln_ai/adapters/vector_store_loaders/test_lancedb_loader.py

@@ -0,0 +1,282 @@
+import os
+import random
+import time
+import uuid
+from dataclasses import dataclass
+
+import pytest
+from pydantic import BaseModel, Field
+
+from kiln_ai.adapters.vector_store.lancedb_adapter import lancedb_construct_from_config
+from kiln_ai.adapters.vector_store_loaders.vector_store_loader import VectorStoreLoader
+from kiln_ai.datamodel.chunk import Chunk, ChunkedDocument
+from kiln_ai.datamodel.datamodel_enums import KilnMimeType
+from kiln_ai.datamodel.embedding import ChunkEmbeddings, Embedding
+from kiln_ai.datamodel.extraction import (
+    Document,
+    Extraction,
+    ExtractionSource,
+    FileInfo,
+    Kind,
+)
+from kiln_ai.datamodel.project import Project
+from kiln_ai.datamodel.rag import RagConfig
+from kiln_ai.datamodel.vector_store import VectorStoreConfig, VectorStoreType
+
+
+@dataclass
+class DocWithChunks:
+    document: Document
+    extraction: Extraction
+    chunked_document: ChunkedDocument
+    chunked_embeddings: ChunkEmbeddings
+
+
+def lorem_ipsum(n: int) -> str:
+    return " ".join(
+        ["Lorem ipsum dolor sit amet, consectetur adipiscing elit." for _ in range(n)]
+    )
+
+
+@pytest.fixture
+def mock_chunks_factory(mock_attachment_factory):
+    def fn(
+        project: Project,
+        rag_config: RagConfig,
+        num_chunks: int = 1,
+        text: str | None = None,
+    ) -> DocWithChunks:
+        doc = Document(
+            id=f"doc_{uuid.uuid4()}",
+            name="Test Document",
+            description="Test Document",
+            original_file=FileInfo(
+                filename="test.pdf",
+                size=100,
+                mime_type="application/pdf",
+                attachment=mock_attachment_factory(KilnMimeType.PDF),
+            ),
+            kind=Kind.DOCUMENT,
+            parent=project,
+        )
+        doc.save_to_file()
+
+        extraction = Extraction(
+            source=ExtractionSource.PROCESSED,
+            extractor_config_id=rag_config.extractor_config_id,
+            output=mock_attachment_factory(KilnMimeType.PDF),
+            parent=doc,
+        )
+        extraction.save_to_file()
+
+        chunks = [
+            Chunk(
+                content=mock_attachment_factory(
+                    KilnMimeType.TXT, text=f"text-{i}: {text or lorem_ipsum(10)}"
+                )
+            )
+            for i in range(num_chunks)
+        ]
+        chunked_document = ChunkedDocument(
+            chunks=chunks,
+            chunker_config_id=rag_config.chunker_config_id,
+            parent=extraction,
+        )
+        chunked_document.save_to_file()
+        chunked_embeddings = ChunkEmbeddings(
+            embeddings=[
+                Embedding(vector=[i + 0.1, i + 0.2, i + 0.3, i + 0.4, i + 0.5])
+                for i in range(num_chunks)
+            ],
+            embedding_config_id=rag_config.embedding_config_id,
+            parent=chunked_document,
+        )
+        chunked_embeddings.save_to_file()
+        return DocWithChunks(
+            document=doc,
+            extraction=extraction,
+            chunked_document=chunked_document,
+            chunked_embeddings=chunked_embeddings,
+        )
+
+    return fn
+
+
+@pytest.fixture
+def mock_project(tmp_path):
+    project = Project(
+        name="Test Project", path=tmp_path / "test_project" / "project.kiln"
+    )
+    project.save_to_file()
+    return project
+
+
+@pytest.fixture
+def rag_config_factory(mock_project):
+    def fn(vector_store_config_id: str) -> RagConfig:
+        rag_config = RagConfig(
+            name="Test Rag Config",
+            parent=mock_project,
+            vector_store_config_id=vector_store_config_id,
+            tool_name="test_tool",
+            tool_description="test_description",
+            extractor_config_id="test_extractor",
+            chunker_config_id="test_chunker",
+            embedding_config_id="test_embedding",
+        )
+        rag_config.save_to_file()
+        return rag_config
+
+    return fn
+
+
+@pytest.fixture
+def vector_store_config_factory(mock_project):
+    def fn(vector_store_type: VectorStoreType) -> VectorStoreConfig:
+        match vector_store_type:
+            case VectorStoreType.LANCE_DB_FTS:
+                vector_store_config = VectorStoreConfig(
+                    name="Test Vector Store Config FTS",
+                    parent=mock_project,
+                    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",
+                    },
+                )
+                vector_store_config.save_to_file()
+                return vector_store_config
+            case VectorStoreType.LANCE_DB_VECTOR:
+                vector_store_config = VectorStoreConfig(
+                    name="Test Vector Store Config KNN",
+                    parent=mock_project,
+                    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,
+                    },
+                )
+                vector_store_config.save_to_file()
+                return vector_store_config
+            case VectorStoreType.LANCE_DB_HYBRID:
+                vector_store_config = VectorStoreConfig(
+                    name="Test Vector Store Config Hybrid",
+                    parent=mock_project,
+                    store_type=VectorStoreType.LANCE_DB_HYBRID,
+                    properties={
+                        "similarity_top_k": 10,
+                        "nprobes": 10,
+                        "overfetch_factor": 20,
+                        "vector_column_name": "vector",
+                        "text_key": "text",
+                        "doc_id_key": "doc_id",
+                    },
+                )
+                vector_store_config.save_to_file()
+                return vector_store_config
+            case _:
+                raise ValueError(f"Invalid vector store type: {vector_store_type}")
+
+    return fn
+
+
+class LanceDBCloudEnvVars(BaseModel):
+    uri: str = Field("LANCE_DB_URI")
+    api_key: str = Field("LANCE_DB_API_KEY")
+    region: str = Field("LANCE_DB_REGION")
+
+
+def lancedb_cloud_env_vars() -> LanceDBCloudEnvVars:
+    lancedb_uri = os.getenv("LANCE_DB_URI")
+    assert lancedb_uri is not None, (
+        "LANCE_DB_URI is not set - test requires lancedb cloud"
+    )
+
+    lancedb_api_key = os.getenv("LANCE_DB_API_KEY")
+    assert lancedb_api_key is not None, (
+        "LANCE_DB_API_KEY is not set - test requires lancedb cloud"
+    )
+
+    lancedb_region = os.getenv("LANCE_DB_REGION")
+    assert lancedb_region is not None, (
+        "LANCE_DB_REGION is not set - test requires lancedb cloud"
+    )
+    return LanceDBCloudEnvVars(
+        uri=lancedb_uri,
+        api_key=lancedb_api_key,
+        region=lancedb_region,
+    )
+
+
+@pytest.mark.parametrize(
+    "vector_store_type",
+    [
+        VectorStoreType.LANCE_DB_FTS,
+        VectorStoreType.LANCE_DB_VECTOR,
+        VectorStoreType.LANCE_DB_HYBRID,
+    ],
+)
+@pytest.mark.paid
+async def test_lancedb_loader_insert_nodes_lancedb_cloud(
+    mock_project,
+    mock_chunks_factory,
+    rag_config_factory,
+    vector_store_type,
+    vector_store_config_factory,
+):
+    lancedb_cloud_config = lancedb_cloud_env_vars()
+
+    vector_store_config = vector_store_config_factory(vector_store_type)
+    rag_config = rag_config_factory(vector_store_config.id)
+
+    # init lancedb store
+    now = time.time()
+    table_name = f"test_lancedb_loader_insert_nodes_{vector_store_type.value}_{now}"
+    lancedb_store = lancedb_construct_from_config(
+        vector_store_config=vector_store_config,
+        uri=lancedb_cloud_config.uri,
+        api_key=lancedb_cloud_config.api_key,
+        region=lancedb_cloud_config.region,
+        table_name=table_name,
+    )
+
+    loader = VectorStoreLoader(
+        project=mock_project,
+        rag_config=rag_config,
+    )
+
+    # create nodes
+    doc_count = 10
+    node_count = 0
+    for i in range(doc_count):
+        nodes_to_add = random.randint(1, 20)
+        # create mock docs, extractions, chunked documents, and chunk embeddings and persist
+        mock_chunks_factory(
+            mock_project,
+            rag_config,
+            num_chunks=nodes_to_add,
+            text=f"Document {i}",
+        )
+        node_count += nodes_to_add
+
+    assert node_count > 0, "No mock nodes were created"
+
+    # insert docs
+    batch_size = 100
+    async for batch in loader.iter_llama_index_nodes(batch_size=batch_size):
+        await lancedb_store.async_add(batch)
+
+    # check if docs are inserted
+    table = lancedb_store.table
+    assert table is not None
+    row_count = table.count_rows()
+    assert row_count == node_count, (
+        f"Expected {node_count} rows (one for each node), got {row_count} instead"
+    )