Skip to content

Latest commit

 

History

History
307 lines (234 loc) · 10.5 KB

File metadata and controls

307 lines (234 loc) · 10.5 KB

rag-query — Query understanding

rag-query (Step 2.6) turns one raw user query into a structured UnderstoodQuery that downstream retrieval, reranking, and packing consume. Four components ship behind their own SPIs; the QueryUnderstandingPipeline orchestrator wires them together and emits the canonical query.understand OpenTelemetry span.

Overview

Component Purpose Targets
Rewriter Alternative phrasings of the same intent Keyword path (BM25)
Decomposer Split multi-hop into sub-questions Per-sub-query fan-out
HyDE Hypothetical-answer paragraph + embedding Dense path
Glossary expansion Domain synonyms / acronyms Keyword path
QueryUnderstandingPipeline Parallel orchestrator Returns UnderstoodQuery

All four components are optional. A pipeline with just a glossary, just HyDE, or all four is valid — missing components contribute their identity output ([] / {} / None).

Usage

from rag_query import (
    QueryUnderstandingPipeline,
    LLMQueryRewriter,
    HeuristicQueryDecomposer,
    InMemoryGlossary,
    HyDEGenerator,
)

pipeline = QueryUnderstandingPipeline(
    rewriter=LLMQueryRewriter(llm),
    decomposer=HeuristicQueryDecomposer(),
    glossary=InMemoryGlossary({"rag": ["retrieval augmented generation"]}),
    hyde=HyDEGenerator(llm, embedder),
    rewrite_n=3,
)

understood = await pipeline.understand(ctx, "what is RAG used for?")

# Hand off to retrieval (Step 2.5):
retriever.retrieve(
    ctx,
    vector=understood.hyde_vector or original_embedding,
    query_text=understood.original_text,   # or fan out across understood.keyword_queries
    graph_seeds=...,
)

Output type — UnderstoodQuery

Frozen Pydantic model. Pass the same instance through retrieval → reranker → packer — every stage may read but never mutate.

Field Type Notes
original_text str Preserved verbatim for audit
rewrites list[str] Excludes the original
sub_queries list[str] Empty = "no decomposition", not "no answer"
expansion_terms dict[str, list[str]] {matched_term_lowercase: [synonym, ...]}
hyde_documents list[str] Hypothetical answer paragraphs
hyde_vector list[float] | None Averaged embedding of hyde_documents
metadata dict[str, Any] elapsed_ms, hyde_embed_model, optional partial_failures list

Convenience property keyword_queries returns [original_text, *rewrites] deduped case-insensitively — handy when fanning out the keyword path.

Rewriter

QueryRewriter

class QueryRewriter(abc.ABC):
    async def rewrite(self, ctx, query_text, *, n=3) -> list[str]: ...

Contract: must not return the original query (case-insensitive); must not return duplicates; may return fewer than n.

LLMQueryRewriter

Prompts an LLM for n paraphrases, parses one-per-line. Robust to common LLM-output quirks — leading numbering (1., 1)), bullets (- , * , ), trailing whitespace, and case-only duplicates are normalised away.

Override the prompt at construction time via system_prompt= (the {n} placeholder is substituted before sending).

Optional cache (Step 2.11 G-06). Pass a generic Cache via cache= to skip the LLM call on repeat inputs:

from rag_core.spi.noop import NoopCache  # or redis-backed Cache
from rag_query.rewriter import LLMQueryRewriter

rewriter = LLMQueryRewriter(
    llm,
    cache=redis_cache,
    cache_ttl_seconds=3600,
    prompt_version="v1",  # bump to invalidate after prompt changes
)

The cache key normalises query text (lowercase + whitespace collapse) and includes (tenant, model_id, prompt_version, n) so independent tenants / model versions / paraphrase counts don't share entries. Corrupt entries heal forward (fall through to a fresh LLM call and overwrite). Empty LLM responses are not cached — a future call retries.

NoopQueryRewriter

Always returns []. Safe default; useful in tests and pipelines that want decomposition / HyDE / expansion but not rewrites.

Decomposer

QueryDecomposer

class QueryDecomposer(abc.ABC):
    async def decompose(self, ctx, query_text) -> list[str]: ...

Contract: [] means "no decomposition applies"; the caller MUST use the original query. Never use [query_text] as a sentinel — that's ambiguous with a successful single-question decomposition.

HeuristicQueryDecomposer

Pure-Python, free, deterministic. Splits on conjunctions (" and ", " then ", " but also ") and sentence terminators (?, ., !). Prefers under-splitting — returns [] when:

  • Query is shorter than min_chars (default 30) — likely single-hop.
  • Splitting produces only one part.
  • Any candidate part is shorter than 3 characters — splitter was probably wrong about boundaries.

LLMQueryDecomposer

LLM-based, robust to numbering/bullets/blanks. Collapses the "LLM echoed the original query" case to [] so the contract holds.

NoopQueryDecomposer

Always returns [].

HyDE

HyDEGenerator

Implements Gao et al. 2022 (https://arxiv.org/abs/2212.10496): LLM-generated hypothetical answer paragraph is often a better dense retrieval target than the (short, vocabulary-mismatched) query itself.

hyde = HyDEGenerator(llm, embedder, n=1)
result = await hyde.generate(ctx, query_text)
# result.documents : list[str]
# result.vector    : list[float] | None  (averaged across n docs)
# result.model     : str | None          (embedder model id)

embedder=None is supported — generate() returns documents but no vector, useful when the caller embeds elsewhere (cache hit, different embedder per tenant).

When n > 1, the LLM is called n times and the resulting vectors are component-wise averaged (the protocol from §2.3 of the paper).

Optional embedding cache (Step 2.11 G-06). Pass an EmbeddingCache via embedding_cache= to skip the embedder call when a hypothetical paragraph has already been embedded for this (tenant, model_id, model_version):

from rag_core.spi.noop import NoopEmbeddingCache

hyde = HyDEGenerator(
    llm, embedder, n=3,
    embedding_cache=redis_embedding_cache,
    embedding_model_version="v1",
)

The cache key normalises hypothetical-doc text the same way the agent loop's own cache does — so HyDE and the agent loop share key namespaces and reuse each other's work when backed by the same store. HyDEResult.cache_hits / .cache_misses surface per-call counters for dashboarding. Wiring a cache without an embedder raises ValueError at construction (the cache has no source of truth to populate).

Glossary expansion

GlossaryExpander

class GlossaryExpander(abc.ABC):
    async def expand(self, ctx, query_text) -> dict[str, list[str]]: ...

Returns {matched_term_lowercase: [synonym, ...]}. Synonyms in glossary insertion order; the caller picks how to fold them into the keyword query (OR-expand, boost original term, etc.).

InMemoryGlossary

Whole-word, case-insensitive lookup over an in-memory dict. Multi-word entries match as phrases (case and surrounding whitespace ignored). Multi-word entries are tried before single-word ones so prefixes don't half-match.

g = InMemoryGlossary({
    "invoice":           ["bill", "statement"],
    "operating system":  ["OS", "platform"],
})

NoopGlossaryExpander

Always returns {}.

QueryUnderstandingPipeline

Async orchestrator. Independent components fan out in parallel via asyncio.gather. Per-component exceptions are caught, logged (query.understand.component_failed), and degraded to that component's identity output — one degraded component must not take retrieval down.

KeyboardInterrupt / SystemExit propagate out of gather rather than being captured.

Span schema — query.understand

Attribute Type Description
rag.tenant_id str From ctx.tenant_id
rag.query.length_chars int Length of input query
rag.query.rewrites_n int Number of rewrites emitted
rag.query.sub_queries_n int Number of sub-questions emitted
rag.query.expansions_n int Number of glossary keys matched
rag.query.hyde_docs_n int HyDE document count
rag.query.hyde_vector_present bool True when an embedder was wired
rag.query.partial_failures_n int Number of components that raised
rag.query.failed_components str Comma-separated component labels (only when failures > 0)
rag.query.elapsed_ms float End-to-end wall time

ragctl understand — smoke command

ragctl understand "what is RAG used for"
ragctl understand "describe SpaceX and explain Falcon 1" --no-hyde
ragctl understand "what is RAG used for" \
    --glossary-term rag \
    --glossary-synonyms "retrieval augmented generation,RAG pipeline"

Uses the noop LLM + noop embedder so it runs anywhere with no credentials. Useful for sanity-checking pipeline wiring and span emission.

Internals

PolicyEngine boundary

The pipeline does not consult PolicyEngine. Query understanding is a pre-retrieval transformation on user-supplied text — it never reads, writes, or egresses stored chunks. The downstream HybridRetriever (Step 2.5) is the canonical read_chunk PDP consumer.

The PolicyEngine coverage linter (tests/policy/coverage.py) explicitly allowlists the three LLM-consuming component files (rewriter.py, decomposer.py, hyde.py) because the gateway (Phase 3) — not the components — will be the LLM-call PDP consumer.

Partial-failure semantics

Component Identity output on failure
Rewriter rewrites = []
Decomposer sub_queries = []
Glossary expansion_terms = {}
HyDE hyde_documents = [], hyde_vector = None

metadata["partial_failures"] lists the labels of components that raised, in alphabetical order.

Extension points

  • New rewriter / decomposer / glossary impl — subclass the matching ABC and pass to QueryUnderstandingPipeline.
  • Tenant-scoped glossary — implement a GlossaryExpander that loads the per-tenant glossary on construction or via ctx.tenant_id.
  • Different HyDE prompt or n — pass system_prompt= / n= to HyDEGenerator.
  • Different fan-out shape (sequential instead of parallel) — subclass QueryUnderstandingPipeline._run_all.