app/main.py

"""FastAPI app: serve UI, POST /ask (retrieve -> generate -> validate), /health.

Importing app.retrieval loads the bge model once at startup.
"""
from __future__ import annotations

from pathlib import Path

from fastapi import FastAPI
from fastapi.responses import FileResponse
from pydantic import BaseModel
from sqlalchemy import text

from app.db import engine
from app.generate import generate_answer, generate_plain_english
from app.retrieval import hybrid_search  # triggers model load at startup
from app.validate import validate_answer

STATIC_DIR = Path(__file__).parent / "static"

app = FastAPI(title="LexRAG", description="Hybrid-search legal RAG over Indian judgments")


class AskRequest(BaseModel):
    question: str


@app.get("/")
def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/health")
def health() -> dict:
    try:
        with engine.connect() as conn:
            conn.execute(text("SELECT 1"))
        return {"status": "ok", "db": "ok"}
    except Exception as e:  # pragma: no cover
        return {"status": "degraded", "db": f"error: {e}"}


@app.post("/ask")
def ask(req: AskRequest) -> dict:
    question = (req.question or "").strip()
    if not question:
        return {"answer": "Please enter a question.", "sources": []}

    sources = hybrid_search(question)
    raw = generate_answer(question, sources)
    result = validate_answer(raw, sources)
    plain = generate_plain_english(result["answer"])

    return {
        "answer": result["answer"],
        "plain_english": plain,
        "removed_citations": result["removed"],
        "sources": [
            {"case_name": s["case_name"], "content": s["content"]} for s in sources
        ],
    }