Spaces:
Runtime error
Runtime error
File size: 3,506 Bytes
0a19322 | 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 | 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()
@asynccontextmanager
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] = []
@app.get("/health")
def health():
return {"status": "ok", "llm_model": settings.llm_model, "models_loaded": _warm.is_set()}
@app.get("/documents")
def documents():
return [public(d) for d in get_registry().all()]
@app.post("/documents")
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}
@app.delete("/documents/{doc_id}")
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}
@app.post("/ask")
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}")
|