Prathamesh1420's picture
Create server.py
a7a28b1 verified
import asyncio
import websockets
import os
from io import BytesIO
from pydub import AudioSegment
from utils import transcribe_audio, get_response_llm, play_text_to_speech, load_whisper
# Load Whisper Model
model = load_whisper()
# Store connected clients
clients = set()
async def handle_audio_stream(websocket):
"""Handles incoming live audio stream from users."""
clients.add(websocket)
try:
while True:
audio_chunk = await websocket.recv()
if not audio_chunk:
break
# Convert received bytes to WAV file
audio = AudioSegment.from_file(BytesIO(audio_chunk))
audio.export("temp_audio_chunk.wav", format="wav")
# Transcribe
text = transcribe_audio(model, "temp_audio_chunk.wav")
if text:
print(f"User: {text}")
# Generate response
response_llm = get_response_llm(user_question=text)
# Play response
play_text_to_speech(text=response_llm)
# Send response to WebSocket client
await websocket.send(response_llm)
except:
pass
finally:
clients.remove(websocket)
async def main():
async with websockets.serve(handle_audio_stream, "0.0.0.0", 8765):
await asyncio.Future() # Keeps the server running
if __name__ == "__main__":
asyncio.run(main())