Realtime Flow

src/agent_session/main.py

@@ -1,8 +1,5 @@
 #!/usr/bin/env python3
-
-"""
-Agent Session for Avurna Flow (Corrected based on Hume Docs)
-"""
+"""Agent Session for Avurna Flow (Fixed HTTP Session Management)"""
 
 import asyncio
 import os
@@ -21,26 +18,56 @@ from livekit.plugins.google import LLM as GoogleLLM
 from livekit.plugins.groq import STT
 from livekit.plugins.hume import TTS, VoiceByName, VoiceProvider
 from livekit.plugins.silero import VAD
-
 from src.agent_session.constants import SYSTEM_PROMPT, GREETING_INSTRUCTIONS
 from src.utils import validate_env_vars
 
-# FastAPI app setup
+# Global HTTP session for the entire application
+_global_http_session: aiohttp.ClientSession | None = None
+
+async def get_http_session() -> aiohttp.ClientSession:
+    """Get or create the global HTTP session"""
+    global _global_http_session
+    if _global_http_session is None or _global_http_session.closed:
+        _global_http_session = aiohttp.ClientSession()
+    return _global_http_session
+
+async def cleanup_http_session():
+    """Clean up the global HTTP session"""
+    global _global_http_session
+    if _global_http_session and not _global_http_session.closed:
+        await _global_http_session.close()
+        _global_http_session = None
+
+# FastAPI app setup with proper lifecycle management
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    async with aiohttp.ClientSession() as session:
-        app.state.http_session = session
-        yield
+    # Startup: Initialize the global HTTP session
+    await get_http_session()
+    print("HTTP session initialized")
+    
+    yield
+    
+    # Shutdown: Clean up the HTTP session
+    await cleanup_http_session()
+    print("HTTP session cleaned up")
 
 app = FastAPI(lifespan=lifespan)
+
 origins = ["*"]
-app.add_middleware(CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"]
+)
 
 class JoinRoomRequest(BaseModel):
     room_name: str
     agent_token: str
 
 async def send_agent_state(room: Room, state: str):
+    """Send agent state to the room"""
     try:
         msg = json.dumps({"type": "agent_state", "state": state})
         await room.local_participant.publish_data(msg)
@@ -52,70 +79,120 @@ class VoiceAssistant(Agent):
     def __init__(self):
         super().__init__(instructions=SYSTEM_PROMPT)
 
-async def entrypoint(ctx: JobContext) -> None:
-    """
-    Configure and run STT, LLM, and TTS in a LiveKit session.
-    """
-    await ctx.connect()
-
+class CustomJobContext:
+    """Custom JobContext that properly manages the room connection"""
+    
+    def __init__(self, room_name: str, agent_token: str):
+        self.room_name = room_name
+        self.agent_token = agent_token
+        self.room = Room()
+        self._connected = False
+    
+    async def connect(self):
+        """Connect to the LiveKit room"""
+        if not self._connected:
+            livekit_url = os.getenv("LIVEKIT_URL")
+            if not livekit_url:
+                raise ValueError("LIVEKIT_URL environment variable not set")
+            
+            await self.room.connect(livekit_url, self.agent_token)
+            self._connected = True
+            print(f"Connected to room: {self.room_name}")
+    
+    async def disconnect(self):
+        """Disconnect from the LiveKit room"""
+        if self._connected:
+            await self.room.disconnect()
+            self._connected = False
+            print(f"Disconnected from room: {self.room_name}")
+
+async def create_agent_session(ctx: CustomJobContext) -> AgentSession:
+    """Create and configure the agent session with proper HTTP session management"""
+    
+    # Get the global HTTP session
+    http_session = await get_http_session()
+    
     # Voice-activity detection + buffering for non-streaming STT
     vad = VAD.load(
-        min_speech_duration=0.1, 
+        min_speech_duration=0.1,
         min_silence_duration=0.5
     )
-
+    
+    # Create TTS with explicit HTTP session
+    tts = TTS(
+        voice=VoiceByName(
+            name="Male English Actor",
+            provider=VoiceProvider.hume,
+        ),
+        instant_mode=True,
+        # Pass the HTTP session to avoid the context error
+        session=http_session
+    )
+    
+    # Create the agent session
     session = AgentSession(
         vad=vad,
         stt=StreamAdapter(
             stt=STT(
-                model="whisper-large-v3-turbo", 
+                model="whisper-large-v3-turbo",
                 language="en",
             ),
             vad=vad,
         ),
         llm=GoogleLLM(
-            model="gemini-2.5-flash", 
+            model="gemini-2.5-flash",
             temperature=0.5,
         ),
-        tts=TTS(
-            voice=VoiceByName(
-                name="Male English Actor",
-                provider=VoiceProvider.hume,
-            ), 
-            instant_mode=True
-        ),
+        tts=tts,
     )
+    
+    return session
 
-    # Start the session and generate greeting
-    await session.start(agent=VoiceAssistant(), room=ctx.room)
-    await session.generate_reply(instructions=GREETING_INSTRUCTIONS)
+async def entrypoint(ctx: CustomJobContext) -> None:
+    """Configure and run STT, LLM, and TTS in a LiveKit session"""
     
-    # Keep the session alive by waiting for it to complete
-    # The session will run until the room is disconnected
     try:
-        # Wait indefinitely - the session manages its own lifecycle
-        await asyncio.Event().wait()
-    except asyncio.CancelledError:
-        pass
+        # Connect to the room
+        await ctx.connect()
+        
+        # Send initial state
+        await send_agent_state(ctx.room, "listening")
+        
+        # Create the agent session
+        session = await create_agent_session(ctx)
+        
+        # Start the session
+        await session.start(agent=VoiceAssistant(), room=ctx.room)
+        
+        # Send greeting state
+        await send_agent_state(ctx.room, "thinking")
+        
+        # Generate greeting
+        await session.generate_reply(instructions=GREETING_INSTRUCTIONS)
+        
+        # Back to listening
+        await send_agent_state(ctx.room, "listening")
+        
+        # Keep the session alive
+        print("Agent session started successfully, waiting for interactions...")
+        
+        # Monitor room connection and keep alive
+        while ctx.room.connection_state.name in ['CONNECTED', 'CONNECTING']:
+            await asyncio.sleep(1)
+            
+        print("Room disconnected, ending agent session")
+        
+    except Exception as e:
+        print(f"Error in agent session: {e}")
+        print(traceback.format_exc())
+        await send_agent_state(ctx.room, "error")
+    finally:
+        # Clean up
+        await ctx.disconnect()
 
-# Custom entrypoint for FastAPI integration
 async def run_agent_with_room(room_name: str, agent_token: str):
-    """
-    Custom entrypoint that connects to a specific room and keeps it alive
-    """
-    livekit_url = os.getenv("LIVEKIT_URL")
-    
-    # Create a mock JobContext for the specific room
-    class MockJobContext:
-        def __init__(self, room_name: str, agent_token: str):
-            self.room_name = room_name
-            self.agent_token = agent_token
-            self.room = Room()
-            
-        async def connect(self):
-            await self.room.connect(livekit_url, self.agent_token)
-    
-    ctx = MockJobContext(room_name, agent_token)
+    """Run the agent in a specific room"""
+    ctx = CustomJobContext(room_name, agent_token)
     
     try:
         await entrypoint(ctx)
@@ -123,7 +200,7 @@ async def run_agent_with_room(room_name: str, agent_token: str):
         print(f"FATAL ERROR in agent session: {e}")
         print(traceback.format_exc())
     finally:
-        await ctx.room.disconnect()
+        await ctx.disconnect()
 
 @app.post("/join-room")
 async def join_room(req: JoinRoomRequest, background_tasks: BackgroundTasks):
@@ -135,6 +212,31 @@ async def join_room(req: JoinRoomRequest, background_tasks: BackgroundTasks):
 async def root():
     return {"status": "avurna_agent_server_online"}
 
+@app.get("/health")
+async def health():
+    """Health check endpoint"""
+    global _global_http_session
+    session_status = "healthy" if _global_http_session and not _global_http_session.closed else "needs_init"
+    return {
+        "status": "healthy",
+        "http_session": session_status,
+        "timestamp": asyncio.get_event_loop().time()
+    }
+
 if __name__ == "__main__":
-    validate_env_vars(["HUME_API_KEY", "LIVEKIT_URL", "LIVEKIT_API_KEY", "LIVEKIT_API_SECRET", "GROQ_API_KEY", "GOOGLE_API_KEY"])
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # Validate required environment variables
+    required_vars = [
+        "HUME_API_KEY", 
+        "LIVEKIT_URL", 
+        "LIVEKIT_API_KEY", 
+        "LIVEKIT_API_SECRET", 
+        "GROQ_API_KEY", 
+        "GOOGLE_API_KEY"
+    ]
+    
+    validate_env_vars(required_vars)
+    
+    print("Starting Avurna Agent Server...")
+    print("Required environment variables validated")
+    
+    uvicorn.run(app, host="0.0.0.0", port=7860)