-
Notifications
You must be signed in to change notification settings - Fork 1.8k
Python: Add Mistral AI embedding client package #5480
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
eavanvalkenburg
merged 5 commits into
microsoft:main
from
daric93:feature/mistral-embedding-client
May 29, 2026
Merged
Changes from all commits
Commits
Show all changes
5 commits
Select commit
Hold shift + click to select a range
5a33077
Python: Add Mistral AI embedding client package
daric93 ef532ba
Address review feedback: fix dimensions check, sort embeddings by ind…
daric93 a200e52
Address review feedback: downgrade to alpha, remove integration tests…
daric93 268563b
Add samples directory for alpha package compliance Per python-package…
daric93 3547d24
Fix ruff formatting in sample file
daric93 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Some comments aren't visible on the classic Files Changed page.
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| # Mistral Package (agent-framework-mistral) | ||
|
|
||
| Integration with Mistral AI for embedding generation. | ||
|
|
||
| ## Main Classes | ||
|
|
||
| - **`MistralEmbeddingClient`** - Embedding client for Mistral AI models | ||
| - **`MistralEmbeddingOptions`** - Options TypedDict for Mistral-specific embedding parameters | ||
| - **`MistralEmbeddingSettings`** - TypedDict settings for Mistral configuration | ||
|
|
||
| ## Usage | ||
|
|
||
| ```python | ||
| from agent_framework_mistral import MistralEmbeddingClient | ||
|
|
||
| # Requires MISTRAL_API_KEY environment variable (or pass api_key= directly) | ||
| client = MistralEmbeddingClient(model="mistral-embed") | ||
| result = await client.get_embeddings(["Hello, world!"]) | ||
| print(result[0].vector) | ||
| ``` | ||
|
|
||
| ## Import Path | ||
|
|
||
| ```python | ||
| from agent_framework_mistral import MistralEmbeddingClient | ||
| ``` | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| MIT License | ||
|
|
||
| Copyright (c) Microsoft Corporation. | ||
|
|
||
| Permission is hereby granted, free of charge, to any person obtaining a copy | ||
| of this software and associated documentation files (the "Software"), to deal | ||
| in the Software without restriction, including without limitation the rights | ||
| to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | ||
| copies of the Software, and to permit persons to whom the Software is | ||
| furnished to do so, subject to the following conditions: | ||
|
|
||
| The above copyright notice and this permission notice shall be included in all | ||
| copies or substantial portions of the Software. | ||
|
|
||
| THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | ||
| IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | ||
| FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | ||
| AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | ||
| LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | ||
| OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | ||
| SOFTWARE |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| # Get Started with Microsoft Agent Framework Mistral AI | ||
|
|
||
| Please install this package: | ||
|
|
||
| ```bash | ||
| pip install agent-framework-mistral --pre | ||
| ``` | ||
|
|
||
| and see the [README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) for more information. | ||
|
|
||
| ## Embedding Client | ||
|
|
||
| The `MistralEmbeddingClient` provides embedding generation using Mistral AI models. | ||
|
|
||
| ### Quick Start | ||
|
|
||
| ```python | ||
| from agent_framework_mistral import MistralEmbeddingClient | ||
|
|
||
| # Using environment variables (MISTRAL_API_KEY, MISTRAL_EMBEDDING_MODEL) | ||
| client = MistralEmbeddingClient() | ||
|
|
||
| # Or passing parameters directly | ||
| client = MistralEmbeddingClient( | ||
| model="mistral-embed", | ||
| api_key="your-api-key", | ||
| ) | ||
|
|
||
| # Generate embeddings | ||
| result = await client.get_embeddings(["Hello, world!", "How are you?"]) | ||
| for embedding in result: | ||
| print(f"Dimensions: {embedding.dimensions}") | ||
| print(f"Vector: {embedding.vector[:5]}...") | ||
| ``` | ||
|
|
||
| ### Configuration | ||
|
|
||
| | Environment Variable | Description | | ||
| |---|---| | ||
| | `MISTRAL_API_KEY` | Your Mistral AI API key | | ||
| | `MISTRAL_EMBEDDING_MODEL` | Embedding model name (e.g., `mistral-embed`) | | ||
| | `MISTRAL_SERVER_URL` | Optional server URL override | |
17 changes: 17 additions & 0 deletions
17
python/packages/mistral/agent_framework_mistral/__init__.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,17 @@ | ||
| # Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| import importlib.metadata | ||
|
|
||
| from ._embedding_client import MistralEmbeddingClient, MistralEmbeddingOptions, MistralEmbeddingSettings | ||
|
|
||
| try: | ||
| __version__ = importlib.metadata.version(__name__) | ||
| except importlib.metadata.PackageNotFoundError: | ||
| __version__ = "0.0.0" # Fallback for development mode | ||
|
|
||
| __all__ = [ | ||
| "MistralEmbeddingClient", | ||
| "MistralEmbeddingOptions", | ||
| "MistralEmbeddingSettings", | ||
| "__version__", | ||
| ] |
250 changes: 250 additions & 0 deletions
250
python/packages/mistral/agent_framework_mistral/_embedding_client.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,250 @@ | ||
| # Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import logging | ||
| import sys | ||
| from collections.abc import Sequence | ||
| from typing import Any, ClassVar, Generic, TypedDict | ||
|
|
||
| from agent_framework import ( | ||
| BaseEmbeddingClient, | ||
| Embedding, | ||
| EmbeddingGenerationOptions, | ||
| GeneratedEmbeddings, | ||
| UsageDetails, | ||
| load_settings, | ||
| ) | ||
| from agent_framework._settings import SecretString | ||
| from agent_framework.observability import EmbeddingTelemetryLayer | ||
| from mistralai.client import Mistral | ||
|
|
||
| if sys.version_info >= (3, 13): | ||
| from typing import TypeVar # type: ignore # pragma: no cover | ||
| else: | ||
| from typing_extensions import TypeVar # type: ignore # pragma: no cover | ||
|
|
||
|
|
||
| logger = logging.getLogger("agent_framework.mistral") | ||
|
|
||
|
|
||
| class MistralEmbeddingOptions(EmbeddingGenerationOptions, total=False): | ||
| """Mistral AI-specific embedding options. | ||
|
|
||
| Extends EmbeddingGenerationOptions with Mistral-specific fields. | ||
|
|
||
| Examples: | ||
| .. code-block:: python | ||
|
|
||
| from agent_framework_mistral import MistralEmbeddingOptions | ||
|
|
||
| options: MistralEmbeddingOptions = { | ||
| "model": "mistral-embed", | ||
| "dimensions": 1024, | ||
| } | ||
| """ | ||
|
|
||
|
|
||
| MistralEmbeddingOptionsT = TypeVar( | ||
| "MistralEmbeddingOptionsT", | ||
| bound=TypedDict, # type: ignore[valid-type] | ||
| default="MistralEmbeddingOptions", | ||
|
daric93 marked this conversation as resolved.
|
||
| covariant=True, | ||
| ) | ||
|
|
||
|
|
||
| class MistralEmbeddingSettings(TypedDict, total=False): | ||
| """Mistral AI embedding settings. | ||
|
|
||
| Fields: | ||
| api_key: Mistral API key. Resolved from ``MISTRAL_API_KEY``. | ||
| embedding_model: Embedding model name. Resolved from ``MISTRAL_EMBEDDING_MODEL``. | ||
| server_url: Optional server URL override. Resolved from ``MISTRAL_SERVER_URL``. | ||
| """ | ||
|
|
||
| api_key: str | None | ||
| embedding_model: str | None | ||
| server_url: str | None | ||
|
|
||
|
|
||
| class RawMistralEmbeddingClient( | ||
| BaseEmbeddingClient[str, list[float], MistralEmbeddingOptionsT], | ||
| Generic[MistralEmbeddingOptionsT], | ||
| ): | ||
| """Raw Mistral AI embedding client without telemetry. | ||
|
|
||
| Keyword Args: | ||
| model: The Mistral embedding model (e.g. "mistral-embed"). | ||
| Can also be set via environment variable ``MISTRAL_EMBEDDING_MODEL``. | ||
| api_key: Mistral API key. Defaults to ``MISTRAL_API_KEY`` environment variable. | ||
| server_url: Optional server URL override. Defaults to ``MISTRAL_SERVER_URL`` | ||
| environment variable, or the Mistral default. | ||
| client: Optional pre-configured ``Mistral`` client instance. | ||
| additional_properties: Additional properties stored on the client instance. | ||
| env_file_path: Path to ``.env`` file for settings. | ||
| env_file_encoding: Encoding for ``.env`` file. | ||
| """ | ||
|
|
||
| INJECTABLE: ClassVar[set[str]] = {"client"} | ||
|
|
||
| def __init__( | ||
| self, | ||
| *, | ||
| model: str | None = None, | ||
| api_key: str | SecretString | None = None, | ||
| server_url: str | None = None, | ||
| client: Mistral | None = None, | ||
| additional_properties: dict[str, Any] | None = None, | ||
| env_file_path: str | None = None, | ||
| env_file_encoding: str | None = None, | ||
| ) -> None: | ||
| """Initialize a raw Mistral AI embedding client.""" | ||
| mistral_settings = load_settings( | ||
| MistralEmbeddingSettings, | ||
| env_prefix="MISTRAL_", | ||
| required_fields=["embedding_model", "api_key"], | ||
| api_key=str(api_key) if isinstance(api_key, SecretString) else api_key, | ||
| embedding_model=model, | ||
| server_url=server_url, | ||
| env_file_path=env_file_path, | ||
| env_file_encoding=env_file_encoding, | ||
| ) | ||
|
|
||
| self.model: str = mistral_settings["embedding_model"] # type: ignore[assignment] | ||
| resolved_api_key: str = mistral_settings["api_key"] # type: ignore[assignment] | ||
| resolved_server_url = mistral_settings.get("server_url") | ||
|
|
||
| if client is not None: | ||
| self.client = client | ||
| else: | ||
| client_kwargs: dict[str, Any] = {"api_key": resolved_api_key} | ||
| if resolved_server_url: | ||
| client_kwargs["server_url"] = resolved_server_url | ||
| self.client = Mistral(**client_kwargs) | ||
|
|
||
| self.server_url = resolved_server_url | ||
| super().__init__(additional_properties=additional_properties) | ||
|
|
||
| def service_url(self) -> str: | ||
| """Get the URL of the service.""" | ||
| return self.server_url or "https://api.mistral.ai" | ||
|
|
||
| async def get_embeddings( | ||
| self, | ||
| values: Sequence[str], | ||
| *, | ||
| options: MistralEmbeddingOptionsT | None = None, | ||
| ) -> GeneratedEmbeddings[list[float], MistralEmbeddingOptionsT]: | ||
| """Call the Mistral AI embeddings API. | ||
|
|
||
| Args: | ||
| values: The text values to generate embeddings for. | ||
| options: Optional embedding generation options. | ||
|
|
||
| Returns: | ||
| Generated embeddings with usage metadata. | ||
|
|
||
| Raises: | ||
| ValueError: If model is not provided or values is empty. | ||
| """ | ||
| if not values: | ||
| return GeneratedEmbeddings([], options=options) | ||
|
|
||
| opts: dict[str, Any] = options or {} # type: ignore | ||
| model = opts.get("model") or self.model | ||
| if not model: | ||
| raise ValueError("model is required") | ||
|
|
||
| kwargs: dict[str, Any] = {"model": model, "inputs": list(values)} | ||
| if "dimensions" in opts: | ||
| kwargs["output_dimension"] = opts["dimensions"] | ||
|
|
||
| response = await self.client.embeddings.create_async(**kwargs) | ||
|
|
||
| embeddings: list[Embedding[list[float]]] = [] | ||
| if response and response.data: | ||
| items = sorted(response.data, key=lambda d: d.index if d.index is not None else 0) | ||
| for item in items: | ||
| vector = list(item.embedding) if item.embedding else [] | ||
| embeddings.append( | ||
| Embedding( | ||
| vector=vector, | ||
| dimensions=len(vector), | ||
| model=response.model or model, | ||
| ) | ||
| ) | ||
|
daric93 marked this conversation as resolved.
|
||
|
|
||
| usage_dict: UsageDetails | None = None | ||
| if response and response.usage: | ||
| usage_dict = { | ||
| "input_token_count": response.usage.prompt_tokens, | ||
| "total_token_count": response.usage.total_tokens, | ||
| } | ||
|
|
||
| return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict) | ||
|
|
||
|
|
||
| class MistralEmbeddingClient( | ||
| EmbeddingTelemetryLayer[str, list[float], MistralEmbeddingOptionsT], | ||
| RawMistralEmbeddingClient[MistralEmbeddingOptionsT], | ||
| Generic[MistralEmbeddingOptionsT], | ||
| ): | ||
| """Mistral AI embedding client with telemetry support. | ||
|
|
||
| Keyword Args: | ||
| model: The Mistral embedding model (e.g. "mistral-embed"). | ||
| Can also be set via environment variable ``MISTRAL_EMBEDDING_MODEL``. | ||
| api_key: Mistral API key. Defaults to ``MISTRAL_API_KEY`` environment variable. | ||
| server_url: Optional server URL override. Defaults to ``MISTRAL_SERVER_URL`` | ||
| environment variable, or the Mistral default. | ||
| client: Optional pre-configured ``Mistral`` client instance. | ||
| otel_provider_name: Optional telemetry provider name override. | ||
| env_file_path: Path to ``.env`` file for settings. | ||
| env_file_encoding: Encoding for ``.env`` file. | ||
|
|
||
| Examples: | ||
| .. code-block:: python | ||
|
|
||
| from agent_framework_mistral import MistralEmbeddingClient | ||
|
|
||
| # Using environment variables | ||
| # Set MISTRAL_API_KEY=your-key | ||
| # Set MISTRAL_EMBEDDING_MODEL=mistral-embed | ||
| client = MistralEmbeddingClient() | ||
|
|
||
| # Or passing parameters directly | ||
| client = MistralEmbeddingClient( | ||
| model="mistral-embed", | ||
| api_key="your-api-key", | ||
| ) | ||
|
|
||
| # Generate embeddings | ||
| result = await client.get_embeddings(["Hello, world!"]) | ||
| print(result[0].vector) | ||
| """ | ||
|
|
||
| OTEL_PROVIDER_NAME: ClassVar[str] = "mistralai" | ||
|
|
||
| def __init__( | ||
| self, | ||
| *, | ||
| model: str | None = None, | ||
| api_key: str | SecretString | None = None, | ||
| server_url: str | None = None, | ||
| client: Mistral | None = None, | ||
| otel_provider_name: str | None = None, | ||
| additional_properties: dict[str, Any] | None = None, | ||
| env_file_path: str | None = None, | ||
| env_file_encoding: str | None = None, | ||
| ) -> None: | ||
| """Initialize a Mistral AI embedding client.""" | ||
| super().__init__( | ||
| model=model, | ||
| api_key=api_key, | ||
| server_url=server_url, | ||
| client=client, | ||
| additional_properties=additional_properties, | ||
| otel_provider_name=otel_provider_name, | ||
| env_file_path=env_file_path, | ||
| env_file_encoding=env_file_encoding, | ||
| ) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
|
|
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.