properly type the output track

Author: tbarbugli

agents-core/vision_agents/core/agents/agents.py

@@ -2,10 +2,9 @@
 import logging
 import time
 import uuid
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import uuid4
 
-import aiortc
 import getstream.models
 from aiortc import VideoStreamTrack
 from getstream.video.rtc import Call
@@ -15,7 +14,7 @@
 from getstream.video.rtc.pb.stream.video.sfu.models.models_pb2 import TrackType
 from ..edge import sfu_events
 from ..edge.events import AudioReceivedEvent, TrackAddedEvent, CallEndedEvent
-from ..edge.types import Connection, Participant, PcmData, User
+from ..edge.types import Connection, Participant, PcmData, User, OutputAudioTrack
 from ..events.manager import EventManager
 from ..llm import events as llm_events
 from ..llm.events import (
@@ -161,7 +160,7 @@ def __init__(
         self._callback_executed = False
         self._track_tasks: Dict[str, asyncio.Task] = {}
         self._connection: Optional[Connection] = None
-        self._audio_track: Optional[aiortc.AudioStreamTrack] = None
+        self._audio_track: Optional[OutputAudioTrack] = None
         self._video_track: Optional[VideoStreamTrack] = None
         self._realtime_connection = None
         self._pc_track_handler_attached: bool = False
@@ -308,15 +307,10 @@ async def on_realtime_agent_speech_transcription(
                 original=event,
             )
 
-        # Listen for TTS audio events and write audio to the output track
         @self.events.subscribe
-        async def _on_tts_audio(event: TTSAudioEvent):
-            try:
-                if self._audio_track and event.audio_data:
-                    track_any = cast(Any, self._audio_track)
-                    await track_any.write(event.audio_data)
-            except Exception as e:
-                self.logger.error(f"Error writing TTS audio to track: {e}")
+        async def _on_tts_audio_write_to_output(event: TTSAudioEvent):
+            if self._audio_track and event and event.audio_data is not None:
+                await self._audio_track.write(event.audio_data)
 
         @self.events.subscribe
         async def on_stt_transcript_event_create_response(event: STTTranscriptEvent):

agents-core/vision_agents/core/edge/edge_transport.py

@@ -1,17 +1,17 @@
 """
 Abstraction for stream vs other services here
 """
+
 import abc
 
 from typing import TYPE_CHECKING, Any, Optional
 
 import aiortc
 from pyee.asyncio import AsyncIOEventEmitter
 
-from vision_agents.core.edge.types import User
+from vision_agents.core.edge.types import User, OutputAudioTrack
 
 if TYPE_CHECKING:
-
     pass
 
 
@@ -31,7 +31,7 @@ async def create_user(self, user: User):
         pass
 
     @abc.abstractmethod
-    def create_audio_track(self):
+    def create_audio_track(self) -> OutputAudioTrack:
         pass
 
     @abc.abstractmethod
@@ -55,6 +55,7 @@ async def create_conversation(self, call: Any, user: User, instructions):
         pass
 
     @abc.abstractmethod
-    def add_track_subscriber(self, track_id: str) -> Optional[aiortc.mediastreams.MediaStreamTrack]:
+    def add_track_subscriber(
+        self, track_id: str
+    ) -> Optional[aiortc.mediastreams.MediaStreamTrack]:
         pass
-

agents-core/vision_agents/core/edge/types.py

@@ -1,5 +1,14 @@
 from dataclasses import dataclass
-from typing import Any, Optional, NamedTuple, Union, Iterator, AsyncIterator
+from typing import (
+    Any,
+    Optional,
+    NamedTuple,
+    Union,
+    Iterator,
+    AsyncIterator,
+    Protocol,
+    runtime_checkable,
+)
 import logging
 
 import numpy as np
@@ -34,6 +43,18 @@ async def close(self):
         pass
 
 
+@runtime_checkable
+class OutputAudioTrack(Protocol):
+    """
+    A protocol describing an output audio track, the actual implementation depends on the edge transported used
+    eg. getstream.video.rtc.audio_track.AudioStreamTrack
+    """
+
+    async def write(self, data: bytes) -> None: ...
+
+    def stop(self) -> None: ...
+
+
 class PcmData(NamedTuple):
     """
     A named tuple representing PCM audio data.

plugins/getstream/vision_agents/plugins/getstream/stream_edge_transport.py

@@ -22,7 +22,7 @@
 
 from vision_agents.core.edge import EdgeTransport, sfu_events
 from vision_agents.plugins.getstream.stream_conversation import StreamConversation
-from vision_agents.core.edge.types import Connection, User
+from vision_agents.core.edge.types import Connection, User, OutputAudioTrack
 from vision_agents.core.events.manager import EventManager
 from vision_agents.core.edge import events
 from vision_agents.core.utils import get_vision_agents_version
@@ -104,7 +104,7 @@ async def _on_track_published(self, event: sfu_events.TrackPublishedEvent):
         track_type_int = event.payload.type  # TrackType enum int from SFU
         expected_kind = self._get_webrtc_kind(track_type_int)
         track_key = (user_id, session_id, track_type_int)
-        is_agent_track = (user_id == self.agent_user_id)
+        is_agent_track = user_id == self.agent_user_id
 
         # First check if track already exists in map (e.g., from previous unpublish/republish)
         if track_key in self._track_map:
@@ -288,7 +288,9 @@ async def on_audio_received(pcm: PcmData, participant: Participant):
         standardize_connection = StreamConnection(connection)
         return standardize_connection
 
-    def create_audio_track(self, framerate: int = 48000, stereo: bool = True):
+    def create_audio_track(
+        self, framerate: int = 48000, stereo: bool = True
+    ) -> OutputAudioTrack:
         return audio_track.AudioStreamTrack(
             framerate=framerate, stereo=stereo
         )  # default to webrtc framerate