dot-inference 2.0.2__py3-none-any.whl

dot_inference/_mistral_compat.py

@@ -0,0 +1,530 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Shared Mistral-native client.
+
+Provides ``MistralClient`` — a class-based LLM client using the native
+Mistral SDK.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import logging
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import httpx
+from pydantic import BaseModel
+
+from dot_inference._client import EmbeddingClient, LLMClient
+from dot_inference._cost_tracker import compute_cost
+from dot_inference._llm_registry import ResolvedModelConfig
+from dot_inference._retry import compute_wait
+from dot_inference._types import (
+    InputJsonDelta,
+    LLMResponse,
+    Message,
+    ReasoningDelta,
+    StreamDone,
+    StreamEvent,
+    StreamUsage,
+    TextDelta,
+    ToolCall,
+    ToolCallDone,
+    ToolDef,
+)
+
+LOGGER = logging.getLogger(__name__)
+# TODO: include the deepika logger
+
+_MAX_ATTEMPTS = 2
+_RETRY_BASE_WAIT = 3  # seconds
+
+# HTTP status codes that should NOT be retried (client-side mistakes).
+_NON_RETRYABLE_STATUS = {400, 401, 403, 404, 422}
+
+
+def _check_mistral_available() -> None:
+    """Raise ``ImportError`` with install instructions if *mistralai* is missing."""
+    try:
+        import mistralai  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "The mistralai package is required for native Mistral support. "
+            "Install it with: pip install dot-inference[mistral]"
+        ) from None
+
+
+def _tools_to_mistral(tools: list[ToolDef]) -> list[dict[str, Any]]:
+    """Convert ``ToolDef`` list to the Mistral tools format.
+
+    Mistral uses the same OpenAI-like structure:
+    ``{"type": "function", "function": {"name": ..., "description": ..., "parameters": ...}}``
+    """
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": t.name,
+                "description": t.description,
+                "parameters": t.parameters,
+            },
+        }
+        for t in tools
+    ]
+
+
+def _normalize_mistral_content(content: Any) -> tuple[str, str]:
+    """Split Mistral message content into ``(text, thinking)``.
+
+    Non-reasoning responses return ``content`` as a plain ``str`` → ``(content, "")``.
+    Reasoning responses (mistralai >= 2.4) return a ``list[ContentChunk]`` mixing
+    ``TextChunk(text=..., type="text")`` (the answer) and
+    ``ThinkChunk(thinking=[...], type="thinking")`` (the chain of thought).
+    Text chunks are concatenated into *text*, thinking chunks into *thinking*;
+    other chunk types (reference/image) carry no ``.text`` and are skipped.
+    """
+    if content is None:
+        return "", ""
+    if isinstance(content, str):
+        return content, ""
+
+    text_parts: list[str] = []
+    thinking_parts: list[str] = []
+    for chunk in content:
+        if getattr(chunk, "type", None) == "thinking":
+            for item in getattr(chunk, "thinking", None) or []:
+                thinking_parts.append(getattr(item, "text", "") or "")
+        else:
+            text_parts.append(getattr(chunk, "text", "") or "")
+    return "".join(text_parts), "".join(thinking_parts)
+
+
+class MistralClient(LLMClient):
+    """LLM client using the native Mistral SDK.
+
+    The underlying ``Mistral`` SDK client is stored as ``self._client``
+    and created once at construction time.
+
+    Requires: ``pip install dot-inference[mistral]``
+    """
+
+    def __init__(
+        self,
+        model: str,
+        api_key: str | None = None,
+        model_config: ResolvedModelConfig | None = None,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        max_tokens: int | None = None,
+        timeout: float = 300,
+    ) -> None:
+        _check_mistral_available()
+
+        from mistralai.client import Mistral
+
+        super().__init__()
+        self._client = Mistral(
+            api_key=api_key,
+            timeout_ms=int(timeout * 1000),
+        )
+        self._model = model
+        self._model_config = model_config
+        self._temperature = temperature
+        self._reasoning_effort = reasoning_effort
+        self._max_tokens = max_tokens
+        self._strict = model_config.supports_strict_json if model_config else True
+
+    def _build_create_kwargs(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None,
+        tools: list[ToolDef] | None,
+        temperature: float | None,
+        reasoning_effort: str | None,
+        tool_choice: str | None = None,
+    ) -> tuple[dict[str, Any], type[BaseModel] | None]:
+        """Build kwargs for Mistral chat completion.
+
+        Returns ``(create_kwargs, response_model)``.
+        """
+        create_kwargs: dict[str, Any] = {
+            "model": self._model,
+            "messages": [m.model_dump(exclude_none=True) for m in messages],
+        }
+
+        is_reasoning = self._model_config.supports_reasoning if self._model_config else False
+        if is_reasoning and reasoning_effort is not None:
+            # Mistral reasoning models (Medium 3.5, Small) take reasoning_effort
+            # natively on the chat-completion request — La Plateforme accepts
+            # "none" (minimal thinking, thinking chunk omitted) and "high"
+            # (full thinking chunk). Temperature is ignored for these.
+            create_kwargs["reasoning_effort"] = reasoning_effort
+        elif not is_reasoning and temperature is not None:
+            create_kwargs["temperature"] = temperature
+
+        response_model: type[BaseModel] | None = None
+
+        self._validate_params(response_format, tools)
+
+        if tools:
+            create_kwargs["tools"] = _tools_to_mistral(tools)
+            create_kwargs["tool_choice"] = tool_choice if tool_choice is not None else "auto"
+        elif response_format is not None:
+            if isinstance(response_format, type) and issubclass(response_format, BaseModel):
+                response_model = response_format
+                schema = response_format.model_json_schema()
+            else:
+                schema = response_format
+
+            create_kwargs["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "Output",
+                    "strict": self._strict,
+                    "schema": schema,
+                },
+            }
+
+        if self._max_tokens is not None:
+            create_kwargs["max_tokens"] = self._max_tokens
+
+        return create_kwargs, response_model
+
+    @staticmethod
+    def _normalize_arguments(arguments: Any) -> str:
+        """Mistral SDK may return arguments as dict instead of str."""
+        return arguments if isinstance(arguments, str) else json.dumps(arguments)
+
+    async def _complete_with_retry(self, create_kwargs: dict[str, Any]) -> Any:
+        """Call ``chat.complete_async`` with transient-error retry."""
+        from mistralai.client.errors.mistralerror import MistralError
+
+        for attempt in range(1, _MAX_ATTEMPTS + 1):
+            try:
+                return await self._client.chat.complete_async(**create_kwargs)
+            except (MistralError, httpx.TransportError) as exc:
+                status = getattr(exc, "status_code", None)
+                if status in _NON_RETRYABLE_STATUS or attempt == _MAX_ATTEMPTS:
+                    raise
+                headers = getattr(exc, "headers", None)
+                retry_after = headers.get("retry-after") if headers is not None else None
+                wait = compute_wait(attempt, _RETRY_BASE_WAIT, retry_after)
+                LOGGER.warning(
+                    "Transient Mistral API error (attempt %d/%d, status %s): %s: %s — retrying in %.1fs",
+                    attempt,
+                    _MAX_ATTEMPTS,
+                    status,
+                    type(exc).__name__,
+                    exc,
+                    wait,
+                )
+                if self.on_retry is not None:
+                    self.on_retry(attempt, _MAX_ATTEMPTS, str(exc))
+                await asyncio.sleep(wait)
+
+    async def call(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None = None,
+        tools: list[ToolDef] | None = None,
+        *,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        tool_choice: str | None = None,
+    ) -> LLMResponse:
+        temperature = self._resolve_temperature(temperature)
+        reasoning_effort = self._resolve_reasoning_effort(reasoning_effort)
+
+        create_kwargs, response_model = self._build_create_kwargs(
+            messages,
+            response_format,
+            tools,
+            temperature,
+            reasoning_effort,
+            tool_choice,
+        )
+
+        response = await self._complete_with_retry(create_kwargs)
+
+        prompt_tokens = response.usage.prompt_tokens if response.usage else None
+        completion_tokens = response.usage.completion_tokens if response.usage else None
+
+        cost = None
+        if self._model_config:
+            cost = compute_cost(
+                response.usage,
+                self._model_config.input_token_cost,
+                self._model_config.output_token_cost,
+            )
+        if response.usage:
+            LOGGER.debug(
+                "tokens: %s in / %s out, cost=$%s",
+                prompt_tokens,
+                completion_tokens,
+                cost or "?",
+            )
+
+        msg = response.choices[0].message
+        # Reasoning models return content as list[ContentChunk]; keep the
+        # answer text only (thinking is display-only and dropped here).
+        text_content, _thinking = _normalize_mistral_content(msg.content)
+
+        tool_calls = None
+        if msg.tool_calls:
+            tool_calls = [
+                ToolCall(
+                    id=tc.id,
+                    name=tc.function.name,
+                    arguments=self._normalize_arguments(tc.function.arguments),
+                )
+                for tc in msg.tool_calls
+            ]
+
+        structured_output = None
+        if response_model is not None:
+            if not text_content:
+                raise RuntimeError("LLM returned empty content with response_format set")
+            structured_output = response_model.model_validate_json(text_content)
+
+        response = LLMResponse(
+            content=text_content or None,
+            tool_calls=tool_calls,
+            structured_output=structured_output,
+            cost=cost,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+        )
+        self.cost_tracker.record(response)
+        return response
+
+    async def stream(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None = None,
+        tools: list[ToolDef] | None = None,
+        *,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        tool_choice: str | None = None,
+    ) -> AsyncGenerator[StreamEvent, None]:
+        temperature = self._resolve_temperature(temperature)
+        reasoning_effort = self._resolve_reasoning_effort(reasoning_effort)
+
+        create_kwargs, response_model = self._build_create_kwargs(
+            messages,
+            response_format,
+            tools,
+            temperature,
+            reasoning_effort,
+            tool_choice,
+        )
+
+        from mistralai.client.errors.mistralerror import MistralError
+
+        content_parts: list[str] = []
+        tool_acc: dict[int, dict[str, str]] = {}
+        finished_tool_calls: list[ToolCall] = []
+        prompt_tokens: int | None = None
+        completion_tokens: int | None = None
+        last_usage: Any = None
+
+        for attempt in range(1, _MAX_ATTEMPTS + 1):
+            has_yielded = False
+            reasoning_seen = False
+            content_parts.clear()
+            tool_acc.clear()
+            finished_tool_calls.clear()
+            prompt_tokens = None
+            completion_tokens = None
+            last_usage = None
+            response = None
+
+            try:
+                response = await self._client.chat.stream_async(**create_kwargs)
+
+                async for event in response:
+                    data = event.data
+
+                    if data.usage:
+                        prompt_tokens = data.usage.prompt_tokens
+                        completion_tokens = data.usage.completion_tokens
+                        last_usage = data.usage
+
+                    if not data.choices:
+                        continue
+
+                    choice = data.choices[0]
+                    delta = choice.delta
+
+                    if delta.content:
+                        # Reasoning models stream content as list[ContentChunk]
+                        # (TextChunk + ThinkChunk); split thinking from answer.
+                        text, thinking = _normalize_mistral_content(delta.content)
+                        if thinking:
+                            if not reasoning_seen:
+                                reasoning_seen = True
+                                LOGGER.info(
+                                    "model '%s' streams reasoning — "
+                                    "emitting reasoning_delta events",
+                                    self._model,
+                                )
+                            has_yielded = True
+                            yield ReasoningDelta(content=thinking)
+                        if text:
+                            content_parts.append(text)
+                            has_yielded = True
+                            yield TextDelta(content=text)
+
+                    if delta.tool_calls:
+                        for tc_delta in delta.tool_calls:
+                            idx = tc_delta.index if hasattr(tc_delta, "index") else 0
+                            if idx not in tool_acc:
+                                tool_acc[idx] = {"id": "", "name": "", "arguments": ""}
+
+                            acc = tool_acc[idx]
+                            tc_id = getattr(tc_delta, "id", None)
+                            tc_fn = getattr(tc_delta, "function", None)
+                            if tc_id:
+                                acc["id"] = tc_id
+                            if tc_fn and getattr(tc_fn, "name", None):
+                                acc["name"] = tc_fn.name
+                            if tc_fn and getattr(tc_fn, "arguments", None):
+                                arg_chunk = self._normalize_arguments(tc_fn.arguments)
+                                acc["arguments"] += arg_chunk
+                                has_yielded = True
+                                yield InputJsonDelta(
+                                    tool_call_index=idx,
+                                    tool_name=acc["name"],
+                                    partial_json=arg_chunk,
+                                )
+
+                    if choice.finish_reason in ("tool_calls", "stop") and tool_acc:
+                        for idx in sorted(tool_acc):
+                            acc = tool_acc[idx]
+                            tc = ToolCall(
+                                id=acc["id"], name=acc["name"], arguments=acc["arguments"]
+                            )
+                            finished_tool_calls.append(tc)
+                            has_yielded = True
+                            yield ToolCallDone(tool_call=tc)
+                        tool_acc.clear()
+
+                break  # Stream completed successfully
+
+            except (MistralError, httpx.TransportError) as exc:
+                status = getattr(exc, "status_code", None)
+                if status in _NON_RETRYABLE_STATUS or has_yielded or attempt == _MAX_ATTEMPTS:
+                    raise
+                headers = getattr(exc, "headers", None)
+                retry_after = headers.get("retry-after") if headers is not None else None
+                wait = compute_wait(attempt, _RETRY_BASE_WAIT, retry_after)
+                LOGGER.warning(
+                    "Transient Mistral API error during stream (attempt %d/%d, status %s): %s: %s — retrying in %.1fs",
+                    attempt,
+                    _MAX_ATTEMPTS,
+                    status,
+                    type(exc).__name__,
+                    exc,
+                    wait,
+                )
+                if self.on_retry is not None:
+                    self.on_retry(attempt, _MAX_ATTEMPTS, str(exc))
+                await asyncio.sleep(wait)
+
+            finally:
+                if response is not None:
+                    aclose = getattr(response, "aclose", None) or getattr(response, "close", None)
+                    if aclose is not None:
+                        result = aclose()
+                        if inspect.isawaitable(result):
+                            await result
+
+        cost = None
+        if self._model_config:
+            cost = compute_cost(
+                last_usage,
+                self._model_config.input_token_cost,
+                self._model_config.output_token_cost,
+            )
+        if prompt_tokens is not None:
+            LOGGER.debug(
+                "tokens: %s in / %s out, cost=$%s",
+                prompt_tokens,
+                completion_tokens,
+                cost or "?",
+            )
+
+        if prompt_tokens is not None or completion_tokens is not None:
+            yield StreamUsage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                cost=cost,
+            )
+
+        content = "".join(content_parts) or None
+        structured_output = None
+        if response_model is not None:
+            if not content:
+                raise RuntimeError("LLM returned empty streamed content with response_format set")
+            structured_output = response_model.model_validate_json(content)
+
+        final_response = LLMResponse(
+            content=content,
+            tool_calls=finished_tool_calls or None,
+            structured_output=structured_output,
+            cost=cost,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+        )
+        self.cost_tracker.record(final_response)
+        yield StreamDone(response=final_response)
+
+
+class MistralEmbeddingClient(EmbeddingClient):
+    """Embedding client using the native Mistral SDK.
+
+    The underlying ``Mistral`` SDK client is stored as ``self._client``
+    and created once at construction time.
+
+    Requires: ``pip install dot-inference[mistral]``
+    """
+
+    def __init__(
+        self,
+        model: str,
+        api_key: str | None = None,
+        batch_size: int = 0,
+        timeout: float = 300,
+    ) -> None:
+        _check_mistral_available()
+
+        from mistralai.client import Mistral
+
+        self._client = Mistral(
+            api_key=api_key,
+            timeout_ms=int(timeout * 1000),
+        )
+        self._model = model
+        self._batch_size = batch_size
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+
+        batches = (
+            [texts[i : i + self._batch_size] for i in range(0, len(texts), self._batch_size)]
+            if self._batch_size > 0
+            else [texts]
+        )
+
+        all_embeddings: list[list[float]] = []
+        for batch in batches:
+            response = await self._client.embeddings.create_async(model=self._model, inputs=batch)
+            items = sorted(response.data, key=lambda x: x.index)
+            all_embeddings.extend(item.embedding for item in items)
+
+        return all_embeddings