api.py

from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from src.rag import answer

app = FastAPI(title="Vietnamese RAG API")


class QueryRequest(BaseModel):
    question: str
    top_k: int = 5


class Source(BaseModel):
    source: str
    page: int
    rrf_score: float
    rerank_score: float


class QueryResponse(BaseModel):
    answer: str
    sources: list[Source]


@app.get("/", response_class=HTMLResponse)
def index():
    with open("static/index.html", encoding="utf-8") as f:
        return f.read()


@app.post("/query", response_model=QueryResponse)
def query(req: QueryRequest):
    result = answer(req.question, top_k=req.top_k)
    return QueryResponse(
        answer=result["answer"],
        sources=[Source(**s) for s in result["sources"]],
    )


@app.get("/health")
def health():
    return {"status": "ok"}