"""Mistral provider for embeddings. Currently supports embeddings only (``mistral-embed``, 1024-dim). Generation can be added later if needed; see ADR-015. """ import logging # mistralai 2.x ships no top-level __init__.py, so `from mistralai import …` # raises ImportError. The canonical public paths are `mistralai.client` (which # re-exports the SDK class via `client/__init__.py`) and `mistralai.client.errors` # (which lazy-loads SDKError). There is no `mistralai.models` subpackage either. from mistralai.client import Mistral from mistralai.client.errors import SDKError from ._retry import retry_on_rate_limit from .base import Provider logger = logging.getLogger(__name__) # Well-known Mistral embedding model dimensions MISTRAL_EMBEDDING_DIMENSIONS: dict[str, int] = { "mistral-embed": 1024, } # Conservative chunk size for batch embeddings. Mistral allows large batches, # but we keep this in line with sibling providers (OpenAI=100, Ollama=32). BATCH_SIZE = 64 _NO_EMBEDDING_MODEL_MSG = "Embedding not supported - no embedding_model configured" def _is_rate_limit(exc: BaseException) -> bool: """True only for HTTP 429 SDKErrors.""" return getattr(exc, "status_code", None) == 429 _retry_429 = retry_on_rate_limit( SDKError, is_rate_limit=_is_rate_limit, provider_name="Mistral" ) class MistralProvider(Provider): """ Mistral provider — embeddings only. Uses the official ``mistralai`` SDK. Lazy dimension detection mirrors the OpenAI provider: known models populate the cached dimension at construction time; unknown models get their dimension detected on the first ``embed()`` call. """ def __init__( self, api_key: str, embedding_model: str | None = "mistral-embed", base_url: str | None = None, ): """ Initialize the Mistral provider. Args: api_key: Mistral API key. embedding_model: Embedding model ID (default: ``mistral-embed``). Pass ``None`` to disable embeddings (the provider will then support no capabilities, which is mostly useful for tests). base_url: Optional base URL override (e.g. proxies, on-prem). """ self.embedding_model = embedding_model self._dimension: int | None = None self.client = Mistral(api_key=api_key, server_url=base_url) if embedding_model and embedding_model in MISTRAL_EMBEDDING_DIMENSIONS: self._dimension = MISTRAL_EMBEDDING_DIMENSIONS[embedding_model] logger.info( "Initialized Mistral provider: base_url=%s, embedding_model=%s, " "dimension=%s", base_url or "default", embedding_model, self._dimension, ) @property def supports_embeddings(self) -> bool: return self.embedding_model is not None @property def supports_generation(self) -> bool: return False @_retry_429 async def embed(self, text: str) -> list[float]: """Generate an embedding for a single text.""" if not self.supports_embeddings: raise NotImplementedError(_NO_EMBEDDING_MODEL_MSG) assert self.embedding_model is not None response = await self.client.embeddings.create_async( model=self.embedding_model, inputs=[text], ) if not response.data or response.data[0].embedding is None: raise RuntimeError( f"Mistral embeddings API returned no embedding for model " f"{self.embedding_model}" ) embedding = response.data[0].embedding if self._dimension is None: self._dimension = len(embedding) logger.info( "Detected embedding dimension: %d for model %s", self._dimension, self.embedding_model, ) return embedding async def embed_batch(self, texts: list[str]) -> list[list[float]]: """Generate embeddings for multiple texts, chunking by ``BATCH_SIZE``.""" embeddings, _ = await self.embed_batch_with_usage(texts) return embeddings async def embed_with_usage(self, text: str) -> tuple[list[float], int]: """Embed one text, reporting the Mistral request's token count.""" embeddings, tokens = await self.embed_batch_with_usage([text]) if not embeddings: raise RuntimeError( f"Mistral embeddings API returned no embedding for model " f"{self.embedding_model}" ) return embeddings[0], tokens async def embed_batch_with_usage( self, texts: list[str] ) -> tuple[list[list[float]], int]: """Embed multiple texts, summing the Mistral-reported token usage. Returns ``(embeddings, total_tokens)`` where ``total_tokens`` is the sum of ``response.usage.total_tokens`` across the ``BATCH_SIZE`` sub- requests (the unit Mistral bills on). Used by the usage-metering hooks to record ``tokens_embedded`` by tokens (Deck #67). """ if not self.supports_embeddings: raise NotImplementedError(_NO_EMBEDDING_MODEL_MSG) if not texts: return [], 0 all_embeddings: list[list[float]] = [] total_tokens = 0 for i in range(0, len(texts), BATCH_SIZE): batch = texts[i : i + BATCH_SIZE] batch_embeddings, batch_tokens = await self._embed_batch_request(batch) all_embeddings.extend(batch_embeddings) total_tokens += batch_tokens if self._dimension is None and batch_embeddings: self._dimension = len(batch_embeddings[0]) logger.info( "Detected embedding dimension: %d for model %s", self._dimension, self.embedding_model, ) return all_embeddings, total_tokens @_retry_429 async def _embed_batch_request( self, batch: list[str] ) -> tuple[list[list[float]], int]: """Single batch request with rate-limit retry. Returns ``(embeddings, token_count)``; ``token_count`` comes from the response's ``usage.total_tokens`` and falls back to a char-based estimate if the API omits usage. """ assert self.embedding_model is not None response = await self.client.embeddings.create_async( model=self.embedding_model, inputs=batch, ) # Defensive: response.data items have Optional fields. Sort by index # (default 0 if missing) and reject None embeddings explicitly. sorted_data = sorted(response.data or [], key=lambda x: x.index or 0) result: list[list[float]] = [] for item in sorted_data: if item.embedding is None: raise RuntimeError( f"Mistral embeddings API returned a null embedding for " f"model {self.embedding_model}" ) result.append(item.embedding) if len(result) != len(batch): raise RuntimeError( f"Mistral embeddings API returned {len(result)} embeddings " f"for {len(batch)} inputs" ) usage = getattr(response, "usage", None) total_tokens = getattr(usage, "total_tokens", None) if usage else None # Guard on numeric type (not just ``is not None``): a real response # gives an int, but test doubles / partial responses can surface a # non-numeric attribute — fall back to the estimate there. tokens = ( round(total_tokens) if isinstance(total_tokens, (int, float)) else self._estimate_tokens(batch) ) return result, tokens def get_dimension(self) -> int: if not self.supports_embeddings: raise NotImplementedError(_NO_EMBEDDING_MODEL_MSG) if self._dimension is None: raise RuntimeError( f"Embedding dimension not detected yet for model " f"{self.embedding_model}. Call embed() first or use a known " "model." ) return self._dimension async def generate(self, prompt: str, max_tokens: int = 500) -> str: raise NotImplementedError( "MistralProvider does not support generation. " "Use OpenAI, Anthropic, or Bedrock for text generation." ) async def close(self) -> None: # The mistralai 2.x client (Speakeasy-generated) does not expose a # public close()/aclose() — only the async-context-manager protocol # (__aenter__/__aexit__). Calling __aexit__ directly is internal API # and brittle across SDK patch versions; the underlying httpx client # is closed during garbage collection, so we leave this as a no-op. return None