Spaces:
Runtime error
Runtime error
| import logging | |
| import threading | |
| from contextlib import asynccontextmanager | |
| from fastapi import BackgroundTasks, FastAPI, File, HTTPException, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel, Field | |
| from . import ingest, store | |
| from .bm25 import bm25 | |
| from .config import settings | |
| from .rag import NoReadyDocuments, answer_question | |
| from .registry import get_registry, public | |
| try: | |
| from groq import RateLimitError | |
| except Exception: # pragma: no cover | |
| RateLimitError = () # type: ignore | |
| logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") | |
| log = logging.getLogger("rag.main") | |
| _warm = threading.Event() | |
| def _warm_up() -> None: | |
| """Load models in the background so the server starts instantly.""" | |
| try: | |
| store.get_embedder() | |
| store.get_qdrant() | |
| ingest.recover_interrupted() | |
| store.get_reranker() | |
| bm25.ensure_fresh() | |
| except Exception: | |
| log.exception("Warm-up failed (the app will retry lazily on first use)") | |
| finally: | |
| _warm.set() | |
| async def lifespan(app: FastAPI): | |
| threading.Thread(target=_warm_up, daemon=True).start() | |
| yield | |
| app = FastAPI(title="RAG Backend", version="0.3.0", lifespan=lifespan) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=[o.strip() for o in settings.cors_origins.split(",") if o.strip()], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| class Turn(BaseModel): | |
| role: str = Field(pattern="^(user|assistant)$") | |
| content: str | |
| class AskRequest(BaseModel): | |
| question: str | |
| doc_ids: list[str] | None = None # None = search every ready document | |
| history: list[Turn] = [] | |
| def health(): | |
| return {"status": "ok", "llm_model": settings.llm_model, "models_loaded": _warm.is_set()} | |
| def documents(): | |
| return [public(d) for d in get_registry().all()] | |
| def upload_document(background: BackgroundTasks, file: UploadFile = File(...)): | |
| data = file.file.read() | |
| try: | |
| entry, duplicate = ingest.submit_document(file.filename or "upload", data) | |
| except ValueError as e: | |
| raise HTTPException(status_code=400, detail=str(e)) | |
| if not duplicate: | |
| background.add_task(ingest.process_document, entry["doc_id"], entry["filename"], data) | |
| return {**public(entry), "duplicate": duplicate} | |
| def delete_document(doc_id: str): | |
| if not ingest.delete_document(doc_id): | |
| raise HTTPException(status_code=404, detail="Document not found.") | |
| return {"deleted": doc_id} | |
| def ask(req: AskRequest): | |
| if not req.question.strip(): | |
| raise HTTPException(status_code=400, detail="Question is empty.") | |
| if not settings.groq_api_key: | |
| raise HTTPException(status_code=500, detail="GROQ_API_KEY is not set in .env") | |
| try: | |
| return answer_question(req.question.strip(), req.doc_ids, [t.model_dump() for t in req.history]) | |
| except NoReadyDocuments: | |
| raise HTTPException(status_code=409, detail="No document is ready yet. Upload one and wait until it finishes processing.") | |
| except RateLimitError: | |
| raise HTTPException(status_code=429, detail="Groq free-tier rate limit reached. Wait a few seconds and try again.") | |
| except Exception as e: # bad key, retired model, network... | |
| log.exception("ask failed") | |
| raise HTTPException(status_code=502, detail=f"LLM error: {e}") | |