mistralai 1.10.1__py3-none-any.whl → 1.12.0__py3-none-any.whl

mistralai/extra/realtime/connection.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import base64
+import json
+from asyncio import CancelledError
+from collections import deque
+from typing import Any, AsyncIterator, Deque, Optional, Union
+
+from pydantic import ValidationError, BaseModel
+
+try:
+    from websockets.asyncio.client import ClientConnection  # websockets >= 13.0
+except ImportError as exc:
+    raise ImportError(
+        "The `websockets` package (>=13.0) is required for real-time transcription. "
+        "Install with: pip install 'mistralai[realtime]'"
+    ) from exc
+
+from mistralai.models import (
+    AudioFormat,
+    RealtimeTranscriptionError,
+    RealtimeTranscriptionSession,
+    RealtimeTranscriptionSessionCreated,
+    RealtimeTranscriptionSessionUpdated,
+    TranscriptionStreamDone,
+    TranscriptionStreamLanguage,
+    TranscriptionStreamSegmentDelta,
+    TranscriptionStreamTextDelta,
+)
+
+
+class UnknownRealtimeEvent(BaseModel):
+    """
+    Forward-compat fallback event:
+    - unknown message type
+    - invalid JSON payload
+    - schema validation failure
+    """
+    type: Optional[str]
+    content: Any
+    error: Optional[str] = None
+
+
+RealtimeEvent = Union[
+    # session lifecycle
+    RealtimeTranscriptionSessionCreated,
+    RealtimeTranscriptionSessionUpdated,
+    # server errors
+    RealtimeTranscriptionError,
+    # transcription events
+    TranscriptionStreamLanguage,
+    TranscriptionStreamSegmentDelta,
+    TranscriptionStreamTextDelta,
+    TranscriptionStreamDone,
+    # forward-compat fallback
+    UnknownRealtimeEvent,
+]
+
+
+_MESSAGE_MODELS: dict[str, Any] = {
+    "session.created": RealtimeTranscriptionSessionCreated,
+    "session.updated": RealtimeTranscriptionSessionUpdated,
+    "error": RealtimeTranscriptionError,
+    "transcription.language": TranscriptionStreamLanguage,
+    "transcription.segment": TranscriptionStreamSegmentDelta,
+    "transcription.text.delta": TranscriptionStreamTextDelta,
+    "transcription.done": TranscriptionStreamDone,
+}
+
+
+def parse_realtime_event(payload: Any) -> RealtimeEvent:
+    """
+    Tolerant parser:
+    - unknown event type -> UnknownRealtimeEvent
+    - validation failures -> UnknownRealtimeEvent (includes error string)
+    - invalid payload -> UnknownRealtimeEvent
+    """
+    if not isinstance(payload, dict):
+        return UnknownRealtimeEvent(
+            type=None, content=payload, error="expected JSON object"
+        )
+
+    msg_type = payload.get("type")
+    if not isinstance(msg_type, str):
+        return UnknownRealtimeEvent(
+            type=None, content=payload, error="missing/invalid 'type'"
+        )
+
+    model_cls = _MESSAGE_MODELS.get(msg_type)
+    if model_cls is None:
+        return UnknownRealtimeEvent(
+            type=msg_type, content=payload, error="unknown event type"
+        )
+    try:
+        parsed = model_cls.model_validate(payload)
+        return parsed
+    except ValidationError as exc:
+        return UnknownRealtimeEvent(type=msg_type, content=payload, error=str(exc))
+
+
+class RealtimeConnection:
+    def __init__(
+        self,
+        websocket: ClientConnection,
+        session: RealtimeTranscriptionSession,
+        *,
+        initial_events: Optional[list[RealtimeEvent]] = None,
+    ) -> None:
+        self._websocket = websocket
+        self._session = session
+        self._audio_format = session.audio_format
+        self._closed = False
+        self._initial_events: Deque[RealtimeEvent] = deque(initial_events or [])
+
+    @property
+    def request_id(self) -> str:
+        return self._session.request_id
+
+    @property
+    def session(self) -> RealtimeTranscriptionSession:
+        return self._session
+
+    @property
+    def audio_format(self) -> AudioFormat:
+        return self._audio_format
+
+    @property
+    def is_closed(self) -> bool:
+        return self._closed
+
+    async def send_audio(
+        self, audio_bytes: Union[bytes, bytearray, memoryview]
+    ) -> None:
+        if self._closed:
+            raise RuntimeError("Connection is closed")
+
+        message = {
+            "type": "input_audio.append",
+            "audio": base64.b64encode(bytes(audio_bytes)).decode("ascii"),
+        }
+        await self._websocket.send(json.dumps(message))
+
+    async def update_session(self, audio_format: AudioFormat) -> None:
+        if self._closed:
+            raise RuntimeError("Connection is closed")
+
+        self._audio_format = audio_format
+        message = {
+            "type": "session.update",
+            "session": {"audio_format": audio_format.model_dump(mode="json")},
+        }
+        await self._websocket.send(json.dumps(message))
+
+    async def end_audio(self) -> None:
+        if self._closed:
+            return
+        await self._websocket.send(json.dumps({"type": "input_audio.end"}))
+
+    async def close(self, *, code: int = 1000, reason: str = "") -> None:
+        if self._closed:
+            return
+        self._closed = True
+        await self._websocket.close(code=code, reason=reason)
+
+    async def __aenter__(self) -> "RealtimeConnection":
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        await self.close()
+
+    def __aiter__(self) -> AsyncIterator[RealtimeEvent]:
+        return self.events()
+
+    async def events(self) -> AsyncIterator[RealtimeEvent]:
+        # replay any handshake/prelude events (including session.created)
+        while self._initial_events:
+            ev = self._initial_events.popleft()
+            self._apply_session_updates(ev)
+            yield ev
+
+        try:
+            async for msg in self._websocket:
+                text = (
+                    msg.decode("utf-8", errors="replace")
+                    if isinstance(msg, (bytes, bytearray))
+                    else msg
+                )
+                try:
+                    data = json.loads(text)
+                except Exception as exc:
+                    yield UnknownRealtimeEvent(
+                        type=None, content=text, error=f"invalid JSON: {exc}"
+                    )
+                    continue
+
+                ev = parse_realtime_event(data)
+                self._apply_session_updates(ev)
+                yield ev
+        except CancelledError:
+            pass
+        finally:
+            await self.close()
+
+    def _apply_session_updates(self, ev: RealtimeEvent) -> None:
+        if isinstance(ev, RealtimeTranscriptionSessionCreated) or isinstance(ev, RealtimeTranscriptionSessionUpdated):
+            self._session = ev.session
+            self._audio_format = ev.session.audio_format

mistralai/extra/realtime/transcription.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from typing import AsyncIterator, Mapping, Optional
+from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
+
+try:
+    from websockets.asyncio.client import (
+        ClientConnection,
+        connect,
+    )  # websockets >= 13.0
+except ImportError as exc:
+    raise ImportError(
+        "The `websockets` package (>=13.0) is required for real-time transcription. "
+        "Install with: pip install 'mistralai[realtime]'"
+    ) from exc
+
+from mistralai import models, utils
+from mistralai.models import (
+    AudioFormat,
+    RealtimeTranscriptionError,
+    RealtimeTranscriptionSession,
+    RealtimeTranscriptionSessionCreated,
+)
+from mistralai.sdkconfiguration import SDKConfiguration
+from mistralai.utils import generate_url, get_security, get_security_from_env
+
+from ..exceptions import RealtimeTranscriptionException, RealtimeTranscriptionWSError
+from .connection import (
+    RealtimeConnection,
+    RealtimeEvent,
+    UnknownRealtimeEvent,
+    parse_realtime_event,
+)
+
+
+class RealtimeTranscription:
+    """Client for realtime transcription over WebSocket (websockets >= 13.0)."""
+
+    def __init__(self, sdk_config: SDKConfiguration) -> None:
+        self._sdk_config = sdk_config
+
+    def _build_url(
+        self,
+        model: str,
+        *,
+        server_url: Optional[str],
+        query_params: Mapping[str, str],
+    ) -> str:
+        if server_url is not None:
+            base_url = utils.remove_suffix(server_url, "/")
+        else:
+            base_url, _ = self._sdk_config.get_server_details()
+
+        url = generate_url(base_url, "/v1/audio/transcriptions/realtime", None)
+
+        parsed = urlparse(url)
+        merged = dict(parse_qsl(parsed.query, keep_blank_values=True))
+        merged["model"] = model
+        merged.update(dict(query_params))
+
+        return urlunparse(parsed._replace(query=urlencode(merged)))
+
+    async def connect(
+        self,
+        model: str,
+        audio_format: Optional[AudioFormat] = None,
+        server_url: Optional[str] = None,
+        timeout_ms: Optional[int] = None,
+        http_headers: Optional[Mapping[str, str]] = None,
+    ) -> RealtimeConnection:
+        if timeout_ms is None:
+            timeout_ms = self._sdk_config.timeout_ms
+
+        security = self._sdk_config.security
+        if security is not None and callable(security):
+            security = security()
+
+        resolved_security = get_security_from_env(security, models.Security)
+
+        headers: dict[str, str] = {}
+        query_params: dict[str, str] = {}
+
+        if resolved_security is not None:
+            security_headers, security_query = get_security(resolved_security)
+            headers |= security_headers
+            for key, values in security_query.items():
+                if values:
+                    query_params[key] = values[-1]
+
+        if http_headers is not None:
+            headers |= dict(http_headers)
+
+        url = self._build_url(model, server_url=server_url, query_params=query_params)
+
+        parsed = urlparse(url)
+        if parsed.scheme == "https":
+            parsed = parsed._replace(scheme="wss")
+        elif parsed.scheme == "http":
+            parsed = parsed._replace(scheme="ws")
+        ws_url = urlunparse(parsed)
+        open_timeout = None if timeout_ms is None else timeout_ms / 1000.0
+        user_agent = self._sdk_config.user_agent
+
+        websocket: Optional[ClientConnection] = None
+        try:
+            websocket = await connect(
+                ws_url,
+                additional_headers=dict(headers),
+                open_timeout=open_timeout,
+                user_agent_header=user_agent,
+            )
+
+            session, initial_events = await _recv_handshake(
+                websocket, timeout_ms=timeout_ms
+            )
+            connection = RealtimeConnection(
+                websocket=websocket,
+                session=session,
+                initial_events=initial_events,
+            )
+
+            if audio_format is not None:
+                await connection.update_session(audio_format)
+
+            return connection
+
+        except RealtimeTranscriptionException:
+            if websocket is not None:
+                await websocket.close()
+            raise
+        except Exception as exc:
+            if websocket is not None:
+                await websocket.close()
+            raise RealtimeTranscriptionException(f"Failed to connect: {exc}") from exc
+
+    async def transcribe_stream(
+        self,
+        audio_stream: AsyncIterator[bytes],
+        model: str,
+        audio_format: Optional[AudioFormat] = None,
+        server_url: Optional[str] = None,
+        timeout_ms: Optional[int] = None,
+        http_headers: Optional[Mapping[str, str]] = None,
+    ) -> AsyncIterator[RealtimeEvent]:
+        """
+        Flow
+          - opens connection
+          - streams audio in background
+          - yields events from the connection
+        """
+        async with await self.connect(
+            model=model,
+            audio_format=audio_format,
+            server_url=server_url,
+            timeout_ms=timeout_ms,
+            http_headers=http_headers,
+        ) as connection:
+
+            async def _send() -> None:
+                async for chunk in audio_stream:
+                    if connection.is_closed:
+                        break
+                    await connection.send_audio(chunk)
+                await connection.end_audio()
+
+            send_task = asyncio.create_task(_send())
+
+            try:
+                async for event in connection:
+                    yield event
+
+                    # stop early (caller still sees the terminating event)
+                    if isinstance(event, RealtimeTranscriptionError):
+                        break
+                    if getattr(event, "type", None) == "transcription.done":
+                        break
+            finally:
+                send_task.cancel()
+                try:
+                    await send_task
+                except asyncio.CancelledError:
+                    pass
+                await connection.close()
+
+
+def _extract_error_message(payload: dict) -> str:
+    err = payload.get("error")
+    if isinstance(err, dict):
+        msg = err.get("message")
+        if isinstance(msg, str):
+            return msg
+        if isinstance(msg, dict):
+            detail = msg.get("detail")
+            if isinstance(detail, str):
+                return detail
+    return "Realtime transcription error"
+
+
+async def _recv_handshake(
+    websocket: ClientConnection,
+    *,
+    timeout_ms: Optional[int],
+) -> tuple[RealtimeTranscriptionSession, list[RealtimeEvent]]:
+    """
+    Read messages until session.created or error.
+    Replay all messages read during handshake as initial events (lossless).
+    """
+    timeout_s = None if timeout_ms is None else timeout_ms / 1000.0
+    deadline = None if timeout_s is None else (time.monotonic() + timeout_s)
+
+    initial_events: list[RealtimeEvent] = []
+
+    def remaining() -> Optional[float]:
+        if deadline is None:
+            return None
+        return max(0.0, deadline - time.monotonic())
+
+    try:
+        while True:
+            raw = await asyncio.wait_for(websocket.recv(), timeout=remaining())
+            text = (
+                raw.decode("utf-8", errors="replace")
+                if isinstance(raw, (bytes, bytearray))
+                else raw
+            )
+
+            try:
+                payload = json.loads(text)
+            except Exception as exc:
+                initial_events.append(
+                    UnknownRealtimeEvent(
+                        type=None, content=text, error=f"invalid JSON: {exc}"
+                    )
+                )
+                continue
+
+            msg_type = payload.get("type") if isinstance(payload, dict) else None
+            if msg_type == "error" and isinstance(payload, dict):
+                parsed = parse_realtime_event(payload)
+                initial_events.append(parsed)
+                if isinstance(parsed, RealtimeTranscriptionError):
+                    raise RealtimeTranscriptionWSError(
+                        _extract_error_message(payload),
+                        payload=parsed,
+                        raw=payload,
+                    )
+                raise RealtimeTranscriptionWSError(
+                    _extract_error_message(payload),
+                    payload=None,
+                    raw=payload,
+                )
+
+            event = parse_realtime_event(payload)
+            initial_events.append(event)
+
+            if isinstance(event, RealtimeTranscriptionSessionCreated):
+                return event.session, initial_events
+
+    except asyncio.TimeoutError as exc:
+        raise RealtimeTranscriptionException(
+            "Timeout waiting for session creation."
+        ) from exc
+    except RealtimeTranscriptionException:
+        raise
+    except Exception as exc:
+        raise RealtimeTranscriptionException(
+            f"Unexpected websocket handshake failure: {exc}"
+        ) from exc