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

kiln_ai/adapters/test_ml_embedding_model_list.py

@@ -1,5 +1,8 @@
+from typing import List
+
 import pytest
 
+from kiln_ai.adapters.embedding.embedding_registry import embedding_adapter_from_type
 from kiln_ai.adapters.ml_embedding_model_list import (
     EmbeddingModelName,
     KilnEmbeddingModel,
@@ -10,24 +13,28 @@ from kiln_ai.adapters.ml_embedding_model_list import (
     get_model_by_name,
 )
 from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+from kiln_ai.datamodel.embedding import EmbeddingConfig
 
 
-class TestEmbeddingModelName:
-    """Test cases for EmbeddingModelName enum"""
-
-    def test_enum_values(self):
-        """Test that enum values are correctly defined"""
-        assert (
-            EmbeddingModelName.openai_text_embedding_3_small
-            == "openai_text_embedding_3_small"
-        )
-        assert (
-            EmbeddingModelName.openai_text_embedding_3_large
-            == "openai_text_embedding_3_large"
-        )
-        assert (
-            EmbeddingModelName.gemini_text_embedding_004 == "gemini_text_embedding_004"
+@pytest.fixture
+def litellm_adapter():
+    adapter = embedding_adapter_from_type(
+        EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name=ModelProviderName.openai,
+            model_name=EmbeddingModelName.openai_text_embedding_3_small,
+            properties={},
         )
+    )
+    return adapter
+
+
+def get_all_embedding_models_and_providers() -> List[tuple[str, str]]:
+    return [
+        (model.name, provider.name)
+        for model in built_in_embedding_models
+        for provider in model.providers
+    ]
 
 
 class TestKilnEmbeddingModelProvider:
@@ -120,222 +127,40 @@ class TestKilnEmbeddingModel:
         assert model.providers[1].name == ModelProviderName.anthropic
 
 
-class TestEmbeddingModelsList:
-    """Test cases for the embedding_models list"""
-
-    def test_embedding_models_not_empty(self):
-        """Test that the embedding_models list is not empty"""
-        assert len(built_in_embedding_models) > 0
-
-    def test_all_models_have_required_fields(self):
-        """Test that all models in the list have required fields"""
-        for model in built_in_embedding_models:
-            assert hasattr(model, "family")
-            assert hasattr(model, "name")
-            assert hasattr(model, "friendly_name")
-            assert hasattr(model, "providers")
-            assert isinstance(model.name, str)
-            assert isinstance(model.friendly_name, str)
-            assert isinstance(model.providers, list)
-            assert len(model.providers) > 0
-
-    def test_all_providers_have_required_fields(self):
-        """Test that all providers in all models have required fields"""
-        for model in built_in_embedding_models:
-            for provider in model.providers:
-                assert hasattr(provider, "name")
-                assert isinstance(provider.name, ModelProviderName)
-
-    def test_model_names_are_unique(self):
-        """Test that all model names in the list are unique"""
-        model_names = [model.name for model in built_in_embedding_models]
-        assert len(model_names) == len(set(model_names))
-
-    def test_specific_models_exist(self):
-        """Test that specific expected models exist in the list"""
-        model_names = [model.name for model in built_in_embedding_models]
-
-        assert EmbeddingModelName.openai_text_embedding_3_small in model_names
-        assert EmbeddingModelName.openai_text_embedding_3_large in model_names
-        assert EmbeddingModelName.gemini_text_embedding_004 in model_names
-
-    def test_openai_embedding_models(self):
-        """Test specific OpenAI embedding models"""
-        openai_models = [
-            model
-            for model in built_in_embedding_models
-            if model.family == KilnEmbeddingModelFamily.openai
-        ]
-
-        assert len(openai_models) >= 2  # Should have at least 2 OpenAI models
-
-        # Check for specific OpenAI models
-        openai_model_names = [model.name for model in openai_models]
-        assert EmbeddingModelName.openai_text_embedding_3_small in openai_model_names
-        assert EmbeddingModelName.openai_text_embedding_3_large in openai_model_names
-
-    def test_gemini_embedding_models(self):
-        """Test specific Gemini embedding models"""
-        gemini_models = [
-            model
-            for model in built_in_embedding_models
-            if model.family == KilnEmbeddingModelFamily.gemini
-        ]
-
-        assert len(gemini_models) >= 1  # Should have at least 1 Gemini model
-
-        # Check for specific Gemini model
-        gemini_model_names = [model.name for model in gemini_models]
-        assert EmbeddingModelName.gemini_text_embedding_004 in gemini_model_names
-
-    def test_openai_text_embedding_3_small_details(self):
-        """Test specific details of OpenAI text-embedding-3-small model"""
-        model = get_model_by_name(EmbeddingModelName.openai_text_embedding_3_small)
-
-        assert model.family == KilnEmbeddingModelFamily.openai
-        assert model.friendly_name == "Text Embedding 3 Small"
-        assert len(model.providers) == 1
-
-        provider = model.providers[0]
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-small"
-        assert provider.n_dimensions == 1536
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
-
-    def test_openai_text_embedding_3_large_details(self):
-        """Test specific details of OpenAI text-embedding-3-large model"""
-        model = get_model_by_name(EmbeddingModelName.openai_text_embedding_3_large)
-
-        assert model.family == KilnEmbeddingModelFamily.openai
-        assert model.friendly_name == "Text Embedding 3 Large"
-        assert len(model.providers) == 1
-
-        provider = model.providers[0]
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-large"
-        assert provider.n_dimensions == 3072
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
-
-    def test_gemini_text_embedding_004_details(self):
-        """Test specific details of Gemini text-embedding-004 model"""
-        model = get_model_by_name(EmbeddingModelName.gemini_text_embedding_004)
-
-        assert model.family == KilnEmbeddingModelFamily.gemini
-        assert model.friendly_name == "Text Embedding 004"
-        assert len(model.providers) == 1
-
-        provider = model.providers[0]
-        assert provider.name == ModelProviderName.gemini_api
-        assert provider.model_id == "text-embedding-004"
-        assert provider.n_dimensions == 768
-        assert provider.max_input_tokens == 2048
-        assert provider.supports_custom_dimensions is False
-
-
 class TestGetModelByName:
-    """Test cases for get_model_by_name function"""
-
-    def test_get_existing_model(self):
-        """Test getting an existing model by name"""
-        model = get_model_by_name(EmbeddingModelName.openai_text_embedding_3_small)
-        assert model.name == EmbeddingModelName.openai_text_embedding_3_small
-        assert model.family == KilnEmbeddingModelFamily.openai
-
-    def test_get_all_existing_models(self):
-        """Test getting all existing models by name"""
-        for model_name in EmbeddingModelName:
-            model = get_model_by_name(model_name)
-            assert model.name == model_name
-
     def test_get_nonexistent_model_raises_error(self):
         """Test that getting a nonexistent model raises ValueError"""
         with pytest.raises(
             ValueError, match="Embedding model nonexistent_model not found"
         ):
-            get_model_by_name("nonexistent_model")
-
-    def test_get_model_with_invalid_enum_value(self):
-        """Test that getting a model with invalid enum value raises ValueError"""
-        with pytest.raises(ValueError, match="Embedding model invalid_enum not found"):
-            get_model_by_name("invalid_enum")
+            get_model_by_name("nonexistent_model")  # type: ignore
 
     @pytest.mark.parametrize(
-        "model_name,expected_family,expected_friendly_name",
-        [
-            (
-                EmbeddingModelName.openai_text_embedding_3_small,
-                KilnEmbeddingModelFamily.openai,
-                "Text Embedding 3 Small",
-            ),
-            (
-                EmbeddingModelName.openai_text_embedding_3_large,
-                KilnEmbeddingModelFamily.openai,
-                "Text Embedding 3 Large",
-            ),
-            (
-                EmbeddingModelName.gemini_text_embedding_004,
-                KilnEmbeddingModelFamily.gemini,
-                "Text Embedding 004",
-            ),
-        ],
+        "model_name",
+        [model.name for model in built_in_embedding_models],
     )
-    def test_parametrized_model_retrieval(
-        self, model_name, expected_family, expected_friendly_name
-    ):
+    def test_model_retrieval(self, model_name):
         """Test retrieving models with parametrized test cases"""
         model = get_model_by_name(model_name)
-        assert model.family == expected_family
-        assert model.friendly_name == expected_friendly_name
+        assert model.family == model.family
+        assert model.friendly_name == model.friendly_name
 
 
 class TestBuiltInEmbeddingModelsFromProvider:
-    """Test cases for built_in_embedding_models_from_provider function"""
-
-    def test_get_existing_provider_for_model(self):
-        """Test getting an existing provider for a model"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.openai,
-            model_name=EmbeddingModelName.openai_text_embedding_3_small,
-        )
+    @pytest.mark.parametrize(
+        "model_name,provider_name", get_all_embedding_models_and_providers()
+    )
+    def test_get_all_existing_models_and_providers(self, model_name, provider_name):
+        provider = built_in_embedding_models_from_provider(provider_name, model_name)
 
         assert provider is not None
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-small"
-        assert provider.n_dimensions == 1536
-
-    def test_get_all_existing_provider_model_combinations(self):
-        """Test getting all existing provider-model combinations"""
-        combinations = [
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_small,
-            ),
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_large,
-            ),
-            (
-                ModelProviderName.gemini_api,
-                EmbeddingModelName.gemini_text_embedding_004,
-            ),
-        ]
-
-        for provider_name, model_name in combinations:
-            provider = built_in_embedding_models_from_provider(
-                provider_name, model_name
-            )
-            assert provider is not None
-            assert provider.name == provider_name
-
-    def test_get_nonexistent_provider_returns_none(self):
-        """Test that getting a nonexistent provider returns None"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.anthropic,  # Not used for embeddings
-            model_name=EmbeddingModelName.openai_text_embedding_3_small,
+        assert provider.name == provider_name
+        assert provider.model_id == provider.model_id
+        assert provider.n_dimensions == provider.n_dimensions
+        assert provider.max_input_tokens == provider.max_input_tokens
+        assert (
+            provider.supports_custom_dimensions == provider.supports_custom_dimensions
         )
-        assert provider is None
 
     def test_get_nonexistent_model_returns_none(self):
         """Test that getting a nonexistent model returns None"""
@@ -353,77 +178,62 @@ class TestBuiltInEmbeddingModelsFromProvider:
         )
         assert provider is None
 
-    def test_get_openai_text_embedding_3_small_provider_details(self):
-        """Test specific details of OpenAI text-embedding-3-small provider"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.openai,
-            model_name=EmbeddingModelName.openai_text_embedding_3_small,
-        )
 
-        assert provider is not None
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-small"
-        assert provider.n_dimensions == 1536
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
+class TestGenerateEmbedding:
+    """Test cases for generate_embedding function"""
 
-    def test_get_openai_text_embedding_3_large_provider_details(self):
-        """Test specific details of OpenAI text-embedding-3-large provider"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.openai,
-            model_name=EmbeddingModelName.openai_text_embedding_3_large,
+    @pytest.mark.parametrize(
+        "model_name,provider_name", get_all_embedding_models_and_providers()
+    )
+    @pytest.mark.paid
+    async def test_generate_embedding(self, model_name, provider_name):
+        """Test generating an embedding"""
+        model_provider = built_in_embedding_models_from_provider(
+            provider_name, model_name
         )
-
-        assert provider is not None
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-large"
-        assert provider.n_dimensions == 3072
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
-
-    def test_get_gemini_text_embedding_004_provider_details(self):
-        """Test specific details of Gemini text-embedding-004 provider"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.gemini_api,
-            model_name=EmbeddingModelName.gemini_text_embedding_004,
+        assert model_provider is not None
+
+        embedding = embedding_adapter_from_type(
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name=provider_name,
+                model_name=model_name,
+                properties={},
+            )
         )
-
-        assert provider is not None
-        assert provider.name == ModelProviderName.gemini_api
-        assert provider.model_id == "text-embedding-004"
-        assert provider.n_dimensions == 768
-        assert provider.max_input_tokens == 2048
-        assert provider.supports_custom_dimensions is False
+        embedding = await embedding.generate_embeddings(["Hello, world!"])
+        assert len(embedding.embeddings) == 1
+        assert len(embedding.embeddings[0].vector) == model_provider.n_dimensions
 
     @pytest.mark.parametrize(
-        "provider_name,model_name,expected_model_id,expected_dimensions",
-        [
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_small,
-                "text-embedding-3-small",
-                1536,
-            ),
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_large,
-                "text-embedding-3-large",
-                3072,
-            ),
-            (
-                ModelProviderName.gemini_api,
-                EmbeddingModelName.gemini_text_embedding_004,
-                "text-embedding-004",
-                768,
-            ),
-        ],
+        "model_name,provider_name", get_all_embedding_models_and_providers()
     )
-    def test_parametrized_provider_retrieval(
-        self, provider_name, model_name, expected_model_id, expected_dimensions
+    @pytest.mark.paid
+    async def test_generate_embedding_with_user_supplied_dimensions(
+        self, model_name, provider_name
     ):
-        """Test retrieving providers with parametrized test cases"""
-        provider = built_in_embedding_models_from_provider(provider_name, model_name)
+        """Test generating an embedding with user supplied dimensions"""
+        model_provider = built_in_embedding_models_from_provider(
+            provider_name=provider_name,
+            model_name=model_name,
+        )
+        assert model_provider is not None
 
-        assert provider is not None
-        assert provider.model_id == expected_model_id
-        assert provider.n_dimensions == expected_dimensions
+        if not model_provider.supports_custom_dimensions:
+            pytest.skip("Model does not support custom dimensions")
+
+        # max dim
+        max_dimensions = model_provider.n_dimensions
+        dimensions_target = max_dimensions // 2
+
+        embedding = embedding_adapter_from_type(
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name=provider_name,
+                model_name=model_name,
+                properties={"dimensions": dimensions_target},
+            )
+        )
+        embedding = await embedding.generate_embeddings(["Hello, world!"])
+        assert len(embedding.embeddings) == 1
+        assert len(embedding.embeddings[0].vector) == dimensions_target

kiln_ai/adapters/test_ml_model_list.py

@@ -360,16 +360,6 @@ def test_uncensored():
         assert provider.suggested_for_uncensored_data_gen
 
 
-def test_multimodal_capable():
-    """Test that multimodal_capable is set correctly"""
-    model = get_model_by_name(ModelName.gpt_4_1)
-    for provider in model.providers:
-        assert provider.multimodal_capable
-        assert provider.supports_doc_extraction
-        assert provider.multimodal_mime_types is not None
-        assert len(provider.multimodal_mime_types) > 0
-
-
 def test_no_empty_multimodal_mime_types():
     """Ensure that multimodal fields are self-consistent as they are interdependent"""
     for model in built_in_models:

kiln_ai/adapters/vector_store/lancedb_adapter.py

@@ -5,12 +5,7 @@ from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Set, TypedDict
 
 from llama_index.core import StorageContext, VectorStoreIndex
-from llama_index.core.schema import (
-    BaseNode,
-    NodeRelationship,
-    RelatedNodeInfo,
-    TextNode,
-)
+from llama_index.core.schema import BaseNode, TextNode
 from llama_index.core.vector_stores.types import (
     VectorStoreQuery as LlamaIndexVectorStoreQuery,
 )
@@ -24,15 +19,19 @@ from kiln_ai.adapters.vector_store.base_vector_store_adapter import (
     SearchResult,
     VectorStoreQuery,
 )
+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.rag import RagConfig
 from kiln_ai.datamodel.vector_store import (
     VectorStoreConfig,
-    VectorStoreType,
     raise_exhaustive_enum_error,
 )
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.env import temporary_env
-from kiln_ai.utils.uuid import string_to_uuid
 
 logger = logging.getLogger(__name__)
 
@@ -48,6 +47,7 @@ class LanceDBAdapter(BaseVectorStoreAdapter):
         self,
         rag_config: RagConfig,
         vector_store_config: VectorStoreConfig,
+        lancedb_vector_store: LanceDBVectorStore | None = None,
     ):
         super().__init__(rag_config, vector_store_config)
         self.config_properties = self.vector_store_config.lancedb_properties
@@ -56,17 +56,15 @@ class LanceDBAdapter(BaseVectorStoreAdapter):
         if vector_store_config.lancedb_properties.nprobes is not None:
             kwargs["nprobes"] = vector_store_config.lancedb_properties.nprobes
 
-        self.lancedb_vector_store = LanceDBVectorStore(
-            mode="create",
-            uri=LanceDBAdapter.lancedb_path_for_config(rag_config),
-            query_type=self.query_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,
-            **kwargs,
+        # allow overriding the vector store with a custom one, useful for user loading into an arbitrary
+        # deployment
+        self.lancedb_vector_store = (
+            lancedb_vector_store
+            or lancedb_construct_from_config(
+                vector_store_config,
+                uri=LanceDBAdapter.lancedb_path_for_config(rag_config),
+            )
         )
-
         self._index = None
 
     @property
@@ -149,7 +147,7 @@ class LanceDBAdapter(BaseVectorStoreAdapter):
 
             chunk_count_for_document = len(chunks)
             deterministic_chunk_ids = [
-                self.compute_deterministic_chunk_id(document_id, chunk_idx)
+                deterministic_chunk_id(document_id, chunk_idx)
                 for chunk_idx in range(chunk_count_for_document)
             ]
 
@@ -176,42 +174,12 @@ class LanceDBAdapter(BaseVectorStoreAdapter):
                 zip(chunks_text, embeddings)
             ):
                 node_batch.append(
-                    TextNode(
-                        id_=deterministic_chunk_ids[chunk_idx],
+                    convert_to_llama_index_node(
+                        document_id=document_id,
+                        chunk_idx=chunk_idx,
+                        node_id=deterministic_chunk_id(document_id, chunk_idx),
                         text=chunk_text,
-                        embedding=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={},
-                            ),
-                        },
+                        vector=embedding.vector,
                     )
                 )
 
@@ -330,10 +298,6 @@ class LanceDBAdapter(BaseVectorStoreAdapter):
                 return []
             raise
 
-    def compute_deterministic_chunk_id(self, 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}"))
-
     async def count_records(self) -> int:
         try:
             table = self.lancedb_vector_store.table
@@ -346,15 +310,7 @@ class LanceDBAdapter(BaseVectorStoreAdapter):
 
     @property
     def query_type(self) -> Literal["fts", "hybrid", "vector"]:
-        match self.vector_store_config.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(self.vector_store_config.store_type)
+        return store_type_to_lancedb_query_type(self.vector_store_config.store_type)
 
     @staticmethod
     def lancedb_path_for_config(rag_config: RagConfig) -> str:
@@ -380,9 +336,7 @@ class LanceDBAdapter(BaseVectorStoreAdapter):
                 kiln_doc_id = row["metadata"]["kiln_doc_id"]
                 if kiln_doc_id not in document_ids:
                     kiln_chunk_idx = row["metadata"]["kiln_chunk_idx"]
-                    record_id = self.compute_deterministic_chunk_id(
-                        kiln_doc_id, kiln_chunk_idx
-                    )
+                    record_id = deterministic_chunk_id(kiln_doc_id, kiln_chunk_idx)
                     rows_to_delete.append(record_id)
 
             if rows_to_delete: