Spaces:
Running
Running
File size: 3,293 Bytes
ed65693 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | 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,
}
|