fix: typing

Author: jean-malo

src/mistralai/extra/realtime/connection.py

@@ -32,6 +32,7 @@
     TranscriptionStreamSegmentDelta,
     TranscriptionStreamTextDelta,
 )
+from mistralai.types import UNSET
 
 
 class UnknownRealtimeEvent(BaseModel):
@@ -41,6 +42,7 @@ class UnknownRealtimeEvent(BaseModel):
     - invalid JSON payload
     - schema validation failure
     """
+
     type: Optional[str]
     content: Any
     error: Optional[str] = None
@@ -61,7 +63,6 @@ class UnknownRealtimeEvent(BaseModel):
     UnknownRealtimeEvent,
 ]
 
-
 _MESSAGE_MODELS: dict[str, Any] = {
     "session.created": RealtimeTranscriptionSessionCreated,
     "session.updated": RealtimeTranscriptionSessionUpdated,
@@ -113,7 +114,6 @@ def __init__(
     ) -> 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 [])
 
@@ -127,7 +127,7 @@ def session(self) -> RealtimeTranscriptionSession:
 
     @property
     def audio_format(self) -> AudioFormat:
-        return self._audio_format
+        return self._session.audio_format
 
     @property
     def is_closed(self) -> bool:
@@ -163,16 +163,13 @@ async def update_session(
         if audio_format is None and target_streaming_delay_ms is None:
             raise ValueError("At least one session field must be provided")
 
-        session_update_data: dict[str, object] = {}
-        if audio_format is not None:
-            self._audio_format = audio_format
-            session_update_data["audio_format"] = audio_format
-        if target_streaming_delay_ms is not None:
-            session_update_data["target_streaming_delay_ms"] = (
-                target_streaming_delay_ms
-            )
         message = RealtimeTranscriptionSessionUpdateMessage(
-            session=RealtimeTranscriptionSessionUpdatePayload(**session_update_data)
+            session=RealtimeTranscriptionSessionUpdatePayload(
+                audio_format=audio_format if audio_format is not None else UNSET,
+                target_streaming_delay_ms=target_streaming_delay_ms
+                if target_streaming_delay_ms is not None
+                else UNSET,
+            )
         )
         await self._websocket.send(message.model_dump_json())
 
@@ -229,6 +226,7 @@ async def events(self) -> AsyncIterator[RealtimeEvent]:
             await self.close()
 
     def _apply_session_updates(self, ev: RealtimeEvent) -> None:
-        if isinstance(ev, RealtimeTranscriptionSessionCreated) or isinstance(ev, RealtimeTranscriptionSessionUpdated):
+        if isinstance(ev, RealtimeTranscriptionSessionCreated) or isinstance(
+            ev, RealtimeTranscriptionSessionUpdated
+        ):
             self._session = ev.session
-            self._audio_format = ev.session.audio_format