Query Rewriting for Code RAG: Expand Symbols and Package Aliases

Query Rewriting for Code RAG: Expand Symbols and Package Aliases

Engineers ask “where’s the payments retry helper?” while the symbol is withStripeBackoff in @acme/payments-core. Vanilla embedding search on the raw question under-recalls; keyword search misses aliases. A thin query rewrite stage expands informal phrasing into canonical symbols, package names, and known aliases before retrieval.

⚡ TL;DR: Maintain an alias/symbol dictionary from the index (exports, package.json names, common nicknames); LLM-rewrite the question into {rewritten, symbols[], packages[]}; retrieve with hybrid lexical+vector using those expansions; log rewrites for eval. Pair with Citation-Required RAG and Code RAG Rerankers.

Build the expansion dictionary from the repo

# build_alias_dict.py
import json, re
from pathlib import Path

def package_names(root: Path) -> dict[str, str]:
    out = {}
    for pkg in root.glob("**/package.json"):
        data = json.loads(pkg.read_text())
        name = data.get("name")
        if name:
            short = name.split("/")[-1]
            out[short] = name
            out[name] = name
    return out

def export_symbols(ts_api_json: Path) -> dict[str, list[str]]:
    """Map lowercase nickname → [ExactSymbol, ...]."""
    api = json.loads(ts_api_json.read_text())
    m: dict[str, list[str]] = {}
    for sym in api["exports"]:
        m.setdefault(sym.lower(), []).append(sym)
        parts = re.findall(r"[A-Z]?[a-z]+|[A-Z]+(?![a-z])", sym)
        if parts:
            m.setdefault(" ".join(p.lower() for p in parts), []).append(sym)
    return m

Refresh on index sync (same webhook cadence as KB sync).

Rewrite stage (cheap model, tight schema)

type Rewrite = {
  rewritten: string;
  symbols: string[];
  packages: string[];
  intent: "locate_symbol" | "explain_behavior" | "how_to" | "debug";
};

const REWRITE_PROMPT = `Expand the developer question for code retrieval.
Use ONLY symbols/packages from the CANDIDATE_HINTS list when relevant.
Return JSON {rewritten, symbols, packages, intent}.
Keep rewritten concise; include exact identifiers.`;

async function rewriteQuery(q: string, hints: string[]): Promise<Rewrite> {
  return bedrockJson({
    model: "amazon.nova-lite-v1:0", // cheap rewriter
    system: REWRITE_PROMPT,
    user: `QUESTION: ${q}\nCANDIDATE_HINTS: ${hints.slice(0, 40).join(", ")}`,
    schema: RewriteSchema,
  });
}
async function retrieve(q: string) {
  const hints = aliasLookup(q); // fuzzy match dictionary
  const rw = await rewriteQuery(q, hints);
  const lexical = [...rw.symbols, ...rw.packages].join(" ");
  return hybridSearch({
    neuralText: rw.rewritten,
    bm25Text: `${q} ${lexical}`,
    filters: inferRepoFilters(rw),
    k: 12,
  });
}

Do not let the rewriter invent symbols not in hints — you amplify hallucinations into retrieval. Constrain to dictionary + retrieved hint list.

Evaluate rewrite lift separately

Ablation hit@10 Notes
Raw query 0.71 Baseline
Dictionary expand only 0.81 Fast, no LLM
LLM rewrite + dictionary 0.91 Best for informal asks
LLM rewrite unconstrained 0.84 Invented symbols hurt

Keep rewrite faithfulness: every symbols[] entry must exist in the dictionary or be dropped.

Latency budget

Rewrites must not dominate TTFT. Cap rewriter tokens, use a small model, and cache rewrites for identical questions (cost-aware caches). For autocomplete, skip LLM rewrite and use dictionary-only expansion.

Closing checklist

Dos
– Build alias/symbol dictionaries from packages + exports
– Constrain LLM rewrites to known identifiers
– Hybrid-search with both raw and expanded terms
– Ablate rewrite lift in the golden eval harness
– Cache rewrites for repeated questions

Donts
– Do not let rewriters mint novel API names
– Do not use your largest model for rewrite
– Do not skip dictionary refresh after big renames
– Do not apply heavy rewrite on every keystroke UI
– Do not fold rewrite bugs into “embedding quality” tickets

Related reading

Last updated on September 11, 2026


Discover more from CheatCoders

Subscribe to get the latest posts sent to your email.

1 Comment

Leave a Reply