[vibe coding] feat: add mistral embedding function support (#153)

chakkk309 · web-flow · commit 2dfa0e68b89f · 2026-01-31T15:00:11.000+08:00
## Summary Fix #135. Integrate [Mistral](https://docs.morphllm.com/api-reference/endpoint/embedding) embedding function ## Solution Description Mistral test result: ``` uv run pytest tests/unit_tests/test_mistral_embedding_function.py -vv ================================================================================ test session starts ================================================================================ platform darwin -- Python 3.11.13, pytest-9.0.2, pluggy-1.6.0 -- /pyseekdb/.venv/bin/python3 cachedir: .pytest_cache rootdir: /pyseekdb configfile: pyproject.toml plugins: anyio-4.12.1 collected 12 items tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_mistral_env PASSED [ 8%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_initialization_with_defaults PASSED [ 16%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_initialization_with_custom_api_key_env PASSED [ 25%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_initialization_with_custom_api_base PASSED [ 33%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_initialization_with_additional_kwargs PASSED [ 41%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_initialization_with_missing_api_key PASSED [ 50%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_dimension_property_for_known_model PASSED [ 58%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_embedding_generation_single_document PASSED [ 66%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_embedding_generation_multiple_documents PASSED [ 75%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_embedding_with_empty_input PASSED [ 83%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_dimension_of_function PASSED [ 91%] tests/unit_tests/test_mistral_embedding_function.py::TestMistralEmbeddingFunction::test_persistence PASSED [100%] ================================================================================ 12 passed in 3.22s =================================================================================  ## Summary by CodeRabbit * **New Features** * Added support for Mistral text embeddings: a new embedding provider is available and can be configured via environment variable and standard embedding settings; supports batch inputs and serialization. * **Tests** * Added comprehensive unit tests for the Mistral embedding provider (guarded for runtime API availability). <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub>
diff --git a/src/pyseekdb/client/embedding_function.py b/src/pyseekdb/client/embedding_function.py
@@ -634,6 +634,7 @@ def _initialize(cls) -> None:
                 CohereEmbeddingFunction,
                 GoogleVertexEmbeddingFunction,
                 JinaEmbeddingFunction,
+                MistralEmbeddingFunction,
                 MorphEmbeddingFunction,
                 OllamaEmbeddingFunction,
                 OpenAIEmbeddingFunction,
@@ -647,6 +648,7 @@ def _initialize(cls) -> None:
             cls._registry["sentence_transformer"] = SentenceTransformerEmbeddingFunction
             cls._registry["openai"] = OpenAIEmbeddingFunction
             cls._registry["qwen"] = QwenEmbeddingFunction
+            cls._registry["mistral"] = MistralEmbeddingFunction
             cls._registry["morph"] = MorphEmbeddingFunction
             cls._registry["siliconflow"] = SiliconflowEmbeddingFunction
             cls._registry["tencent_hunyuan"] = TencentHunyuanEmbeddingFunction
diff --git a/src/pyseekdb/utils/embedding_functions/__init__.py b/src/pyseekdb/utils/embedding_functions/__init__.py
@@ -10,6 +10,7 @@
 from .google_vertex_embedding_function import GoogleVertexEmbeddingFunction
 from .jina_embedding_function import JinaEmbeddingFunction
 from .litellm_base_embedding_function import LiteLLMBaseEmbeddingFunction
+from .mistral_embedding_function import MistralEmbeddingFunction
 from .morph_embedding_function import MorphEmbeddingFunction
 from .ollama_embedding_function import OllamaEmbeddingFunction
 from .openai_base_embedding_function import OpenAIBaseEmbeddingFunction
@@ -28,6 +29,7 @@
     "GoogleVertexEmbeddingFunction",
     "JinaEmbeddingFunction",
     "LiteLLMBaseEmbeddingFunction",
+    "MistralEmbeddingFunction",
     "MorphEmbeddingFunction",
     "OllamaEmbeddingFunction",
     "OpenAIBaseEmbeddingFunction",
diff --git a/src/pyseekdb/utils/embedding_functions/mistral_embedding_function.py b/src/pyseekdb/utils/embedding_functions/mistral_embedding_function.py
@@ -0,0 +1,153 @@
+import warnings
+from typing import Any
+
+from pyseekdb.client.embedding_function import Documents, Embeddings
+from pyseekdb.utils.embedding_functions.openai_base_embedding_function import (
+    OpenAIBaseEmbeddingFunction,
+)
+
+# Known Mistral embedding model dimensions
+# Source: https://docs.mistral.ai/capabilities/embeddings/text_embeddings
+_MISTRAL_MODEL_DIMENSIONS = {
+    "mistral-embed": 1024,
+}
+
+
+class MistralEmbeddingFunction(OpenAIBaseEmbeddingFunction):
+    """
+    A convenient embedding function for Mistral text embedding models.
+
+    This class provides a simplified interface to Mistral text embeddings using the
+    OpenAI-compatible API.
+
+    Note: The embeddings API only accepts the model name and input texts.
+
+    For more information about Mistral embeddings, see:
+    https://docs.mistral.ai/capabilities/embeddings/text_embeddings
+
+    Example:
+        pip install pyseekdb openai
+
+    .. code-block:: python
+        import pyseekdb
+        from pyseekdb.utils.embedding_functions import MistralEmbeddingFunction
+
+        # Using Mistral text embedding model
+        # Set MISTRAL_API_KEY environment variable first
+        ef = MistralEmbeddingFunction(model_name="mistral-embed")
+
+        # Using with additional parameters
+        ef = MistralEmbeddingFunction(
+            model_name="mistral-embed",
+            timeout=30,
+            max_retries=3
+        )
+
+        db = pyseekdb.Client(path="./seekdb.db")
+        collection = db.create_collection(name="my_collection", embedding_function=ef)
+        # Add documents
+        collection.add(ids=["1", "2"], documents=["Hello world", "How are you?"], metadatas=[{"id": 1}, {"id": 2}])
+        # Query using semantic search
+        results = collection.query("How are you?", n_results=1)
+        print(results)
+
+    """
+
+    def __init__(
+        self,
+        model_name: str = "mistral-embed",
+        api_key_env: str | None = None,
+        api_base: str | None = None,
+        dimensions: int | None = None,
+        **kwargs: Any,
+    ):
+        """Initialize MistralEmbeddingFunction.
+
+        Args:
+            model_name (str, optional): Name of the Mistral embedding model.
+                Defaults to "mistral-embed".
+            api_key_env (str, optional): Name of the environment variable containing the Mistral API key.
+                Defaults to "MISTRAL_API_KEY" if not provided.
+            api_base (str, optional): Base URL for the Mistral API endpoint.
+                Defaults to "https://api.mistral.ai/v1" if not provided.
+            dimensions (int, optional): This parameter is not supported by the Mistral embeddings API.
+                If provided, a warning will be issued and the parameter will be ignored.
+            **kwargs: Additional arguments to pass to the OpenAI client.
+                Common options include:
+                - timeout: Request timeout in seconds
+                - max_retries: Maximum number of retries
+                - See https://github.com/openai/openai-python for more options
+        """
+        if dimensions is not None:
+            warnings.warn(
+                "The dimensions parameter is not supported by Mistral embeddings. "
+                "The provided dimensions parameter will be ignored.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        super().__init__(
+            model_name=model_name,
+            api_key_env=api_key_env,
+            api_base=api_base,
+            dimensions=None,
+            **kwargs,
+        )
+
+    def _get_default_api_base(self) -> str:
+        return "https://api.mistral.ai/v1"
+
+    def _get_default_api_key_env(self) -> str:
+        return "MISTRAL_API_KEY"
+
+    def _get_model_dimensions(self) -> dict[str, int]:
+        return _MISTRAL_MODEL_DIMENSIONS
+
+    @staticmethod
+    def name() -> str:
+        return "mistral"
+
+    def __call__(self, documents: Documents) -> Embeddings:
+        """Generate embeddings for the given documents using Mistral's input parameter."""
+        if isinstance(documents, str):
+            documents = [documents]
+
+        if not documents:
+            return []
+
+        request_params = {
+            "model": self.model_name,
+            "input": documents,
+        }
+
+        response = self._client.embeddings.create(**request_params)
+        embeddings = [item.embedding for item in response.data]
+
+        if len(embeddings) != len(documents):
+            raise ValueError(f"Expected {len(documents)} embeddings but got {len(embeddings)} from API")
+
+        return embeddings
+
+    def get_config(self) -> dict[str, Any]:
+        return super().get_config()
+
+    @staticmethod
+    def build_from_config(config: dict[str, Any]) -> "MistralEmbeddingFunction":
+        model_name = config.get("model_name")
+        if model_name is None:
+            raise ValueError("Missing required field 'model_name' in configuration")
+
+        api_key_env = config.get("api_key_env")
+        api_base = config.get("api_base")
+        dimensions = config.get("dimensions")
+        client_kwargs = config.get("client_kwargs", {})
+        if not isinstance(client_kwargs, dict):
+            raise TypeError(f"client_kwargs must be a dictionary, but got {client_kwargs}")
+
+        return MistralEmbeddingFunction(
+            model_name=model_name,
+            api_key_env=api_key_env,
+            api_base=api_base,
+            dimensions=dimensions,
+            **client_kwargs,
+        )
diff --git a/tests/unit_tests/test_mistral_embedding_function.py b/tests/unit_tests/test_mistral_embedding_function.py
@@ -0,0 +1,188 @@
+"""
+Unit tests for MistralEmbeddingFunction.
+
+Tests Mistral embedding function initialization, embedding generation, and dimension detection.
+Uses real API calls - requires MISTRAL_API_KEY environment variable to be set.
+
+To run this test manually:
+    pytest tests/unit_tests/test_mistral_embedding_function.py -v -s
+    # Or with environment variable:
+    MISTRAL_API_KEY=your-key pytest tests/unit_tests/test_mistral_embedding_function.py -v -s
+"""
+
+import importlib.util
+import os
+
+import pytest
+
+from pyseekdb.client.embedding_function import dimension_of
+from pyseekdb.utils.embedding_functions import MistralEmbeddingFunction
+
+from .test_utils import env_guard
+
+
+def is_openai_available() -> bool:
+    """
+    Check if openai is available for testing.
+
+    Returns:
+        True if openai is available, False otherwise.
+    """
+    return importlib.util.find_spec("openai") is not None
+
+
+# Skip this test by default - it requires external API access and API keys
+@pytest.mark.skipif(
+    not os.environ.get("MISTRAL_API_KEY") or not is_openai_available(),
+    reason="MISTRAL_API_KEY environment variable must be set",
+)
+class TestMistralEmbeddingFunction:
+    """Test MistralEmbeddingFunction - skipped by default, requires manual execution"""
+
+    def test_mistral_env(self):
+        """Test if openai package is installed and required environment variables are set."""
+        if not is_openai_available():
+            print("openai package is not installed")
+            raise AssertionError("openai package is not installed")
+
+        if not os.environ.get("MISTRAL_API_KEY"):
+            print("MISTRAL_API_KEY environment variable is not set")
+            raise AssertionError("MISTRAL_API_KEY environment variable is not set")
+
+    def test_initialization_with_defaults(self):
+        """Test MistralEmbeddingFunction initialization with default values"""
+        print("\nTesting MistralEmbeddingFunction initialization with defaults")
+
+        # Check if openai is available and env vars are set
+        self.test_mistral_env()
+
+        ef = MistralEmbeddingFunction()
+
+        assert ef is not None
+        assert ef.model_name == "mistral-embed"
+        assert ef.api_key_env == "MISTRAL_API_KEY"
+        assert ef.api_base == "https://api.mistral.ai/v1"
+        assert ef._dimensions_param is None
+        print(f"   Model name: {ef.model_name}")
+        print(f"   API key env: {ef.api_key_env}")
+        print(f"   API base: {ef.api_base}")
+
+    def test_initialization_with_custom_api_key_env(self):
+        """Test MistralEmbeddingFunction initialization with custom API key env"""
+        print("\nTesting MistralEmbeddingFunction initialization with custom API key env")
+
+        self.test_mistral_env()
+
+        custom_key_env = "CUSTOM_MISTRAL_KEY"
+        if not os.environ.get(custom_key_env):
+            os.environ[custom_key_env] = "your-custom-key"
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed", api_key_env=custom_key_env)
+        assert ef.api_key_env == custom_key_env
+        print(f"   Custom API key env: {ef.api_key_env}")
+
+    def test_initialization_with_custom_api_base(self):
+        """Test MistralEmbeddingFunction initialization with custom API base"""
+        print("\nTesting MistralEmbeddingFunction initialization with custom API base")
+
+        self.test_mistral_env()
+
+        custom_base = "https://api.mistral.ai/v1"
+        ef = MistralEmbeddingFunction(model_name="mistral-embed", api_base=custom_base)
+        assert ef.api_base == custom_base
+        print(f"   Custom API base: {ef.api_base}")
+
+    def test_initialization_with_additional_kwargs(self):
+        """Test MistralEmbeddingFunction initialization with additional kwargs"""
+        print("\nTesting MistralEmbeddingFunction initialization with kwargs")
+
+        self.test_mistral_env()
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed", timeout=30, max_retries=3)
+        assert ef._client_kwargs.get("timeout") == 30
+        assert ef._client_kwargs.get("max_retries") == 3
+        print(f"   Client kwargs: {ef._client_kwargs}")
+
+    def test_initialization_with_missing_api_key(self):
+        """Test MistralEmbeddingFunction initialization with missing API key"""
+        print("\nTesting MistralEmbeddingFunction initialization with missing API key")
+
+        with env_guard(MISTRAL_API_KEY=None):
+            try:
+                MistralEmbeddingFunction(model_name="mistral-embed")
+                raise AssertionError("Expected ValueError for missing API key")
+            except ValueError as e:
+                print(f"   Caught expected error: {e}")
+
+    def test_dimension_property_for_known_model(self):
+        """Test MistralEmbeddingFunction dimension property for known model"""
+        print("\nTesting MistralEmbeddingFunction dimension property for known model")
+
+        self.test_mistral_env()
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed")
+        assert ef.dimension == 1024
+        print(f"   Model mistral-embed dimension: {ef.dimension}")
+
+    def test_embedding_generation_single_document(self):
+        """Test MistralEmbeddingFunction embedding generation (single document)"""
+        print("\nTesting MistralEmbeddingFunction embedding generation (single document)")
+
+        self.test_mistral_env()
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed")
+        embeddings = ef("Hello world")
+        assert len(embeddings) == 1
+        assert len(embeddings[0]) > 0
+        print(f"   Embedding dimension: {len(embeddings[0])}")
+
+    def test_embedding_generation_multiple_documents(self):
+        """Test MistralEmbeddingFunction embedding generation (multiple documents)"""
+        print("\nTesting MistralEmbeddingFunction embedding generation (multiple documents)")
+
+        self.test_mistral_env()
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed")
+        embeddings = ef(["Hello world", "How are you?"])
+        assert len(embeddings) == 2
+        assert len(embeddings[0]) > 0
+        print(f"   Embedding dimension: {len(embeddings[0])}")
+
+    def test_embedding_with_empty_input(self):
+        """Test MistralEmbeddingFunction with empty input"""
+        print("\nTesting MistralEmbeddingFunction with empty input")
+
+        self.test_mistral_env()
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed")
+        embeddings = ef([])
+        assert embeddings == []
+        print("   Empty input returned empty embeddings")
+
+    def test_dimension_of_function(self):
+        """Test dimension_of function with MistralEmbeddingFunction"""
+        print("\nTesting dimension_of function with MistralEmbeddingFunction")
+
+        self.test_mistral_env()
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed")
+        dim = dimension_of(ef)
+        assert dim == 1024
+        print(f"   dimension_of returned: {dim}")
+
+    def test_persistence(self):
+        """Test persistence for MistralEmbeddingFunction"""
+        print("\nTesting MistralEmbeddingFunction persistence")
+
+        self.test_mistral_env()
+
+        assert MistralEmbeddingFunction.name() == "mistral"
+
+        ef = MistralEmbeddingFunction(model_name="mistral-embed")
+        config = ef.get_config()
+
+        restored_ef = MistralEmbeddingFunction.build_from_config(config)
+        assert isinstance(restored_ef, MistralEmbeddingFunction)
+        assert restored_ef.model_name == ef.model_name
+        assert restored_ef.api_key_env == ef.api_key_env
+        assert restored_ef.api_base == ef.api_base
