ObrishBlesson/multi_model_final

1# === main.py ===
2from fastapi import FastAPI, UploadFile, File, HTTPException, Body
3from fastapi.middleware.cors import CORSMiddleware
4from fastapi.responses import JSONResponse
5from fastapi.staticfiles import StaticFiles
6import os
7import uuid
8
9# NOTE: This used to say "rag_utils". Keep it as "utils" to match the file name below.
10from rag_utils import (
11    extract_text_from_pdfs,
12    chunk_text,
13    save_vectors_simple,
14    transcribe_with_assemblyai,
15    classify_question,
16    get_chat_response,
17    friendly_agent,
18    tts_generate_audio,         # NEW: ElevenLabs TTS
19)
20
21app = FastAPI()
22
23app.add_middleware(
24    CORSMiddleware,
25    allow_origins=["*"],
26    allow_methods=["*"],
27    allow_headers=["*"],
28)
29
30TEMP_DIR = "temp_audio"
31AUDIO_OUT_DIR = "generated_audio"
32os.makedirs(TEMP_DIR, exist_ok=True)
33os.makedirs(AUDIO_OUT_DIR, exist_ok=True)
34
35# Serve generated audio files statically at /audio/...
36# Docs pattern for static files in FastAPI. 
37# Ref: FastAPI StaticFiles docs. :contentReference[oaicite:0]{index=0}
38app.mount("/audio", StaticFiles(directory=AUDIO_OUT_DIR), name="audio")
39
40
41# === PDF Upload ===
42@app.post("/upload-pdf")
43async def upload_pdf(file: UploadFile = File(...)):
44    if not file.filename.lower().endswith(".pdf"):
45        raise HTTPException(status_code=400, detail="Only PDF files are supported.")
46
47    try:
48        text = extract_text_from_pdfs([file.file])
49        print("Text extracted")
50        chunks = chunk_text(text)
51        print("chunk")
52        save_vectors_simple(chunks)
53        print("vector")
54        return {"message": "PDF processed and stored in vector DB."}
55    except Exception as e:
56        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")
57
58
59# === Voice chat (audio in -> text out) — existing behavior kept
60@app.post("/voice-chat")
61async def voice_chat(audio: UploadFile = File(...)):
62    try:
63        uid = str(uuid.uuid4())
64        input_audio_path = os.path.join(TEMP_DIR, f"{uid}.wav")
65        with open(input_audio_path, "wb") as f:
66            f.write(await audio.read())
67
68        question_text = transcribe_with_assemblyai(input_audio_path)
69        route = classify_question(question_text)
70
71        if route == "pdf":
72            answer_text = get_chat_response(question_text)
73        else:
74            answer_text = friendly_agent(question_text)
75
76        return JSONResponse({"transcribed": question_text, "answer": answer_text})
77    except Exception as e:
78        raise HTTPException(status_code=500, detail=str(e))
79
80
81# === NEW: Voice QA with TTS (audio in -> audio out)
82@app.post("/voice-chat-audio")
83async def voice_chat_audio(audio: UploadFile = File(...)):
84    """
85    Accepts an audio question, answers (using PDF RAG or friendly agent), and returns an MP3 URL.
86    """
87    try:
88        uid = str(uuid.uuid4())
89        input_audio_path = os.path.join(TEMP_DIR, f"{uid}.wav")
90        with open(input_audio_path, "wb") as f:
91            f.write(await audio.read())
92
93        # STT with AssemblyAI
94        question_text = transcribe_with_assemblyai(input_audio_path)
95
96        # Route question
97        route = classify_question(question_text)
98        if route == "pdf":
99            answer_text = get_chat_response(question_text)
100        else:
101            answer_text = friendly_agent(question_text)
102
103        # TTS with ElevenLabs -> save MP3 under /audio
104        filename, rel_url = tts_generate_audio(answer_text, out_dir=AUDIO_OUT_DIR)
105        return JSONResponse({
106            "transcribed": question_text,
107            "answer": answer_text,
108            "audio_url": f"/audio/{filename}"
109        })
110    except Exception as e:
111        raise HTTPException(status_code=500, detail=str(e))
112
113
114# === Text chat (text in -> text out) — existing
115@app.get("/chat")
116async def text_chat(question: str):
117    try:
118        route = classify_question(question)
119        if route == "pdf":
120            answer_text = get_chat_response(question)
121        else:
122            answer_text = friendly_agent(question)
123        return JSONResponse({"question": question, "answer": answer_text})
124    except Exception as e:
125        raise HTTPException(status_code=500, detail=str(e))
126
127
128# === NEW: Text-to-Speech endpoint (text in -> audio out)
129@app.post("/tts")
130async def tts(text: str = Body(..., embed=True)):
131    """
132    Convert text to MP3 via ElevenLabs. Returns relative URL to /audio/<file>.
133    """
134    try:
135        filename, rel_url = tts_generate_audio(text, out_dir=AUDIO_OUT_DIR)
136        return JSONResponse({"audio_url": f"/audio/{filename}"})
137    except Exception as e:
138        raise HTTPException(status_code=500, detail=str(e))
139