Improve transcription demo

Author: AHGIJMKLKKZNPJKQR

examples/transcription/main.py

@@ -1,11 +1,10 @@
-import asyncio
 from contextlib import asynccontextmanager
 from typing import Annotated
 
 from fastapi import Depends, FastAPI
 from transcription.notifier import make_notifier
 from transcription.room import RoomService, fishjam
-from transcription.transcription import transcription_session_factory
+from transcription.worker import async_worker
 
 from fishjam import PeerOptions, SubscribeOptions
 
@@ -20,19 +19,14 @@ def get_room_service():
 
 @asynccontextmanager
 async def lifespan(_app: FastAPI):
-    async with (
-        transcription_session_factory() as session_factory,
-        asyncio.TaskGroup() as tg,
-    ):
+    async with async_worker() as worker:
         global _room_service
-        _room_service = RoomService(session_factory)
+        _room_service = RoomService(worker)
         notifier = make_notifier(_room_service)
-        notifier_task = tg.create_task(notifier.connect())
+        worker.run_in_background(notifier.connect())
 
         yield
 
-        notifier_task.cancel()
-
 
 app = FastAPI(lifespan=lifespan)
 

examples/transcription/pyproject.toml

@@ -3,11 +3,8 @@ name = "transcription"
 version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
-requires-python = ">=3.12"
-dependencies = [
-    "fastapi[standard]==0.116.0",
-    "fishjam-server-sdk",
-]
+requires-python = ">=3.10"
+dependencies = ["fastapi[standard]==0.116.0", "fishjam-server-sdk"]
 
 [tool.uv.sources]
 fishjam-server-sdk = { workspace = true }

examples/transcription/transcription/agent.py

@@ -1,61 +1,55 @@
 import asyncio
-from typing import Callable
 
 from fishjam.agent import Agent, AgentResponseTrackData
+from transcription.worker import BackgroundWorker
 
 from .transcription import TranscriptionSession
 
 
 class TranscriptionAgent:
-    def __init__(
-        self,
-        room_id: str,
-        agent: Agent,
-        on_text: Callable[[str, str], None],
-        session_factory: Callable[[str], TranscriptionSession],
-    ):
+    def __init__(self, room_id: str, agent: Agent, worker: BackgroundWorker):
         self._room_id = room_id
         self._agent = agent
-        self._peers: dict[str, TranscriptionSession] = {}
-        self._on_text = on_text
-        self._session_factory = session_factory
-        self._leave_event = asyncio.Event()
-        self.done = False
+        self._sessions: dict[str, TranscriptionSession] = {}
+        self._disconnect = asyncio.Event()
+        self._worker = worker
 
         @agent.on_track_data
         def _(track_data: AgentResponseTrackData):
-            if track_data.peer_id not in self._peers:
+            if track_data.peer_id not in self._sessions:
                 return
-            self._peers[track_data.peer_id].transcribe(track_data.data)
+            self._sessions[track_data.peer_id].transcribe(track_data.data)
 
-        print(f"Created agent for room {room_id}")
-        self._task = asyncio.create_task(self._run_agent())
-
-    async def _run_agent(self):
-        print(f"Connecting agent to room {self._room_id}")
+    async def _start(self):
         async with self._agent:
             print(f"Agent connected to room {self._room_id}")
-            await self._leave_event.wait()
+            await self._disconnect.wait()
+        self._disconnect.clear()
         print(f"Agent disconnected from room {self._room_id}")
 
+    def _stop(self):
+        self._disconnect.set()
+
+    def _handle_transcription(self, peer_id: str, text: str):
+        print(f"Peer {peer_id} in room {self._room_id} said: {text}")
+
     def on_peer_enter(self, peer_id: str):
-        if peer_id in self._peers:
+        if peer_id in self._sessions:
             return
 
-        print(f"Starting transcription session for peer {peer_id}")
+        if len(self._sessions) == 0:
+            self._worker.run_in_background(self._start())
 
-        self._peers[peer_id] = self._session_factory(peer_id)
+        session = TranscriptionSession(lambda t: self._handle_transcription(peer_id, t))
+        self._sessions[peer_id] = session
+        self._worker.run_in_background(session.start(peer_id))
 
     def on_peer_leave(self, peer_id: str):
-        if peer_id not in self._peers:
+        if peer_id not in self._sessions:
             return
 
-        print(f"Ending transcription session for peer {peer_id}")
-
-        session = self._peers[peer_id]
-        session.end()
-        self._peers.pop(peer_id)
+        self._sessions[peer_id].end()
+        self._sessions.pop(peer_id)
 
-        if len(self._peers) == 0:
-            self._leave_event.set()
-            self.done = True
+        if len(self._sessions) == 0:
+            self._stop()

examples/transcription/transcription/notifier.py

@@ -11,7 +11,6 @@ def make_notifier(room_service: RoomService):
 
     @notifier.on_server_notification
     def _(notification: AllowedNotification):
-        print(f"Received notification {notification}")
         match notification:
             case ServerMessagePeerConnected(peer_id=peer_id, room_id=room_id):
                 handle_peer_connected(peer_id, room_id)
@@ -20,18 +19,11 @@ def _(notification: AllowedNotification):
                 handle_peer_disconnected(peer_id, room_id)
 
     def handle_peer_connected(peer_id: str, room_id: str):
-        if room_id != room_service.room_id:
-            return
-
-        agent = room_service.create_agent()
-        agent.on_peer_enter(peer_id)
+        if room_id == room_service.room.id:
+            room_service.agent.on_peer_enter(peer_id)
 
     def handle_peer_disconnected(peer_id: str, room_id: str):
-        if room_id != room_service.room_id:
-            return
-
-        agent = room_service.get_agent()
-        if agent:
-            agent.on_peer_leave(peer_id)
+        if room_id == room_service.room.id:
+            room_service.agent.on_peer_leave(peer_id)
 
     return notifier

examples/transcription/transcription/room.py

@@ -1,11 +1,9 @@
-from typing import Callable
-
-from fishjam import FishjamClient, FishjamNotifier, Room
+from fishjam import FishjamClient, Room
 from fishjam.errors import NotFoundError
+from transcription.worker import BackgroundWorker
 
 from .agent import TranscriptionAgent
 from .config import FISHJAM_ID, FISHJAM_TOKEN, FISHJAM_URL
-from .transcription import TranscriptionSession
 
 fishjam = FishjamClient(
     FISHJAM_ID,
@@ -15,49 +13,27 @@
 
 
 class RoomService:
-    def __init__(
-        self,
-        session_factory: Callable[[Callable[[str], None]], TranscriptionSession],
-    ):
-        self.room_id = fishjam.create_room().id
-        self._agent: TranscriptionAgent | None = None
-        self._notifier = FishjamNotifier(
-            FISHJAM_ID,
-            FISHJAM_TOKEN,
-            fishjam_url=FISHJAM_URL,
-        )
-
-        def _make_session(peer_id: str):
-            return session_factory(lambda t: self._handle_transcription(peer_id, t))
-
-        self._session_factory = _make_session
+    def __init__(self, worker: BackgroundWorker):
+        self._worker = worker
+        self._create_room()
 
     def get_room(self) -> Room:
         try:
-            return fishjam.get_room(self.room_id)
+            self.room = fishjam.get_room(self.room.id)
         except NotFoundError:
-            self.clear()
-            room = fishjam.create_room()
-            self.room_id = room.id
-            return room
-
-    def clear(self):
-        fishjam.delete_room(self.room_id)
-
-    def _handle_transcription(self, peer_id: str, text: str):
-        print(f"Peer {peer_id} in room {self.room_id} said: {text}")
-
-    def create_agent(self):
-        if not self._agent or self._agent.done:
-            self._agent = TranscriptionAgent(
-                self.room_id,
-                fishjam.create_agent(self.room_id),
-                self._handle_transcription,
-                self._session_factory,
-            )
-        return self._agent
+            self._create_room()
+        return self.room
+
+    def _create_room(self):
+        self.room = fishjam.create_room()
+        self._create_agent()
+
+    def _create_agent(self):
+        self.agent = TranscriptionAgent(
+            self.room.id,
+            fishjam.create_agent(self.room.id),
+            self._worker,
+        )
 
     def get_agent(self):
-        if self._agent.done:
-            self._agent = None
-        return self._agent
+        return self.agent

examples/transcription/transcription/transcription.py

@@ -1,4 +1,4 @@
-from asyncio import Event, Queue, Task, TaskGroup
+from asyncio import Event, Queue, TaskGroup
 from contextlib import asynccontextmanager
 from typing import Callable
 
@@ -17,7 +17,7 @@ def __init__(self, on_text: Callable[[str], None]):
         self._model = TRANSCRIPTION_MODEL
         self._on_text = on_text
 
-    async def start(self):
+    async def start(self, peer_id: str):
         async with self._gemini.aio.live.connect(
             model=self._model,
             config=TRANSCRIPTION_CONFIG,
@@ -26,12 +26,13 @@ async def start(self):
                 send_task = tg.create_task(self._send_loop(session))
                 recv_task = tg.create_task(self._recv_loop(session))
 
-                print("Started transcription session")
+                print(f"Started transcription session for peer {peer_id}")
 
                 await self._end_event.wait()
 
                 send_task.cancel()
                 recv_task.cancel()
+            print(f"Stopped transcription session for peer {peer_id}")
 
     def transcribe(self, audio: bytes):
         self._audio_queue.put_nowait(audio)
@@ -48,27 +49,16 @@ async def _send_loop(self, session: AsyncSession):
 
     async def _recv_loop(self, session: AsyncSession):
         while True:
+            acc = ""
             async for res in session.receive():
                 if (
                     (content := res.server_content)
                     and (transcription := content.input_transcription)
                     and (text := transcription.text)
                 ):
-                    self._on_text(text)
+                    acc += text
 
+            if acc:
+                self._on_text(acc)
 
-@asynccontextmanager
-async def transcription_session_factory():
-    async with TaskGroup() as tg:
-        sessions: list[TranscriptionSession] = []
-
-        def make_session(on_text: Callable[[str], None]):
-            session = TranscriptionSession(on_text)
-            sessions.append(session)
-            tg.create_task(session.start())
             return session
-
-        yield make_session
-
-        for session in sessions:
-            session.end()

examples/transcription/transcription/worker.py

@@ -0,0 +1,30 @@
+from asyncio import Task, TaskGroup
+from contextlib import asynccontextmanager
+from typing import Any, Coroutine
+
+
+class BackgroundWorker:
+    def __init__(self, tg: TaskGroup) -> None:
+        self._tg = tg
+        self._tasks: set[Task[None]] = set()
+
+    def run_in_background(self, coro: Coroutine[Any, Any, None]):
+        task = self._tg.create_task(coro)
+        task.add_done_callback(self._remove_task)
+        self._tasks.add(task)
+
+    def _remove_task(self, task: Task[None]):
+        self._tasks.discard(task)
+
+    def cleanup(self):
+        for task in self._tasks:
+            task.cancel()
+        self._tasks = set()
+
+
+@asynccontextmanager
+async def async_worker():
+    async with TaskGroup() as tg:
+        worker = BackgroundWorker(tg)
+        yield worker
+        worker.cleanup()