app.py

import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["OMP_NUM_THREADS"] = "1"
import tempfile
import uuid
import torch
import subprocess
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse

# pyrefly: ignore [missing-import]
try:
    from faster_whisper import WhisperModel
except ImportError:
    print("WARNING: faster_whisper is not installed. STT will not work.")
    WhisperModel = None

from groq import Groq
from dotenv import load_dotenv

# Auto-clone sooktam2 repository (code only) if it's missing or empty
sooktam_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sooktam2")
if not os.path.exists(sooktam_path) or not os.listdir(sooktam_path):
    print("Cloning sooktam2 repository (code only)...")
    env = os.environ.copy()
    env["GIT_LFS_SKIP_SMUDGE"] = "1"
    try:
        subprocess.run(["git", "clone", "https://huggingface.co/bharatgenai/sooktam2", sooktam_path], env=env, check=True)
        print("Cloning complete.")
    except Exception as e:
        print(f"Error cloning sooktam2: {e}")

# Import TTS Engine
from tts_engine import synthesize_speech

# 1. Create a clean FastAPI instance first
app = FastAPI(title="Voice Assistant API")

# 2. Setup CORS for the React frontend (explicit origins to support credentials)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5173", "http://127.0.0.1:5173"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load Environment Variables
load_dotenv()
api_key = os.environ.get("GROQ_API_KEY")

# Initialize models conditionally to avoid hanging indefinitely if they fail
print("Initializing Whisper STT...")
try:
    if WhisperModel is None:
        raise Exception("faster_whisper module not found")
        
    if os.environ.get("SPACE_ID"):
        device = "cpu"
        compute_type = "int8"
    else:
        device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
        compute_type = "float16" if device == "cuda" else "int8"
    stt_model = WhisperModel("large-v3-turbo", device=device, compute_type=compute_type)
    print(f"Whisper STT loaded on {device}")
except Exception as e:
    print(f"Error loading Whisper: {e}")
    stt_model = None

print("Initializing Groq...")
groq_client = Groq(api_key=api_key) if api_key else None

if not groq_client:
    print("WARNING: GROQ_API_KEY is not set. LLM calls will fail or return mock data.")

sessions_history = {}
DEFAULT_SYSTEM_PROMPT = {"role": "system", "content": "You are a close friend chatting in highly colloquial, casual spoken Tamil ( The user is speaking in casual Tamil. You MUST reply exclusively in extremely casual, colloquial Tamil exactly how two friends speak on the streets of tamilnadu. DO NOT use formal or textbook Tamil. Use casual words like 'machan', 'da', 'machi'. CRITICAL: Write EVERYTHING strictly in pure Tamil script. DO NOT use English letters or punctuation marks (no commas or periods).\n\nIMPORTANT GRAMMAR RULES:\n- Never invent weird mashup words.\n- 'I will do it for you' = 'பண்ணி தரேன்' or 'பண்ணி குடுக்கறேன்' (NEVER use 'பண்ணிக்கிடுக்கேன்').\n- 'Understood' = 'புரிஞ்சது'.\n\nEXAMPLES:\nUser: எனக்கு டிக்கெட் புக் பண்ணி தருவியா\nYou: ஓகே மச்சி நான் இப்பவே புக் பண்ணி தரேன் கவலைப்படாத\n\nKeep answers extremely brief (maximum 10 to 15 words). Speak naturally like a human friend."}

@app.get("/")
def read_root():
    return {"status": "Voice API is running"}

@app.post("/api/stt")
async def process_stt(audio: UploadFile = File(...)):
    """Receives audio file from frontend and returns transcribed text"""
    if not stt_model:
        return JSONResponse(status_code=500, content={"error": "STT Model not loaded on server."})

    try:
        original_filename = audio.filename or "recording.webm"
        _, ext = os.path.splitext(original_filename)
        if not ext:
            ext = ".webm"
        
        temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
        content = await audio.read()
        
        print(f"[STT] Received audio: {len(content)} bytes, filename: {original_filename}, ext: {ext}")
        
        if len(content) == 0:
            return JSONResponse(status_code=400, content={"error": "Received empty audio file"})
        
        temp_audio.write(content)
        temp_audio.close()

        segments, _ = stt_model.transcribe(
            temp_audio.name, 
            beam_size=1, 
            language="ta",
            vad_filter=True,
            vad_parameters=dict(min_silence_duration_ms=500)
        )
        text = "".join([s.text for s in segments]).strip()
        
        print(f"[STT] Transcribed: '{text}'")
        os.unlink(temp_audio.name)
        return {"text": text}
    except Exception as e:
        print(f"[STT] Error: {e}")
        return JSONResponse(status_code=500, content={"error": str(e)})

@app.post("/api/llm")
async def process_llm(data: dict):
    """Receives text from STT and returns Groq LLM response"""
    text = data.get("text", "")
    session_id = data.get("session_id", "default")
    
    if not text:
        return {"response": ""}
        
    if not groq_client:
        return {"response": "GROQ_API_KEY not found. (Mock response)"}
        
    if session_id not in sessions_history:
        sessions_history[session_id] = [DEFAULT_SYSTEM_PROMPT.copy()]
        
    chat_history = sessions_history[session_id]
    chat_history.append({"role": "user", "content": text})
    
    if len(chat_history) > 11:
        sessions_history[session_id] = [chat_history[0]] + chat_history[-10:]
        chat_history = sessions_history[session_id]

    try:
        import re
        response = groq_client.chat.completions.create(
            model='openai/gpt-oss-120b',
            messages=chat_history
        )
        llm_output = response.choices[0].message.content.strip()
        llm_output = re.sub(r'<think>.*?</think>', '', llm_output, flags=re.DOTALL).strip()
        
        chat_history.append({"role": "assistant", "content": llm_output})
        return {"response": llm_output}
    except Exception as e:
        chat_history.pop()
        print(f"Groq API Error: {e}")
        return {"response": f"Error calling LLM: {str(e)}"}

@app.post("/api/tts")
async def process_tts(data: dict, background_tasks: BackgroundTasks):
    """Receives LLM response text and returns synthesized audio file"""
    text = data.get("text", "")
    if not text:
        raise HTTPException(status_code=400, detail="No text provided")
        
    try:
        output_filename = f"response_{uuid.uuid4().hex}.wav"
        out_path = synthesize_speech(text, output_filename=output_filename)
        
        if out_path and os.path.exists(out_path):
            background_tasks.add_task(os.remove, out_path)
            return FileResponse(
                path=out_path, 
                media_type="audio/wav", 
                filename=output_filename
            )
        else:
            raise HTTPException(status_code=500, detail="TTS Engine failed to generate audio")
    except Exception as e:
        print(f"TTS API Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 3. Create Gradio Blocks to satisfy Hugging Face Space runner
import gradio as gr
with gr.Blocks() as demo:
    gr.Markdown("# Shabdham Voice Assistant API")
    gr.Markdown("The backend API is running. Connect your Vercel frontend here.")

# 4. Mount Gradio inside FastAPI at root
app = gr.mount_gradio_app(app, demo, path="/")