feat: add query understanding layer for better retrieval

Pre-retrieval LLM step that reformulates user queries into
search-optimized forms (expanded keywords, resolved pronouns,
domain-aware terms) and optionally asks clarification questions.

- New src/query_engine.py with understand_query() + JSON parsing
- New prompt template in src/prompts.py
- CLI: query understanding + conversation history + clarification flow
- Web UI: query understanding + pending_clarification state
- Config: query_understanding.enabled/max_history/max_clarifications
- 10 new tests (59 total, all passing)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

app_cli.py

@@ -9,6 +9,7 @@ from rich.text import Text
 
 from src.config_loader import load_config, get_api_key
 from src.ingest import ingest_documents
+from src.query_engine import understand_query
 from src.retriever import retrieve
 from src.verifier import verify_and_respond
 from src.llm import list_models
@@ -214,6 +215,7 @@ def main() -> None:
     state: dict = {
         "last_retrieval": None,
         "web_search_override": None,
+        "conversation_history": [],
     }
 
     while True:
@@ -254,7 +256,6 @@ def main() -> None:
             continue
 
         # ── RAG pipeline ─────────────────────────────────────────────────
-        query = user_input
 
         # Apply web search override if set
         effective_cfg = cfg
@@ -266,24 +267,83 @@ def main() -> None:
                 "enabled": state["web_search_override"],
             }
 
-        # Retrieve
+        # ── Query understanding ─────────────────────────────────────────
+        qu_cfg = effective_cfg.get("query_understanding", {})
+        qu_enabled = qu_cfg.get("enabled", True)
+        max_clarifications = qu_cfg.get("max_clarifications", 1)
+
+        original_query = user_input
+        search_query = user_input
+
+        if qu_enabled:
+            with console.status("[bold blue]Understanding your question...[/bold blue]"):
+                try:
+                    qu_result = understand_query(
+                        user_input, effective_cfg,
+                        state.get("conversation_history", []),
+                    )
+                except Exception as e:
+                    console.print(f"[yellow]Query understanding failed, using raw query: {e}[/yellow]")
+                    qu_result = {"action": "search", "search_query": user_input, "original_query": user_input}
+
+            # Handle clarification
+            clarification_rounds = 0
+            while qu_result["action"] == "clarify" and clarification_rounds < max_clarifications:
+                console.print(f"\n[bold yellow]Clarification needed:[/bold yellow] {qu_result['clarification_question']}")
+                try:
+                    clarification = input("You> ").strip()
+                except (EOFError, KeyboardInterrupt):
+                    clarification = ""
+                if not clarification:
+                    break
+                combined = f"{user_input} — {clarification}"
+                with console.status("[bold blue]Understanding your question...[/bold blue]"):
+                    try:
+                        qu_result = understand_query(
+                            combined, effective_cfg,
+                            state.get("conversation_history", []),
+                        )
+                    except Exception:
+                        qu_result = {"action": "search", "search_query": combined, "original_query": user_input}
+                clarification_rounds += 1
+
+            # After max clarification rounds, force search
+            if qu_result["action"] == "clarify":
+                qu_result["action"] = "search"
+                # search_query fallback is already set in _parse_qu_result
+
+            search_query = qu_result.get("search_query", user_input)
+            original_query = qu_result.get("original_query", user_input)
+
+            # Show reformulated query if different from original
+            if search_query != user_input:
+                console.print(f"[dim]Searching for: \"{search_query}\"[/dim]")
+
+        # ── Retrieve ────────────────────────────────────────────────────
         with console.status("[bold blue]Searching knowledge base...[/bold blue]"):
             try:
-                retrieval_result = retrieve(query, effective_cfg)
+                retrieval_result = retrieve(search_query, effective_cfg)
             except Exception as e:
                 console.print(f"[red]Retrieval error: {e}[/red]")
                 continue
 
         state["last_retrieval"] = retrieval_result
 
-        # Generate + verify
+        # ── Generate + verify ───────────────────────────────────────────
         with console.status("[bold blue]Generating response...[/bold blue]"):
             try:
-                result = verify_and_respond(query, retrieval_result, effective_cfg)
+                result = verify_and_respond(original_query, retrieval_result, effective_cfg)
             except Exception as e:
                 console.print(f"[red]Generation error: {e}[/red]")
                 continue
 
+        # ── Update conversation history ─────────────────────────────────
+        history = state.get("conversation_history", [])
+        max_history = qu_cfg.get("max_history", 6)
+        history.append({"role": "user", "content": user_input})
+        history.append({"role": "assistant", "content": result["response"]})
+        state["conversation_history"] = history[-max_history:]
+
         # ── Display response ─────────────────────────────────────────────
         response_text = result["response"]
 

app_web.py

@@ -4,6 +4,7 @@ import streamlit as st
 
 from src.config_loader import load_config, get_api_key
 from src.ingest import get_chroma_collection, ingest_documents
+from src.query_engine import understand_query
 from src.retriever import retrieve
 from src.verifier import verify_and_respond
 from src.llm import list_models
@@ -19,6 +20,8 @@ def init_session():
         st.session_state.cfg = load_config()
     if "last_retrieval" not in st.session_state:
         st.session_state.last_retrieval = None
+    if "pending_clarification" not in st.session_state:
+        st.session_state.pending_clarification = None
 
 
 # ── Sidebar ──────────────────────────────────────────────────────────────────
@@ -133,11 +136,56 @@ def render_chat():
         with st.chat_message("user"):
             st.markdown(prompt)
 
+        # ── Check if this is a clarification response ───────────────────
+        pending = st.session_state.pending_clarification
+        if pending is not None:
+            # This prompt is the user's clarification answer
+            combined = f"{pending} — {prompt}"
+            original_query = pending
+            st.session_state.pending_clarification = None
+        else:
+            combined = prompt
+            original_query = prompt
+
+        # ── Query understanding ─────────────────────────────────────────
+        qu_cfg = cfg.get("query_understanding", {})
+        qu_enabled = qu_cfg.get("enabled", True)
+        max_history = qu_cfg.get("max_history", 6)
+
+        search_query = combined
+
+        if qu_enabled:
+            history = [
+                {"role": m["role"], "content": m["content"]}
+                for m in st.session_state.messages[-max_history:]
+            ]
+            try:
+                qu_result = understand_query(combined, cfg, history)
+            except Exception:
+                qu_result = {"action": "search", "search_query": combined, "original_query": original_query}
+
+            if qu_result["action"] == "clarify" and pending is None:
+                # Ask clarification — store original query, show question
+                st.session_state.pending_clarification = original_query
+                clarification_msg = f"**Before I search, could you clarify?** {qu_result['clarification_question']}"
+                st.session_state.messages.append(
+                    {"role": "assistant", "content": clarification_msg}
+                )
+                with st.chat_message("assistant"):
+                    st.markdown(clarification_msg)
+                return
+
+            search_query = qu_result.get("search_query", combined)
+
         # Generate assistant response
         with st.chat_message("assistant"):
             with st.status("Searching knowledge base...", expanded=True) as status:
+                # Show reformulated query if different
+                if search_query != original_query:
+                    status.update(label=f'Searching for: "{search_query}"...')
+
                 # Retrieval
-                retrieval_result = retrieve(prompt, cfg)
+                retrieval_result = retrieve(search_query, cfg)
                 st.session_state.last_retrieval = retrieval_result
 
                 n_local = len(retrieval_result.get("db_results", []))
@@ -146,8 +194,8 @@ def render_chat():
                     label=f"Found {n_local} local + {n_web} web sources. Generating response..."
                 )
 
-                # Verification and response generation
-                result = verify_and_respond(prompt, retrieval_result, cfg)
+                # Verification and response generation (uses original query)
+                result = verify_and_respond(original_query, retrieval_result, cfg)
 
                 # Update status based on verification outcome
                 if result.get("refused"):

setup.py

@@ -56,6 +56,11 @@ def generate_config(
             "backend": "semantic_scholar",
             "max_results": 5,
         },
+        "query_understanding": {
+            "enabled": True,
+            "max_history": 6,
+            "max_clarifications": 1,
+        },
         "verification": {
             "enabled": True,
             "max_iterations": 3,

src/prompts.py

@@ -3,6 +3,7 @@ Prompt builder with anti-hallucination guardrails and verification prompts.
 
 Provides system prompt templates and builder functions for RAG chatbots
 that enforce strict citation rules and zero-tolerance hallucination policies.
+Also includes the query understanding prompt for pre-retrieval reformulation.
 """
 
 SYSTEM_PROMPT_TEMPLATE = """\
@@ -120,6 +121,77 @@ If there are no errors, return an empty "errors" list,
 """
 
 
+QUERY_UNDERSTANDING_PROMPT_TEMPLATE = """\
+You are a query reformulation assistant for a {domain} research chatbot.
+
+Your job is to analyze the user's question and decide:
+1. SEARCH -- if the question is clear enough, rewrite it as an optimized search query
+2. CLARIFY -- if the question is genuinely ambiguous, ask ONE clarification question
+
+RULES for SEARCH:
+- Expand abbreviations and acronyms (e.g., "HR" -> "human rights")
+- Resolve pronouns using conversation history (e.g., "they" -> the entity from prior turn)
+- Add relevant domain-specific keywords that would match document content
+- Keep the reformulated query concise (15-40 words)
+- Preserve the user's original intent -- do NOT change what they're asking about
+- If the query is already clear and specific, return it with minimal changes
+
+RULES for CLARIFY:
+- ONLY ask for clarification when the question is genuinely ambiguous
+  (could mean 2+ clearly different things)
+- Ask ONE specific, short question
+- Provide 2-3 options when possible
+- Do NOT ask for clarification on clear questions -- even if they are broad
+- A broad question is NOT the same as an ambiguous question
+
+CONVERSATION HISTORY:
+{history}
+
+USER'S QUESTION: {query}
+
+Respond in JSON:
+{{
+  "action": "search" or "clarify",
+  "search_query": "reformulated search query (only if action is search)",
+  "clarification_question": "question to ask (only if action is clarify)",
+  "reasoning": "one sentence explaining your choice"
+}}
+"""
+
+
+def build_query_understanding_prompt(
+    query: str, domain: str, history: list[dict],
+) -> str:
+    """Build the query understanding prompt for pre-retrieval reformulation.
+
+    Args:
+        query: The user's raw question.
+        domain: Knowledge domain from config.
+        history: Recent conversation messages for context resolution.
+
+    Returns:
+        Formatted query understanding prompt string.
+    """
+    if history:
+        history_lines = []
+        for msg in history:
+            role = msg.get("role", "user").capitalize()
+            content = msg.get("content", "")
+            # Truncate long assistant responses to save tokens
+            if role == "Assistant" and len(content) > 200:
+                content = content[:200] + "..."
+            history_lines.append(f"{role}: {content}")
+        history_str = "\n".join(history_lines)
+    else:
+        history_str = "(No prior conversation)"
+
+    return QUERY_UNDERSTANDING_PROMPT_TEMPLATE.format(
+        domain=domain,
+        history=history_str,
+        query=query,
+    )
+
+
 def build_prompt(context: str, bot_name: str, domain: str) -> str:
     """Build the system prompt with anti-hallucination guardrails.
 

src/query_engine.py

@@ -0,0 +1,119 @@
+"""Query understanding layer: reformulates user queries for better retrieval."""
+
+import json
+import re
+
+from src.llm import generate
+from src.prompts import build_query_understanding_prompt
+
+
+def understand_query(
+    user_query: str,
+    cfg: dict,
+    conversation_history: list[dict] | None = None,
+) -> dict:
+    """Analyze and reformulate a user query for better retrieval.
+
+    Uses the LLM to either:
+    - SEARCH: Rewrite the query into a search-optimized form
+    - CLARIFY: Ask the user a clarification question
+
+    Args:
+        user_query: Raw user input.
+        cfg: App config (provides domain, LLM settings).
+        conversation_history: Recent messages for pronoun/reference resolution.
+
+    Returns:
+        Dict with keys:
+        - ``action``: ``"search"`` or ``"clarify"``
+        - ``search_query``: Reformulated query (always present)
+        - ``original_query``: The raw user input
+        - ``clarification_question``: Question to ask (only if action is "clarify")
+    """
+    qu_cfg = cfg.get("query_understanding", {})
+
+    # Disabled: pass through raw query
+    if not qu_cfg.get("enabled", True):
+        return {
+            "action": "search",
+            "search_query": user_query,
+            "original_query": user_query,
+        }
+
+    domain = cfg.get("chatbot", {}).get("domain", "research")
+    max_history = qu_cfg.get("max_history", 6)
+
+    # Trim conversation history
+    history = conversation_history or []
+    if len(history) > max_history:
+        history = history[-max_history:]
+
+    prompt = build_query_understanding_prompt(user_query, domain, history)
+
+    try:
+        raw = generate(
+            "You are a query reformulation assistant. Return only JSON.",
+            prompt,
+            cfg,
+            max_tokens=256,
+        )
+        result = _parse_qu_result(raw, user_query)
+    except Exception:
+        # LLM failure: fall back to raw query
+        result = {
+            "action": "search",
+            "search_query": user_query,
+            "original_query": user_query,
+        }
+
+    return result
+
+
+def _parse_qu_result(raw: str, original_query: str) -> dict:
+    """Parse the LLM's JSON response into a structured result.
+
+    Falls back to raw query passthrough if parsing fails.
+    """
+    parsed = None
+
+    # Try direct parse
+    try:
+        parsed = json.loads(raw)
+    except (json.JSONDecodeError, TypeError):
+        pass
+
+    # Regex fallback: find first { ... } block
+    if parsed is None:
+        match = re.search(r"\{[\s\S]*\}", raw)
+        if match:
+            try:
+                parsed = json.loads(match.group())
+            except (json.JSONDecodeError, TypeError):
+                pass
+
+    # Unparseable: fall back to raw query
+    if not isinstance(parsed, dict) or "action" not in parsed:
+        return {
+            "action": "search",
+            "search_query": original_query,
+            "original_query": original_query,
+        }
+
+    action = parsed.get("action", "search")
+
+    if action == "clarify":
+        return {
+            "action": "clarify",
+            "search_query": original_query,  # fallback if clarification is skipped
+            "original_query": original_query,
+            "clarification_question": parsed.get(
+                "clarification_question", "Could you be more specific?"
+            ),
+        }
+
+    # action == "search"
+    return {
+        "action": "search",
+        "search_query": parsed.get("search_query", original_query),
+        "original_query": original_query,
+    }

tests/test_prompts.py

@@ -36,3 +36,24 @@ def test_build_verification_prompt():
     assert "citation" in prompt.lower()
     assert "JSON" in prompt or "json" in prompt
     assert "error_count" in prompt
+
+
+def test_build_query_understanding_prompt():
+    from src.prompts import build_query_understanding_prompt
+    prompt = build_query_understanding_prompt(
+        "What happened in China?",
+        "human rights",
+        [{"role": "user", "content": "prior question"}],
+    )
+    assert "human rights" in prompt
+    assert "What happened in China?" in prompt
+    assert "prior question" in prompt
+    assert "SEARCH" in prompt or "search" in prompt
+    assert "CLARIFY" in prompt or "clarify" in prompt
+
+
+def test_build_query_understanding_prompt_no_history():
+    from src.prompts import build_query_understanding_prompt
+    prompt = build_query_understanding_prompt("test query", "research", [])
+    assert "No prior conversation" in prompt
+    assert "test query" in prompt

tests/test_query_engine.py

@@ -0,0 +1,110 @@
+import json
+import pytest
+from unittest.mock import patch
+
+
+def _make_cfg(enabled=True, max_history=6, max_clarifications=1):
+    """Build a minimal config dict for query engine tests."""
+    return {
+        "chatbot": {"name": "TestBot", "domain": "human rights research"},
+        "llm": {"provider": "openai", "model": "gpt-4o", "temperature": 0.0, "max_tokens": 2048},
+        "api_keys": {"openai": "fake-key"},
+        "query_understanding": {
+            "enabled": enabled,
+            "max_history": max_history,
+            "max_clarifications": max_clarifications,
+        },
+    }
+
+
+def _mock_generate_search(system_prompt, user_message, cfg, **kwargs):
+    """Mock LLM that always returns a search reformulation."""
+    return json.dumps({
+        "action": "search",
+        "search_query": "reformulated test query about human rights",
+        "reasoning": "expanded for better retrieval",
+    })
+
+
+def _mock_generate_clarify(system_prompt, user_message, cfg, **kwargs):
+    """Mock LLM that always returns a clarification request."""
+    return json.dumps({
+        "action": "clarify",
+        "clarification_question": "Are you asking about X or Y?",
+        "reasoning": "ambiguous query",
+    })
+
+
+def _mock_generate_error(system_prompt, user_message, cfg, **kwargs):
+    """Mock LLM that raises an exception."""
+    raise RuntimeError("LLM unavailable")
+
+
+def test_understand_query_reformulates():
+    from src.query_engine import understand_query
+    with patch("src.query_engine.generate", _mock_generate_search):
+        result = understand_query("test query", _make_cfg())
+    assert result["action"] == "search"
+    assert result["search_query"] == "reformulated test query about human rights"
+    assert result["original_query"] == "test query"
+
+
+def test_understand_query_clarifies():
+    from src.query_engine import understand_query
+    with patch("src.query_engine.generate", _mock_generate_clarify):
+        result = understand_query("AI", _make_cfg())
+    assert result["action"] == "clarify"
+    assert "clarification_question" in result
+    assert result["original_query"] == "AI"
+    # search_query should still be present as fallback
+    assert "search_query" in result
+
+
+def test_understand_query_preserves_original():
+    from src.query_engine import understand_query
+    original = "What happened in China?"
+    with patch("src.query_engine.generate", _mock_generate_search):
+        result = understand_query(original, _make_cfg())
+    assert result["original_query"] == original
+
+
+def test_understand_query_with_history():
+    from src.query_engine import understand_query
+    history = [
+        {"role": "user", "content": "Tell me about China"},
+        {"role": "assistant", "content": "China has various human rights issues..."},
+    ]
+    with patch("src.query_engine.generate", _mock_generate_search):
+        result = understand_query("what about 2023?", _make_cfg(), history)
+    assert result["action"] == "search"
+
+
+def test_understand_query_disabled():
+    from src.query_engine import understand_query
+    result = understand_query("test query", _make_cfg(enabled=False))
+    assert result["action"] == "search"
+    assert result["search_query"] == "test query"
+    assert result["original_query"] == "test query"
+
+
+def test_understand_query_fallback_on_error():
+    from src.query_engine import understand_query
+    with patch("src.query_engine.generate", _mock_generate_error):
+        result = understand_query("test query", _make_cfg())
+    assert result["action"] == "search"
+    assert result["search_query"] == "test query"
+
+
+def test_parse_qu_result_malformed_json():
+    from src.query_engine import _parse_qu_result
+    result = _parse_qu_result("not json at all", "original query")
+    assert result["action"] == "search"
+    assert result["search_query"] == "original query"
+
+
+def test_parse_qu_result_embedded_json():
+    from src.query_engine import _parse_qu_result
+    raw = 'Here is the result: {"action": "search", "search_query": "better query", "reasoning": "test"} end'
+    result = _parse_qu_result(raw, "original")
+    assert result["action"] == "search"
+    assert result["search_query"] == "better query"

tests/test_setup.py

@@ -22,6 +22,10 @@ def test_generate_config_yaml():
     assert cfg["retrieval"]["chunk_size"] == 1000
     assert cfg["retrieval"]["chunk_overlap"] == 100
     assert cfg["embeddings"]["provider"] == "local"
+    # Query understanding config
+    assert cfg["query_understanding"]["enabled"] is True
+    assert cfg["query_understanding"]["max_history"] == 6
+    assert cfg["query_understanding"]["max_clarifications"] == 1
 
 def test_generate_env_file():
     from setup import generate_env