diff --git a/CONTEXT.md b/CONTEXT.md index 7d0a11c0..38e7f396 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -31,3 +31,19 @@ behaviour behind a small interface), **seam** (a boundary you can substitute at) telemetry assembly and the eval harness use one source of truth; the eval package imports from here, never the reverse. See [ADR 0003](docs/adr/0003-telemetry-ownership.md). +- **Retriever** — the seam (Protocol) every retrieval strategy hides behind: + `retrieve(query, top_k) -> list[SearchResult]`. Implementations either conform + directly (`DenseRetriever`, `BM25HybridRetriever`) or *compose* an inner + Retriever (`RerankingRetriever` over-fetches then cross-encodes; + `MultiQueryRetriever` fans rewritten queries out and dedups). Live in + `src/retrieval/`; selected for production by `build_retriever`. See + [ADR 0004](docs/adr/0004-retriever-seam-and-query-engine.md). +- **QueryEngine** (`src/query_engine/`) — the deep module owning retrieve→generate + for both the sync (`ask`) and streaming (`ask_stream`) paths: one Markdown + answer prompt, filename-prefixed context, an optional refusal gate, and + telemetry assembly — all in one place. Both `RAGBackend` and the eval harness + call it, so eval measures the shipped pipeline. See + [ADR 0004](docs/adr/0004-retriever-seam-and-query-engine.md). +- **RefusalHandler** — an answerability gate (not a Retriever): refuses when the + top-1 similarity is below a threshold (or nothing was retrieved). Applied + inside the QueryEngine; off by default in production. diff --git a/docs/adr/0004-retriever-seam-and-query-engine.md b/docs/adr/0004-retriever-seam-and-query-engine.md new file mode 100644 index 00000000..ef044a63 --- /dev/null +++ b/docs/adr/0004-retriever-seam-and-query-engine.md @@ -0,0 +1,104 @@ +# ADR 0004 — Retriever seam and the shared QueryEngine + +- **Status:** Accepted +- **Sequencing:** Step 4 of the RAG architecture deepening spec ([issue #16](https://github.com/elkaix/rag-document-qa/issues/16)); resolves the Retriever-seam and shared-QueryEngine map tickets. +- **Date:** 2026-07-12 + +## Context + +Two coupled problems, both rooted in there being no seam between retrieval and +generation: + +1. **Proven retrieval levers couldn't ship.** Hybrid BM25, cross-encoder + reranking, query rewriting, and refusal handling existed only in + `src/eval/`, with no production interface to activate them. +2. **Three diverged answer prompts, and eval measured a different pipeline.** + The sync query path used a *plain* answer prompt; the streaming path used a + *Markdown* one; the eval harness carried a *third* copy ("Answer **the** + question… say so **clearly**") and joined context with a bare newline join + (no `[filename]` prefix) and different chunking defaults (512/64 vs 500/50). + Eval therefore scored a pipeline that was not the one served, and telemetry + assembly was duplicated across four backend sites. + +## Decision + +**A `Retriever` seam.** A runtime-checkable Protocol — +`retrieve(query, top_k) -> list[SearchResult]` — with adapters that either +conform directly or compose an inner Retriever: + +- `DenseRetriever` wraps the vector store (the default). +- `BM25HybridRetriever` conforms directly. +- `RerankingRetriever` composes an inner Retriever: over-fetches, then + re-scores with a cross-encoder. +- `MultiQueryRetriever` composes an inner Retriever: fans rewritten queries out, + unions, and dedups by chunk_id keeping each chunk's best score. + +The four eval-proven levers were **promoted from `src/eval/` to a core +`src/retrieval/` package** (via `git mv`, no shims — same dependency-direction +fix as ADR 0003), so production can activate them without importing eval. + +**A deep `QueryEngine` module** owns retrieve→generate for both paths behind a +small interface (`ask` sync, `ask_stream` streaming). The two are separate +methods sharing prompt/context/telemetry helpers — only streaming runs the +planning pass, so sync keeps its single LLM call. The engine owns: + +- The **single answer prompt**: the Markdown one, for every path (the sync path + adopts it — a deliberate, product-improving change: the frontend renderer + expects Markdown, and eval must measure the shipped prompt). +- **Filename-prefixed context** everywhere (the eval bare join is retired). +- **Telemetry assembly** once, from provider-reported `Usage` (ADR 0003). +- An **optional refusal gate**, off by default. The gate is checked *before* the + no-documents branch, so an empty retrieval is itself an answerability signal + the gate may act on. + +**Production selects a strategy by config** (`RETRIEVER_STRATEGY`, default +`dense`) through a `build_retriever` factory: `dense` and `reranked` are wired; +`hybrid` and `multi_query` are recognised but **deferred** (see Consequences). +`RAGBackend` delegates `query`/`query_with_telemetry`/`stream_query` to the +engine and owns only conversation persistence. + +**The eval harness converges onto the engine.** `EvalPipeline` composes its +levers into one Retriever behind the seam and delegates retrieve→generate to a +`QueryEngine`; its divergent prompt/context/top-k copies are deleted. +`EvalConfig` chunking/top-k/model defaults now derive from `src/config.py` +(single source of truth) — production ingestion reads the same constants. The +values are unchanged (config holds production's actual 512/64), so this is pure +single-sourcing with no behaviour delta on either side. A parity test pins eval +to the shipped prompt and context builders. + +## Consequences + +- **Behaviour preserved for production:** `test_backend.py` and + `test_backend_telemetry.py` pass unchanged — the facade contract + (`query`/`query_with_telemetry`/`stream_query` shapes, the streaming event + protocol, the empty-store path) is intact. The sync answer becoming Markdown + is invisible to those tests (dummy LLM) and is the intended product change. +- **Streaming persistence** now happens at one point (the terminal `result` + event), gated on non-empty results — the empty-store conversation path + persists nothing, as before, and conversation writes concentrate ahead of the + step-5 ConversationStore extraction. The only residual difference from the old + code is an untested mid-generation-crash edge (old: a dangling user message; + now: nothing) — "don't persist a half-failed turn" is the more defensible + behaviour. +- **Eval now measures the shipped pipeline.** Deliberate, spec-accepted changes: + eval telemetry granularity collapses from per-lever stages + (`rewrite`/`rerank`/`refusal_check`) to the engine's `retrieve`/`generate`; + the dead `rewriter_cost_usd` field is dropped (verified: zero readers); and + multi-query dedup shifts from first-seen to best-score-and-truncate (the + *shipped* `MultiQueryRetriever` semantics) — a retrieval-metric shift that is + correct-by-definition once eval measures production. And because the gate is + checked before the no-documents branch, an eval run with an **empty index and + no refusal handler** now returns the no-documents sentinel instead of + generating from empty context (the old eval path) — latent, since eval always + ingests before querying. +- **Deferred, with signal:** `hybrid` needs a live BM25 corpus kept in sync with + ingestion/deletion (a genuinely new feature), and `multi_query`'s production + wiring is held so it lands deliberately; `build_retriever` raises a clear + error pointing here rather than silently falling back. When hybrid *is* + enabled, `BM25HybridRetriever` emits empty `metadata`/`doc_id` (its corpus is + `chunk_id -> text`), so citations degrade — acceptable while the lever is off + by default. +- **New coverage:** contract tests across every adapter; engine tests with a + fake Retriever and fake LLM asserting sync and streaming issue identical + answer instructions; a factory strategy→type test; and the eval↔production + parity test. diff --git a/src/backend.py b/src/backend.py index 12a46223..aa14ee92 100644 --- a/src/backend.py +++ b/src/backend.py @@ -32,7 +32,6 @@ import hashlib import logging import tempfile -import time import uuid from datetime import datetime, timezone from pathlib import Path @@ -43,10 +42,14 @@ from .api.schemas.telemetry import StageTelemetry from .config import ( + CHUNK_OVERLAP, + CHUNK_SIZE, DEFAULT_MODEL, EVAL_MODEL, MAX_TITLE_LENGTH, REASONING_MODEL, + RERANK_OVER_FETCH_N, + RETRIEVER_STRATEGY, SLIDING_WINDOW_SIZE, TOP_K_RESULTS, ) @@ -56,18 +59,34 @@ evaluate_context_precision, evaluate_faithfulness, ) -from .llm_handler import LLMHandler, Usage -from .telemetry.pricing import cost_usd +from .llm_handler import LLMHandler from .models.conversation import Conversation from .models.document import DocumentRecord from .models.evaluation import MessageEvaluation from .models.message import Message, MessageSource -from .observability import get_tracer -from .vector_store import ChromaVectorStore +from .query_engine import QueryEngine, StreamResult +from .retrieval import build_retriever +from .vector_store import ChromaVectorStore, SearchResult logger = logging.getLogger(__name__) +def _source_dict(result: SearchResult) -> dict[str, Any]: + """Shape one retrieved chunk into the source-citation dict the API returns. + + Returns the fields common to both query paths; the synchronous path attaches + ``chunk_index`` on top (the streaming path historically omits it). The two + shapes unify under a typed SourceInfo in step 7 of issue #16. + """ + return { + "doc_id": result.doc_id, + "chunk_id": result.chunk_id, + "filename": result.metadata.get("filename"), + "score": round(result.score, 4), + "excerpt": result.content[:300], + } + + class RAGBackend: """Stateful RAG facade that persists data across requests and restarts. @@ -106,10 +125,12 @@ def __init__(self, engine: Engine, collection: Any) -> None: # TRADE-OFF: Recursive chunking gives better retrieval quality than # fixed-size because it respects paragraph/sentence boundaries. - # 512-char chunks with 64-char overlap is a balanced default. + # SINGLE SOURCE: chunk size/overlap come from config (CHUNK_SIZE/ + # CHUNK_OVERLAP) so the eval harness benchmarks production's + # actual chunking, not a drifted copy (issue #16, step 4c). self.chunker = TextChunker( - chunk_size=512, - chunk_overlap=64, + chunk_size=CHUNK_SIZE, + chunk_overlap=CHUNK_OVERLAP, strategy="recursive", ) @@ -136,9 +157,25 @@ def __init__(self, engine: Engine, collection: Any) -> None: # while strong enough to catch factual errors. self.eval_llm = LLMHandler(model=EVAL_MODEL, max_tokens=4096) + # PATTERN: The QueryEngine owns retrieve->generate for both the sync and + # streaming paths. The Retriever is selected from config + # (RETRIEVER_STRATEGY) behind the seam, so a validated eval + # chain is promoted to production by configuration, not a + # rewrite. Default "dense" preserves current behaviour. + self.query_engine = QueryEngine( + retriever=build_retriever( + RETRIEVER_STRATEGY, self.vector_store, + rerank_over_fetch_n=RERANK_OVER_FETCH_N, + ), + llm=self.llm, + reasoning_llm=self.reasoning_llm, + top_k=TOP_K_RESULTS, + ) + logger.info( - "RAGBackend initialised (engine=%s, answer_model=%s, reasoning_model=%s, eval_model=%s)", - engine.url, DEFAULT_MODEL, REASONING_MODEL, EVAL_MODEL, + "RAGBackend initialised (engine=%s, answer_model=%s, reasoning_model=%s, " + "eval_model=%s, retriever=%s)", + engine.url, DEFAULT_MODEL, REASONING_MODEL, EVAL_MODEL, RETRIEVER_STRATEGY, ) # ------------------------------------------------------------------ # @@ -324,122 +361,41 @@ def query_with_telemetry( ) -> tuple[dict[str, Any], StageTelemetry]: """Run a full RAG query and return per-stage observability data. - Identical to query() in output, but also returns a StageTelemetry - object with retrieve_ms, generate_ms, prompt_tokens, completion_tokens, - and cost_usd. The route layer (Task 5) calls this method so the REST - response can include telemetry without changing the query() contract. - - WHY a sibling instead of modifying query(): - Existing tests assert on result["answer"] and result["sources"] from - query(). Changing query() to return a tuple would break them silently - at dict-access time. The sibling keeps the tested contract intact. + Delegates retrieve->generate to the shared QueryEngine, then shapes the + {answer, sources, confidence} dict the API layer expects. Identical + output to query(), plus a StageTelemetry (retrieve/generate timing and + provider-reported token cost). RAG Pipeline Position: - Question -> [RETRIEVE (traced)] -> [GENERATE (traced)] -> (answer + telemetry) + Question -> [QueryEngine: retrieve -> generate -> telemetry] -> (answer + telemetry) Args: question: Natural language question from the user. top_k: Number of chunks to retrieve (default from config). - model: LLM model override (creates a new handler if different). + model: LLM model override. Returns: - Tuple of (result_dict, StageTelemetry). result_dict has the same - shape as query(): {answer, sources, confidence}. + Tuple of (result_dict, StageTelemetry). result_dict has the shape + {answer, sources, confidence}. """ - k = top_k or TOP_K_RESULTS - tracer = get_tracer() - - # ---- PHASE 1: Retrieval (timed + traced) ---------------------------- - t_retrieve_start = time.perf_counter() - with tracer.start_as_current_span("rag.retrieve") as retrieve_span: - retrieve_span.set_attribute("top_k", k) - retrieve_span.set_attribute("question_len", len(question)) - results = self.vector_store.query(query_text=question, top_k=k) - retrieve_span.set_attribute("results_count", len(results)) - retrieve_ms = (time.perf_counter() - t_retrieve_start) * 1000 - - if not results: - # PATTERN: Early return with zero telemetry — no LLM call was made. - return ( - { - "answer": "No documents indexed yet. Please upload documents first.", - "sources": [], - "confidence": 0.0, - }, - StageTelemetry( - retrieve_ms=round(retrieve_ms, 2), - generate_ms=0.0, - prompt_tokens=0, - completion_tokens=0, - cost_usd=0.0, - ), - ) - - # Build context string from retrieved chunks - context = "\n\n".join( - f"[{r.metadata.get('filename', 'unknown')}] {r.content}" - for r in results - ) - - # Generate answer (create a per-query handler if model differs) - handler = self.llm - if model and model != self.llm.model: - handler = LLMHandler(model=model) - - # Build the answer prompt (system + user). Passing it through - # generate_with_usage means telemetry counts the exact tokens the - # provider billed — no reconstruction of a proxy string. - answer_system_prompt = ( - "You are a helpful assistant. Answer the user's question based solely on the " - "provided context. If the context does not contain enough information, say so." - ) - answer_user_prompt = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:" - - # ---- PHASE 2: Generation (timed + traced) --------------------------- - t_generate_start = time.perf_counter() - with tracer.start_as_current_span("rag.generate") as generate_span: - generate_span.set_attribute("model", handler.model) - answer, prompt_tokens, completion_tokens = handler.generate_with_usage( - answer_user_prompt, system_prompt=answer_system_prompt - ) - generate_span.set_attribute("answer_len", len(answer)) - generate_ms = (time.perf_counter() - t_generate_start) * 1000 + results, answer, telemetry = self.query_engine.ask(question, top_k=top_k, model=model) sources = [ - { - "doc_id": r.doc_id, - "chunk_id": r.chunk_id, - "filename": r.metadata.get("filename"), - "score": round(r.score, 4), - "excerpt": r.content[:300], - "chunk_index": r.metadata.get("chunk_index"), - } + {**_source_dict(r), "chunk_index": r.metadata.get("chunk_index")} for r in results ] - # PATTERN: Confidence = clamped average of top-3 similarity scores. + # PATTERN: Confidence = clamped average of top-3 similarity scores; 0.0 + # when there are no results (empty index or a refusal). top_scores = [r.score for r in results[: min(3, len(results))]] - confidence = max(0.0, min(1.0, sum(top_scores) / len(top_scores))) - - # ---- Telemetry assembly --------------------------------------------- - # Token counts come from the provider-reported usage above (or the - # adapter's local count fallback) — never a reconstructed prompt. - total_cost = cost_usd(handler.model, prompt_tokens, completion_tokens) - - telemetry = StageTelemetry( - retrieve_ms=round(retrieve_ms, 2), - generate_ms=round(generate_ms, 2), - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - cost_usd=total_cost, + confidence = ( + round(max(0.0, min(1.0, sum(top_scores) / len(top_scores))), 4) + if top_scores + else 0.0 ) return ( - { - "answer": answer, - "sources": sources, - "confidence": round(confidence, 4), - }, + {"answer": answer, "sources": sources, "confidence": confidence}, telemetry, ) @@ -474,234 +430,64 @@ def stream_query( WHY telemetry covers the answer pass only (not reasoning): The reasoning pass uses a separate model (REASONING_MODEL) with its own - cost. Telemetry here tracks the user-visible answer generation; mixing - two model costs into one StageTelemetry would confuse the "per-query - cost" display. Reasoning cost is a separate concern. + cost. Telemetry here tracks the user-visible answer generation; the + engine drops the reasoning pass's usage so the "per-query cost" display + reflects one model's spend. + + WHY persistence is deferred to the terminal event: the QueryEngine owns + retrieve->generate and yields a terminal ("result", ...) carrying the + chunks + telemetry. This facade owns only conversation persistence — it + saves the user and assistant messages together, and only when retrieval + produced results, so an empty index leaves nothing persisted (as before) + and all conversation writes concentrate at one point. Yields: Tuples as described above. """ - k = top_k or TOP_K_RESULTS - tracer = get_tracer() - - # ---- PHASE 0: Retrieval (timed + traced) -------------------------------- - yield ("status", "Searching indexed documents...") - t_retrieve_start = time.perf_counter() - with tracer.start_as_current_span("rag.retrieve") as retrieve_span: - retrieve_span.set_attribute("top_k", k) - retrieve_span.set_attribute("question_len", len(question)) - results = self.vector_store.query(query_text=question, top_k=k) - retrieve_span.set_attribute("results_count", len(results)) - retrieve_ms = (time.perf_counter() - t_retrieve_start) * 1000 - - if not results: - yield ("status", "No indexed documents — nothing to retrieve.") - yield ("token", "No documents indexed yet. Please upload documents first.") - yield ("done", {"sources": []}) - # PATTERN: Emit zero telemetry even on early return so the route - # layer always gets a telemetry event it can forward. - yield ("telemetry", StageTelemetry( - retrieve_ms=round(retrieve_ms, 2), - generate_ms=0.0, - prompt_tokens=0, - completion_tokens=0, - cost_usd=0.0, - ).model_dump()) - return - - # WHY: Summarise retrieval in one status line so the user can see which - # files contributed without inspecting the sources panel yet. - filenames = [r.metadata.get("filename", "unknown") for r in results] - unique_files = sorted({f for f in filenames if f}) - file_summary = ", ".join(unique_files[:3]) - if len(unique_files) > 3: - file_summary += f" (+{len(unique_files) - 3} more)" - yield ( - "status", - f"Retrieved {len(results)} chunk(s) across {len(unique_files)} file(s): {file_summary}", - ) - - # Build context from retrieved chunks (shared by reasoning + answer) - context = "\n\n".join( - f"[{r.metadata.get('filename', 'unknown')}] {r.content}" - for r in results - ) - - handler = self.llm - if model and model != self.llm.model: - handler = LLMHandler(model=model) - - sources = [ - { - "doc_id": r.doc_id, - "chunk_id": r.chunk_id, - "filename": r.metadata.get("filename"), - "score": round(r.score, 4), - "excerpt": r.content[:300], - } - for r in results - ] - - # ---- PHASE 1: Reasoning pass (chain-of-thought) ------------------------ - # WHY a dedicated reasoning prompt: Asking the model to "think first, - # answer later" in a single call is brittle — formatting drifts between - # providers. A separate short call with a focused system prompt gives - # deterministic reasoning tokens we can stream as their own event type. - # - # WHY a dedicated reasoning model: The CoT output is short, throwaway - # scaffolding. Running it through the user's (potentially premium) - # answer model doubles cost for no quality gain. self.reasoning_llm is - # cached to REASONING_MODEL (default: gpt-5-nano) independent of the - # answer model — so premium answers stay cheap to "think" about. - yield ("status", f"Analyzing retrieved context ({self.reasoning_llm.model})...") - - # WHY a RESEARCH-PLAN style prompt (not raw CoT): exposing raw generated - # chain-of-thought is a documented product risk — the model may - # verbalise uncertain or incorrect intermediate beliefs that users - # mistake for confident answers. Instead we ask for a concise, - # outcome-oriented *reasoning summary* that describes the plan - # without asserting factual conclusions. Still streams live so - # the UI's two-step feel (plan → answer) is preserved. - reasoning_system = ( - "You are the planning step of a retrieval-augmented Q&A system. " - "In 3-5 concise sentences, summarise how you will construct the " - "answer using the retrieved excerpts. Cover:\n" - "1) What the user is asking, resolving any ambiguity explicitly.\n" - "2) Which excerpts are most relevant and the gist of their support.\n" - "3) Any gaps or conflicts the reader should be aware of.\n" - "4) The shape of the answer you will give next.\n" - "Stay factual and outcome-oriented — describe the plan, do not " - "verbalise stream-of-consciousness reasoning. No markdown headings, " - "no bullet lists, no preamble. Do NOT produce the final answer." - ) - reasoning_user = ( - f"Context:\n{context}\n\nQuestion: {question}\n\n" - "Reasoning plan (summary only, do not answer):" - ) - - try: - for item in self.reasoning_llm.stream_response( - reasoning_user, system_prompt=reasoning_system - ): - # The reasoning pass reports its own usage, but telemetry covers - # the answer pass only — so drop the reasoning terminal Usage. - if isinstance(item, Usage): - continue - yield ("reasoning", item) - except Exception as exc: - # PATTERN: Reasoning is best-effort — a failure here must not block - # the final answer. Log, emit a terse status, continue to the answer. - logger.warning("Reasoning pass failed: %s", exc) - yield ("status", "Reasoning unavailable — skipping to answer.") - - # ---- PHASE 2: Answer pass (timed + traced) ----------------------------- - yield ("status", "Composing answer...") - - system_prompt = ( - "You are a helpful assistant. Answer the user's question based solely on the " - "provided context. If the context does not contain enough information, say so.\n\n" - "Format your response using Markdown for readability:\n" - "- Use ## for main sections and ### for sub-sections (max 3 levels)\n" - "- Use **bold** for key terms and important concepts\n" - "- Use bullet points (-) for lists of related items\n" - "- Use numbered lists (1.) for sequential steps\n" - "- Use `inline code` for technical terms, parameters, or commands\n" - "- Use fenced code blocks (```language) for code snippets\n" - "- Use > blockquotes for notable quotes from the context\n" - "- Keep paragraphs short (2-3 sentences max)\n" - "- Add blank lines between sections for visual breathing room\n" - "Do NOT use # (h1) headings. Start directly with content or ## sections." - ) - user_prompt = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:" - - # Accumulate full response for persistence; capture the answer pass's - # reported usage (yielded as the stream's terminal event) for telemetry. - full_response: list[str] = [] - answer_usage: Usage | None = None - - t_generate_start = time.perf_counter() - - if conversation_id: - # PHASE 1: Save user message BEFORE streaming + # The sliding window is prior *completed* pairs; the current question is + # unpaired, so computing the window before persisting the user message is + # behaviour-identical to computing it after (the old ordering). + history = self._get_sliding_window(conversation_id) if conversation_id else [] + + answer_parts: list[str] = [] + result: StreamResult | None = None + for event_type, data in self.query_engine.ask_stream( + question, top_k=top_k, model=model, history=history + ): + if isinstance(data, StreamResult): + result = data # internal terminal event — consumed, not forwarded + continue + if event_type == "token": + answer_parts.append(data) + yield (event_type, data) + + # The engine always emits exactly one terminal StreamResult (both the + # normal and the no-results/refusal paths), so this binds every time. + assert result is not None, "QueryEngine.ask_stream emitted no terminal result" + results = result.results + sources = [_source_dict(r) for r in results] + + if conversation_id and results: self._save_message(conversation_id, "user", question) - - # PHASE 2: Load sliding window (completed pairs only — excludes - # the just-saved user message because it has no assistant - # reply yet). - window = self._get_sliding_window(conversation_id) - - # PHASE 3: Build messages list for multi-turn generation - messages = [{"role": "system", "content": system_prompt}] - messages.extend(window) - messages.append({"role": "user", "content": user_prompt}) - - # PHASE 4: Stream via messages API (multi-turn aware) - with tracer.start_as_current_span("rag.generate") as generate_span: - generate_span.set_attribute("model", handler.model) - generate_span.set_attribute("has_conversation", True) - for item in handler.stream_messages(messages): - if isinstance(item, Usage): - answer_usage = item - continue - full_response.append(item) - yield ("token", item) - generate_span.set_attribute("answer_len", sum(len(t) for t in full_response)) - - generate_ms = (time.perf_counter() - t_generate_start) * 1000 - - # PHASE 5: Save assistant message + sources - assistant_content = "".join(full_response) assistant_msg_id = self._save_message( - conversation_id, "assistant", assistant_content, - model=handler.model, sources=sources, + conversation_id, "assistant", "".join(answer_parts), + model=result.model, sources=sources, ) - - # PHASE 6: Auto-title on first message + # WHY: Auto-title on the first turn so the sidebar shows something + # meaningful; message_id + conversation_id let the frontend + # update local state without re-fetching. self._auto_title(conversation_id, question) - - # WHY: Include message_id and conversation_id in the done event - # so the frontend can update its local state (add the new - # message to the conversation without re-fetching). yield ("done", { "sources": sources, "message_id": assistant_msg_id, "conversation_id": conversation_id, }) - else: - # No conversation — simple single-turn streaming - with tracer.start_as_current_span("rag.generate") as generate_span: - generate_span.set_attribute("model", handler.model) - generate_span.set_attribute("has_conversation", False) - for item in handler.stream_response(user_prompt, system_prompt=system_prompt): - if isinstance(item, Usage): - answer_usage = item - continue - full_response.append(item) - yield ("token", item) - generate_span.set_attribute("answer_len", sum(len(t) for t in full_response)) - - generate_ms = (time.perf_counter() - t_generate_start) * 1000 - yield ("done", {"sources": sources}) - # ---- Telemetry assembly (after done, additive) ----------------------- - # WHY after done: the done event is what the client waits for to show - # sources. Telemetry is a secondary signal — emit it last. - # - # Usage is the answer pass's reported figure (provider-reported where the - # provider supplies it, adapter-counted otherwise), carried on the - # stream's terminal event — no prompt reconstruction. - usage = answer_usage or Usage(prompt_tokens=0, completion_tokens=0) - total_cost = cost_usd(handler.model, usage.prompt_tokens, usage.completion_tokens) - - yield ("telemetry", StageTelemetry( - retrieve_ms=round(retrieve_ms, 2), - generate_ms=round(generate_ms, 2), - prompt_tokens=usage.prompt_tokens, - completion_tokens=usage.completion_tokens, - cost_usd=total_cost, - ).model_dump()) + # Telemetry last (additive): the done event is what the client waits on + # for sources; telemetry is a secondary signal. + yield ("telemetry", result.telemetry.model_dump()) # ------------------------------------------------------------------ # # Document management # @@ -1400,11 +1186,13 @@ def _get_sliding_window( assistant message. A user message with no assistant reply is NOT a complete pair and must be excluded. - WHY exclude unpaired: The current user question is saved BEFORE streaming - starts (phase 1 of stream_query). If we included it in the window, the - LLM would see the question twice — once in the window and once as the - explicit user turn appended by stream_query. This causes the model to - repeat itself or get confused. + WHY exclude unpaired: the window is built from PRIOR completed pairs and + fed to the answer prompt, which appends the current question as its own + user turn. A dangling user message — e.g. a prior turn whose generation + failed before its assistant reply was persisted — must not leak in, or + the LLM would see a question with no answer and may repeat or get + confused. (Since step 4c the current turn is persisted only AFTER + generation, so the live question is never in the window regardless.) Args: conversation_id: UUID of the conversation. diff --git a/src/config.py b/src/config.py index 833dfe01..a736652a 100644 --- a/src/config.py +++ b/src/config.py @@ -54,14 +54,29 @@ # PATTERN: Seed the PRNG for reproducible TF-IDF splits and test fixtures. RANDOM_SEED: int = 42 -# TRADE-OFF: 500-char chunks balance context (enough text for meaning) vs. -# precision (small enough for specific retrieval). Overlap of 50 +# TRADE-OFF: 512-char chunks balance context (enough text for meaning) vs. +# precision (small enough for specific retrieval). Overlap of 64 # prevents answers that straddle chunk boundaries from being lost. -CHUNK_SIZE: int = 500 -CHUNK_OVERLAP: int = 50 +# SINGLE SOURCE: production ingestion (RAGBackend) and the eval harness both +# read these — so "baseline" eval measures production's chunking +# (issue #16, step 4c). These are production's actual shipped values. +CHUNK_SIZE: int = 512 +CHUNK_OVERLAP: int = 64 TOP_K_RESULTS: int = 5 +# WHY a strategy selector: the Retriever seam (ADR 0004) makes dense, reranked, +# hybrid, and multi-query retrieval interchangeable. Production picks one +# here, so an eval-validated chain is promoted by config, not by a rewrite. +# "dense" is the behaviour-preserving default; "reranked" is wired; +# "hybrid"/"multi_query" are recognised but deferred (see build_retriever). +RETRIEVER_STRATEGY: str = "dense" + +# WHY 20: the reranked strategy over-fetches this many dense candidates before +# the cross-encoder narrows them — wide enough to give the precise +# reranker real choice, matching the eval-tuned default. +RERANK_OVER_FETCH_N: int = 20 + # --------------------------------------------------------------------------- # API server # --------------------------------------------------------------------------- diff --git a/src/eval/config.py b/src/eval/config.py index e9ebd0f0..3d17d8d9 100644 --- a/src/eval/config.py +++ b/src/eval/config.py @@ -21,6 +21,19 @@ import yaml from pydantic import BaseModel, ConfigDict, Field +# WHY import production config: eval "baseline" runs must benchmark the pipeline +# users actually get, so chunking/top-k/model defaults derive from the single +# source of truth (src/config.py) rather than drifting as independent literals +# (issue #16, step 4c). An explicit YAML value still overrides any default. +from src.config import ( + CHUNK_OVERLAP, + CHUNK_SIZE, + DEFAULT_MODEL, + EVAL_MODEL, + REASONING_MODEL, + TOP_K_RESULTS, +) + class ChunkerCfg(BaseModel): """Chunking strategy configuration. @@ -31,8 +44,8 @@ class ChunkerCfg(BaseModel): """ strategy: Literal["fixed", "recursive", "semantic"] = "recursive" - chunk_size: int = 512 - chunk_overlap: int = 64 + chunk_size: int = CHUNK_SIZE + chunk_overlap: int = CHUNK_OVERLAP class RetrieverCfg(BaseModel): @@ -42,7 +55,7 @@ class RetrieverCfg(BaseModel): Pipeline position: QUERYING step — Embeddings → [Retriever] → Top-K chunks. """ - top_k: int = 5 + top_k: int = TOP_K_RESULTS class GeneratorCfg(BaseModel): @@ -53,9 +66,9 @@ class GeneratorCfg(BaseModel): Pipeline position: QUERYING step — Chunks → [Generator] → Answer. """ - model: str = "gpt-5-mini" + model: str = DEFAULT_MODEL # WHY: reasoning_model is optional — None disables the CoT pre-pass. - reasoning_model: str | None = "gpt-4.1-nano" + reasoning_model: str | None = REASONING_MODEL class EmbedderCfg(BaseModel): @@ -155,7 +168,7 @@ class EvalCfg(BaseModel): """ datasets: list[Literal["squad_v2_dev_200", "ml_papers_v1"]] - judge_model: str = "gpt-4.1-mini" + judge_model: str = EVAL_MODEL bootstrap_n: int = 1000 permutation_n: int = 10000 seed: int = 42 diff --git a/src/eval/pipeline_factory.py b/src/eval/pipeline_factory.py index 2bed4747..d0128700 100644 --- a/src/eval/pipeline_factory.py +++ b/src/eval/pipeline_factory.py @@ -12,11 +12,10 @@ - Ephemeral Chroma collection per (config, dataset) so two concurrent runs cannot pollute each other's vectors. Random suffix on the collection name guards against collisions. - - Per-stage timings via time.perf_counter() so the runner can record - p50/p95/p99 latency at aggregation time. - - Token counting: tiktoken if available, word-count×1.3 fallback — - eval should not hard-fail because a tokenizer for a new model - isn't installed. + - retrieve->generate is delegated to the shared QueryEngine (issue #16, + step 4c): the levers become a composed Retriever behind the seam, and the + prompt / context / telemetry come from production's one module — so eval + measures the pipeline that is actually shipped, not a hand-copied twin. - Test doubles (DummyLLM) inject via *_override params; production uses LLMHandler(model_name). @@ -30,7 +29,6 @@ from __future__ import annotations import logging -import time import uuid from dataclasses import dataclass, field from typing import Any @@ -38,10 +36,19 @@ import chromadb from src.document_loader import TextChunker -from src.telemetry.tokens import count_tokens from src.eval.config import EvalConfig from src.eval.schemas import EvalQuestion from src.llm_handler import LLMHandler +from src.query_engine import QueryEngine +from src.retrieval import ( + CrossEncoderReranker, + DenseRetriever, + MultiQueryRetriever, + QueryRewriter, + RefusalHandler, + RerankingRetriever, + Retriever, +) from src.vector_store import ChromaVectorStore, SearchResult logger = logging.getLogger(__name__) @@ -71,11 +78,12 @@ class EvalPipeline: config: EvalConfig dataset_name: str - # Phase 2 additions — None when the corresponding lever is off. - hybrid_retriever: object | None = None # BM25HybridRetriever or None - reranker: object | None = None # CrossEncoderReranker or None - rewriter: object | None = None # QueryRewriter or None - refusal_handler: object | None = None # RefusalHandler or None + # Phase 2 additions — None when the corresponding lever is off. Composed into + # a single Retriever by _get_engine(); the refusal gate is passed to the engine. + hybrid_retriever: Retriever | None = None # BM25HybridRetriever, set during ingest + reranker: CrossEncoderReranker | None = None + rewriter: QueryRewriter | None = None + refusal_handler: RefusalHandler | None = None # Private: needed for teardown() — ChromaVectorStore doesn't own the client. # WHY not reach into vector_store._collection._client: that would couple us @@ -83,6 +91,10 @@ class EvalPipeline: _client: chromadb.ClientAPI = field(repr=False, default=None) # type: ignore[assignment] _collection_name: str = field(repr=False, default="") + # Lazily-built QueryEngine, cached after the first query(). Deferred because + # the hybrid retriever is only assembled during ingest() (it needs the corpus). + _engine: QueryEngine | None = field(repr=False, default=None) + def ingest(self, questions: list[EvalQuestion]) -> None: """Upsert question contexts into the vector store. @@ -205,111 +217,67 @@ def _ingest_ml_papers(self) -> None: ) def query(self, question: str) -> tuple[list[SearchResult], str, dict]: - """Retrieve relevant chunks and generate an answer with timing + cost telemetry. + """Retrieve chunks and generate an answer via the shared QueryEngine. - Phase 2 pipeline steps: rewrite → retrieve (hybrid or dense) → rerank → - refusal gate → generate. Each step is a no-op when the corresponding - config lever is off, preserving backward compatibility with Phase 1 callers. + The levers become a composed Retriever behind the seam (hybrid-or-dense, + wrapped in multi-query and reranking adapters as configured); the engine + applies the refusal gate, builds the shipped prompt + context, and + assembles telemetry. This is the convergence that makes eval measure the + production pipeline (issue #16, step 4c). Args: question: Natural language question from the eval set. Returns: Tuple of (chunks, answer, telemetry). telemetry keys: - timings_ms: dict of stage→ms for rewrite, retrieve, rerank, - refusal_check, generate + timings_ms: {"retrieve": float, "generate": float} tokens: {"prompt": int, "completion": int} - cost_usd: float (generator side) - rewriter_cost_usd: float (rewriter side, 0.0 when disabled) + cost_usd: float """ - from src.telemetry import pricing + results, answer, stage = self._get_engine().ask(question) + return results, answer, { + "timings_ms": {"retrieve": stage.retrieve_ms, "generate": stage.generate_ms}, + "tokens": {"prompt": stage.prompt_tokens, "completion": stage.completion_tokens}, + "cost_usd": stage.cost_usd, + } + + def _get_engine(self) -> QueryEngine: + """Build (once) and return the QueryEngine composed from the configured levers. - timings: dict[str, float] = {} - rewriter_cost = 0.0 + Built lazily because the hybrid retriever is only available after ingest(). + The lever composition is the Retriever seam used as designed: multi-query + wraps the base retriever, reranking wraps that. + """ + if self._engine is not None: + return self._engine - # ---- Rewrite (lever 2e) ----------------------------------------------- - t = time.perf_counter() + base = self.hybrid_retriever or DenseRetriever(self.vector_store) + retriever = base if self.rewriter is not None: - queries, rewriter_cost, _, _ = self.rewriter.expand(question) - else: - queries = [question] - timings["rewrite"] = (time.perf_counter() - t) * 1000.0 - - # ---- Retrieve --------------------------------------------------------- - # WHY use rerank_top_n for initial fetch when a reranker is active: - # the reranker needs a wider candidate pool to re-score before final_top_k. - top_k_initial = ( - self.config.pipeline.reranker.rerank_top_n - if self.reranker is not None else self.config.pipeline.retriever.top_k - ) - t = time.perf_counter() - if self.hybrid_retriever is not None: - seen: dict[str, SearchResult] = {} - for q in queries: - for r in self.hybrid_retriever.retrieve(q, top_k=top_k_initial): - if r.chunk_id not in seen: - seen[r.chunk_id] = r - results = list(seen.values()) - else: - seen = {} - for q in queries: - for r in self.vector_store.query(query_text=q, top_k=top_k_initial): - if r.chunk_id not in seen: - seen[r.chunk_id] = r - results = list(seen.values()) - timings["retrieve"] = (time.perf_counter() - t) * 1000.0 - - # ---- Rerank (lever 2d) ------------------------------------------------ - t = time.perf_counter() + retriever = MultiQueryRetriever(inner=retriever, rewriter=self.rewriter) if self.reranker is not None: - results = self.reranker.rerank( - question, results, - final_top_k=self.config.pipeline.reranker.final_top_k, + retriever = RerankingRetriever( + inner=retriever, reranker=self.reranker, + over_fetch_n=self.config.pipeline.reranker.rerank_top_n, ) - else: - results = results[: self.config.pipeline.retriever.top_k] - timings["rerank"] = (time.perf_counter() - t) * 1000.0 - - # ---- Refusal gate (lever 2g) ------------------------------------------ - t = time.perf_counter() - if self.refusal_handler is not None and self.refusal_handler.should_refuse(results): - chunks, answer = self.refusal_handler.refuse_response() - timings["refusal_check"] = (time.perf_counter() - t) * 1000.0 - return chunks, answer, { - "timings_ms": timings, - "tokens": {"prompt": 0, "completion": 0}, - "cost_usd": 0.0, - "rewriter_cost_usd": rewriter_cost, - } - timings["refusal_check"] = (time.perf_counter() - t) * 1000.0 - - # ---- Generate --------------------------------------------------------- - context = "\n\n".join(r.content for r in results) - system_prompt = ( - "You are a helpful assistant. Answer the question based solely on the " - "provided context. If the context does not contain enough information, " - "say so clearly." - ) - user_prompt = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:" - # WHY count both: the LLM sees system_prompt + user_prompt as prompt tokens. - full_prompt_text = system_prompt + "\n" + user_prompt - model = self.config.pipeline.generator.model - - t = time.perf_counter() - answer = self.llm.generate(user_prompt, system_prompt=system_prompt) - timings["generate"] = (time.perf_counter() - t) * 1000.0 - - # ---- Token counting + cost estimation --------------------------------- - prompt_tokens = count_tokens(full_prompt_text, model) - completion_tokens = count_tokens(answer, model) - cost = pricing.cost_usd(model, prompt_tokens, completion_tokens) - return results, answer, { - "timings_ms": timings, - "tokens": {"prompt": prompt_tokens, "completion": completion_tokens}, - "cost_usd": cost, - "rewriter_cost_usd": rewriter_cost, - } + # When reranking is on, the final count is the reranker's final_top_k + # (it over-fetches the wider rerank_top_n first); otherwise it is top_k. + top_k = ( + self.config.pipeline.reranker.final_top_k + if self.reranker is not None + else self.config.pipeline.retriever.top_k + ) + # reasoning_llm is unused on the sync ask() path (the eval harness never + # streams), so the answer LLM stands in for the constructor requirement. + self._engine = QueryEngine( + retriever=retriever, + llm=self.llm, + reasoning_llm=self.llm, + top_k=top_k, + refusal=self.refusal_handler, + ) + return self._engine def teardown(self) -> None: """Delete the ephemeral Chroma collection and release the client reference. @@ -454,7 +422,7 @@ def _build_hybrid_retriever(cfg, vector_store, documents: dict[str, str]): """ if not cfg.enabled: return None - from src.eval.retrievers import BM25HybridRetriever + from src.retrieval import BM25HybridRetriever return BM25HybridRetriever( vector_store=vector_store, documents=documents, bm25_top_k=cfg.bm25_top_k, dense_top_k=cfg.dense_top_k, rrf_k=cfg.rrf_k, @@ -472,7 +440,7 @@ def _build_reranker(cfg): """ if cfg.model is None: return None - from src.eval.retrievers import CrossEncoderReranker + from src.retrieval import CrossEncoderReranker return CrossEncoderReranker() @@ -489,7 +457,7 @@ def _build_rewriter(cfg, llm): """ if cfg.model is None: return None - from src.eval.transforms import QueryRewriter + from src.retrieval import QueryRewriter return QueryRewriter(model=cfg.model, max_expansions=cfg.max_expansions, llm=llm) @@ -504,7 +472,7 @@ def _build_refusal(cfg): """ if not cfg.enabled: return None - from src.eval.transforms import RefusalHandler + from src.retrieval import RefusalHandler return RefusalHandler( enabled=True, similarity_threshold=cfg.similarity_threshold, no_answer_text=cfg.no_answer_text, diff --git a/src/eval/retrievers/__init__.py b/src/eval/retrievers/__init__.py deleted file mode 100644 index 9417a9e5..00000000 --- a/src/eval/retrievers/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Phase 2 retriever package — hybrid sparse/dense retrieval and reranking.""" - -from src.eval.retrievers.bm25_hybrid import BM25HybridRetriever -from src.eval.retrievers.reranker import CrossEncoderReranker - -__all__ = ["BM25HybridRetriever", "CrossEncoderReranker"] diff --git a/src/eval/retrievers/reranker.py b/src/eval/retrievers/reranker.py deleted file mode 100644 index a79d53da..00000000 --- a/src/eval/retrievers/reranker.py +++ /dev/null @@ -1,62 +0,0 @@ -"""CrossEncoderReranker — re-scores retrieval candidates with a cross-encoder model. - -Pipeline position: - Retriever top-N → [CrossEncoderReranker] → top-K → Refusal / Generator - -Phase 2 lever 2d. Cross-encoders (single-tower models that consume both -the query and a candidate together) typically outperform bi-encoder retrieval -in precision at the cost of latency. We use ms-marco-MiniLM-L-6-v2 — small -enough to run on CPU in milliseconds per pair, trained on MS MARCO so the -ranking signal transfers well to general-domain QA. -""" - -from __future__ import annotations - -from src.vector_store import SearchResult - - -class CrossEncoderReranker: - """Wraps sentence-transformers CrossEncoder to re-score retrieval candidates.""" - - MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2" - - def __init__(self) -> None: - from sentence_transformers import CrossEncoder - - self._model = CrossEncoder(self.MODEL_NAME) - - def rerank( - self, - query: str, - candidates: list[SearchResult], - final_top_k: int, - ) -> list[SearchResult]: - """Re-score candidates against the query and return top-K reranked. - - Args: - query: Original user query. - candidates: Pre-retrieved chunks (typically top-N from a base retriever). - final_top_k: How many to keep after reranking. - - Returns: - Top-K SearchResult ordered by descending cross-encoder score. The - original `score` field is *replaced* with the cross-encoder score so - downstream consumers reading `result.score` get the more precise signal. - """ - if not candidates: - return [] - pairs = [(query, c.content) for c in candidates] - scores = self._model.predict(pairs) - scored = sorted( - zip(candidates, scores), key=lambda t: t[1], reverse=True, - )[:final_top_k] - return [ - SearchResult( - doc_id=c.doc_id, - chunk_id=c.chunk_id, - content=c.content, - score=float(s), - metadata=c.metadata, - ) - for c, s in scored - ] diff --git a/src/eval/transforms/__init__.py b/src/eval/transforms/__init__.py deleted file mode 100644 index 1dc285e7..00000000 --- a/src/eval/transforms/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Phase 2 transforms — pre/post pipeline hooks (rewriter, refusal handler).""" - -from src.eval.transforms.query_rewriter import QueryRewriter -from src.eval.transforms.refusal_handler import RefusalHandler - -__all__ = ["QueryRewriter", "RefusalHandler"] diff --git a/src/eval/transforms/query_rewriter.py b/src/eval/transforms/query_rewriter.py deleted file mode 100644 index 8040ab98..00000000 --- a/src/eval/transforms/query_rewriter.py +++ /dev/null @@ -1,107 +0,0 @@ -"""QueryRewriter — LLM-based query expansion with token/cost capture. - -Pipeline position: - user query → [QueryRewriter] → {q, q', q''} → Retriever → ... - -Phase 2 lever 2e. Expansion gives the retriever multiple lexical/semantic -formulations of the same intent, which raises recall on questions where the -original phrasing diverges from the corpus phrasing. We use a tiny model -(gpt-4.1-nano) because the task is cheap and we don't want this lever to -dominate the cost ledger. -""" - -from __future__ import annotations - -import json -import logging -import re -from typing import Protocol - -from src.telemetry import pricing - -logger = logging.getLogger(__name__) - - -class _LLMHandler(Protocol): - """Structural type for any object exposing generate_with_usage.""" - - def generate_with_usage( - self, prompt: str, system_prompt: str | None = None, - ) -> tuple[str, int, int]: ... - - -class QueryRewriter: - """Expands one user query into up to N alternative phrasings via an LLM.""" - - SYSTEM_PROMPT = ( - "You rewrite user search queries into alternative phrasings that preserve " - "the original intent but vary surface form. Respond ONLY with a JSON " - "array of strings — no prose, no code fences." - ) - - def __init__( - self, - model: str | None, - max_expansions: int, - llm: _LLMHandler | None, - ) -> None: - """Configure the rewriter. - - Args: - model: LLM model name. None disables rewriting (pass-through). - max_expansions: Cap on the number of alternative phrasings to return. - llm: Object exposing generate_with_usage(prompt, system_prompt). Required - if model is not None. - """ - self._model = model - self._max_expansions = max_expansions - self._llm = llm - - def expand(self, query: str) -> tuple[list[str], float, int, int]: - """Expand `query` into up to N+1 unique phrasings. - - Returns: - (queries, cost_usd, prompt_tokens, completion_tokens). The original - query is always the first element. When `model is None`, returns - ([query], 0.0, 0, 0) and skips the LLM call. - """ - if self._model is None: - return [query], 0.0, 0, 0 - if self._llm is None: - raise ValueError("QueryRewriter has model set but no llm handler provided.") - - user_prompt = ( - f'Original query: "{query}"\n\n' - f"Return a JSON array of up to {self._max_expansions} alternative " - f"phrasings of this query. Do NOT include the original." - ) - raw, p_t, c_t = self._llm.generate_with_usage( - user_prompt, system_prompt=self.SYSTEM_PROMPT, - ) - cost = pricing.cost_usd(self._model, p_t, c_t) - - expansions = self._parse_expansions(raw) - # Always lead with original; dedupe; cap at original + max_expansions. - ordered: list[str] = [query] - for alt in expansions: - if alt and alt not in ordered: - ordered.append(alt) - if len(ordered) >= self._max_expansions + 1: - break - return ordered, cost, p_t, c_t - - @staticmethod - def _parse_expansions(raw: str) -> list[str]: - """Strip code fences and parse the JSON array; return [] on failure.""" - stripped = re.sub(r"^```(?:json)?\s*", "", raw.strip()) - stripped = re.sub(r"\s*```$", "", stripped).strip() - try: - parsed = json.loads(stripped) - except json.JSONDecodeError: - logger.warning( - "QueryRewriter got non-JSON response — falling back to [query] only." - ) - return [] - if not isinstance(parsed, list): - return [] - return [str(item) for item in parsed if isinstance(item, str)] diff --git a/src/query_engine/__init__.py b/src/query_engine/__init__.py new file mode 100644 index 00000000..1728b857 --- /dev/null +++ b/src/query_engine/__init__.py @@ -0,0 +1,13 @@ +"""QueryEngine package — the shared retrieve->generate module (issue #16, step 4). + +`QueryEngine` owns retrieval, prompt construction, generation, and telemetry +assembly behind a small interface (`ask`, `ask_stream`). Both the production +`RAGBackend` facade and the eval harness call it, so eval measures the shipped +pipeline. See [ADR 0004](../../docs/adr/0004-retriever-seam-and-query-engine.md). +""" + +from __future__ import annotations + +from src.query_engine.engine import QueryEngine, StreamResult + +__all__ = ["QueryEngine", "StreamResult"] diff --git a/src/query_engine/engine.py b/src/query_engine/engine.py new file mode 100644 index 00000000..bef30a9b --- /dev/null +++ b/src/query_engine/engine.py @@ -0,0 +1,249 @@ +"""QueryEngine — the one module that owns retrieve->generate. + +RAG Pipeline Position: + Query -> [QUERYENGINE] -> Answer + Retriever -> (refusal gate) -> prompt -> LLM -> telemetry + +What concept it teaches: + A *deep* module: a small interface (`ask`, `ask_stream`) hiding retrieval, + prompt construction, generation, and telemetry assembly. Both the production + facade and the eval harness call it, so a measured improvement is an + improvement in the shipped system — there is no second retrieve->generate + implementation to drift from. + +Design Decisions (full rationale in ADR 0004): + - The `Retriever` is injected (constructor injection), so dense / hybrid / + reranked / multi-query retrieval are swapped by configuration, not code. + - `ask` (sync) and `ask_stream` (streaming) are two methods sharing the same + prompt / context / telemetry helpers, NOT one body: only streaming runs the + extra planning pass, so the sync path keeps its single LLM call. + - The refusal gate is optional and off by default, preserving production. +""" + +from __future__ import annotations + +import logging +import time +from typing import Iterator + +from src.api.schemas.telemetry import StageTelemetry +from src.llm_handler import LLMHandler, Usage +from src.observability import get_tracer +from src.query_engine import telemetry as telemetry_asm +from src.query_engine.prompt import ( + ANSWER_SYSTEM_PROMPT, + NO_DOCUMENTS_ANSWER, + REASONING_SYSTEM_PROMPT, + build_answer_user_prompt, + build_context, + build_reasoning_user_prompt, +) +from src.query_engine.streaming import ( + StreamEvent, + StreamResult, + build_answer_messages, + retrieval_summary, +) +from src.retrieval import RefusalHandler, Retriever +from src.vector_store import SearchResult + +logger = logging.getLogger(__name__) + + +class QueryEngine: + """Owns retrieve->generate for both the sync and the streaming query paths.""" + + def __init__( + self, + retriever: Retriever, + llm: LLMHandler, + reasoning_llm: LLMHandler, + top_k: int, + refusal: RefusalHandler | None = None, + ) -> None: + """Wire the engine to its retrieval and generation collaborators. + + Args: + retriever: The retrieval strategy (dense by default). Injected so it + is interchangeable behind the Retriever seam. + llm: The answer-generation handler. + reasoning_llm: The (cheaper) handler for the streaming planning pass. + top_k: Default number of chunks to retrieve. + refusal: Optional answerability gate; when None there is no gate. + """ + self._retriever = retriever + self._llm = llm + self._reasoning_llm = reasoning_llm + self._top_k = top_k + self._refusal = refusal + + def _handler_for(self, model: str | None) -> LLMHandler: + """Return the default handler, or a per-query one if `model` differs.""" + if model and model != self._llm.model: + return LLMHandler(model=model) + return self._llm + + def _refusal_text(self, results: list[SearchResult]) -> str | None: + """Return the no-answer text if the gate refuses, else None. + + A local binding (not ``self._refusal``) so the None-narrowing is visible + to the type checker without a cast. + """ + gate = self._refusal + if gate is not None and gate.should_refuse(results): + return gate.refuse_response()[1] + return None + + # ------------------------------------------------------------------ # + # Synchronous path # + # ------------------------------------------------------------------ # + + def ask( + self, question: str, top_k: int | None = None, model: str | None = None, + ) -> tuple[list[SearchResult], str, StageTelemetry]: + """Retrieve, generate, and assemble telemetry for one question. + + Args: + question: The user's natural-language question. + top_k: Chunks to retrieve (defaults to the engine's configured top_k). + model: Optional per-query answer-model override. + + Returns: + (results, answer, telemetry). On an empty index or a refusal, results + is empty and telemetry has zero generation fields (no LLM call). + """ + k = top_k or self._top_k + tracer = get_tracer() + + start = time.perf_counter() + with tracer.start_as_current_span("rag.retrieve") as span: + span.set_attribute("top_k", k) + span.set_attribute("question_len", len(question)) + results = self._retriever.retrieve(question, top_k=k) + span.set_attribute("results_count", len(results)) + retrieve_ms = (time.perf_counter() - start) * 1000 + + # The refusal gate is checked BEFORE the no-documents branch: an empty + # retrieval is itself an answerability signal the gate is entitled to + # act on (should_refuse([]) is True when enabled). Production leaves the + # gate off, so it always falls through to the no-documents notice. + refusal_text = self._refusal_text(results) + if refusal_text is not None: + return [], refusal_text, telemetry_asm.zero(retrieve_ms) + if not results: + return [], NO_DOCUMENTS_ANSWER, telemetry_asm.zero(retrieve_ms) + + context = build_context(results) + handler = self._handler_for(model) + + gen_start = time.perf_counter() + with tracer.start_as_current_span("rag.generate") as span: + span.set_attribute("model", handler.model) + answer, prompt_tokens, completion_tokens = handler.generate_with_usage( + build_answer_user_prompt(context, question), + system_prompt=ANSWER_SYSTEM_PROMPT, + ) + span.set_attribute("answer_len", len(answer)) + generate_ms = (time.perf_counter() - gen_start) * 1000 + + usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens) + return results, answer, telemetry_asm.assemble( + retrieve_ms, generate_ms, handler.model, usage + ) + + # ------------------------------------------------------------------ # + # Streaming path # + # ------------------------------------------------------------------ # + + def ask_stream( + self, + question: str, + top_k: int | None = None, + model: str | None = None, + history: list[dict[str, str]] | None = None, + ) -> Iterator[StreamEvent]: + """Stream retrieve -> plan -> answer as typed events. + + Yields ("status"|"reasoning"|"token", str) display events, then one + terminal ("result", StreamResult). The facade consumes the terminal event + to persist the turn and emit its own done/telemetry events. + + Args: + question: The user's question. + top_k: Chunks to retrieve (defaults to the configured top_k). + model: Optional per-query answer-model override. + history: Prior conversation turns (role/content dicts). When present, + the answer pass runs multi-turn; when empty, single-turn. + """ + k = top_k or self._top_k + history = history or [] + tracer = get_tracer() + + yield ("status", "Searching indexed documents...") + start = time.perf_counter() + with tracer.start_as_current_span("rag.retrieve") as span: + span.set_attribute("top_k", k) + span.set_attribute("question_len", len(question)) + results = self._retriever.retrieve(question, top_k=k) + span.set_attribute("results_count", len(results)) + retrieve_ms = (time.perf_counter() - start) * 1000 + + refusal_text = self._refusal_text(results) + if refusal_text is not None: + yield ("token", refusal_text) + yield ("result", StreamResult([], telemetry_asm.zero(retrieve_ms), self._llm.model)) + return + if not results: + yield ("status", "No indexed documents — nothing to retrieve.") + yield ("token", NO_DOCUMENTS_ANSWER) + yield ("result", StreamResult([], telemetry_asm.zero(retrieve_ms), self._llm.model)) + return + + yield ("status", retrieval_summary(results)) + context = build_context(results) + handler = self._handler_for(model) + + yield from self._stream_reasoning(context, question) + + yield ("status", "Composing answer...") + answer_usage: Usage | None = None + gen_start = time.perf_counter() + with tracer.start_as_current_span("rag.generate") as span: + span.set_attribute("model", handler.model) + span.set_attribute("has_conversation", bool(history)) + stream = ( + handler.stream_messages(build_answer_messages(history, context, question)) + if history + else handler.stream_response( + build_answer_user_prompt(context, question), + system_prompt=ANSWER_SYSTEM_PROMPT, + ) + ) + answer_len = 0 + for item in stream: + if isinstance(item, Usage): + answer_usage = item + continue + answer_len += len(item) + yield ("token", item) + span.set_attribute("answer_len", answer_len) + generate_ms = (time.perf_counter() - gen_start) * 1000 + + usage = answer_usage or Usage(prompt_tokens=0, completion_tokens=0) + telemetry = telemetry_asm.assemble(retrieve_ms, generate_ms, handler.model, usage) + yield ("result", StreamResult(results, telemetry, handler.model)) + + def _stream_reasoning(self, context: str, question: str) -> Iterator[StreamEvent]: + """Stream the planning pass; a failure here must not block the answer.""" + yield ("status", f"Analyzing retrieved context ({self._reasoning_llm.model})...") + try: + for item in self._reasoning_llm.stream_response( + build_reasoning_user_prompt(context, question), + system_prompt=REASONING_SYSTEM_PROMPT, + ): + if isinstance(item, Usage): + continue # reasoning usage is out of telemetry scope (ADR 0003) + yield ("reasoning", item) + except Exception as exc: # noqa: BLE001 — best-effort; degrade to answer. + logger.warning("Reasoning pass failed: %s", exc) + yield ("status", "Reasoning unavailable — skipping to answer.") diff --git a/src/query_engine/prompt.py b/src/query_engine/prompt.py new file mode 100644 index 00000000..9a8baeb3 --- /dev/null +++ b/src/query_engine/prompt.py @@ -0,0 +1,82 @@ +"""Prompt and context assembly — the single source of answer instructions. + +RAG Pipeline Position: + retrieved chunks + question -> [PROMPT] -> (system, user) -> LLM + +Design Decision: + Before step 4 the answer system prompt existed in three diverged copies (a + plain one on the sync query path, a Markdown one on the streaming path, and a + third in the eval harness). This module makes the **Markdown** prompt the one + answer prompt for every path — it is what the frontend's renderer expects, + and the eval harness must measure the shipped prompt. Context is + filename-prefixed everywhere so citations survive. See ADR 0004. +""" + +from __future__ import annotations + +from src.vector_store import SearchResult + +# The single answer system prompt for sync, streaming, and eval paths. +ANSWER_SYSTEM_PROMPT = ( + "You are a helpful assistant. Answer the user's question based solely on the " + "provided context. If the context does not contain enough information, say so.\n\n" + "Format your response using Markdown for readability:\n" + "- Use ## for main sections and ### for sub-sections (max 3 levels)\n" + "- Use **bold** for key terms and important concepts\n" + "- Use bullet points (-) for lists of related items\n" + "- Use numbered lists (1.) for sequential steps\n" + "- Use `inline code` for technical terms, parameters, or commands\n" + "- Use fenced code blocks (```language) for code snippets\n" + "- Use > blockquotes for notable quotes from the context\n" + "- Keep paragraphs short (2-3 sentences max)\n" + "- Add blank lines between sections for visual breathing room\n" + "Do NOT use # (h1) headings. Start directly with content or ## sections." +) + +# The planning-pass prompt (streaming only). It asks for an outcome-oriented +# reasoning summary, not raw chain-of-thought, so users are not shown uncertain +# intermediate beliefs they might mistake for the answer. +REASONING_SYSTEM_PROMPT = ( + "You are the planning step of a retrieval-augmented Q&A system. " + "In 3-5 concise sentences, summarise how you will construct the " + "answer using the retrieved excerpts. Cover:\n" + "1) What the user is asking, resolving any ambiguity explicitly.\n" + "2) Which excerpts are most relevant and the gist of their support.\n" + "3) Any gaps or conflicts the reader should be aware of.\n" + "4) The shape of the answer you will give next.\n" + "Stay factual and outcome-oriented — describe the plan, do not " + "verbalise stream-of-consciousness reasoning. No markdown headings, " + "no bullet lists, no preamble. Do NOT produce the final answer." +) + +# Shown (and returned) when the index has no documents to retrieve from. +NO_DOCUMENTS_ANSWER = "No documents indexed yet. Please upload documents first." + + +def build_context(results: list[SearchResult]) -> str: + """Join retrieved chunks into a filename-prefixed context block. + + Args: + results: Retrieved chunks, best first. + + Returns: + A ``"[filename] content"`` block per chunk, separated by blank lines — + the prefix is what lets the model (and downstream citations) attribute + each passage to its source document. + """ + return "\n\n".join( + f"[{r.metadata.get('filename', 'unknown')}] {r.content}" for r in results + ) + + +def build_answer_user_prompt(context: str, question: str) -> str: + """Assemble the answer-pass user message from context and question.""" + return f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:" + + +def build_reasoning_user_prompt(context: str, question: str) -> str: + """Assemble the planning-pass user message (summary only, no answer).""" + return ( + f"Context:\n{context}\n\nQuestion: {question}\n\n" + "Reasoning plan (summary only, do not answer):" + ) diff --git a/src/query_engine/streaming.py b/src/query_engine/streaming.py new file mode 100644 index 00000000..2d38b25b --- /dev/null +++ b/src/query_engine/streaming.py @@ -0,0 +1,51 @@ +"""Streaming support for the QueryEngine — the terminal payload and pure helpers. + +The engine's `ask_stream` yields display events (status/reasoning/token) plus one +terminal `("result", StreamResult)` so the facade can persist the turn and emit +its own done/telemetry events without the engine knowing about conversations. +This module holds that payload type, the event alias, and the small pure helpers +the streaming path uses. (Step 6 of issue #16 will add the shared streaming event +vocabulary here.) +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from src.api.schemas.telemetry import StageTelemetry +from src.query_engine.prompt import ANSWER_SYSTEM_PROMPT, build_answer_user_prompt +from src.vector_store import SearchResult + + +@dataclass(frozen=True) +class StreamResult: + """The terminal payload of `ask_stream`, carried on the ("result", ...) event.""" + + results: list[SearchResult] + telemetry: StageTelemetry + model: str + + +# One streamed event: a label plus either a display string or the terminal result. +StreamEvent = tuple[str, "str | StreamResult"] + + +def retrieval_summary(results: list[SearchResult]) -> str: + """Return a one-liner naming the files that contributed to the retrieval.""" + unique_files = sorted( + {r.metadata.get("filename", "unknown") for r in results if r.metadata.get("filename")} + ) + summary = ", ".join(unique_files[:3]) + if len(unique_files) > 3: + summary += f" (+{len(unique_files) - 3} more)" + return f"Retrieved {len(results)} chunk(s) across {len(unique_files)} file(s): {summary}" + + +def build_answer_messages( + history: list[dict[str, str]], context: str, question: str, +) -> list[dict[str, str]]: + """Build the multi-turn messages list: system + prior turns + current user.""" + messages: list[dict[str, str]] = [{"role": "system", "content": ANSWER_SYSTEM_PROMPT}] + messages.extend(history) + messages.append({"role": "user", "content": build_answer_user_prompt(context, question)}) + return messages diff --git a/src/query_engine/telemetry.py b/src/query_engine/telemetry.py new file mode 100644 index 00000000..a8ed5fc6 --- /dev/null +++ b/src/query_engine/telemetry.py @@ -0,0 +1,54 @@ +"""Telemetry assembly — turn a generation's Usage into a StageTelemetry payload. + +RAG Pipeline Position: + (retrieve_ms, generate_ms, model, Usage) -> [TELEMETRY] -> StageTelemetry + +Design Decision: + Before step 4 the StageTelemetry construction was duplicated across four call + sites in the backend. The QueryEngine assembles it once, here, from the + provider-reported ``Usage`` (ADR 0003) and the core pricing table (ADR 0003) + — never from a reconstructed prompt. +""" + +from __future__ import annotations + +from src.api.schemas.telemetry import StageTelemetry +from src.llm_handler import Usage +from src.telemetry.pricing import cost_usd + + +def assemble( + retrieve_ms: float, generate_ms: float, model: str, usage: Usage, +) -> StageTelemetry: + """Build a StageTelemetry from stage timings and provider-reported usage. + + Args: + retrieve_ms: Retrieval wall time in milliseconds. + generate_ms: Generation wall time in milliseconds. + model: The model whose pricing applies to ``usage``. + usage: Provider-reported (or adapter-counted) token counts. + + Returns: + StageTelemetry with cost priced from ``model`` and ``usage``. + """ + return StageTelemetry( + retrieve_ms=round(retrieve_ms, 2), + generate_ms=round(generate_ms, 2), + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + cost_usd=cost_usd(model, usage.prompt_tokens, usage.completion_tokens), + ) + + +def zero(retrieve_ms: float) -> StageTelemetry: + """Telemetry for a path that never called the LLM (no docs, or refusal). + + Retrieval time is real; every generation field is zero. + """ + return StageTelemetry( + retrieve_ms=round(retrieve_ms, 2), + generate_ms=0.0, + prompt_tokens=0, + completion_tokens=0, + cost_usd=0.0, + ) diff --git a/src/retrieval/__init__.py b/src/retrieval/__init__.py new file mode 100644 index 00000000..16c6d108 --- /dev/null +++ b/src/retrieval/__init__.py @@ -0,0 +1,40 @@ +"""Retrieval package — the Retriever seam and its adapters (issue #16, step 4). + +One interface, `Retriever` (`retrieve(query, top_k) -> list[SearchResult]`), with +adapters that either conform directly or compose an inner Retriever: + +- `DenseRetriever` — dense vector search (the default). +- `BM25HybridRetriever` — sparse (BM25) + dense fused by RRF; conforms directly. +- `RerankingRetriever` — composes an inner Retriever, over-fetches, re-scores with + a cross-encoder (`CrossEncoderReranker`). +- `MultiQueryRetriever` — composes an inner Retriever, fans out rewritten queries + (`QueryRewriter`), dedups. + +`RefusalHandler` is not a Retriever — it is an answerability gate the QueryEngine +applies after retrieval. These modules were promoted from `src/eval/` so +production can activate the eval-proven levers by configuration. See +[ADR 0004](../../docs/adr/0004-retriever-seam-and-query-engine.md). +""" + +from __future__ import annotations + +from src.retrieval.base import Retriever +from src.retrieval.dense import DenseRetriever +from src.retrieval.factory import build_retriever +from src.retrieval.hybrid import BM25HybridRetriever, reciprocal_rank_fusion +from src.retrieval.query_rewriter import MultiQueryRetriever, QueryRewriter +from src.retrieval.refusal_handler import RefusalHandler +from src.retrieval.reranker import CrossEncoderReranker, RerankingRetriever + +__all__ = [ + "Retriever", + "DenseRetriever", + "build_retriever", + "BM25HybridRetriever", + "reciprocal_rank_fusion", + "CrossEncoderReranker", + "RerankingRetriever", + "QueryRewriter", + "MultiQueryRetriever", + "RefusalHandler", +] diff --git a/src/retrieval/base.py b/src/retrieval/base.py new file mode 100644 index 00000000..ac791805 --- /dev/null +++ b/src/retrieval/base.py @@ -0,0 +1,45 @@ +"""Retriever seam — the one interface every retrieval strategy hides behind. + +RAG Pipeline Position: + Query -> [RETRIEVER] -> list[SearchResult] -> QueryEngine -> Answer + ^^^^^^^^^ + This module defines the *seam*: a single `retrieve(query, top_k)` interface + that dense, hybrid, reranked, and multi-query retrieval all present. The + QueryEngine (step 4b) depends only on this Protocol, so a retrieval strategy + validated offline in the eval harness is promoted to production by + *configuration*, not by a code change. + +Design Decision: + A `Protocol` (not an ABC) per the project standard — retrieval strategies + conform structurally without inheriting, and `@runtime_checkable` lets tests + assert conformance with `isinstance`. `SearchResult` stays the shared result + type (defined in `vector_store`) so no adapter invents its own shape. +""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from src.vector_store import SearchResult + + +@runtime_checkable +class Retriever(Protocol): + """Anything that turns a query into ranked chunks. + + Implementations either *conform* directly (a dense store wrapper, the BM25 + hybrid retriever) or *compose* an inner Retriever (reranking, multi-query), + always presenting this same interface outward. + """ + + def retrieve(self, query: str, top_k: int = 5) -> list[SearchResult]: + """Return up to `top_k` chunks most relevant to `query`. + + Args: + query: Natural-language query. + top_k: Maximum number of results to return, best first. + + Returns: + SearchResult list ordered by descending relevance (possibly empty). + """ + ... diff --git a/src/retrieval/dense.py b/src/retrieval/dense.py new file mode 100644 index 00000000..f7b085c2 --- /dev/null +++ b/src/retrieval/dense.py @@ -0,0 +1,45 @@ +"""DenseRetriever — the default adapter: pure dense vector search. + +RAG Pipeline Position: + Query -> [DenseRetriever -> ChromaVectorStore] -> list[SearchResult] + +This is the baseline retrieval strategy and the QueryEngine's default. It is a +thin adapter that presents the `Retriever` interface over `ChromaVectorStore`, +whose own `query(query_text=..., top_k=...)` already returns `SearchResult`s. + +Design Decision: + The adapter exists (rather than passing the store directly) so the seam is + uniform: dense, hybrid, reranked, and multi-query all expose `retrieve()`. + The store's `query` keyword (`query_text=`) is an implementation detail this + adapter hides behind the Protocol's positional `query`. +""" + +from __future__ import annotations + +from src.vector_store import ChromaVectorStore, SearchResult + + +class DenseRetriever: + """Dense retrieval over a Chroma collection, behind the Retriever seam.""" + + def __init__(self, vector_store: ChromaVectorStore) -> None: + """Wrap a vector store. + + Args: + vector_store: The dense index to query. The adapter never reaches + past its public `query` method. + """ + self._vector_store = vector_store + + def retrieve(self, query: str, top_k: int = 5) -> list[SearchResult]: + """Return the `top_k` nearest chunks to `query` by cosine similarity. + + Args: + query: Natural-language query text. + top_k: Number of results to return. + + Returns: + SearchResult list ordered by descending similarity (empty if the + index has no documents). + """ + return self._vector_store.query(query_text=query, top_k=top_k) diff --git a/src/retrieval/factory.py b/src/retrieval/factory.py new file mode 100644 index 00000000..0f8b3f21 --- /dev/null +++ b/src/retrieval/factory.py @@ -0,0 +1,65 @@ +"""Production retriever selection — turn a config strategy into a Retriever. + +RAG Pipeline Position: + config strategy -> [build_retriever] -> Retriever -> QueryEngine + +This is the seam through which a validated eval chain is promoted to production +by *configuration* (ADR 0004). `RAGBackend` reads `RETRIEVER_STRATEGY` from +config and calls this factory; the QueryEngine never knows which strategy it got. + +Wiring status (step 4): + - ``dense`` — wired (the default; the behaviour-preserving anchor). + - ``reranked`` — wired (composes dense + a cross-encoder; cost-clean). + - ``multi_query`` — deferred to step 4c (its rewriter-cost surfacing lands + with its production wiring, so it never ships a path that silently drops + cost). + - ``hybrid`` — deferred: BM25 needs a live corpus kept in sync with + ingestion/deletion, a genuinely new feature beyond this refactor. +""" + +from __future__ import annotations + +from src.retrieval.base import Retriever +from src.retrieval.dense import DenseRetriever +from src.retrieval.reranker import CrossEncoderReranker, RerankingRetriever +from src.vector_store import ChromaVectorStore + +_DEFERRED = { + "hybrid": "needs a live BM25 corpus synced with ingestion (a new feature)", + "multi_query": "lands in step 4c with its rewriter-cost surfacing", +} + + +def build_retriever( + strategy: str, + vector_store: ChromaVectorStore, + rerank_over_fetch_n: int = 20, +) -> Retriever: + """Construct the production Retriever for a config strategy. + + Args: + strategy: One of ``dense`` or ``reranked`` (wired), or ``hybrid`` / + ``multi_query`` (recognised but deferred — see module docstring). + vector_store: The dense index every strategy is built over. + rerank_over_fetch_n: Candidate width the reranked strategy over-fetches. + + Returns: + A Retriever conforming to the seam. + + Raises: + ValueError: If the strategy is unknown, or recognised but not yet wired + for production. + """ + dense = DenseRetriever(vector_store) + if strategy == "dense": + return dense + if strategy == "reranked": + return RerankingRetriever( + inner=dense, reranker=CrossEncoderReranker(), over_fetch_n=rerank_over_fetch_n, + ) + if strategy in _DEFERRED: + raise ValueError( + f"Retriever strategy {strategy!r} is validated in the eval harness but " + f"not yet wired for production ({_DEFERRED[strategy]}); see ADR 0004." + ) + raise ValueError(f"Unknown retriever strategy: {strategy!r}") diff --git a/src/eval/retrievers/bm25_hybrid.py b/src/retrieval/hybrid.py similarity index 100% rename from src/eval/retrievers/bm25_hybrid.py rename to src/retrieval/hybrid.py diff --git a/src/retrieval/query_rewriter.py b/src/retrieval/query_rewriter.py new file mode 100644 index 00000000..6f0de84f --- /dev/null +++ b/src/retrieval/query_rewriter.py @@ -0,0 +1,169 @@ +"""Multi-query expansion — the LLM rewriter and the Retriever adapter that composes it. + +Pipeline position: + user query → [MultiQueryRetriever → QueryRewriter] → {q, q', q''} → inner Retriever → union + +Two collaborators live here: + +- `QueryRewriter` expands one user query into alternative phrasings via a tiny LLM + (gpt-4.1-nano — cheap, so this lever doesn't dominate the cost ledger). Expansion + raises recall when the user's phrasing diverges from the corpus phrasing. +- `MultiQueryRetriever` presents the `Retriever` interface by *composing* an inner + Retriever: it fans the expansions out, unions the results, and dedups by + chunk_id keeping each chunk's best score. The "compose rather than conform" + adapter from ADR 0004. + +The rewriter reports its own token cost, but the pure `Retriever` interface has +no cost channel, so it is dropped at this seam. This was deliberate in step 4c: +the eval harness's old `rewriter_cost_usd` field had zero readers (verified), so +convergence dropped it rather than plumb a cost path nothing consumed. A cost +channel can be added if a consumer ever needs multi-query spend broken out. +""" + +from __future__ import annotations + +import json +import logging +import re +from typing import Protocol + +from src.retrieval.base import Retriever +from src.telemetry import pricing +from src.vector_store import SearchResult + +logger = logging.getLogger(__name__) + + +class _LLMHandler(Protocol): + """Structural type for any object exposing generate_with_usage.""" + + def generate_with_usage( + self, prompt: str, system_prompt: str | None = None, + ) -> tuple[str, int, int]: ... + + +class QueryRewriter: + """Expands one user query into up to N alternative phrasings via an LLM.""" + + SYSTEM_PROMPT = ( + "You rewrite user search queries into alternative phrasings that preserve " + "the original intent but vary surface form. Respond ONLY with a JSON " + "array of strings — no prose, no code fences." + ) + + def __init__( + self, + model: str | None, + max_expansions: int, + llm: _LLMHandler | None, + ) -> None: + """Configure the rewriter. + + Args: + model: LLM model name. None disables rewriting (pass-through). + max_expansions: Cap on the number of alternative phrasings to return. + llm: Object exposing generate_with_usage(prompt, system_prompt). Required + if model is not None. + """ + self._model = model + self._max_expansions = max_expansions + self._llm = llm + + def expand(self, query: str) -> tuple[list[str], float, int, int]: + """Expand `query` into up to N+1 unique phrasings. + + Returns: + (queries, cost_usd, prompt_tokens, completion_tokens). The original + query is always the first element. When `model is None`, returns + ([query], 0.0, 0, 0) and skips the LLM call. + """ + if self._model is None: + return [query], 0.0, 0, 0 + if self._llm is None: + raise ValueError("QueryRewriter has model set but no llm handler provided.") + + user_prompt = ( + f'Original query: "{query}"\n\n' + f"Return a JSON array of up to {self._max_expansions} alternative " + f"phrasings of this query. Do NOT include the original." + ) + raw, p_t, c_t = self._llm.generate_with_usage( + user_prompt, system_prompt=self.SYSTEM_PROMPT, + ) + cost = pricing.cost_usd(self._model, p_t, c_t) + + expansions = self._parse_expansions(raw) + # Always lead with original; dedupe; cap at original + max_expansions. + ordered: list[str] = [query] + for alt in expansions: + if alt and alt not in ordered: + ordered.append(alt) + if len(ordered) >= self._max_expansions + 1: + break + return ordered, cost, p_t, c_t + + @staticmethod + def _parse_expansions(raw: str) -> list[str]: + """Strip code fences and parse the JSON array; return [] on failure.""" + stripped = re.sub(r"^```(?:json)?\s*", "", raw.strip()) + stripped = re.sub(r"\s*```$", "", stripped).strip() + try: + parsed = json.loads(stripped) + except json.JSONDecodeError: + logger.warning( + "QueryRewriter got non-JSON response — falling back to [query] only." + ) + return [] + if not isinstance(parsed, list): + return [] + return [str(item) for item in parsed if isinstance(item, str)] + + +class _Rewriter(Protocol): + """Structural type for a query expander (the one collaborator we inject).""" + + def expand(self, query: str) -> tuple[list[str], float, int, int]: ... + + +class MultiQueryRetriever: + """Retriever adapter: fan an inner Retriever out over rewritten queries. + + Presents `retrieve(query, top_k)` while delegating expansion to a rewriter and + candidate generation to an inner Retriever — so multi-query retrieval is + interchangeable with any other strategy behind the same seam. + """ + + def __init__(self, inner: Retriever, rewriter: _Rewriter) -> None: + """Compose an inner Retriever with a query expander. + + Args: + inner: The Retriever run once per expanded query. + rewriter: Produces the alternative phrasings (original query first). + """ + self._inner = inner + self._rewriter = rewriter + + def retrieve(self, query: str, top_k: int = 5) -> list[SearchResult]: + """Retrieve for every expansion, union, dedup by chunk_id, rank best-first. + + Args: + query: The original user query. + top_k: Number of results to return after the union is ranked. Each + expansion is itself retrieved at `top_k` before the union. + + Returns: + Up to `top_k` SearchResults ordered by descending score. When a chunk + surfaces under several expansions, its highest score wins (dense + similarities share the embedding space, so they compare directly). + """ + # expand() also reports rewrite cost/tokens; the pure Retriever interface + # has no cost channel, so only the query list is used here (see step 4c). + expansions, *_cost_and_tokens = self._rewriter.expand(query) + best: dict[str, SearchResult] = {} + for expansion in expansions: + for result in self._inner.retrieve(expansion, top_k=top_k): + current = best.get(result.chunk_id) + if current is None or result.score > current.score: + best[result.chunk_id] = result + ranked = sorted(best.values(), key=lambda r: r.score, reverse=True) + return ranked[:top_k] diff --git a/src/eval/transforms/refusal_handler.py b/src/retrieval/refusal_handler.py similarity index 100% rename from src/eval/transforms/refusal_handler.py rename to src/retrieval/refusal_handler.py diff --git a/src/retrieval/reranker.py b/src/retrieval/reranker.py new file mode 100644 index 00000000..9fc3f3a6 --- /dev/null +++ b/src/retrieval/reranker.py @@ -0,0 +1,120 @@ +"""Cross-encoder reranking — the re-scorer and the Retriever adapter that composes it. + +Pipeline position: + Retriever top-N → [RerankingRetriever → CrossEncoderReranker] → top-K → Generator + +Two collaborators live here: + +- `CrossEncoderReranker` re-scores a candidate list. Cross-encoders (single-tower + models that consume the query and a candidate together) typically outperform + bi-encoder retrieval in precision at the cost of latency. We use + ms-marco-MiniLM-L-6-v2 — small enough to run on CPU in milliseconds per pair, + trained on MS MARCO so the ranking signal transfers to general-domain QA. +- `RerankingRetriever` presents the `Retriever` interface by *composing* an inner + Retriever: it over-fetches a wide candidate set, then narrows via the reranker. + This is the "compose rather than conform" adapter from ADR 0004. +""" + +from __future__ import annotations + +from typing import Protocol + +from src.retrieval.base import Retriever +from src.vector_store import SearchResult + + +class CrossEncoderReranker: + """Wraps sentence-transformers CrossEncoder to re-score retrieval candidates.""" + + MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2" + + def __init__(self) -> None: + from sentence_transformers import CrossEncoder + + self._model = CrossEncoder(self.MODEL_NAME) + + def rerank( + self, + query: str, + candidates: list[SearchResult], + final_top_k: int, + ) -> list[SearchResult]: + """Re-score candidates against the query and return top-K reranked. + + Args: + query: Original user query. + candidates: Pre-retrieved chunks (typically top-N from a base retriever). + final_top_k: How many to keep after reranking. + + Returns: + Top-K SearchResult ordered by descending cross-encoder score. The + original `score` field is *replaced* with the cross-encoder score so + downstream consumers reading `result.score` get the more precise signal. + """ + if not candidates: + return [] + pairs = [(query, c.content) for c in candidates] + scores = self._model.predict(pairs) + scored = sorted( + zip(candidates, scores), key=lambda t: t[1], reverse=True, + )[:final_top_k] + return [ + SearchResult( + doc_id=c.doc_id, + chunk_id=c.chunk_id, + content=c.content, + score=float(s), + metadata=c.metadata, + ) + for c, s in scored + ] + + +class _Reranker(Protocol): + """Structural type for a candidate re-scorer (the one collaborator we inject).""" + + def rerank( + self, query: str, candidates: list[SearchResult], final_top_k: int, + ) -> list[SearchResult]: ... + + +class RerankingRetriever: + """Retriever adapter: over-fetch from an inner Retriever, then cross-encode. + + Presents `retrieve(query, top_k)` while delegating candidate generation to an + inner Retriever and precision re-scoring to a reranker — so reranking is + interchangeable with any other retrieval strategy behind the same seam. + """ + + def __init__( + self, + inner: Retriever, + reranker: _Reranker, + over_fetch_n: int, + ) -> None: + """Compose an inner Retriever with a candidate re-scorer. + + Args: + inner: The Retriever that produces the initial candidate set. + reranker: The cross-encoder re-scorer applied to those candidates. + over_fetch_n: How many candidates to pull from `inner` before + reranking. Wider than the final `top_k` so the precise reranker + has real choice; the eval-tuned default is 20. + """ + self._inner = inner + self._reranker = reranker + self._over_fetch_n = over_fetch_n + + def retrieve(self, query: str, top_k: int = 5) -> list[SearchResult]: + """Over-fetch `over_fetch_n` candidates, rerank, return the top `top_k`. + + Args: + query: Natural-language query. + top_k: Final number of results after reranking. + + Returns: + The reranked top-`top_k` SearchResults (empty if the inner retriever + found nothing). + """ + candidates = self._inner.retrieve(query, top_k=self._over_fetch_n) + return self._reranker.rerank(query, candidates, final_top_k=top_k) diff --git a/tests/test_eval_config.py b/tests/test_eval_config.py index 78073810..360e2d2d 100644 --- a/tests/test_eval_config.py +++ b/tests/test_eval_config.py @@ -46,11 +46,15 @@ def test_full_construction(self): assert cfg.eval.datasets == ["squad_v2_dev_200"] def test_defaults_applied_when_omitted(self): + from src.config import CHUNK_SIZE + d = _baseline_dict() del d["pipeline"]["chunker"]["chunk_size"] del d["eval"]["bootstrap_n"] cfg = EvalConfig.model_validate(d) - assert cfg.pipeline.chunker.chunk_size == 512 # default + # Step 4c: the chunk-size default now derives from production config + # (single source of truth), not a hard-coded eval literal. + assert cfg.pipeline.chunker.chunk_size == CHUNK_SIZE assert cfg.eval.bootstrap_n == 1000 # default def test_missing_required_field_raises(self): diff --git a/tests/test_eval_integration.py b/tests/test_eval_integration.py index c23c591b..19074c1a 100644 --- a/tests/test_eval_integration.py +++ b/tests/test_eval_integration.py @@ -18,6 +18,8 @@ class DummyLLM: """Returns canned data — JSON for judges, plain for generation.""" + model = "gpt-4.1-nano" # engine reads .model for spans + cost pricing + def generate(self, prompt: str, system_prompt: str | None = None) -> str: if "JSON" in (system_prompt or "") or '"score"' in prompt or '"is_refusal"' in prompt: return json.dumps({ @@ -26,6 +28,12 @@ def generate(self, prompt: str, system_prompt: str | None = None) -> str: }) return "" + def generate_with_usage( + self, prompt: str, system_prompt: str | None = None + ) -> tuple[str, int, int]: + text = self.generate(prompt, system_prompt) + return text, max(1, len(prompt.split())), len(text.split()) + def _make_config(name: str, top_k: int) -> EvalConfig: return EvalConfig.model_validate({ diff --git a/tests/test_eval_pipeline_factory.py b/tests/test_eval_pipeline_factory.py index 6e8783d7..3fae47d1 100644 --- a/tests/test_eval_pipeline_factory.py +++ b/tests/test_eval_pipeline_factory.py @@ -15,12 +15,19 @@ class DummyLLM: """Returns a fixed answer; tracks calls.""" def __init__(self, answer: str = ""): self.answer = answer + self.model = "gpt-4.1-nano" # engine reads .model for spans + cost pricing self.calls: list[tuple[str, str | None]] = [] def generate(self, prompt: str, system_prompt: str | None = None) -> str: self.calls.append((prompt, system_prompt)) return self.answer + def generate_with_usage( + self, prompt: str, system_prompt: str | None = None + ) -> tuple[str, int, int]: + self.calls.append((prompt, system_prompt)) + return self.answer, max(1, len(prompt.split())), len(self.answer.split()) + def _baseline_config() -> EvalConfig: return EvalConfig.model_validate({ diff --git a/tests/test_eval_pipeline_factory_phase2.py b/tests/test_eval_pipeline_factory_phase2.py index e4f00d63..30ac1350 100644 --- a/tests/test_eval_pipeline_factory_phase2.py +++ b/tests/test_eval_pipeline_factory_phase2.py @@ -24,6 +24,7 @@ @pytest.fixture def stub_llm(): class _S: + model = "gpt-4.1-nano" # engine reads .model for spans + cost pricing def generate(self, prompt, system_prompt=None): return "stub answer" def generate_with_usage(self, prompt, system_prompt=None): @@ -67,11 +68,13 @@ def test_phase2_query_with_refusal_short_circuits(stub_llm): llm_override=stub_llm, judge_llm_override=stub_llm, ) try: - # Empty index → retrieval returns [] → handler refuses. + # Empty index → retrieval returns [] → handler refuses. Post-convergence + # (step 4c) the engine owns telemetry, so timings_ms is {retrieve, generate} + # rather than the old per-lever stages (refusal is applied inside the engine). chunks, answer, telemetry = pipeline.query("what is x?") assert chunks == [] assert answer == cfg.pipeline.refusal_handler.no_answer_text - assert "refusal_check" in telemetry["timings_ms"] + assert "retrieve" in telemetry["timings_ms"] finally: pipeline.teardown() diff --git a/tests/test_eval_production_parity.py b/tests/test_eval_production_parity.py new file mode 100644 index 00000000..1553d1ab --- /dev/null +++ b/tests/test_eval_production_parity.py @@ -0,0 +1,83 @@ +"""Eval<->production parity — the regression guard for prompt/context drift. + +The whole point of step 4c (issue #16) is that the eval harness measures the +*shipped* pipeline. Before convergence, the eval pipeline carried its own copy +of the answer prompt (worded differently) and joined context without filename +prefixes, so eval scored a pipeline that was not the one served. This test pins +the eval path to the single shipped prompt + context builders in +`src.query_engine.prompt` — the same ones the production RAGBackend uses. If +either drifts, this fails. +""" + +from __future__ import annotations + +from src.eval.config import EvalConfig +from src.eval.pipeline_factory import build_pipeline +from src.eval.schemas import EvalQuestion +from src.query_engine.prompt import ( + ANSWER_SYSTEM_PROMPT, + build_answer_user_prompt, + build_context, +) + + +class _RecordingLLM: + """Captures exactly the (system, user) instructions the answer pass receives.""" + + model = "gpt-4.1-nano" + + def __init__(self) -> None: + self.system: str | None = None + self.user: str | None = None + + def generate(self, prompt: str, system_prompt: str | None = None) -> str: + return "{}" # judge path — unused for answer capture + + def generate_with_usage(self, prompt, system_prompt=None): + self.system = system_prompt + self.user = prompt + return "captured", 5, 2 + + +def _baseline_config() -> EvalConfig: + return EvalConfig.model_validate({ + "name": "parity", "description": "", + "pipeline": { + "chunker": {"strategy": "recursive", "chunk_size": 256, "chunk_overlap": 32}, + "retriever": {"top_k": 3}, + "generator": {"model": "gpt-4.1-nano", "reasoning_model": None}, + }, + "eval": { + "datasets": ["squad_v2_dev_200"], "judge_model": "gpt-4.1-nano", + "bootstrap_n": 100, "permutation_n": 100, "seed": 7, + }, + }) + + +def test_eval_pipeline_issues_the_shipped_prompt_and_context(): + """Eval sends the production ANSWER_SYSTEM_PROMPT and filename-prefixed context.""" + recorder = _RecordingLLM() + pipeline = build_pipeline( + _baseline_config(), "squad_v2_dev_200", + llm_override=recorder, judge_llm_override=_RecordingLLM(), + ) + try: + pipeline.ingest([ + EvalQuestion( + id="q1", question="What is the capital of France?", + gold_answer="Paris", gold_chunk_ids=["q1"], + metadata={"context": "Paris is the capital of France.", "title": "t"}, + ) + ]) + results, _answer, _telemetry = pipeline.query("What is the capital of France?") + + # Eval uses the ONE shipped answer prompt — not a reworded eval copy. + assert recorder.system == ANSWER_SYSTEM_PROMPT + # ...and the ONE shipped context + user builders, byte-for-byte. A bare + # join or a divergent template would make this inequality fail. + expected_user = build_answer_user_prompt( + build_context(results), "What is the capital of France?" + ) + assert recorder.user == expected_user + finally: + pipeline.teardown() diff --git a/tests/test_eval_runner.py b/tests/test_eval_runner.py index 52733173..cdae266a 100644 --- a/tests/test_eval_runner.py +++ b/tests/test_eval_runner.py @@ -17,6 +17,7 @@ class DummyLLM: """Returns canned answers / canned JSON for any prompt.""" def __init__(self, answer: str = "", judge_payload: dict | None = None): self.answer = answer + self.model = "gpt-4.1-nano" # engine reads .model for spans + cost pricing self.judge_payload = judge_payload or { "score": 1.0, "claims": [], "chunks": [], "factual_match": 1.0, "is_refusal": False, "reasoning": "ok", @@ -29,6 +30,12 @@ def generate(self, prompt: str, system_prompt: str | None = None) -> str: return json.dumps(self.judge_payload) return self.answer + def generate_with_usage( + self, prompt: str, system_prompt: str | None = None + ) -> tuple[str, int, int]: + text = self.generate(prompt, system_prompt) + return text, max(1, len(prompt.split())), len(text.split()) + def _baseline_config() -> EvalConfig: return EvalConfig.model_validate({ diff --git a/tests/test_query_engine.py b/tests/test_query_engine.py new file mode 100644 index 00000000..2ee4fda5 --- /dev/null +++ b/tests/test_query_engine.py @@ -0,0 +1,241 @@ +"""Tests for QueryEngine — the shared retrieve->generate module (issue #16, step 4b). + +RAG Pipeline Position: + Query -> [QUERYENGINE: Retriever -> prompt -> LLM -> telemetry] -> Answer + +What concept it teaches: + One deep module owns retrieve->generate for BOTH the synchronous and the + streaming path, so the answer prompt, context format, and telemetry + assembly exist in exactly one place. These tests drive the engine with a + fake Retriever and a fake LLM (its two injected seams) and assert the + contract the RAGBackend facade and the eval harness both depend on — + including that sync and streaming issue *identical* answer instructions. +""" + +from __future__ import annotations + +from src.llm_handler import Usage +from src.query_engine import QueryEngine +from src.query_engine.prompt import ANSWER_SYSTEM_PROMPT, NO_DOCUMENTS_ANSWER +from src.vector_store import SearchResult + + +# --------------------------------------------------------------------------- # +# Fakes at the two engine seams # +# --------------------------------------------------------------------------- # + +class _FakeRetriever: + def __init__(self, results: list[SearchResult]): + self._results = results + self.calls: list[tuple[str, int]] = [] + + def retrieve(self, query: str, top_k: int = 5) -> list[SearchResult]: + self.calls.append((query, top_k)) + return list(self._results)[:top_k] + + +class _FakeLLM: + """Records the (system, user) instructions it is handed; scripts its output.""" + + def __init__(self, model: str = "fake-answer", answer: str = "hello world", p: int = 7, c: int = 3): + self.model = model + self._answer = answer + self._p, self._c = p, c + self.seen_system: list[str | None] = [] + self.seen_user: list[str] = [] + self.seen_messages: list[list[dict]] = [] + + def generate_with_usage(self, prompt, system_prompt=None): + self.seen_system.append(system_prompt) + self.seen_user.append(prompt) + return self._answer, self._p, self._c + + def stream_response(self, prompt, system_prompt=None): + self.seen_system.append(system_prompt) + self.seen_user.append(prompt) + for tok in self._answer.split(): + yield tok + yield Usage(prompt_tokens=self._p, completion_tokens=self._c) + + def stream_messages(self, messages): + self.seen_messages.append(messages) + for tok in self._answer.split(): + yield tok + yield Usage(prompt_tokens=self._p, completion_tokens=self._c) + + +def _sr(chunk_id: str, content: str, score: float, filename: str = "doc.txt") -> SearchResult: + return SearchResult( + chunk_id=chunk_id, content=content, score=score, + metadata={"filename": filename, "chunk_index": 0}, doc_id="d1", + ) + + +def _engine(results, answer_llm=None, reasoning_llm=None, refusal=None, top_k=5): + return QueryEngine( + retriever=_FakeRetriever(results), + llm=answer_llm or _FakeLLM(), + reasoning_llm=reasoning_llm or _FakeLLM(model="fake-reason", answer="plan step"), + top_k=top_k, + refusal=refusal, + ) + + +# --------------------------------------------------------------------------- # +# Slice 1 — ask() happy path: retrieve -> prompt -> generate -> telemetry # +# --------------------------------------------------------------------------- # + +def test_ask_returns_results_answer_and_telemetry(): + llm = _FakeLLM(answer="Paris is the capital.", p=11, c=4) + engine = _engine([_sr("c1", "Paris is the capital of France.", 0.9)], answer_llm=llm) + + results, answer, telemetry = engine.ask("What is the capital of France?") + + assert [r.chunk_id for r in results] == ["c1"] + assert answer == "Paris is the capital." + assert telemetry.prompt_tokens == 11 + assert telemetry.completion_tokens == 4 + assert telemetry.retrieve_ms >= 0.0 + assert telemetry.generate_ms >= 0.0 + assert telemetry.cost_usd >= 0.0 + + +def test_ask_uses_the_markdown_answer_prompt_and_filename_context(): + llm = _FakeLLM() + engine = _engine([_sr("c1", "Body text.", 0.8, filename="paper.pdf")], answer_llm=llm) + + engine.ask("Q?") + + # The single markdown answer prompt is the system instruction. + assert llm.seen_system == [ANSWER_SYSTEM_PROMPT] + # Context is filename-prefixed, not a bare join. + assert "[paper.pdf] Body text." in llm.seen_user[0] + assert llm.seen_user[0].endswith("Question: Q?\n\nAnswer:") + + +def test_ask_top_k_defaults_and_overrides(): + retriever_results = [_sr(f"c{i}", f"t{i}", 0.5) for i in range(10)] + engine = _engine(retriever_results, top_k=5) + + engine.ask("q") + engine.ask("q", top_k=3) + + # First call used the engine default (5); second used the override (3). + assert engine._retriever.calls == [("q", 5), ("q", 3)] + + +# --------------------------------------------------------------------------- # +# Slice 2 — ask() no-documents and refusal gate skip generation # +# --------------------------------------------------------------------------- # + +def test_ask_with_no_documents_returns_zero_generation_telemetry(): + llm = _FakeLLM() + engine = _engine([], answer_llm=llm) + + results, answer, telemetry = engine.ask("anything") + + assert results == [] + assert answer == NO_DOCUMENTS_ANSWER + assert telemetry.generate_ms == 0.0 + assert telemetry.prompt_tokens == 0 + assert telemetry.cost_usd == 0.0 + # No LLM call was made. + assert llm.seen_user == [] + + +def test_ask_with_refusal_gate_short_circuits_before_generation(): + from src.retrieval import RefusalHandler + + llm = _FakeLLM() + gate = RefusalHandler(enabled=True, similarity_threshold=0.5, no_answer_text="I don't know.") + # Top score 0.3 < threshold 0.5 -> refuse. + engine = _engine([_sr("c1", "weakly related", 0.3)], answer_llm=llm, refusal=gate) + + results, answer, telemetry = engine.ask("q") + + assert results == [] + assert answer == "I don't know." + assert telemetry.generate_ms == 0.0 + assert telemetry.prompt_tokens == 0 + assert llm.seen_user == [] # generation skipped + + +def test_ask_refusal_gate_fires_on_empty_index_before_no_documents_notice(): + """An enabled gate treats an empty retrieval as unanswerable (refuse, not 'no docs').""" + from src.retrieval import RefusalHandler + + gate = RefusalHandler(enabled=True, similarity_threshold=0.5, no_answer_text="Cannot answer.") + engine = _engine([], refusal=gate) + + results, answer, telemetry = engine.ask("q") + + assert results == [] + assert answer == "Cannot answer." # not NO_DOCUMENTS_ANSWER + assert telemetry.prompt_tokens == 0 + + +# --------------------------------------------------------------------------- # +# Slice 3 — ask_stream events + sync/stream instruction parity # +# --------------------------------------------------------------------------- # + +def _event_types(events): + return [t for t, _ in events] + + +def test_ask_stream_emits_status_reasoning_token_then_result(): + engine = _engine([_sr("c1", "body", 0.9)]) + events = list(engine.ask_stream("q")) + + types = _event_types(events) + assert "status" in types + assert "reasoning" in types + assert "token" in types + # The terminal event is the internal ("result", StreamResult). + last_type, last_data = events[-1] + assert last_type == "result" + assert [r.chunk_id for r in last_data.results] == ["c1"] + assert last_data.telemetry.completion_tokens == 3 + # reasoning precedes the first answer token. + assert types.index("reasoning") < types.index("token") + + +def test_ask_stream_no_documents_yields_notice_and_empty_result(): + engine = _engine([]) + events = list(engine.ask_stream("q")) + + assert ("token", NO_DOCUMENTS_ANSWER) in events + last_type, last_data = events[-1] + assert last_type == "result" + assert last_data.results == [] + assert last_data.telemetry.generate_ms == 0.0 + + +def test_sync_and_streaming_issue_identical_answer_instructions(): + """The spec's headline guarantee: one prompt, regardless of path.""" + results = [_sr("c1", "shared body", 0.9)] + sync_llm = _FakeLLM() + stream_llm = _FakeLLM() + + _engine(results, answer_llm=sync_llm).ask("same question") + list(_engine(results, answer_llm=stream_llm).ask_stream("same question")) + + # Same system prompt AND same user prompt across both paths. + assert sync_llm.seen_system[0] == stream_llm.seen_system[0] == ANSWER_SYSTEM_PROMPT + assert sync_llm.seen_user[0] == stream_llm.seen_user[0] + + +def test_ask_stream_with_history_uses_multi_turn_messages(): + stream_llm = _FakeLLM() + engine = _engine([_sr("c1", "body", 0.9)], answer_llm=stream_llm) + history = [ + {"role": "user", "content": "prior q"}, + {"role": "assistant", "content": "prior a"}, + ] + + list(engine.ask_stream("now q", history=history)) + + messages = stream_llm.seen_messages[0] + assert messages[0] == {"role": "system", "content": ANSWER_SYSTEM_PROMPT} + assert messages[1:3] == history + assert messages[-1]["role"] == "user" + assert messages[-1]["content"].endswith("Question: now q\n\nAnswer:") diff --git a/tests/test_retrieval_adapters.py b/tests/test_retrieval_adapters.py new file mode 100644 index 00000000..4dc80dad --- /dev/null +++ b/tests/test_retrieval_adapters.py @@ -0,0 +1,223 @@ +"""Contract tests for the Retriever seam (issue #16, step 4a). + +RAG Pipeline Position: + Query -> [RETRIEVER] -> list[SearchResult] -> Generator + +What concept it teaches: + A `Retriever` Protocol lets dense, hybrid, reranked, and multi-query + retrieval be interchangeable behind one interface — `retrieve(query, top_k) + -> list[SearchResult]`. These tests assert every adapter honours that + contract, so the QueryEngine (step 4b) can accept any of them by injection. + +Why fakes for the composing adapters: + RerankingRetriever and MultiQueryRetriever compose an *inner* Retriever plus + an injected re-scorer/rewriter. Their behaviour under test is the + composition wiring (over-fetch, delegate, dedup) — not the ML model inside + the reranker or the LLM inside the rewriter. Faking those injected + collaborators keeps the contract test deterministic and fast; the real + CrossEncoderReranker / QueryRewriter have their own dedicated tests. +""" + +from __future__ import annotations + +import chromadb +import pytest + +from src.retrieval import Retriever +from src.vector_store import ChromaVectorStore, SearchResult + + +def _sr(chunk_id: str, content: str, score: float) -> SearchResult: + return SearchResult( + chunk_id=chunk_id, content=content, score=score, metadata={}, doc_id="" + ) + + +class _FakeRetriever: + """A Retriever that returns a scripted list and records the top_k asked for.""" + + def __init__(self, results_by_query: dict[str, list[SearchResult]]): + self._by_query = results_by_query + self.calls: list[tuple[str, int]] = [] + + def retrieve(self, query: str, top_k: int = 5) -> list[SearchResult]: + self.calls.append((query, top_k)) + return list(self._by_query.get(query, []))[:top_k] + + +# --------------------------------------------------------------------------- # +# Slice 1 — Retriever Protocol + DenseRetriever # +# --------------------------------------------------------------------------- # + +def _chroma_store() -> ChromaVectorStore: + client = chromadb.EphemeralClient() + coll = client.get_or_create_collection( + name="test_dense", metadata={"hnsw:space": "cosine"} + ) + coll.upsert( + ids=["d1", "d2", "d3"], + documents=[ + "Paris is the capital of France.", + "Cats are small carnivorous mammals.", + "Airplanes have fixed wings and jet engines.", + ], + metadatas=[{"filename": "geo.txt"}, {"filename": "animals.txt"}, {"filename": "air.txt"}], + ) + return ChromaVectorStore(collection=coll) + + +def test_dense_retriever_conforms_to_protocol(): + """DenseRetriever satisfies the runtime-checkable Retriever Protocol.""" + from src.retrieval import DenseRetriever + + retriever = DenseRetriever(_chroma_store()) + assert isinstance(retriever, Retriever) + + +def test_dense_retriever_returns_search_results_from_store(): + """retrieve() delegates to the vector store and returns ranked SearchResults.""" + from src.retrieval import DenseRetriever + + retriever = DenseRetriever(_chroma_store()) + out = retriever.retrieve("What is the capital of France?", top_k=2) + + assert len(out) == 2 + assert all(isinstance(r, SearchResult) for r in out) + # The geography chunk is the obvious top hit. + assert out[0].chunk_id == "d1" + assert out[0].metadata["filename"] == "geo.txt" + + +# --------------------------------------------------------------------------- # +# Slice 2 — BM25HybridRetriever conforms directly # +# --------------------------------------------------------------------------- # + +def test_hybrid_retriever_conforms_to_protocol(): + """BM25HybridRetriever already exposes retrieve() — it conforms directly.""" + from src.retrieval import BM25HybridRetriever + + store = _chroma_store() + retriever = BM25HybridRetriever( + vector_store=store, + documents={"d1": "Paris is the capital of France."}, + ) + assert isinstance(retriever, Retriever) + + +# --------------------------------------------------------------------------- # +# Slice 3 — RerankingRetriever composes inner + reranker (over-fetch) # +# --------------------------------------------------------------------------- # + +class _FakeReranker: + """Records the candidates + final_top_k it received; reverses then truncates.""" + + def __init__(self) -> None: + self.seen_candidates: list[SearchResult] = [] + self.seen_final_top_k: int | None = None + + def rerank(self, query, candidates, final_top_k): + self.seen_candidates = candidates + self.seen_final_top_k = final_top_k + return list(reversed(candidates))[:final_top_k] + + +def test_reranking_retriever_conforms_to_protocol(): + from src.retrieval import RerankingRetriever + + inner = _FakeRetriever({}) + adapter = RerankingRetriever(inner=inner, reranker=_FakeReranker(), over_fetch_n=20) + assert isinstance(adapter, Retriever) + + +def test_reranking_retriever_over_fetches_then_reranks_to_top_k(): + """It fetches `over_fetch_n` from the inner retriever, then reranks to `top_k`.""" + from src.retrieval import RerankingRetriever + + candidates = [_sr(f"c{i}", f"text {i}", 0.5) for i in range(8)] + inner = _FakeRetriever({"q": candidates}) + reranker = _FakeReranker() + adapter = RerankingRetriever(inner=inner, reranker=reranker, over_fetch_n=8) + + out = adapter.retrieve("q", top_k=3) + + # Inner was asked for the wide candidate set, not top_k. + assert inner.calls == [("q", 8)] + # Reranker received those candidates and the final top_k. + assert len(reranker.seen_candidates) == 8 + assert reranker.seen_final_top_k == 3 + # Output is the reranker's reordered, truncated result. + assert [r.chunk_id for r in out] == ["c7", "c6", "c5"] + + +# --------------------------------------------------------------------------- # +# Slice 4 — MultiQueryRetriever fans out expansions, dedups # +# --------------------------------------------------------------------------- # + +class _FakeRewriter: + """Returns a scripted expansion list (the QueryRewriter.expand contract).""" + + def __init__(self, expansions: list[str]) -> None: + self._expansions = expansions + + def expand(self, query: str) -> tuple[list[str], float, int, int]: + return self._expansions, 0.0, 0, 0 + + +def _scored(chunk_id: str, score: float) -> SearchResult: + return SearchResult( + chunk_id=chunk_id, content=chunk_id, score=score, metadata={}, doc_id="" + ) + + +def test_multi_query_retriever_conforms_to_protocol(): + from src.retrieval import MultiQueryRetriever + + adapter = MultiQueryRetriever(inner=_FakeRetriever({}), rewriter=_FakeRewriter(["q"])) + assert isinstance(adapter, Retriever) + + +def test_multi_query_fans_out_dedups_and_ranks_best_first(): + """Expansions are retrieved, deduped by chunk_id (keeping the best score), ranked.""" + from src.retrieval import MultiQueryRetriever + + inner = _FakeRetriever({ + "q": [_scored("c1", 0.9), _scored("c2", 0.5)], + "q2": [_scored("c3", 0.8), _scored("c2", 0.7)], + }) + adapter = MultiQueryRetriever(inner=inner, rewriter=_FakeRewriter(["q", "q2"])) + + out = adapter.retrieve("q", top_k=2) + + # Both expansions were retrieved. + assert {c[0] for c in inner.calls} == {"q", "q2"} + # c2 was deduped to its higher score (0.7), the union ranked best-first, + # then truncated to top_k: c1(0.9), c3(0.8) win over c2(0.7). + assert [r.chunk_id for r in out] == ["c1", "c3"] + + +# --------------------------------------------------------------------------- # +# Factory — config strategy -> Retriever type # +# --------------------------------------------------------------------------- # + +def test_build_retriever_dense_is_the_default_strategy(): + from src.retrieval import DenseRetriever, build_retriever + + retriever = build_retriever("dense", _chroma_store()) + assert isinstance(retriever, DenseRetriever) + + +def test_build_retriever_reranked_composes_dense_and_a_reranker(monkeypatch): + """reranked wires a RerankingRetriever without loading the real model here.""" + from src.retrieval import RerankingRetriever, build_retriever + + monkeypatch.setattr("src.retrieval.factory.CrossEncoderReranker", _FakeReranker) + retriever = build_retriever("reranked", _chroma_store(), rerank_over_fetch_n=15) + assert isinstance(retriever, RerankingRetriever) + + +@pytest.mark.parametrize("strategy", ["hybrid", "multi_query", "totally-bogus"]) +def test_build_retriever_rejects_unwired_or_unknown_strategies(strategy): + from src.retrieval import build_retriever + + with pytest.raises(ValueError): + build_retriever(strategy, _chroma_store()) diff --git a/tests/test_eval_retriever_bm25_hybrid.py b/tests/test_retrieval_hybrid.py similarity index 93% rename from tests/test_eval_retriever_bm25_hybrid.py rename to tests/test_retrieval_hybrid.py index 34e798cc..cd3b15ed 100644 --- a/tests/test_eval_retriever_bm25_hybrid.py +++ b/tests/test_retrieval_hybrid.py @@ -11,7 +11,7 @@ def _sr(chunk_id: str, content: str, score: float) -> SearchResult: def test_rrf_fusion_asymmetric_inputs(): """RRF on A=[a,b,c,d], B=[d,a] with rrf_k=60 yields fused order a, d, b, c.""" - from src.eval.retrievers.bm25_hybrid import reciprocal_rank_fusion + from src.retrieval.hybrid import reciprocal_rank_fusion A = ["a", "b", "c", "d"] B = ["d", "a"] fused = reciprocal_rank_fusion([A, B], rrf_k=60) @@ -21,7 +21,7 @@ def test_rrf_fusion_asymmetric_inputs(): def test_hybrid_retrieve_returns_top_k(): """End-to-end: hybrid retriever combines BM25 and Chroma results into top-K.""" import chromadb - from src.eval.retrievers.bm25_hybrid import BM25HybridRetriever + from src.retrieval.hybrid import BM25HybridRetriever from src.vector_store import ChromaVectorStore client = chromadb.EphemeralClient() diff --git a/tests/test_eval_transform_rewriter.py b/tests/test_retrieval_query_rewriter.py similarity index 94% rename from tests/test_eval_transform_rewriter.py rename to tests/test_retrieval_query_rewriter.py index 48dc2707..95513a87 100644 --- a/tests/test_eval_transform_rewriter.py +++ b/tests/test_retrieval_query_rewriter.py @@ -20,7 +20,7 @@ def generate_with_usage(self, prompt: str, system_prompt: str | None = None def test_no_model_passthrough(): """When model is None, expand returns [query] unchanged with zero cost.""" - from src.eval.transforms import QueryRewriter + from src.retrieval import QueryRewriter rw = QueryRewriter(model=None, max_expansions=3, llm=None) queries, cost, p_t, c_t = rw.expand("What is RAG?") assert queries == ["What is RAG?"] @@ -31,7 +31,7 @@ def test_no_model_passthrough(): def test_expansion_returns_dedup_list_and_cost(): """With a real model name and stub LLM, expand returns deduped expansions + cost.""" - from src.eval.transforms import QueryRewriter + from src.retrieval import QueryRewriter stub = _StubLLM( response='["What does RAG stand for?", "Define retrieval augmented generation", ' '"What is RAG?"]', @@ -53,7 +53,7 @@ def test_expansion_returns_dedup_list_and_cost(): def test_malformed_llm_response_falls_back_to_passthrough(): """If the LLM returns non-JSON, expand returns [query] and logs a warning.""" - from src.eval.transforms import QueryRewriter + from src.retrieval import QueryRewriter stub = _StubLLM(response="not json at all", prompt_tokens=50, completion_tokens=10) rw = QueryRewriter(model="gpt-4.1-nano", max_expansions=3, llm=stub) queries, cost, _, _ = rw.expand("What is RAG?") diff --git a/tests/test_eval_transform_refusal.py b/tests/test_retrieval_refusal.py similarity index 85% rename from tests/test_eval_transform_refusal.py rename to tests/test_retrieval_refusal.py index 9cf1d6f6..1f3e9c1d 100644 --- a/tests/test_eval_transform_refusal.py +++ b/tests/test_retrieval_refusal.py @@ -11,28 +11,28 @@ def _sr(score: float, chunk_id: str = "d1") -> SearchResult: def test_refuses_when_top1_below_threshold(): - from src.eval.transforms import RefusalHandler + from src.retrieval import RefusalHandler h = RefusalHandler(enabled=True, similarity_threshold=0.35, no_answer_text="I don't know.") assert h.should_refuse([_sr(0.20), _sr(0.10)]) is True def test_does_not_refuse_when_top1_above_threshold(): - from src.eval.transforms import RefusalHandler + from src.retrieval import RefusalHandler h = RefusalHandler(enabled=True, similarity_threshold=0.35, no_answer_text="I don't know.") assert h.should_refuse([_sr(0.50), _sr(0.10)]) is False def test_refuses_on_empty_candidates(): - from src.eval.transforms import RefusalHandler + from src.retrieval import RefusalHandler h = RefusalHandler(enabled=True, similarity_threshold=0.35, no_answer_text="I don't know.") assert h.should_refuse([]) is True def test_disabled_handler_never_refuses(): - from src.eval.transforms import RefusalHandler + from src.retrieval import RefusalHandler h = RefusalHandler(enabled=False, similarity_threshold=0.35, no_answer_text="I don't know.") assert h.should_refuse([_sr(0.0)]) is False @@ -40,7 +40,7 @@ def test_disabled_handler_never_refuses(): def test_refuse_response_returns_text_and_no_chunks(): - from src.eval.transforms import RefusalHandler + from src.retrieval import RefusalHandler h = RefusalHandler(enabled=True, similarity_threshold=0.35, no_answer_text="I cannot answer.") chunks, answer = h.refuse_response() diff --git a/tests/test_eval_retriever_reranker.py b/tests/test_retrieval_reranker_model.py similarity index 95% rename from tests/test_eval_retriever_reranker.py rename to tests/test_retrieval_reranker_model.py index 135ff083..cccceabd 100644 --- a/tests/test_eval_retriever_reranker.py +++ b/tests/test_retrieval_reranker_model.py @@ -14,7 +14,7 @@ def _sr(chunk_id: str, content: str, score: float, metadata: dict | None = None) @pytest.fixture(scope="module") def reranker(): - from src.eval.retrievers.reranker import CrossEncoderReranker + from src.retrieval.reranker import CrossEncoderReranker return CrossEncoderReranker()