diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index ebd93e3dbd..5a7318c3bf 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -1641,3 +1641,30 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str: if not model.startswith("gemini/"): model = f"gemini/{model}" return model + + +class MistralEmbeddingParameters(BaseEmbeddingParameters): + """See https://docs.litellm.ai/docs/providers/mistral.""" + + api_key: str + embed_batch_size: int | None = None + + @staticmethod + def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: + metadata_copy = {**adapter_metadata} + metadata_copy["model"] = MistralEmbeddingParameters.validate_model(metadata_copy) + + return MistralEmbeddingParameters(**metadata_copy).model_dump() + + @staticmethod + def validate_model(adapter_metadata: dict[str, "Any"]) -> str: + raw_model = adapter_metadata.get("model") + model = raw_model.strip() if isinstance(raw_model, str) else "" + if not model: + raise ValueError( + "The 'model' field is required for the Mistral embedding adapter. " + "Example: 'mistral-embed'" + ) + if not model.startswith("mistral/"): + model = f"mistral/{model}" + return model diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/__init__.py b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/__init__.py index 18f240f04e..3be59f61c3 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/__init__.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/__init__.py @@ -4,6 +4,7 @@ from unstract.sdk1.adapters.embedding1.azure_openai import AzureOpenAIEmbeddingAdapter from unstract.sdk1.adapters.embedding1.bedrock import AWSBedrockEmbeddingAdapter from unstract.sdk1.adapters.embedding1.gemini import GeminiEmbeddingAdapter +from unstract.sdk1.adapters.embedding1.mistral import MistralEmbeddingAdapter from unstract.sdk1.adapters.embedding1.nvidia_build import NvidiaBuildEmbeddingAdapter from unstract.sdk1.adapters.embedding1.ollama import OllamaEmbeddingAdapter from unstract.sdk1.adapters.embedding1.openai import OpenAIEmbeddingAdapter @@ -22,6 +23,7 @@ "AzureOpenAIEmbeddingAdapter", "AWSBedrockEmbeddingAdapter", "GeminiEmbeddingAdapter", + "MistralEmbeddingAdapter", "NvidiaBuildEmbeddingAdapter", "OpenAIEmbeddingAdapter", "OpenAICompatibleEmbeddingAdapter", diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/mistral.py b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/mistral.py new file mode 100644 index 0000000000..3440c1f498 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/mistral.py @@ -0,0 +1,40 @@ +from typing import Any + +from unstract.sdk1.adapters.base1 import BaseAdapter, MistralEmbeddingParameters +from unstract.sdk1.adapters.enums import AdapterTypes + + +class MistralEmbeddingAdapter(MistralEmbeddingParameters, BaseAdapter): + @staticmethod + def get_id() -> str: + return "mistral|28a9d784-bc33-46a4-804a-436058e1354f" + + @staticmethod + def get_metadata() -> dict[str, Any]: + return { + "name": "Mistral", + "version": "1.0.0", + "adapter": MistralEmbeddingAdapter, + "description": "Mistral embedding adapter", + "is_active": True, + } + + @staticmethod + def get_name() -> str: + return "Mistral" + + @staticmethod + def get_description() -> str: + return "Mistral embedding adapter" + + @staticmethod + def get_provider() -> str: + return "mistral" + + @staticmethod + def get_icon() -> str: + return "/icons/adapter-icons/Mistral%20AI.png" + + @staticmethod + def get_adapter_type() -> AdapterTypes: + return AdapterTypes.EMBEDDING diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/mistral.json b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/mistral.json new file mode 100644 index 0000000000..9fab2fac89 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/mistral.json @@ -0,0 +1,38 @@ +{ + "title": "Mistral Embedding", + "type": "object", + "required": [ + "adapter_name", + "api_key", + "model" + ], + "properties": { + "adapter_name": { + "type": "string", + "title": "Name", + "default": "", + "description": "Provide a unique name for this adapter instance. Example: mistral-emb-1" + }, + "model": { + "type": "string", + "title": "Model", + "default": "mistral-embed", + "description": "Provide the name of the model. The mistral/ prefix will be added automatically if omitted. Recommended: mistral-embed." + }, + "api_key": { + "type": "string", + "title": "API Key", + "default": "", + "format": "password", + "description": "API Key for Mistral embeddings" + }, + "timeout": { + "type": "number", + "minimum": 0, + "multipleOf": 1, + "title": "Timeout", + "default": 240, + "description": "Timeout in seconds" + } + } +} diff --git a/unstract/sdk1/tests/test_mistral_embedding.py b/unstract/sdk1/tests/test_mistral_embedding.py new file mode 100644 index 0000000000..6569d0ec78 --- /dev/null +++ b/unstract/sdk1/tests/test_mistral_embedding.py @@ -0,0 +1,139 @@ +import json + +import pytest +from unstract.sdk1.adapters.embedding1.mistral import MistralEmbeddingAdapter +from unstract.sdk1.adapters.enums import AdapterTypes + + +class TestMistralEmbeddingAdapter: + def test_adapter_registration(self) -> None: + from unstract.sdk1.adapters.embedding1 import adapters + + mistral_ids = [k for k in adapters if "mistral" in k.lower()] + assert len(mistral_ids) == 1 + + def test_get_id_format(self) -> None: + adapter_id = MistralEmbeddingAdapter.get_id() + assert adapter_id.startswith("mistral|") + # Standard UUID-4 with hyphens is 36 characters + uuid_part = adapter_id.split("|")[1] + assert len(uuid_part) == 36 + + def test_get_adapter_type(self) -> None: + assert MistralEmbeddingAdapter.get_adapter_type() == AdapterTypes.EMBEDDING + + def test_get_name(self) -> None: + assert MistralEmbeddingAdapter.get_name() == "Mistral" + + def test_get_provider(self) -> None: + assert MistralEmbeddingAdapter.get_provider() == "mistral" + + def test_json_schema_loads(self) -> None: + schema = json.loads(MistralEmbeddingAdapter.get_json_schema()) + assert isinstance(schema, dict) + assert "title" in schema + assert "properties" in schema + assert schema["title"] == "Mistral Embedding" + + def test_json_schema_required_fields(self) -> None: + schema = json.loads(MistralEmbeddingAdapter.get_json_schema()) + assert set(schema["required"]) == {"adapter_name", "api_key", "model"} + + def test_json_schema_omits_batch_size(self) -> None: + # embed_batch_size is an inert client-side hint and is not exposed. + schema = json.loads(MistralEmbeddingAdapter.get_json_schema()) + assert "embed_batch_size" not in schema["properties"] + + def test_json_schema_api_key_password_format(self) -> None: + schema = json.loads(MistralEmbeddingAdapter.get_json_schema()) + assert schema["properties"]["api_key"]["format"] == "password" + + def test_json_schema_model_default(self) -> None: + schema = json.loads(MistralEmbeddingAdapter.get_json_schema()) + assert schema["properties"]["model"]["default"] == "mistral-embed" + + def test_validate_model_adds_prefix(self) -> None: + meta = {"model": "mistral-embed", "api_key": "test"} + result = MistralEmbeddingAdapter.validate_model(meta) + assert result == "mistral/mistral-embed" + + def test_validate_model_idempotent(self) -> None: + meta = {"model": "mistral/mistral-embed", "api_key": "test"} + result = MistralEmbeddingAdapter.validate_model(meta) + assert result == "mistral/mistral-embed" + + def test_validate_model_does_not_mutate_input(self) -> None: + meta = {"model": "mistral-embed", "api_key": "test"} + MistralEmbeddingAdapter.validate_model(meta) + assert meta["model"] == "mistral-embed" + + def test_validate_does_not_mutate_input(self) -> None: + meta = {"model": "mistral-embed", "api_key": "test-key"} + original_model = meta["model"] + MistralEmbeddingAdapter.validate(meta) + assert meta["model"] == original_model + + def test_validate_model_empty_string_raises(self) -> None: + meta = {"model": "", "api_key": "test"} + with pytest.raises(ValueError, match="model.*required"): + MistralEmbeddingAdapter.validate_model(meta) + + def test_validate_model_whitespace_only_raises(self) -> None: + meta = {"model": " ", "api_key": "test"} + with pytest.raises(ValueError, match="model.*required"): + MistralEmbeddingAdapter.validate_model(meta) + + def test_validate_model_none_raises(self) -> None: + meta = {"model": None, "api_key": "test"} + with pytest.raises(ValueError, match="model.*required"): + MistralEmbeddingAdapter.validate_model(meta) + + def test_validate_model_missing_key_raises(self) -> None: + meta = {"api_key": "test"} + with pytest.raises(ValueError, match="model.*required"): + MistralEmbeddingAdapter.validate_model(meta) + + def test_validate_empty_model_raises(self) -> None: + meta = {"model": "", "api_key": "test-key"} + with pytest.raises(ValueError, match="model.*required"): + MistralEmbeddingAdapter.validate(meta) + + def test_validate_none_model_raises(self) -> None: + meta = {"model": None, "api_key": "test-key"} + with pytest.raises(ValueError, match="model.*required"): + MistralEmbeddingAdapter.validate(meta) + + def test_validate_missing_api_key_raises(self) -> None: + from pydantic import ValidationError + + meta = {"model": "mistral/mistral-embed"} + with pytest.raises(ValidationError): + MistralEmbeddingAdapter.validate(meta) + + def test_validate_calls_validate_model(self) -> None: + meta = {"model": "mistral-embed", "api_key": "test-key"} + validated = MistralEmbeddingAdapter.validate(meta) + assert validated["model"] == "mistral/mistral-embed" + + def test_validate_strips_extra_fields(self) -> None: + meta = { + "model": "mistral/mistral-embed", + "api_key": "test-key", + "adapter_name": "my-adapter", + "unknown_field": "should_be_dropped", + } + validated = MistralEmbeddingAdapter.validate(meta) + assert "adapter_name" not in validated + assert "unknown_field" not in validated + + def test_validate_includes_base_fields(self) -> None: + meta = {"model": "mistral/mistral-embed", "api_key": "test-key"} + validated = MistralEmbeddingAdapter.validate(meta) + assert "timeout" in validated + assert "max_retries" in validated + + def test_metadata(self) -> None: + metadata = MistralEmbeddingAdapter.get_metadata() + assert metadata["name"] == "Mistral" + assert metadata["is_active"] is True + assert metadata["adapter"] is MistralEmbeddingAdapter