from contextlib import asynccontextmanager from pathlib import Path from fastapi import FastAPI, File, HTTPException, UploadFile from pydantic import BaseModel, Field from app.retriever import RETRIEVAL_MODES, UPLOAD_DIR, DocumentStore store = DocumentStore() @asynccontextmanager async def lifespan(app: FastAPI): UPLOAD_DIR.mkdir(parents=True, exist_ok=True) store.load() yield app = FastAPI( title="Semantic Document Retrieval API", description=( "Hybrid retrieval API combining BM25 sparse search, FAISS dense search, " "entropy-weighted adaptive fusion, and cross-encoder reranking. " "Supports 7 retrieval modes for ablation comparison." ), version="1.0.0", lifespan=lifespan, ) class QueryRequest(BaseModel): question: str = Field(..., min_length=3, examples=["What does the document say about AI?"]) top_k: int = Field(default=3, ge=1, le=20) mode: str = Field( default="hybrid_calibrated_rerank", description="Retrieval mode. See /modes for available options.", ) @app.get("/") async def root() -> dict[str, str]: return { "name": "Semantic Document Retrieval API", "version": "1.0.0", "docs": "/docs", "health": "/health", "modes": "/modes", } @app.get("/health") async def health() -> dict[str, object]: return { "status": "ok", "index_loaded": store.vector_store is not None, "bm25_loaded": store.bm25_index.bm25 is not None, "corpus_cdfs_loaded": store.corpus_cdfs is not None, "documents": sorted(store.documents), } @app.get("/modes") async def list_modes() -> dict[str, object]: """List all available retrieval modes with descriptions.""" return { "modes": RETRIEVAL_MODES, "default": "hybrid_calibrated_rerank", } @app.post("/upload") async def upload_document(file: UploadFile = File(...)) -> dict[str, object]: if not file.filename: raise HTTPException(status_code=400, detail="Upload must include a filename.") suffix = Path(file.filename).suffix.lower() if suffix not in {".pdf", ".txt", ".md"}: raise HTTPException(status_code=400, detail="Only PDF, TXT, and MD files are supported.") safe_name = Path(file.filename).name file_path = UPLOAD_DIR / safe_name content = await file.read() file_path.write_bytes(content) try: stats = store.add_file(file_path) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc return { "message": "Document indexed successfully.", **stats, } @app.post("/query") async def query_documents(request: QueryRequest) -> dict[str, object]: if request.mode not in RETRIEVAL_MODES: raise HTTPException( status_code=400, detail=f"Unknown mode '{request.mode}'. Available: {list(RETRIEVAL_MODES)}", ) result = store.answer(request.question, request.top_k, request.mode) return { "question": request.question, **result, } @app.get("/documents") async def list_documents() -> dict[str, object]: return { "documents": sorted(store.documents), "index_ready": store.vector_store is not None, }