File size: 2,732 Bytes
c6c1581 03b3d27 c6c1581 03b3d27 c6c1581 03b3d27 c6c1581 03b3d27 c6c1581 03b3d27 c6c1581 03b3d27 c6c1581 03b3d27 c6c1581 | 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 | """FastAPI app exposing keyword and similarity search over the ACL Anthology corpus."""
import logging
import os
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from app.config import get_settings
from app.schemas import PaperResult, SimilarityResult, HealthResponse
from common.db import init_db
from common.vector_index import load_index
from sync.embeddings_client import EmbeddingsClient
from sync.hf_dataset_store import download_snapshot
from app.search_service import SearchService
logger = logging.getLogger(__name__)
app = FastAPI(title="ACL Anthology Search API")
app.state.startup_ok = False
app.state.search_service = None
app.state.last_synced_at = None
@app.on_event("startup")
def load_snapshot() -> None:
settings = get_settings()
try:
state = download_snapshot(settings.hf_repo_id, settings.data_dir, settings.hf_token)
conn = init_db(os.path.join(settings.data_dir, "papers.db"))
index = load_index(os.path.join(settings.data_dir, "index.faiss"))
embeddings_client = EmbeddingsClient(base_url=settings.embedding_base_url, api_key=settings.embedding_api_key)
app.state.search_service = SearchService(conn=conn, index=index, embeddings_client=embeddings_client)
app.state.last_synced_at = state.get("last_synced_at")
app.state.startup_ok = True
except Exception:
logger.exception("Failed to load index/DB snapshot on startup")
app.state.startup_ok = False
@app.middleware("http")
async def block_when_not_ready(request, call_next):
if not app.state.startup_ok and request.url.path != "/health":
return JSONResponse(status_code=503, content={"detail": "service not ready"})
return await call_next(request)
@app.get("/health", response_model=HealthResponse)
def health():
if not app.state.startup_ok:
raise HTTPException(status_code=503, detail="index/db not loaded")
return HealthResponse(
status="ok",
index_size=app.state.search_service.index.ntotal,
last_synced_at=getattr(app.state, "last_synced_at", None),
)
@app.get("/search/keyword", response_model=list[PaperResult])
def search_keyword(q: str, limit: int = 10):
return app.state.search_service.keyword_search(q, limit=limit)
@app.get("/search/similarity", response_model=list[SimilarityResult])
def search_similarity(q: str, k: int = 10):
return app.state.search_service.similarity_search(q, k=k)
@app.get("/paper/{paper_id}", response_model=PaperResult)
def get_paper(paper_id: str):
paper = app.state.search_service.get_paper(paper_id)
if paper is None:
raise HTTPException(status_code=404, detail="paper not found")
return paper
|