evanderpool/rag-knowledge-base
0
1"""2query.py — Answer generation layer for RAG Knowledge Base Builder.3 4The final stage of the RAG pipeline:5 1. Embed the user's question via the same model used at ingest time6 2. Retrieve the top-k most semantically similar chunks from ChromaDB7 3. Apply a token budget — only include chunks that fit in the prompt8 4. Build a structured prompt with numbered context blocks + citation instructions9 5. Send to Groq (Llama 3.3 70B) — free tier, ~500 tok/s throughput10 6. Return a typed QueryResult with answer and deduplicated source list11 12Why Groq + Llama 3.3 70B:13 - Free API with generous rate limits for development and demos14 - 70B model produces coherent, faithful answers from dense document context15 - 128k context window — comfortably holds 5-10 retrieved chunks + system prompt16 17Faithfulness constraint: sources are extracted from the retrieved chunks18(not from model output) — the model cannot hallucinate a source that wasn't19in the retrieved context, even if it hallucinates the citation text.20"""21from __future__ import annotations22 23import logging24import os25from dataclasses import dataclass, field26 27from groq import Groq28 29from embedder import KnowledgeBase, RetrievedChunk30 31# Re-export so app.py only needs to import from query32__all__ = [33 "QueryConfig", "QueryResult", "QueryEngine",34 "_extract_sources", # used by app.py to build source list after streaming35 "_apply_token_budget",36]37 38logger = logging.getLogger(__name__)39 40 41# ── Config & Result Types ──────────────────────────────────────────────────────42 43@dataclass44class QueryConfig:45 """46 Controls LLM behavior and retrieval parameters.47 Tune top_k and max_context_tokens together — more chunks = more context48 but larger prompts. 5 chunks at ~1000 chars each ≈ 1250 tokens of context.49 """50 model: str = "llama-3.3-70b-versatile"51 temperature: float = 0.1 # low temp keeps answers factual and consistent52 max_tokens: int = 1500 # answer length limit — increased for dense documents53 top_k: int = 10 # chunks retrieved per query — increased for coverage54 max_context_tokens: int = 6000 # hard cap on context — headroom for 10+ chunks55 56 57@dataclass58class QueryResult:59 """Structured answer with full provenance for display and logging."""60 question: str61 answer: str62 sources: list[str] # deduplicated: ["policy.pdf (p.4)", "policy.pdf (p.7)"]63 chunks_used: int # how many chunks fit within the token budget64 model: str65 66 67# ── Prompt Engineering ─────────────────────────────────────────────────────────68 69_SYSTEM_PROMPT = """\70You are a precise, helpful document Q&A assistant.71 72Rules you must follow:731. Answer ONLY using the provided context. Do not use prior knowledge or make assumptions beyond what is written.742. If the context does not contain enough information to answer the question, respond with exactly:75 "I don't have enough information in the provided documents to answer that question."763. Cite your sources inline using the format [Source: filename, page N].77 For plain text files with no page number, cite as [Source: filename].784. Lead with the direct answer, then support it with relevant details from the context.795. Be concise. Do not pad your response with filler phrases.80"""81 82 83def _build_context_block(chunks: list[RetrievedChunk]) -> str:84 """85 Format retrieved chunks into a numbered context block for the prompt.86 Numbered references help the model cite accurately and help us debug87 which chunks were used when answers are wrong.88 """89 lines: list[str] = ["Context:\n"]90 for i, chunk in enumerate(chunks, start=1):91 if chunk.page_number > 0:92 header = f"[{i}] Source: {chunk.source}, Page {chunk.page_number} (relevance: {chunk.score:.2f})"93 else:94 header = f"[{i}] Source: {chunk.source} (relevance: {chunk.score:.2f})"95 lines.append(f"{header}\n{chunk.text}")96 return "\n\n".join(lines)97 98 99def _extract_sources(chunks: list[RetrievedChunk]) -> list[str]:100 """101 Build a deduplicated, human-readable source list from the retrieved chunks.102 Sources come from the actual chunks used — not parsed from model output —103 so the model cannot fabricate a citation that wasn't in the context.104 """105 seen: set[tuple[str, int]] = set()106 sources: list[str] = []107 for chunk in chunks:108 key = (chunk.source, chunk.page_number)109 if key not in seen:110 seen.add(key)111 if chunk.page_number > 0:112 sources.append(f"{chunk.source} (p.{chunk.page_number})")113 else:114 sources.append(chunk.source)115 return sources116 117 118def _apply_token_budget(119 chunks: list[RetrievedChunk], max_tokens: int120) -> list[RetrievedChunk]:121 """122 Greedily include chunks until the token budget is exhausted.123 Uses len(text)//4 as a token estimate — same heuristic as loader.py.124 Chunks arrive sorted by relevance (highest first), so the most relevant125 context is always included before less relevant chunks are cut.126 """127 selected: list[RetrievedChunk] = []128 used = 0129 for chunk in chunks:130 estimate = len(chunk.text) // 4131 if used + estimate > max_tokens:132 break133 selected.append(chunk)134 used += estimate135 return selected136 137 138# ── QueryEngine ────────────────────────────────────────────────────────────────139 140class QueryEngine:141 """142 Orchestrates retrieval and generation for a single knowledge base.143 144 Example:145 engine = QueryEngine(kb)146 result = engine.ask("What is the vacation policy?")147 print(result.answer)148 for src in result.sources:149 print(f" • {src}")150 """151 152 def __init__(153 self,154 kb: KnowledgeBase,155 config: QueryConfig | None = None,156 api_key: str | None = None,157 ) -> None:158 self._kb = kb159 self._cfg = config or QueryConfig()160 self._client = Groq(api_key=api_key or os.environ["GROQ_API_KEY"])161 162 def ask(self, question: str, stream: bool = False) -> QueryResult:163 """164 Answer a natural language question using retrieved document context.165 166 Args:167 question: The user's question in plain English.168 stream: If True, stream the answer to stdout while building it.169 Useful for the CLI — makes long answers feel responsive.170 171 Returns:172 QueryResult with answer, sources, and metadata.173 174 Raises:175 RuntimeError: If the knowledge base is empty.176 groq.APIError: If the Groq API call fails.177 """178 question = question.strip()179 if not question:180 raise ValueError("Question cannot be empty.")181 182 # Retrieve + apply token budget183 raw_chunks = self._kb.retrieve(question, top_k=self._cfg.top_k)184 chunks = _apply_token_budget(raw_chunks, self._cfg.max_context_tokens)185 186 logger.info(187 "Retrieved %d chunks (%d within token budget) for: '%s'",188 len(raw_chunks), len(chunks), question,189 )190 191 context_block = _build_context_block(chunks)192 messages = [193 {"role": "system", "content": _SYSTEM_PROMPT},194 {"role": "user", "content": f"{context_block}\n\n---\nQuestion: {question}"},195 ]196 197 if stream:198 answer = self._stream_answer(messages)199 else:200 response = self._client.chat.completions.create(201 model=self._cfg.model,202 messages=messages,203 temperature=self._cfg.temperature,204 max_tokens=self._cfg.max_tokens,205 )206 answer = response.choices[0].message.content.strip()207 208 return QueryResult(209 question=question,210 answer=answer,211 sources=_extract_sources(chunks),212 chunks_used=len(chunks),213 model=self._cfg.model,214 )215 216 def prepare_context(217 self, question: str218 ) -> tuple[list[dict], list[RetrievedChunk]]:219 """220 Build the prompt messages and return the retrieved chunks separately.221 Called by app.py so Streamlit can stream tokens itself via stream_tokens().222 """223 question = question.strip()224 raw_chunks = self._kb.retrieve(question, top_k=self._cfg.top_k)225 chunks = _apply_token_budget(raw_chunks, self._cfg.max_context_tokens)226 messages = [227 {"role": "system", "content": _SYSTEM_PROMPT},228 {229 "role": "user",230 "content": f"{_build_context_block(chunks)}\n\n---\nQuestion: {question}",231 },232 ]233 return messages, chunks234 235 def stream_tokens(self, messages: list[dict]):236 """237 Generator that yields raw token strings from the Groq stream.238 Designed for st.write_stream() in app.py — Streamlit consumes the generator239 and handles display; the caller gets the full answer back from st.write_stream.240 """241 with self._client.chat.completions.create(242 model=self._cfg.model,243 messages=messages,244 temperature=self._cfg.temperature,245 max_tokens=self._cfg.max_tokens,246 stream=True,247 ) as stream:248 for chunk in stream:249 yield chunk.choices[0].delta.content or ""250 251 def _stream_answer(self, messages: list[dict]) -> str:252 """Stream the response to stdout token-by-token. Returns the full answer string."""253 answer_parts: list[str] = []254 with self._client.chat.completions.create(255 model=self._cfg.model,256 messages=messages,257 temperature=self._cfg.temperature,258 max_tokens=self._cfg.max_tokens,259 stream=True,260 ) as stream:261 for chunk in stream:262 delta = chunk.choices[0].delta.content or ""263 print(delta, end="", flush=True)264 answer_parts.append(delta)265 print() # newline after stream ends266 return "".join(answer_parts)267 268 269# ── CLI ────────────────────────────────────────────────────────────────────────270 271def _print_result(result: QueryResult) -> None:272 """Format and print a QueryResult to stdout."""273 print(f"\nAnswer:\n{result.answer}\n")274 if result.sources:275 print("Sources:")276 for src in result.sources:277 print(f" • {src}")278 print(f"\n[{result.chunks_used} chunks used · {result.model}]\n")279 280 281def main() -> None:282 import sys283 from dotenv import load_dotenv284 285 load_dotenv()286 logging.basicConfig(level=logging.WARNING, format="%(levelname)s: %(message)s")287 288 kb = KnowledgeBase()289 290 if kb.size == 0:291 print("Knowledge base is empty.")292 print("Ingest a document first: python ingest.py <file.pdf>")293 sys.exit(1)294 295 docs = kb.list_documents()296 print(f"\nKnowledge base: {kb.size} chunks · {len(docs)} document(s)")297 for doc in docs:298 print(f" • {doc['source']} — {doc['chunks']} chunks, {doc['pages']} pages")299 300 engine = QueryEngine(kb)301 print("\nType your question (or 'quit' to exit):\n")302 303 while True:304 try:305 question = input("Q: ").strip()306 except (KeyboardInterrupt, EOFError):307 print("\nGoodbye.")308 break309 310 if not question:311 continue312 if question.lower() in {"quit", "exit", "q"}:313 break314 315 result = engine.ask(question, stream=True)316 _print_result(result)317 318 319if __name__ == "__main__":320 main()321