Spaces:
Sleeping
Sleeping
File size: 1,851 Bytes
e7a3876 | 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 | """
The complete RAG pipeline in one place.
RAGChain.query() is the single entrypoint for the app and API.
It orchestrates: retrieve → build prompt → generate → return with sources.
This is the class you demo in interviews.
"""
from src.retrieval.retriever import Retriever
from src.generation.prompt_builder import build_prompt
from src.generation.llm_client import LLMClient
from src.utils.config import config
from src.utils.logger import logger
class RAGChain:
def __init__(self):
self.retriever = Retriever()
self.llm = LLMClient()
def query(self, question: str, top_k: int = None) -> dict:
"""
Full RAG pipeline: question in, answer + sources out.
Returns:
{
"question": str,
"answer": str,
"sources": [{"source": str, "page": int, "score": float}],
"chunks_used": int,
}
"""
logger.info(f"Query received: '{question}'")
# Step 1: Retrieve relevant chunks
chunks = self.retriever.retrieve(question, top_k=top_k or config.top_k)
# Step 2: Build the grounded prompt
prompt = build_prompt(question, chunks)
# Step 3: Generate answer
logger.info("Sending to LLM...")
answer = self.llm.generate(prompt)
# Step 4: Package sources for attribution
sources = [
{"source": c["source"], "page": c["page"], "score": c["score"]}
for c in chunks
]
result = {
"question": question,
"answer": answer,
"sources": sources,
"chunks_used": len(chunks),
}
logger.info(f"Answer generated. Sources: {[s['source'] for s in sources]}")
return result
|