Realtime Flow

src/agent_session/main.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
 """
-Agent Session for Avurna Flow (Final Corrected Version with HTTP Session Management)
+Agent Session for Avurna Flow (Final Corrected Version with TARGETED HTTP Session Management)
 """
 import os
 import json
 import traceback
-import aiohttp # 1. Import aiohttp
+import aiohttp
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, BackgroundTasks, Request
 from pydantic import BaseModel
@@ -24,14 +24,11 @@ from livekit.plugins.silero import VAD
 from src.agent_session.constants import SYSTEM_PROMPT, GREETING_INSTRUCTIONS
 from src.utils import validate_env_vars
 
-# 2. Create a lifespan manager for the http session
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Create a single, shared session for the entire application
     async with aiohttp.ClientSession() as session:
         app.state.http_session = session
         yield
-    # The session is automatically closed when the app shuts down
 
 app = FastAPI(lifespan=lifespan)
 origins = ["*"]
@@ -47,16 +44,30 @@ class VoiceAssistant(Agent):
 
 async def send_agent_state(room: Room, state: str):
     # This function is correct
-    # ...
-    pass # Omitted for brevity
+    try:
+        msg = json.dumps({"type": "agent_state", "state": state})
+        await room.local_participant.publish_data(msg)
+        print(f"DEBUG: Sent agent state: {state}")
+    except Exception as e:
+        print(f"DEBUG: Error publishing agent state: {e}")
 
 class LLMStateWrapper(LLM):
     # This class is correct
-    # ...
-    pass # Omitted for brevity
+    def __init__(self, llm: LLM, room: Room):
+        super().__init__()
+        self._llm = llm
+        self._room = room
+
+    @asynccontextmanager
+    async def chat(self, **kwargs):
+        await send_agent_state(self._room, "thinking")
+        try:
+            async with self._llm.chat(**kwargs) as stream:
+                yield stream
+        finally:
+            await send_agent_state(self._room, "listening")
 
 
-# 3. Modify run_agent_session to accept the http_session
 async def run_agent_session(room_name: str, agent_token: str, http_session: aiohttp.ClientSession):
     livekit_url = os.getenv("LIVEKIT_URL")
     room = Room()
@@ -69,25 +80,23 @@ async def run_agent_session(room_name: str, agent_token: str, http_session: aioh
         await send_agent_state(room, "listening")
 
         print("DEBUG: 3. Initializing plugins with shared http_session...")
-        # 4. Pass the session to every plugin that makes an external HTTP call
-        llm_wrapper = LLMStateWrapper(
-            llm=GoogleLLM(model="gemini-1.5-flash", session=http_session), 
-            room=room
-        )
-        vad = VAD.load(min_speech_duration=0.1)
-        stt = StreamAdapter(
-            stt=STT(model="whisper-large-v3-turbo", session=http_session), 
-            vad=vad
-        )
-        tts = TTS(
-            voice=VoiceByName(name="Tiktok Fashion Influencer"), 
-            instant_mode=True, 
-            session=http_session
-        )
+        
+        # --- THE DEFINITIVE FIX: Apply the session ONLY where needed ---
+
+        # Google LLM does NOT take a session argument.
+        google_llm = GoogleLLM(model="gemini-1.5-flash")
+
+        # Hume TTS and Groq STT DO require the session.
+        tts = TTS(voice=VoiceByName(name="Tiktok Fashion Influencer"), instant_mode=True, session=http_session)
+        stt = StreamAdapter(stt=STT(model="whisper-large-v3-turbo", session=http_session), vad=VAD.load(min_speech_duration=0.1))
+
+        # Wrap the correctly initialized Google LLM
+        llm_wrapper = LLMStateWrapper(llm=google_llm, room=room)
+
         print("DEBUG: 4. Plugins initialized.")
 
         print("DEBUG: 5. Creating AgentSession...")
-        session = AgentSession(vad=vad, stt=stt, llm=llm_wrapper, tts=tts)
+        session = AgentSession(vad=stt.vad, stt=stt, llm=llm_wrapper, tts=tts)
         print("DEBUG: 6. AgentSession created.")
 
         print("DEBUG: 7. Starting session...")
@@ -106,7 +115,7 @@ async def run_agent_session(room_name: str, agent_token: str, http_session: aioh
         print(f"DEBUG: Agent session for room {room_name} is ending. Cleaning up.")
         await room.disconnect()
 
-# 5. Modify the endpoint to retrieve the session from the app state
+
 @app.post("/join-room")
 async def join_room(req: JoinRoomRequest, request: Request, background_tasks: BackgroundTasks):
     http_session = request.app.state.http_session