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

kiln_ai/adapters/embedding/litellm_embedding_adapter.py

@@ -0,0 +1,199 @@
+from functools import cached_property
+from typing import Any, Dict, List, Tuple
+
+import litellm
+from litellm.types.utils import EmbeddingResponse
+from pydantic import BaseModel, Field
+
+from kiln_ai.adapters.embedding.base_embedding_adapter import (
+    BaseEmbeddingAdapter,
+    Embedding,
+    EmbeddingResult,
+)
+from kiln_ai.adapters.ml_embedding_model_list import (
+    KilnEmbeddingModelProvider,
+    built_in_embedding_models_from_provider,
+)
+from kiln_ai.adapters.provider_tools import LiteLlmCoreConfig
+from kiln_ai.datamodel.embedding import EmbeddingConfig
+from kiln_ai.utils.litellm import get_litellm_provider_info
+
+# litellm enforces a limit, documented here:
+# https://docs.litellm.ai/docs/embedding/supported_embedding
+# but some providers impose lower limits that LiteLLM does not know about
+# for example, Gemini currently has a limit of 100 inputs per request
+MAX_BATCH_SIZE = 100
+
+
+class EmbeddingOptions(BaseModel):
+    dimensions: int | None = Field(
+        default=None,
+        description="The number of dimensions to return for embeddings. Some models support requesting vectors of different dimensions.",
+    )
+
+
+def validate_map_to_embeddings(
+    response: EmbeddingResponse,
+    expected_embedding_count: int,
+) -> List[Embedding]:
+    # LiteLLM has an Embedding type in litellm.types.utils, but the EmbeddingResponse data has a list of untyped dicts,
+    # which can be dangerous especially if we upgrade litellm, so we do some sanity checks here
+    if not isinstance(response, EmbeddingResponse):
+        raise RuntimeError(f"Expected EmbeddingResponse, got {type(response)}.")
+
+    list_to_validate = response.data
+    if len(list_to_validate) != expected_embedding_count:
+        raise RuntimeError(
+            f"Expected the number of embeddings in the response to be {expected_embedding_count}, got {len(list_to_validate)}."
+        )
+
+    validated_vectors: List[Tuple[list[float], int]] = []
+    for embedding_dict in list_to_validate:
+        object_type = embedding_dict.get("object")
+        if object_type != "embedding":
+            raise RuntimeError(
+                f"Embedding response data has an unexpected shape. Property 'object' is not 'embedding'. Got {object_type}."
+            )
+
+        embedding_property_value = embedding_dict.get("embedding")
+        if embedding_property_value is None:
+            raise RuntimeError(
+                "Embedding response data has an unexpected shape. Property 'embedding' is None in response data item."
+            )
+        if not isinstance(embedding_property_value, list):
+            raise RuntimeError(
+                f"Embedding response data has an unexpected shape. Property 'embedding' is not a list. Got {type(embedding_property_value)}."
+            )
+
+        index_property_value = embedding_dict.get("index")
+        if index_property_value is None:
+            raise RuntimeError(
+                "Embedding response data has an unexpected shape. Property 'index' is None in response data item."
+            )
+        if not isinstance(index_property_value, int):
+            raise RuntimeError(
+                f"Embedding response data has an unexpected shape. Property 'index' is not an integer. Got {type(index_property_value)}."
+            )
+
+        validated_vectors.append((embedding_property_value, index_property_value))
+
+    # sort by index, in place - the data should already be sorted by index,
+    # but litellm docs are not explicit about this
+    validated_vectors.sort(key=lambda x: x[1])
+
+    return [
+        Embedding(vector=embedding_vector) for embedding_vector, _ in validated_vectors
+    ]
+
+
+class LitellmEmbeddingAdapter(BaseEmbeddingAdapter):
+    def __init__(
+        self, embedding_config: EmbeddingConfig, litellm_core_config: LiteLlmCoreConfig
+    ):
+        super().__init__(embedding_config)
+
+        self.litellm_core_config = litellm_core_config
+
+    async def _generate_embeddings(self, input_texts: List[str]) -> EmbeddingResult:
+        # batch the requests
+        batches: List[List[str]] = []
+        for i in range(0, len(input_texts), MAX_BATCH_SIZE):
+            batches.append(input_texts[i : i + MAX_BATCH_SIZE])
+
+        # generate embeddings for each batch
+        results: List[EmbeddingResult] = []
+        for batch in batches:
+            batch_response = await self._generate_embeddings_for_batch(batch)
+            results.append(batch_response)
+
+        # merge the results
+        combined_embeddings: List[Embedding] = []
+        combined_usage = None
+
+        # we prefer returning None overall usage if any of the results is missing usage
+        # better than returning a misleading usage
+        all_have_usage = all(result.usage is not None for result in results)
+        if all_have_usage:
+            combined_usage = litellm.Usage(
+                prompt_tokens=0, total_tokens=0, completion_tokens=0
+            )
+            for result in results:
+                if result.usage is not None:
+                    combined_usage.prompt_tokens += result.usage.prompt_tokens
+                    combined_usage.total_tokens += result.usage.total_tokens
+                    combined_usage.completion_tokens += result.usage.completion_tokens
+
+        for result in results:
+            combined_embeddings.extend(result.embeddings)
+
+        return EmbeddingResult(
+            embeddings=combined_embeddings,
+            usage=combined_usage,
+        )
+
+    async def _generate_embeddings_for_batch(
+        self, input_texts: List[str]
+    ) -> EmbeddingResult:
+        if len(input_texts) > MAX_BATCH_SIZE:
+            raise ValueError(
+                f"Too many input texts, max batch size is {MAX_BATCH_SIZE}, got {len(input_texts)}"
+            )
+
+        completion_kwargs: Dict[str, Any] = {}
+        if self.litellm_core_config.additional_body_options:
+            completion_kwargs.update(self.litellm_core_config.additional_body_options)
+
+        if self.litellm_core_config.base_url:
+            completion_kwargs["api_base"] = self.litellm_core_config.base_url
+
+        if self.litellm_core_config.default_headers:
+            completion_kwargs["default_headers"] = (
+                self.litellm_core_config.default_headers
+            )
+
+        response = await litellm.aembedding(
+            model=self.litellm_model_id,
+            input=input_texts,
+            **self.build_options().model_dump(exclude_none=True),
+            **completion_kwargs,
+        )
+
+        validated_embeddings = validate_map_to_embeddings(
+            response, expected_embedding_count=len(input_texts)
+        )
+
+        return EmbeddingResult(
+            embeddings=validated_embeddings,
+            usage=response.usage,
+        )
+
+    def build_options(self) -> EmbeddingOptions:
+        dimensions = self.embedding_config.properties.get("dimensions", None)
+        if dimensions is not None:
+            if not isinstance(dimensions, int) or dimensions <= 0:
+                raise ValueError("Dimensions must be a positive integer")
+
+        return EmbeddingOptions(
+            dimensions=dimensions,
+        )
+
+    @cached_property
+    def model_provider(self) -> KilnEmbeddingModelProvider:
+        provider = built_in_embedding_models_from_provider(
+            self.embedding_config.model_provider_name, self.embedding_config.model_name
+        )
+        if provider is None:
+            raise ValueError(
+                f"Embedding model {self.embedding_config.model_name} not found in the list of built-in models"
+            )
+        return provider
+
+    @cached_property
+    def litellm_model_id(self) -> str:
+        provider_info = get_litellm_provider_info(self.model_provider)
+        if provider_info.is_custom and self.litellm_core_config.base_url is None:
+            raise ValueError(
+                f"Provider {self.model_provider.name.value} must have an explicit base URL"
+            )
+
+        return provider_info.litellm_model_id

kiln_ai/adapters/embedding/test_base_embedding_adapter.py

@@ -0,0 +1,283 @@
+import pytest
+
+from kiln_ai.adapters.embedding.base_embedding_adapter import (
+    BaseEmbeddingAdapter,
+    Embedding,
+    EmbeddingResult,
+)
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+from kiln_ai.datamodel.embedding import EmbeddingConfig
+
+
+class MockEmbeddingAdapter(BaseEmbeddingAdapter):
+    """Concrete implementation of BaseEmbeddingAdapter for testing purposes."""
+
+    async def _generate_embeddings(self, text_inputs: list[str]) -> EmbeddingResult:
+        # Simple test implementation that returns mock embeddings
+        embeddings = []
+        for i, _ in enumerate(text_inputs):
+            embeddings.append(Embedding(vector=[0.1 * (i + 1)] * 3))
+
+        return EmbeddingResult(embeddings=embeddings)
+
+
+class MockEmbeddingAdapterWithUsage(BaseEmbeddingAdapter):
+    """Concrete implementation that includes usage information."""
+
+    async def _generate_embeddings(self, text_inputs: list[str]) -> EmbeddingResult:
+        from litellm import Usage
+
+        embeddings = []
+        for i, _ in enumerate(text_inputs):
+            embeddings.append(Embedding(vector=[0.1 * (i + 1)] * 3))
+
+        usage = Usage(
+            prompt_tokens=len(text_inputs) * 10, total_tokens=len(text_inputs) * 10
+        )
+        return EmbeddingResult(embeddings=embeddings, usage=usage)
+
+
+@pytest.fixture
+def mock_embedding_config():
+    """Create a mock embedding config for testing."""
+    return EmbeddingConfig(
+        name="test-embedding",
+        model_provider_name=ModelProviderName.openai,
+        model_name="openai_text_embedding_3_small",
+        properties={},
+    )
+
+
+@pytest.fixture
+def test_adapter(mock_embedding_config):
+    """Create a test adapter instance."""
+    return MockEmbeddingAdapter(mock_embedding_config)
+
+
+@pytest.fixture
+def test_adapter_with_usage(mock_embedding_config):
+    """Create a test adapter instance that includes usage information."""
+    return MockEmbeddingAdapterWithUsage(mock_embedding_config)
+
+
+class TestEmbedding:
+    """Test the Embedding model."""
+
+    def test_creation(self):
+        """Test creating a Embedding with a vector."""
+        vector = [0.1, 0.2, 0.3]
+        embedding = Embedding(vector=vector)
+        assert embedding.vector == vector
+
+    def test_empty_vector(self):
+        """Test creating a Embedding with an empty vector."""
+        embedding = Embedding(vector=[])
+        assert embedding.vector == []
+
+    def test_large_vector(self):
+        """Test creating a Embedding with a large vector."""
+        vector = [0.1] * 1536
+        embedding = Embedding(vector=vector)
+        assert len(embedding.vector) == 1536
+        assert all(v == 0.1 for v in embedding.vector)
+
+
+class TestEmbeddingResult:
+    """Test the EmbeddingResult model."""
+
+    def test_creation_with_embeddings(self):
+        """Test creating an EmbeddingResult with embeddings."""
+        embeddings = [
+            Embedding(vector=[0.1, 0.2, 0.3]),
+            Embedding(vector=[0.4, 0.5, 0.6]),
+        ]
+        result = EmbeddingResult(embeddings=embeddings)
+        assert result.embeddings == embeddings
+        assert result.usage is None
+
+    def test_creation_with_usage(self):
+        """Test creating an EmbeddingResult with usage information."""
+        from litellm import Usage
+
+        embeddings = [Embedding(vector=[0.1, 0.2, 0.3])]
+        usage = Usage(prompt_tokens=10, total_tokens=10)
+        result = EmbeddingResult(embeddings=embeddings, usage=usage)
+        assert result.embeddings == embeddings
+        assert result.usage == usage
+
+    def test_empty_embeddings(self):
+        """Test creating an EmbeddingResult with empty embeddings."""
+        result = EmbeddingResult(embeddings=[])
+        assert result.embeddings == []
+        assert result.usage is None
+
+
+class TestBaseEmbeddingAdapter:
+    """Test the BaseEmbeddingAdapter abstract base class."""
+
+    def test_init(self, mock_embedding_config):
+        """Test successful initialization of the adapter."""
+        adapter = MockEmbeddingAdapter(mock_embedding_config)
+        assert adapter.embedding_config == mock_embedding_config
+        assert adapter.embedding_config.name == "test-embedding"
+        assert adapter.embedding_config.model_provider_name == ModelProviderName.openai
+        assert adapter.embedding_config.model_name == "openai_text_embedding_3_small"
+        assert adapter.embedding_config.properties == {}
+
+    def test_cannot_instantiate_abstract_class(self, mock_embedding_config):
+        """Test that BaseEmbeddingAdapter cannot be instantiated directly."""
+        with pytest.raises(TypeError):
+            BaseEmbeddingAdapter(mock_embedding_config)
+
+    async def test_generate_embeddings_empty_list(self, test_adapter):
+        """Test embed method with empty text list."""
+        result = await test_adapter.generate_embeddings([])
+        assert result.embeddings == []
+        assert result.usage is None
+
+    async def test_generate_embeddings_single_text(self, test_adapter):
+        """Test embed method with a single text."""
+        result = await test_adapter.generate_embeddings(["hello world"])
+        assert len(result.embeddings) == 1
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+        assert result.usage is None
+
+    async def test_generate_embeddings_multiple_texts(self, test_adapter):
+        """Test embed method with multiple texts."""
+        texts = ["hello world", "my name is john", "i like to eat apples"]
+        result = await test_adapter.generate_embeddings(texts)
+        assert len(result.embeddings) == 3
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+        assert result.embeddings[1].vector == [0.2, 0.2, 0.2]
+        assert result.embeddings[2].vector == pytest.approx([0.3, 0.3, 0.3])
+        assert result.usage is None
+
+    async def test_generate_embeddings_with_usage(self, test_adapter_with_usage):
+        """Test embed method with usage information."""
+        texts = ["hello", "world"]
+        result = await test_adapter_with_usage.generate_embeddings(texts)
+        assert len(result.embeddings) == 2
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+        assert result.embeddings[1].vector == [0.2, 0.2, 0.2]
+        assert result.usage is not None
+        assert result.usage.prompt_tokens == 20
+        assert result.usage.total_tokens == 20
+
+    async def test_generate_embeddings_with_none_input(self, test_adapter):
+        """Test embed method with None input (should be treated as empty list)."""
+        result = await test_adapter.generate_embeddings(None)  # type: ignore
+        assert result.embeddings == []
+        assert result.usage is None
+
+    async def test_generate_embeddings_with_whitespace_only_texts(self, test_adapter):
+        """Test embed method with texts containing only whitespace."""
+        texts = ["   ", "\n", "\t"]
+        result = await test_adapter.generate_embeddings(texts)
+        assert len(result.embeddings) == 3
+        # Should still generate embeddings for whitespace-only texts
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+        assert result.embeddings[1].vector == [0.2, 0.2, 0.2]
+        assert result.embeddings[2].vector == pytest.approx([0.3, 0.3, 0.3])
+
+    async def test_generate_embeddings_with_duplicate_texts(self, test_adapter):
+        """Test embed method with duplicate texts."""
+        texts = ["hello", "hello", "world"]
+        result = await test_adapter.generate_embeddings(texts)
+        assert len(result.embeddings) == 3
+        # Each text should get its own embedding, even if duplicate
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+        assert result.embeddings[1].vector == [0.2, 0.2, 0.2]
+        assert result.embeddings[2].vector == pytest.approx([0.3, 0.3, 0.3])
+
+
+class TestBaseEmbeddingAdapterEdgeCases:
+    """Test edge cases for BaseEmbeddingAdapter."""
+
+    async def test_generate_embeddings_with_very_long_text(self, test_adapter):
+        """Test embed method with very long text."""
+        long_text = "a" * 10000
+        result = await test_adapter.generate_embeddings([long_text])
+        assert len(result.embeddings) == 1
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+
+    async def test_generate_embeddings_with_special_characters(self, test_adapter):
+        """Test embed method with special characters."""
+        texts = ["hello\nworld", "test\twith\ttabs", "unicode: 🚀🌟"]
+        result = await test_adapter.generate_embeddings(texts)
+        assert len(result.embeddings) == 3
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+        assert result.embeddings[1].vector == [0.2, 0.2, 0.2]
+        assert result.embeddings[2].vector == pytest.approx([0.3, 0.3, 0.3])
+
+    async def test_generate_embeddings_with_empty_strings(self, test_adapter):
+        """Test embed method with empty strings."""
+        texts = ["", "", ""]
+        result = await test_adapter.generate_embeddings(texts)
+        assert len(result.embeddings) == 3
+        assert result.embeddings[0].vector == [0.1, 0.1, 0.1]
+        assert result.embeddings[1].vector == [0.2, 0.2, 0.2]
+        assert result.embeddings[2].vector == pytest.approx([0.3, 0.3, 0.3])
+
+    def test_embedding_config_properties(self, mock_embedding_config):
+        """Test that embedding config properties are accessible."""
+        mock_embedding_config.properties = {"dimensions": 1536, "normalize": True}
+        adapter = MockEmbeddingAdapter(mock_embedding_config)
+        assert adapter.embedding_config.properties["dimensions"] == 1536
+        assert adapter.embedding_config.properties["normalize"] is True
+
+
+class TestBaseEmbeddingAdapterIntegration:
+    """Integration tests for BaseEmbeddingAdapter."""
+
+    async def test_generate_embeddings_method_calls_abstract_method(
+        self, mock_embedding_config
+    ):
+        """Test that embed method properly calls the abstract _embed method."""
+
+        # Create a mock adapter that tracks if _embed was called
+        class MockAdapter(BaseEmbeddingAdapter):
+            def __init__(self, config):
+                super().__init__(config)
+                self._generate_embeddings_called = False
+                self._generate_embeddings_args = None
+
+            async def _generate_embeddings(
+                self, text_inputs: list[str]
+            ) -> EmbeddingResult:
+                self._generate_embeddings_called = True
+                self._generate_embeddings_args = text_inputs
+                return EmbeddingResult(embeddings=[])
+
+        adapter = MockAdapter(mock_embedding_config)
+        texts = ["hello", "world"]
+
+        result = await adapter.generate_embeddings(texts)
+
+        assert adapter._generate_embeddings_called
+        assert adapter._generate_embeddings_args == texts
+        assert result.embeddings == []
+        assert result.usage is None
+
+    async def test_generate_embeddings_empty_list_does_not_call_abstract_method(
+        self, mock_embedding_config
+    ):
+        """Test that embed method with empty list does not call _embed."""
+
+        class MockAdapter(BaseEmbeddingAdapter):
+            def __init__(self, config):
+                super().__init__(config)
+                self._generate_embeddings_called = False
+
+            async def _generate_embeddings(
+                self, text_inputs: list[str]
+            ) -> EmbeddingResult:
+                self._generate_embeddings_called = True
+                return EmbeddingResult(embeddings=[])
+
+        adapter = MockAdapter(mock_embedding_config)
+
+        result = await adapter.generate_embeddings([])
+
+        assert not adapter._generate_embeddings_called
+        assert result.embeddings == []
+        assert result.usage is None

kiln_ai/adapters/embedding/test_embedding_registry.py

@@ -0,0 +1,166 @@
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai.adapters.embedding.embedding_registry import embedding_adapter_from_type
+from kiln_ai.adapters.embedding.litellm_embedding_adapter import LitellmEmbeddingAdapter
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.provider_tools import LiteLlmCoreConfig
+from kiln_ai.datamodel.embedding import EmbeddingConfig
+
+
+@pytest.fixture
+def mock_provider_configs():
+    with patch("kiln_ai.utils.config.Config.shared") as mock_config:
+        mock_config.return_value.open_ai_api_key = "test-openai-key"
+        mock_config.return_value.gemini_api_key = "test-gemini-key"
+        yield mock_config
+
+
+def test_embedding_adapter_from_type(mock_provider_configs):
+    """Test basic embedding adapter creation with valid config."""
+    embedding_config = EmbeddingConfig(
+        name="test-embedding",
+        model_provider_name=ModelProviderName.gemini_api,
+        model_name="text-embedding-003",
+        properties={"dimensions": 768},
+    )
+
+    adapter = embedding_adapter_from_type(embedding_config)
+
+    assert isinstance(adapter, LitellmEmbeddingAdapter)
+    assert adapter.embedding_config.model_name == "text-embedding-003"
+    assert adapter.embedding_config.model_provider_name == ModelProviderName.gemini_api
+
+
+@patch(
+    "kiln_ai.adapters.embedding.embedding_registry.lite_llm_core_config_for_provider"
+)
+def test_embedding_adapter_from_type_uses_litellm_core_config(
+    mock_get_litellm_core_config,
+):
+    """Test that embedding adapter receives auth details from provider_tools."""
+    mock_litellm_core_config = LiteLlmCoreConfig(
+        base_url="https://test.com",
+        additional_body_options={"api_key": "test-key"},
+        default_headers={},
+    )
+    mock_get_litellm_core_config.return_value = mock_litellm_core_config
+
+    embedding_config = EmbeddingConfig(
+        name="test-embedding",
+        model_provider_name=ModelProviderName.openai,
+        model_name="text-embedding-3-small",
+        properties={"dimensions": 1536},
+    )
+
+    adapter = embedding_adapter_from_type(embedding_config)
+
+    assert isinstance(adapter, LitellmEmbeddingAdapter)
+    assert adapter.litellm_core_config == mock_litellm_core_config
+    mock_get_litellm_core_config.assert_called_once()
+
+
+def test_embedding_adapter_from_type_invalid_provider():
+    """Test that invalid model provider names raise a clear error."""
+    # Create a valid config first, then test the enum conversion logic
+    embedding_config = EmbeddingConfig(
+        name="test-embedding",
+        model_provider_name=ModelProviderName.openai,
+        model_name="some-model",
+        properties={"dimensions": 768},
+    )
+
+    # Mock the ModelProviderName constructor to simulate an invalid provider
+    with patch(
+        "kiln_ai.adapters.embedding.embedding_registry.ModelProviderName"
+    ) as mock_enum:
+        mock_enum.side_effect = ValueError("Invalid provider")
+
+        with pytest.raises(
+            ValueError,
+            match="Unsupported model provider name: openai",
+        ):
+            embedding_adapter_from_type(embedding_config)
+
+
+def test_embedding_adapter_from_type_no_config_found(mock_provider_configs):
+    """Test that missing provider configuration raises an error."""
+    with patch(
+        "kiln_ai.adapters.embedding.embedding_registry.lite_llm_core_config_for_provider"
+    ) as mock_lite_llm_core_config_for_provider:
+        mock_lite_llm_core_config_for_provider.return_value = None
+
+        embedding_config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name=ModelProviderName.openai,
+            model_name="text-embedding-3-small",
+            properties={"dimensions": 1536},
+        )
+
+        with pytest.raises(
+            ValueError, match="No configuration found for core provider:"
+        ):
+            embedding_adapter_from_type(embedding_config)
+
+
+@pytest.mark.parametrize(
+    "provider_name",
+    [
+        ModelProviderName.openai,
+        ModelProviderName.gemini_api,
+    ],
+)
+def test_embedding_adapter_from_type_different_providers(
+    provider_name, mock_provider_configs
+):
+    """Test that different providers work correctly."""
+    embedding_config = EmbeddingConfig(
+        name="test-embedding",
+        model_provider_name=provider_name,
+        model_name="test-model",
+        properties={"dimensions": 768},
+    )
+
+    adapter = embedding_adapter_from_type(embedding_config)
+
+    assert isinstance(adapter, LitellmEmbeddingAdapter)
+    assert adapter.embedding_config.model_provider_name == provider_name
+
+
+def test_embedding_adapter_from_type_with_description(mock_provider_configs):
+    """Test embedding adapter creation with description."""
+    embedding_config = EmbeddingConfig(
+        name="test-embedding",
+        description="Test embedding configuration",
+        model_provider_name=ModelProviderName.openai,
+        model_name="text-embedding-3-small",
+        properties={"dimensions": 1536},
+    )
+
+    adapter = embedding_adapter_from_type(embedding_config)
+
+    assert isinstance(adapter, LitellmEmbeddingAdapter)
+    assert adapter.embedding_config.description == "Test embedding configuration"
+
+
+def test_embedding_adapter_from_type_with_additional_properties(
+    mock_provider_configs,
+):
+    """Test embedding adapter creation with additional properties."""
+    embedding_config = EmbeddingConfig(
+        name="test-embedding",
+        model_provider_name=ModelProviderName.openai,
+        model_name="text-embedding-3-small",
+        properties={
+            "dimensions": 1536,
+            "batch_size": 100,
+            "max_retries": 3,
+        },
+    )
+
+    adapter = embedding_adapter_from_type(embedding_config)
+
+    assert isinstance(adapter, LitellmEmbeddingAdapter)
+    assert adapter.embedding_config.properties["batch_size"] == 100
+    assert adapter.embedding_config.properties["max_retries"] == 3