File size: 3,086 Bytes
f2ea5fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from contextlib import asynccontextmanager
from pydantic import BaseModel
from core.backend import AIBackend
import uvicorn, json, os
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from services.stt import StreamingSTT
from services.tts import text_to_speech_stream
from fastapi.staticfiles import StaticFiles
chatbot_obj = AIBackend()
@asynccontextmanager
async def lifespan(app: FastAPI):
await chatbot_obj.async_setup()
yield
if chatbot_obj.conn:
await chatbot_obj.conn.close()
app = FastAPI(lifespan=lifespan)
class UserRequest(BaseModel):
user_id: str
user_query: str
@app.post("/chat")
async def chat(request: UserRequest):
stream = await chatbot_obj.main(
user_id=request.user_id,
user_query=request.user_query,
)
return StreamingResponse(stream, media_type="text/event-stream")
@app.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
await websocket.accept()
try:
while True:
# receive frontend message
data = await websocket.receive_text()
payload = json.loads(data)
user_id = payload["user_id"]
user_query = payload["user_query"]
# stream AI response
stream = await chatbot_obj.main(
user_id=user_id,
user_query=user_query
)
async for chunk in stream:
await websocket.send_text(chunk)
# notify frontend response finished
await websocket.send_text("[[END]]")
except WebSocketDisconnect:
print("Client disconnected")
@app.websocket("/ws/voice")
async def voice_ws(websocket: WebSocket):
await websocket.accept()
stt = StreamingSTT()
try:
while True:
message = await websocket.receive()
# 🎤 AUDIO INPUT
if "bytes" in message:
audio_chunk = message["bytes"]
stt.add_audio(audio_chunk)
text = stt.transcribe_if_ready()
if not text:
continue
await websocket.send_text(f"[STT]: {text}")
# 🤖 LLM STREAM
stream = chatbot_obj.main(
user_id="voice_user",
user_query=text
)
full_response = ""
async for token in stream:
full_response += token
await websocket.send_text(f"[LLM]: {token}")
# 🔊 TTS STREAM
async for audio_chunk in text_to_speech_stream(full_response):
await websocket.send_bytes(audio_chunk)
await websocket.send_text("[END]")
except WebSocketDisconnect:
print("Voice client disconnected")
if __name__ == "__main__":
uvicorn.run("app:app", host="127.0.0.1", port=8679, reload=True) |