src/internal/rtc/rtc_call.py

import fastapi
from fastapi.middleware.cors import CORSMiddleware
from fastrtc import ReplyOnPause, Stream, AlgoOptions, SileroVadOptions, get_cloudflare_turn_credentials_async, get_cloudflare_turn_credentials
from src.utils.audio_helper import audio_to_bytes, resample_audio
from dotenv import load_dotenv
from src.internal.rag.chat_template import get_chat_template
from src.internal.tts.base_tts import TTS
from src.internal.stt.base_stt import STT
from src.internal.agents import Agent, AgentRequest, BankingCRUDAgent, CSAgent
import logging
import time
import platform
import socket
import os
import numpy as np
import asyncio
import asyncio
from src.config.constant import HF_TOKEN
import re

load_dotenv()
logging.basicConfig(level=logging.INFO)

class RTCHandler:
    def __init__(self, agent : Agent , stt: STT, tts : TTS):
        self.agent = agent
        self.stt = stt
        self.tts = tts
        self.full_response = ""
        self.stream = None
        self.app = None

        self._setup_webrtc_ip()

    def _setup_webrtc_ip(self):
        """Setup WebRTC IP for Windows"""
        if platform.system() == 'Windows':
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try:
                s.connect(('8.8.8.8', 80))
                local_ip = s.getsockname()[0]
            except Exception:
                local_ip = '127.0.0.1'
            finally:
                s.close()
            os.environ['WEBRTC_IP'] = local_ip

    
    def echo(self, audio):
            try:
                chat_memory = []
                stt_time = time.time()
                logging.info("Performing STT")

                transcription = self.stt.transcribe(audio_to_bytes(audio))
                prompt = transcription
                if prompt == "":
                    logging.info("STT returned empty string")
                    return

                logging.info(f"STT response: {transcription}")
                logging.info(f"STT took {time.time() - stt_time} seconds")

                llm_time = time.time()

                self.full_response = ""
                router_agent_request = None
                if(type(self.agent) is BankingCRUDAgent):
                    router_agent_request = AgentRequest(
                    chat_memory = chat_memory,
                    prompt_template = {
                        "api_banking_template" : get_chat_template("api_banking"),
                        "responder_template" : get_chat_template("responder_banking")
                    },
                    question = prompt
                )
                else:
                    router_agent_request = AgentRequest(
                        chat_memory = chat_memory,
                        prompt_template = get_chat_template("customer_service"),
                        question = prompt
                    )

                async def stream_text_to_audio():
                    chunk_size = 1024
                    text_buffer = ""
                    async for stream_data in self.agent.get_result(router_agent_request):
                        if stream_data["type"] == "chunk":
                            chunk = stream_data["data"]["chunk"]
                            self.full_response += chunk
                            text_buffer += chunk
                            
                            if re.search(r'[.,?;!]', chunk):
                                try:
                                    audio_buffer_gen =  await self.tts.generate_audio_buffer(text_buffer)
                                    audio_buffer = audio_buffer_gen[0]
                                    resampled = resample_audio(audio_buffer)
                                    for i in range(0, len(resampled), chunk_size):
                                        yield (24000, resampled[i:i + chunk_size])
                                    no_buffer = 0
                                    text_buffer = ""
                                except Exception as e:
                                    logging.error(f"TTS generation failed for chunk: {e}")
                                    continue
                                    
                        elif stream_data["type"] == "metadata":
                            setup_time = stream_data['data']['setup_time']
                            print(f"\nSetup completed in {setup_time:.2f}s")
                            
                        elif stream_data["type"] == "complete":
                            total_time = stream_data['data']['total_time']
                            print(f"\nTotal time: {total_time:.2f}s")
                            break

                
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                
                try:
                    async_gen = stream_text_to_audio()
                    while True:
                        try:
                            chunk = loop.run_until_complete(async_gen.__anext__())
                            yield chunk
                        except StopAsyncIteration:
                            break
                finally:
                    loop.close()
                if(len(chat_memory) >= 15):
                    chat_memory = []
                chat_memory.append({"role": "assistant", "content": self.full_response + " "})
                logging.info(f"LLM response: {self.full_response}")
                logging.info(f"LLM took {time.time() - llm_time} seconds")

            except Exception as e:
                logging.error(f"Error in echo function: {e}")
                error_audio = np.zeros(24000, dtype=np.float32)
                yield (24000, error_audio)

    def reset_conversation(self):
        logging.info("Resetting chat")
        self.messages = [{"role": "system", "content": self.sys_prompt}]
        self.full_response = ""

    def create_stream(self):
        try:
            async def get_credentials():
                return await get_cloudflare_turn_credentials_async(hf_token=HF_TOKEN)
            self.stream = Stream(
                rtc_configuration=get_credentials,
                server_rtc_configuration=get_cloudflare_turn_credentials(ttl=360_000),
                handler = ReplyOnPause(
                    self.echo,
                    algo_options=AlgoOptions(
                        audio_chunk_duration=0.5,
                        started_talking_threshold=0.1,
                        speech_threshold=0.03
                    ),
                    model_options=SileroVadOptions(
                        threshold=0.90,
                        min_speech_duration_ms=250,
                        min_silence_duration_ms=2000,
                        speech_pad_ms=400,
                        max_speech_duration_s=15
                    )
                ),
                modality="audio",
                mode="send-receive",
                ui_args={"title": "Sakura A.I Customer Service"},
            )
            return self.stream
        except Exception as e:
            logging.error(f"Error creating stream: {e}")
            raise

    def create_fastapi_app(self):
        try:
            self.app = fastapi.FastAPI()
            self.app.add_middleware(
                CORSMiddleware,
                allow_origins=["*"],
                allow_credentials=True,
                allow_methods=["*"],
                allow_headers=["*"],
            )

            if not self.stream:
                self.create_stream()
            self.stream.mount(self.app)

            @self.app.get("/reset")
            async def reset():
                try:
                    self.reset_conversation()
                    return {"status": "success"}
                except Exception as e:
                    logging.error(f"Error in reset endpoint: {e}")
                    return {"status": "error", "message": str(e)}

            @self.app.get("/status")
            async def status():
                try:
                    return {
                        "status": "running",
                        "messages_count": len(self.messages),
                        "last_response": self.full_response
                    }
                except Exception as e:
                    logging.error(f"Error in status endpoint: {e}")
                    return {"status": "error", "message": str(e)}

            return self.app
        except Exception as e:
            logging.error(f"Error creating FastAPI app: {e}")
            raise

    def start_server(self, host: str = "0.0.0.0", port: int = 7862):
        import uvicorn
        if not self.app:
            self.create_fastapi_app()
        logging.info(f"Starting server on {host}:{port}")
        try:
            uvicorn.run(self.app, host=host, port=port, log_level="info")
        except Exception as e:
            logging.error(f"Error starting server: {e}")
            raise
    def launch_ui(self, browser: bool = True, port = 7860):
        try:
            if not self.stream:
                self.create_stream()
            if not self.app:
                self.create_fastapi_app()
            logging.info("Launching RTC UI...")
            self.stream.ui.launch(self.app,
                                  server_name="0.0.0.0",
                                  server_port=port,
				  share = True	
                                  )
        except Exception as e:
            logging.error(f"Error launching UI: {e}")
            raise

    def get_conversation_history(self):
        return self.messages.copy()

    def get_last_response(self):
        return self.full_response