File size: 1,931 Bytes
f67f20b 3813085 f67f20b | 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 | from fastapi import APIRouter
from src.schemas.models import QueryRequest
from src.services.retrieval import Retriever
from src.services.generator import generate_answer
from src.services.query_rewriter import rewrite_query
router = APIRouter()
retriever = Retriever()
@router.post("/query")
async def query_documents(request: QueryRequest):
original_query = request.query
session_id = request.session_id
# -----------------------------
# QUERY REWRITING
# -----------------------------
rewritten_query = rewrite_query(original_query)
# -----------------------------
# RETRIEVAL
# -----------------------------
results = retriever.retrieve(
session_id=session_id,
query=rewritten_query,
top_k=5,
documents=request.documents,
)
# -----------------------------
# OUT-OF-SYLLABUS REJECTION
# -----------------------------
if not results:
return {
"query": original_query,
"rewritten_query": rewritten_query,
"message": ("The uploaded material " "does not cover this topic."),
}
# -----------------------------
# GENERATION
# -----------------------------
answer = generate_answer(context_chunks=results, question=rewritten_query)
# -----------------------------
# CITATIONS
# -----------------------------
citations = []
seen = set()
for result in results:
chunk = result["chunk"]
key = (chunk.doc_name, chunk.page)
if key in seen:
continue
seen.add(key)
citations.append({"document": chunk.doc_name, "page": chunk.page})
# -----------------------------
# FINAL RESPONSE
# -----------------------------
return {
"query": original_query,
"rewritten_query": rewritten_query,
"answer": answer,
"citations": citations,
}
|