main_stream_draft.py

import asyncio
import base64
import json
import logging
import queue
import threading
from fastapi import FastAPI, WebSocket, Request, WebSocketDisconnect
from fastapi.responses import Response
from twilio.twiml.voice_response import VoiceResponse, Connect, Stream
from dotenv import load_dotenv

from database import Database
from config import Config
from ai_agent import SwiftLogisticsAgent
from streaming_services import StreamingSpeechService

# Load environment variables
load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="SwiftLogistics Voice AI - Streaming")

config = Config()
# Re-use existing DB and Agent setup
db = Database(config.database_url)
speech_service = StreamingSpeechService(config.google_credentials_path)
ai_agent = SwiftLogisticsAgent(db)

@app.on_event("startup")
async def startup_event():
    await db.initialize()
    logger.info("Database initialized")

@app.on_event("shutdown")
async def shutdown_event():
    await db.close()

@app.post("/webhook/voice")
async def handle_voice_call(request: Request):
    """Initial TwiML response to start the WebSocket stream"""
    response = VoiceResponse()
    
    # 1. Say a quick greeting (optional, or let the stream handle it)
    response.say("Welcome to AeroLogistics. Connecting you.")
    
    # 2. Connect to the stream
    connect = Connect()
    stream = Stream(url=f"wss://{request.headers.get('host')}/stream")
    connect.append(stream)
    response.append(connect)
    
    return Response(content=str(response), media_type="application/xml")

@app.websocket("/stream")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    logger.info("WebSocket connected")
    
    # Queues for communicating between WebSocket (async) and Google Speech (sync thread)
    audio_queue = queue.Queue()
    transcript_queue = asyncio.Queue()
    
    # State for this call
    call_sid = None
    stream_sid = None
    conversation_state = 'greeting'
    context = {}
    
    # Create a loop reference for thread-safe enqueuing
    loop = asyncio.get_event_loop()
    
    def on_transcript(text):
        """Callback to put transcript into async queue"""
        # This runs in the background thread
        if text:
            # Schedule putting the item in the async queue
            asyncio.run_coroutine_threadsafe(transcript_queue.put(text), loop)

    def speech_recognition_thread():
        """Runs the blocking Google Streaming Recognize loop"""
        try:
            # We iterate over the generator
            for transcript in asyncio.run(speech_service.transcribe_stream_wrapper(audio_queue)):
                on_transcript(transcript)
        except Exception as e:
            logger.error(f"Error in speech thread: {e}")

    # We need a wrapper in streaming_services because transcribe_stream is async generator but uses blocking calls?
    # Actually my previous streaming_services.py defined transcribe_stream as async but used blocking calls.
    # It's better to run it entirely in a thread.
    # Let's fix streaming_services usage here.
    
    # Start the speech recognition thread
    recognition_thread = threading.Thread(target=run_speech_recognition, args=(audio_queue, on_transcript))
    recognition_thread.start()

    try:
        # Main loop: Interleave receiving audio and processing transcripts
        
        # We need a concurrent consumer for transcripts so we don't block reading websocket
        transcript_task = asyncio.create_task(process_transcripts(websocket, transcript_queue, context, conversation_state))
        
        while True:
            # Wait for message from Twilio
            data = await websocket.receive_text()
            message = json.loads(data)
            
            if message['event'] == 'start':
                call_sid = message['start']['callSid']
                stream_sid = message['start']['streamSid']
                logger.info(f"Stream started: {stream_sid}")
                
            elif message['event'] == 'media':
                # Decode audio and put in queue for Google
                # message['media']['payload'] is base64 mulaw
                chunk = base64.b64decode(message['media']['payload'])
                audio_queue.put(chunk)
                
            elif message['event'] == 'stop':
                logger.info("Stream stopped")
                audio_queue.put(None) # Signal thread to stop
                break
                
    except WebSocketDisconnect:
        logger.info("WebSocket disconnected")
        audio_queue.put(None) 
    except Exception as e:
        logger.error(f"Error in websocket loop: {e}")
        audio_queue.put(None)
    finally:
        transcript_task.cancel()
        recognition_thread.join()

def run_speech_recognition(audio_queue, callback):
    """
    Helper to run the sync google blocking code. 
    We need to instantiate a fresh generator here because of thread boundaries.
    """
    try:
        # Create a generator that pulls from the queue
        def request_generator():
            while True:
                chunk = audio_queue.get()
                if chunk is None:
                    return
                # Wrap in StreamingRecognizeRequest
                yield speech.StreamingRecognizeRequest(audio_content=chunk)

        service = speech_service # Use the global one
        responses = service.speech_client.streaming_recognize(
            config=service.streaming_config,
            requests=request_generator()
        )
        
        for response in responses:
            if not response.results:
                continue
            result = response.results[0]
            if not result.alternatives:
                continue
            if result.is_final:
                transcript = result.alternatives[0].transcript
                callback(transcript)
                
    except Exception as e:
        logger.error(f"Error in run_speech_recognition: {e}")

async def process_transcripts(websocket, queue, context, conversation_state):
    """Consumer loop for transcripts"""
    try:
        while True:
            text = await queue.get()
            logger.info(f"User said: {text}")
            
            # --- AI LOGIC ---
            # Reuse the existing agent logic
            # We need to handle state updates. Since state is local variable in websocket_endpoint,
            # we'd need a class or mutable dict. 
            # For this MVP, we just use the mutable 'context' dict and assume 'next_state' is tracked there 
            # or passed back. 
            
            # WORKAROUND: The database/agent state needs a persistent reference.
            # In the real code we'd need to update the outer variables.
            # For now, let's treat conversation_state is passed in purely for reading? 
            # No, it changes. We need a StateManager.
            
            # Let's create a simple session object
            # Note: We can't easily update 'conversation_state' variable in outer scope.
            # We will use context['state'] to track it.
            
            current_state = context.get('state', 'greeting')
            
            response = await ai_agent.process_message(text, current_state, context)
            
            # Update state
            context['state'] = response['next_state']
            context.update(response['context'])
            
            message_text = response['message']
            
            # --- TTS ---
            # Convert text to audio
            audio_bytes = await speech_service.text_to_speech_stream(message_text)
            
            # --- SEND TO TWILIO ---
            # Base64 encode
            payload = base64.b64encode(audio_bytes).decode('utf-8')
            
            msg = {
                "event": "media",
                "streamSid": "STREAM_SID_PLACEHOLDER", # We need the streamSid!
                "media": {
                    "payload": payload
                }
            }
            # We need streamSid. It's in the main loop variables. 
            # We should pass a "send_audio" callback to this consumer function.
            await websocket.send_json(msg)
            
            # Also send a "mark" event if needed to handle interruption
            
    except asyncio.CancelledError:
        pass
    except Exception as e:
        logger.error(f"Error processing transcript: {e}")

# This file currently relies on 'stream_sid' being available in process_transcripts
# which is hard due to scoping.
# I will refactor `websocket_endpoint` to be a class or cleaner function.