mcp-nextcloud/nextcloud_mcp_server/providers/mistral.py

"""Mistral provider for embeddings.

Currently supports embeddings only (``mistral-embed``, 1024-dim). Generation
can be added later if needed; see ADR-015.
"""

import logging
from functools import wraps

import anyio
from mistralai.client import Mistral
from mistralai.client.errors.sdkerror import SDKError

from .base import Provider

logger = logging.getLogger(__name__)

MAX_RETRIES = 5
INITIAL_RETRY_DELAY = 2.0
MAX_RETRY_DELAY = 60.0


def retry_on_rate_limit(func):
    """Retry on Mistral 429 (rate limit) responses with exponential backoff."""

    @wraps(func)
    async def wrapper(*args, **kwargs):
        retry_delay = INITIAL_RETRY_DELAY
        last_error: Exception | None = None

        for attempt in range(1, MAX_RETRIES + 1):
            try:
                return await func(*args, **kwargs)
            except SDKError as e:
                # SDKError carries a status_code attribute populated from the
                # raw response. Only 429 is retryable here.
                status = getattr(e, "status_code", None)
                if status != 429:
                    raise
                last_error = e
                if attempt < MAX_RETRIES:
                    logger.warning(
                        "Mistral rate limit hit (attempt %d/%d), retrying in %.1fs...",
                        attempt,
                        MAX_RETRIES,
                        retry_delay,
                    )
                    await anyio.sleep(retry_delay)
                    retry_delay = min(retry_delay * 2, MAX_RETRY_DELAY)

        logger.error("Mistral rate limit exceeded after %d attempts", MAX_RETRIES)
        raise last_error  # type: ignore[misc]

    return wrapper


# 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


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_on_rate_limit
    async def embed(self, text: str) -> list[float]:
        """Generate an embedding for a single text."""
        if not self.supports_embeddings:
            raise NotImplementedError(
                "Embedding not supported - no embedding_model configured"
            )

        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``."""
        if not self.supports_embeddings:
            raise NotImplementedError(
                "Embedding not supported - no embedding_model configured"
            )

        if not texts:
            return []

        all_embeddings: list[list[float]] = []
        for i in range(0, len(texts), BATCH_SIZE):
            batch = texts[i : i + BATCH_SIZE]
            batch_embeddings = await self._embed_batch_request(batch)
            all_embeddings.extend(batch_embeddings)

            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

    @retry_on_rate_limit
    async def _embed_batch_request(self, batch: list[str]) -> list[list[float]]:
        """Single batch request with rate-limit retry."""
        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"
            )
        return result

    def get_dimension(self) -> int:
        if not self.supports_embeddings:
            raise NotImplementedError(
                "Embedding not supported - no embedding_model configured"
            )

        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 Mistral SDK manages its own httpx client lifecycle; close it
        # via the SDK's context-manager hook if present, otherwise no-op.
        close = getattr(self.client, "__aexit__", None)
        if close is not None:
            try:
                await close(None, None, None)
            except Exception:  # pragma: no cover - best-effort cleanup
                logger.debug("Mistral client close raised; ignoring", exc_info=True)