Spaces:
Sleeping
Sleeping
File size: 3,666 Bytes
a4538e5 551503c a4538e5 551503c a4538e5 6b27270 a4538e5 | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import os
from typing import List, Optional
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from dotenv import load_dotenv
from src.search import RAGSearch, RetrievalResult
import uvicorn
load_dotenv()
# Global variable for RAG system
rag_search: Optional[RAGSearch] = None
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup
global rag_search
try:
persist_dir = os.getenv("PERSIST_DIR", "faiss_store")
embedding_model = os.getenv("EMBEDDING_MODEL", "all-MiniLM-L6-v2")
llm_model = os.getenv("LLM_MODEL", "llama-3.1-8b-instant")
rag_search = RAGSearch(
persist_dir=persist_dir,
embedding_model=embedding_model,
llm_model=llm_model,
)
print("[INFO] RAG system loaded successfully")
except Exception as e:
print(f"[ERROR] Failed to load RAG system: {e}")
raise
yield # Application runs here
# Shutdown (cleanup if needed)
print("[INFO] Shutting down RAG system")
# -------------------------
# FastAPI App
# -------------------------
# root_path is needed for HF Spaces reverse proxy to serve /docs correctly
app = FastAPI(
title="RAG Question Answering API",
description="FAISS + SentenceTransformers + Groq LLM",
version="2.0.0",
lifespan=lifespan,
root_path=os.getenv("ROOT_PATH", ""),
)
# CORS for React/Node clients
cors_origins = os.getenv("CORS_ORIGINS", "*").split(",")
app.add_middleware(
CORSMiddleware,
allow_origins=[o.strip() for o in cors_origins] if cors_origins else ["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# -------------------------
# Request / Response Models
# -------------------------
class SourceItem(BaseModel):
index: int
distance: float
text: Optional[str] = None
class QueryRequest(BaseModel):
query: str
top_k: int = 3
class QueryResponse(BaseModel):
query: str
answer: str
sources: List[SourceItem]
# -------------------------
# Routes
# -------------------------
@app.get("/")
def root():
return {"message": "RAG API is running. Go to /docs"}
@app.get("/health")
def health():
if not rag_search:
return {"ready": False}
meta_count = len(rag_search.vectorstore.metadata) if rag_search.vectorstore else 0
return {
"ready": True,
"persist_dir": rag_search.vectorstore.persist_dir,
"documents_indexed": meta_count,
"embedding_model": rag_search.embedding_model,
"llm_model": rag_search.llm_model,
}
@app.post("/query", response_model=QueryResponse)
def query_rag(payload: QueryRequest):
if not rag_search:
raise HTTPException(status_code=503, detail="RAG system not ready")
try:
# Retrieve and summarize
sources: List[RetrievalResult] = rag_search.retrieve(payload.query, top_k=payload.top_k)
answer: str = rag_search.summarize(payload.query, sources)
# Map sources for response
resp_sources = [
SourceItem(index=s.index, distance=float(s.distance), text=s.text)
for s in sources
]
return QueryResponse(query=payload.query, answer=answer, sources=resp_sources)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# -------------------------
# Run locally
# -------------------------
if __name__ == "__main__":
uvicorn.run(
"app:app",
host="0.0.0.0",
port=int(os.getenv("PORT", "7860")),
reload=True
) |