weflayr.sdk.mistralai.client

  1"""Weflayr instrumented wrapper for the Mistral AI SDK.
  2
  3Connector: **Mistral AI**
  4Provider:  https://mistral.ai
  5
  6Available connectors
  7--------------------
  8- :class:`Mistral` — top-level client (drop-in for ``mistralai.Mistral``)
  9
 10Available methods
 11-----------------
 12Via ``client.chat``:
 13
 14- :meth:`Chat.complete`       — synchronous chat completion with telemetry
 15- :meth:`Chat.complete_async` — async chat completion with telemetry
 16
 17Example::
 18
 19    from weflayr.sdk.mistralai.client import Mistral
 20
 21    client = Mistral(api_key="sk-...")
 22    response = client.chat.complete(
 23        model="mistral-small-latest",
 24        messages=[{"role": "user", "content": "Hello"}],
 25    )
 26"""
 27
 28from __future__ import annotations
 29
 30from typing import Any
 31
 32from mistralai.client import Mistral as _Mistral
 33from mistralai.client.chat import Chat as _Chat
 34
 35from weflayr.sdk.helpers import INTAKE_URL, track_sync, track_async
 36
 37
 38def _usage(response) -> dict:
 39    """Extract token usage from a Mistral chat response.
 40
 41    Args:
 42        response: A ``ChatCompletionResponse`` returned by the Mistral SDK.
 43
 44    Returns:
 45        A dict with ``prompt_tokens`` and ``completion_tokens`` (both ``int | None``).
 46    """
 47    usage = getattr(response, "usage", None)
 48    return {
 49        "prompt_tokens": getattr(usage, "prompt_tokens", None),
 50        "completion_tokens": getattr(usage, "completion_tokens", None),
 51    }
 52
 53
 54class Chat(_Chat):
 55    """Instrumented Mistral chat client.
 56
 57    Wraps every completion call with before/after telemetry events sent to
 58    the Weflayr intake API.
 59
 60    Args:
 61        *args: Forwarded to the upstream ``mistralai.client.chat.Chat``.
 62        intake_url: Weflayr intake API base URL.
 63            Defaults to :data:`~weflayr.sdk.helpers.INTAKE_URL`.
 64        **kwargs: Forwarded to the upstream ``mistralai.client.chat.Chat``.
 65    """
 66
 67    def __init__(self, *args: Any, intake_url: str = INTAKE_URL, **kwargs: Any) -> None:
 68        super().__init__(*args, **kwargs)
 69        self._intake_url = intake_url
 70
 71    def complete(self, **kwargs: Any):
 72        """Send a synchronous chat completion request with telemetry.
 73
 74        Fires a ``chat.complete.before`` event to the intake API, calls the
 75        upstream Mistral API, then fires a ``chat.complete.after`` event with
 76        timing and token usage.
 77
 78        Args:
 79            model (str): Mistral model identifier (e.g. ``"mistral-small-latest"``).
 80            messages (list[dict]): Conversation history in OpenAI message format.
 81            tags (dict, optional): Arbitrary key/value metadata forwarded to
 82                both telemetry events. Stripped before the upstream call.
 83            **kwargs: Any other kwargs accepted by ``mistralai`` ``chat.complete()``.
 84
 85        Returns:
 86            ``ChatCompletionResponse``: The upstream Mistral response, unmodified.
 87        """
 88        tags = kwargs.pop("tags", {})
 89        return track_sync(
 90            url=self._intake_url,
 91            call="chat.complete",
 92            before={"model": kwargs.get("model"), "message_count": len(kwargs.get("messages", [])), "tags": tags},
 93            fn=lambda: super(Chat, self).complete(**kwargs),
 94            after_extra=_usage,
 95        )
 96
 97    async def complete_async(self, **kwargs: Any):
 98        """Send an async chat completion request with telemetry.
 99
100        Async equivalent of :meth:`complete`. Fires ``chat.complete_async.before``
101        and ``chat.complete_async.after`` events.
102
103        Args:
104            model (str): Mistral model identifier (e.g. ``"mistral-small-latest"``).
105            messages (list[dict]): Conversation history in OpenAI message format.
106            tags (dict, optional): Arbitrary key/value metadata forwarded to
107                both telemetry events. Stripped before the upstream call.
108            **kwargs: Any other kwargs accepted by ``mistralai`` ``chat.complete_async()``.
109
110        Returns:
111            ``ChatCompletionResponse``: The upstream Mistral response, unmodified.
112        """
113        tags = kwargs.pop("tags", {})
114        return await track_async(
115            url=self._intake_url,
116            call="chat.complete_async",
117            before={"model": kwargs.get("model"), "message_count": len(kwargs.get("messages", [])), "tags": tags},
118            fn=lambda: super(Chat, self).complete_async(**kwargs),
119            after_extra=_usage,
120        )
121
122
123class Mistral(_Mistral):
124    """Drop-in replacement for ``mistralai.client.Mistral`` with Weflayr telemetry.
125
126    Intercepts the ``chat`` attribute to inject an instrumented :class:`Chat`
127    instance. All other attributes and behaviours are inherited unchanged from
128    the upstream client.
129
130    Args:
131        api_key (str): Your Mistral API key.
132        intake_url (str, optional): Weflayr intake API base URL.
133            Defaults to :data:`~weflayr.sdk.helpers.INTAKE_URL`.
134        **kwargs: Forwarded unchanged to ``mistralai.client.Mistral``.
135
136    Attributes:
137        chat (:class:`Chat`): Instrumented chat client exposing
138            :meth:`~Chat.complete` and :meth:`~Chat.complete_async`.
139
140    Example::
141
142        client = Mistral(api_key="sk-...")
143        client.chat.complete(model="mistral-small-latest", messages=[...])
144    """
145
146    def __init__(self, *args: Any, intake_url: str = INTAKE_URL, **kwargs: Any) -> None:
147        super().__init__(*args, **kwargs)
148        self._intake_url = intake_url
149
150    def __getattr__(self, name: str) -> Any:
151        instance = super().__getattr__(name)
152        if name == "chat":
153            instance = Chat(self.sdk_configuration, parent_ref=self, intake_url=self._intake_url)
154            object.__setattr__(self, name, instance)
155        return instance