Update app.py

app.py

@@ -1,116 +1,4 @@
-# rentbot/app.py
-import os
-import base64
-import json
-import asyncio
-import numpy as np
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-from dotenv import load_dotenv
-
-from audio_utils import ulaw_to_pcm16
-from stt_handler import transcribe_audio_chunk
-from llm_handler import get_llm_response
-from tts_handler import text_to_speech_stream
-from tool_handler import execute_tool_call
-
-# Load environment variables
-load_dotenv()
-
-app = FastAPI()
-
-# Configuration
-SILENCE_THRESHOLD_SECONDS = 0.7
-AUDIO_RATE = 8000  # Hz for Twilio media streams
-AUDIO_BUFFER_SIZE = int(SILENCE_THRESHOLD_SECONDS * AUDIO_RATE)
-
-# In-memory session storage (for demonstration)
-# In production, use Redis or another persistent store.
-sessions = {}
-
-@app.websocket("/rentbot")
-async def websocket_endpoint(ws: WebSocket):
-    await ws.accept()
-    stream_sid = None
-    audio_buffer = np.array([], dtype=np.int16)
-    
-    try:
-        # Initial greeting
-        # We need a stream_sid to send audio, so we wait for the first 'start' message
-        # A more robust solution might send a pre-recorded greeting or handle this flow differently.
-        
-        async for message in ws.iter_text():
-            data = json.loads(message)
-
-            if data['event'] == 'start':
-                stream_sid = data['start']['streamSid']
-                sessions[stream_sid] = {
-                    "messages": [{"role": "system", "content": os.getenv("SYSTEM_PROMPT")}],
-                    "processing_task": None
-                }
-                print(f"New stream started: {stream_sid}")
-
-                # Send an initial greeting
-                initial_greeting = "Hi! I'm RentBot, your leasing assistant. How can I help you today?"
-                sessions[stream_sid]["messages"].append({"role": "assistant", "content": initial_greeting})
-                
-                async def send_initial_greeting():
-                    tts_iterator = text_to_speech_stream(iter([initial_greeting]))
-                    async for audio_chunk in tts_iterator:
-                        payload = base64.b64encode(audio_chunk).decode('utf-8')
-                        await ws.send_json({
-                            "event": "media",
-                            "streamSid": stream_sid,
-                            "media": {"payload": payload}
-                        })
-                    # Mark the end of the bot's speech
-                    await ws.send_json({"event": "mark", "streamSid": stream_sid, "mark": {"name": "bot_turn_end"}})
-
-                asyncio.create_task(send_initial_greeting())
-
-
-            elif data['event'] == 'media':
-                if not stream_sid: continue
-
-                # Decode the base64 µ-law audio and add to buffer
-                chunk_ulaw = base64.b64decode(data['media']['payload'])
-                chunk_pcm = ulaw_to_pcm16(chunk_ulaw)
-                audio_buffer = np.append(audio_buffer, chunk_pcm)
-
-                # If buffer is full (indicating continuous speech), process it
-                if len(audio_buffer) >= AUDIO_BUFFER_SIZE:
-                    # If a task is already running, let it finish
-                    if sessions[stream_sid]["processing_task"] and not sessions[stream_sid]["processing_task"].done():
-                        continue # Skip starting a new task
-
-                    # Start processing in a background task
-                    task = asyncio.create_task(process_user_audio(ws, stream_sid, audio_buffer))
-                    sessions[stream_sid]["processing_task"] = task
-                    audio_buffer = np.array([], dtype=np.int16) # Reset buffer
-
-            elif data['event'] == 'mark':
-                # This indicates the user has likely paused. Let's process any remaining audio.
-                if not stream_sid: continue
-                
-                if len(audio_buffer) > 1000: # Heuristic: process if there's meaningful audio left
-                    if not (sessions[stream_sid]["processing_task"] and not sessions[stream_sid]["processing_task"].done()):
-                        task = asyncio.create_task(process_user_audio(ws, stream_sid, audio_buffer))
-                        sessions[stream_sid]["processing_task"] = task
-                        audio_buffer = np.array([], dtype=np.int16) # Reset buffer
-
-            elif data['event'] == 'stop':
-                print(f"Stream stopped: {stream_sid}")
-                break
-
-    except WebSocketDisconnect:
-        print(f"WebSocket disconnected for stream {stream_sid}")
-    except Exception as e:
-        print(f"An error occurred: {e}")
-    finally:
-        if stream_sid and stream_sid in sessions:
-            if sessions[stream_sid]["processing_task"]:
-                sessions[stream_sid]["processing_task"].cancel()
-            del sessions[stream_sid]
-        print(f"Session cleaned up for stream {stream_sid}")
+# In rentbot/app.py, find the process_user_audio function and replace it with this:
 
 async def process_user_audio(ws: WebSocket, stream_sid: str, audio_chunk: np.ndarray):
     """The main logic loop: STT -> LLM -> (Tool/TTS)"""
@@ -125,45 +13,37 @@ async def process_user_audio(ws: WebSocket, stream_sid: str, audio_chunk: np.nda
     print(f"[{stream_sid}] User said: {user_text}")
     sessions[stream_sid]["messages"].append({"role": "user", "content": user_text})
 
-    # 2. Text-to-Assistant Reply (LLM)
-    llm_response_generator = get_llm_response(sessions[stream_sid]["messages"])
-    
-    # We need to accumulate the text to know if there's a tool call
-    text_chunks_for_tts = []
-    assistant_message = None
-    tool_calls = None
-
-    async def llm_logic_handler():
-        nonlocal assistant_message, tool_calls
-        # This async generator will yield text chunks and then return the final message and tool calls
-        # This is a slightly advanced pattern to handle both streaming and final return values
-        class ResponseHandler:
-            def __init__(self, generator):
-                self._generator = generator
-                self.final_result = None
-
-            async def __aiter__(self):
-                return self
+    # This queue will hold text chunks from the LLM for the TTS to process.
+    tts_queue = asyncio.Queue()
 
-            async def __anext__(self):
-                try:
-                    return await self._generator.__anext__()
-                except StopAsyncIteration as e:
-                    self.final_result = e.value
-                    raise
+    # This async function is the "handler" that get_llm_response will call.
+    async def llm_chunk_handler(chunk):
+        await tts_queue.put(chunk)
 
-        response_handler = ResponseHandler(llm_response_generator)
+    # This async generator reads from the queue for the TTS stream.
+    async def tts_text_iterator():
+        while True:
+            chunk = await tts_queue.get()
+            if chunk is None:  # Use None as a sentinel to end the stream
+                break
+            yield chunk
 
-        async for chunk in response_handler:
-            text_chunks_for_tts.append(chunk)
-        
-        assistant_message, tool_calls = response_handler.final_result
+    # 2. Start LLM and TTS tasks concurrently
+    llm_task = asyncio.create_task(
+        get_llm_response(sessions[stream_sid]["messages"], llm_chunk_handler)
+    )
+    tts_task = asyncio.create_task(
+        stream_and_send_audio(ws, stream_sid, tts_text_iterator())
+    )
 
-    # Run the LLM logic and TTS streaming concurrently
-    llm_task = asyncio.create_task(llm_logic_handler())
-    tts_task = asyncio.create_task(stream_and_send_audio(ws, stream_sid, (chunk for chunk in text_chunks_for_tts)))
+    # Wait for the LLM to finish and get the final message and tool calls
+    assistant_message, tool_calls = await llm_task
 
-    await asyncio.gather(llm_task, tts_task) # Wait for both to complete
+    # Now that the LLM is done, signal the TTS to finish by putting None on the queue
+    await tts_queue.put(None)
+    
+    # Wait for the TTS to finish sending all audio
+    await tts_task
 
     # Add the full assistant response to history
     if assistant_message and assistant_message.get("content"):
@@ -171,47 +51,39 @@ async def process_user_audio(ws: WebSocket, stream_sid: str, audio_chunk: np.nda
 
     # 3. Handle Tool Calls if any
     if tool_calls:
-        # Add the assistant's request for a tool call to history
-        sessions[stream_sid]["messages"].append(assistant_message)
-        
-        for tool_call in tool_calls:
+        for tool_call_data in tool_calls:
+            # Recreate the tool_call object structure for execute_tool_call
+            tool_call = type('ToolCall', (), {
+                'id': tool_call_data['id'], 
+                'function': type('Function', (), tool_call_data['function'])
+            })
+            
             print(f"[{stream_sid}] Executing tool: {tool_call.function.name}")
             tool_result_message = execute_tool_call(tool_call)
             sessions[stream_sid]["messages"].append(tool_result_message)
 
         # 4. Get a final response from the LLM after executing the tool
-        final_response_generator = get_llm_response(sessions[stream_sid]["messages"])
+        final_tts_queue = asyncio.Queue()
+        async def final_llm_chunk_handler(chunk):
+            await final_tts_queue.put(chunk)
+            
+        async def final_tts_iterator():
+            while True:
+                chunk = await final_tts_queue.get()
+                if chunk is None:
+                    break
+                yield chunk
+
+        final_llm_task = asyncio.create_task(
+            get_llm_response(sessions[stream_sid]["messages"], final_llm_chunk_handler)
+        )
+        final_tts_task = asyncio.create_task(
+            stream_and_send_audio(ws, stream_sid, final_tts_iterator())
+        )
         
-        final_text_chunks = []
-        async for chunk in final_response_generator:
-            final_text_chunks.append(chunk)
+        final_assistant_message, _ = await final_llm_task
+        await final_tts_queue.put(None)
+        await final_tts_task
 
-        # Add final response to history
-        final_assistant_message, _ = await final_response_generator
         if final_assistant_message:
-            sessions[stream_sid]["messages"].append(final_assistant_message)
-
-        # Stream the final response audio
-        await stream_and_send_audio(ws, stream_sid, iter(final_text_chunks))
-
-
-async def stream_and_send_audio(ws: WebSocket, stream_sid: str, text_iterator):
-    """Stream text to TTS and send the resulting audio back over the WebSocket."""
-    async for audio_chunk in text_to_speech_stream(text_iterator):
-        if audio_chunk:
-            payload = base64.b64encode(audio_chunk).decode('utf-8')
-            await ws.send_json({
-                "event": "media",
-                "streamSid": stream_sid,
-                "media": {"payload": payload}
-            })
-    
-    # Mark the end of the bot's turn to let Twilio know it can listen for the user again
-    await ws.send_json({"event": "mark", "streamSid": stream_sid, "mark": {"name": "bot_turn_end"}})
-    print(f"[{stream_sid}] Finished sending bot's audio turn.")
-
-
-if __name__ == "__main__":
-    import uvicorn
-    print("Starting RentBot server...")
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+            sessions[stream_sid]["messages"].append(final_assistant_message)