[PR #5171] [CLOSED] feat: LLM-based query rewriting for contextual RAG #5335

Closed
opened 2026-06-05 15:21:00 -04:00 by yindo · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/Mintplex-Labs/anything-llm/pull/5171
Author: @Alminc91
Created: 3/8/2026
Status: Closed

Base: masterHead: feat/query-rewriting-for-contextual-rag


📝 Commits (3)

  • 21a3059 feat: LLM-based query rewriting for contextual RAG in multi-turn conversations
  • 78590fd improve: selective rewriting with UNCHANGED signal and 3-layer output validation
  • a657017 improve: add universal script fallback for CJK/Thai in output validation

📊 Changes

5 files changed (+164 additions, -6 deletions)

View changed files

📝 server/utils/chats/apiChatHandler.js (+15 -2)
📝 server/utils/chats/embed.js (+8 -1)
📝 server/utils/chats/openaiCompatible.js (+15 -2)
📝 server/utils/chats/stream.js (+8 -1)
server/utils/helpers/chat/queryRewriter.js (+118 -0)

📄 Description

Problem

In multi-turn conversations, follow-up queries lack context for vector search, causing RAG to retrieve irrelevant or no results.

Example conversation:

  1. User: "What German courses do you offer?" → RAG finds German courses
  2. Bot: "Here are some German courses... Are you looking for A1 or B1?"
  3. User: "Yes, B1" → RAG searches for "Yes, B1" → finds nothing
  4. User: "Is this course also available on weekends?" → RAG searches literally → nothing

The chat history is already loaded and sent to the LLM for response generation — but it is not used for the vector search. The raw user query goes directly into performSimilaritySearch, which fails for context-dependent follow-ups.

This has been reported multiple times:

  • #4352 (Contextual RAG)
  • #4071 (RAG retrieval misses chat history reference)
  • #4072 (follow-up)

Solution

A lightweight LLM-based query rewriter that runs before vector search. It takes the user's message + recent chat history and rewrites it into a standalone search query.

How it works

  1. Trigger: Only when chat history exists AND query is ≤12 words (configurable). First messages and long self-contained queries skip rewriting entirely.
  2. Rewrite: Sends a short prompt (~550 tokens total) with the last 2 turn-pairs to the already-loaded LLM. Assistant messages are truncated to 150 words to keep input small.
  3. Use: The rewritten query goes to performSimilaritySearch. The original message is still used for the LLM chat completion and stored in chat history.
  4. Fallback: On any error, the original query is used unchanged. Zero risk of breaking existing behavior.

Selective rewriting with UNCHANGED signal

Self-contained queries (ones that already contain their own subject/topic) are not rewritten. The LLM responds with a single token UNCHANGED instead of reproducing the full query, which provides two benefits:

  • Latency: 1 output token (~125ms) vs. reproducing the query (5-15 tokens, ~300ms) — roughly 40% faster for queries that don't need rewriting
  • Correctness: No risk of the LLM subtly altering a self-contained query (e.g., changing "What yoga courses are there?" to "Which yoga courses are available?" — same meaning, different search vector)

3-layer output validation (model-agnostic)

LLM output is non-deterministic. Different models (GPT-4, Mistral, Llama, Qwen, local models) may produce meta-responses like "no rewrite needed" or "The query is already self-contained" instead of clean output. If such text reaches vector search, it returns nothing — a silent failure.

The validation guarantees that rewriteQueryForSearch() only ever returns the original query or a valid rewritten query, never meta-text:

Layer Check Catches
1 startsWith("UNCHANGED") Explicit signal from compliant LLMs (fast path, 1 token)
2 Verbatim string comparison LLM copies the query exactly (fallback for less capable models)
3 Content word overlap Valid rewrites share topic words (>3 chars) with the conversation context. Meta-responses like "no rewrite needed" do not.

Layer 3 is the core safety net — it works across all languages and models without keyword lists. A valid rewrite necessarily contains subject/topic words from the conversation ("German courses", "B1", "weekends"), while meta-responses contain only generic words ("rewrite", "needed", "query") that never appear in the conversation context.

Universal script support: For non-space-separated scripts (CJK, Thai, etc.), word-level matching is supplemented with character-level overlap (characters with charCode > 127). This ensures the validation works for Chinese, Japanese, Korean, Thai, and other languages where spaces don't delimit words.

Production results

Tested in production across 65 RAG chatbot instances. Actual rewrite logs:

[QueryRewrite] "Yes, B1 please"                            → "German courses B1 integration course"      (339ms)
[QueryRewrite] "Is this course also available on weekends?" → "Integration courses German B1 on weekends" (258ms)
[QueryRewrite] "What about regular German courses?"         → "German courses B1 on weekends"             (219ms)
[QueryRewrite] "Let's search for Spanish beginner courses!" → "UNCHANGED"                                 (125ms)  ← self-contained, no rewrite
[QueryRewrite] "Maybe I'll try B1 after all!"               → "Spanish courses B1 advanced"               (369ms)
[QueryRewrite] "Is that the only one???"                    → "Spanish courses B1 available"               (332ms)

Overhead: 125–370ms per query (~550 tokens in, 1-15 tokens out). Self-contained queries cost ~125ms (1 token). The rewrite uses the same LLM that's already loaded for chat — no additional model required.

Configuration

Three environment variables, no database changes:

Variable Default Description
DISABLE_QUERY_REWRITING false Kill switch to disable entirely
QUERY_REWRITE_WORD_THRESHOLD 12 Queries above N words skip rewriting
QUERY_REWRITE_MAX_HISTORY 2 Number of turn-pairs used for context

Changes

File Change
server/utils/helpers/chat/queryRewriter.js New — rewrite logic, 3-layer validation, universal script support
server/utils/chats/embed.js Import + rewrite before vector search
server/utils/chats/stream.js Import + rewrite before vector search
server/utils/chats/apiChatHandler.js Import + rewrite before vector search (2 handlers)
server/utils/chats/openaiCompatible.js Import + rewrite before vector search (2 handlers)

All 6 performSimilaritySearch call sites across all chat interfaces (embed widget, workspace UI, API sync, API stream, OpenAI-compatible sync/stream) are covered.

Why this approach

Alternative Why not
String concatenation (keywords from history) Fragile, unreliable for real conversations
Embed history + query together Dilutes the embedding signal
Conditional on word count only Misses context-dependent longer queries ("Is this course available on weekends?")
Always rewrite Wastes compute on first messages and long self-contained queries
Trust LLM output directly Different models produce unpredictable meta-text — needs validation

The LLM-based approach with a word-count threshold is the industry standard (LangChain's create_history_aware_retriever, Open WebUI's search query generation). This implementation goes further with selective rewriting (UNCHANGED signal) and model-agnostic output validation (3-layer check). One file, ~115 lines of logic, no dependencies.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/Mintplex-Labs/anything-llm/pull/5171 **Author:** [@Alminc91](https://github.com/Alminc91) **Created:** 3/8/2026 **Status:** ❌ Closed **Base:** `master` ← **Head:** `feat/query-rewriting-for-contextual-rag` --- ### 📝 Commits (3) - [`21a3059`](https://github.com/Mintplex-Labs/anything-llm/commit/21a3059698b05702737e53a6b9cf83a70bbd67fc) feat: LLM-based query rewriting for contextual RAG in multi-turn conversations - [`78590fd`](https://github.com/Mintplex-Labs/anything-llm/commit/78590fd3e789b8ff95bda18ad5eedd5952a5289e) improve: selective rewriting with UNCHANGED signal and 3-layer output validation - [`a657017`](https://github.com/Mintplex-Labs/anything-llm/commit/a657017ed0e8cb841e4fde8d7d718f68f0284041) improve: add universal script fallback for CJK/Thai in output validation ### 📊 Changes **5 files changed** (+164 additions, -6 deletions) <details> <summary>View changed files</summary> 📝 `server/utils/chats/apiChatHandler.js` (+15 -2) 📝 `server/utils/chats/embed.js` (+8 -1) 📝 `server/utils/chats/openaiCompatible.js` (+15 -2) 📝 `server/utils/chats/stream.js` (+8 -1) ➕ `server/utils/helpers/chat/queryRewriter.js` (+118 -0) </details> ### 📄 Description ## Problem In multi-turn conversations, follow-up queries lack context for vector search, causing RAG to retrieve irrelevant or no results. **Example conversation:** 1. User: "What German courses do you offer?" → RAG finds German courses ✅ 2. Bot: "Here are some German courses... Are you looking for A1 or B1?" 3. User: "Yes, B1" → RAG searches for "Yes, B1" → finds **nothing** ❌ 4. User: "Is this course also available on weekends?" → RAG searches literally → **nothing** ❌ The chat history is already loaded and sent to the LLM for response generation — but it is **not used for the vector search**. The raw user query goes directly into `performSimilaritySearch`, which fails for context-dependent follow-ups. This has been reported multiple times: - #4352 (Contextual RAG) - #4071 (RAG retrieval misses chat history reference) - #4072 (follow-up) ## Solution A lightweight LLM-based query rewriter that runs **before** vector search. It takes the user's message + recent chat history and rewrites it into a standalone search query. ### How it works 1. **Trigger**: Only when chat history exists AND query is ≤12 words (configurable). First messages and long self-contained queries skip rewriting entirely. 2. **Rewrite**: Sends a short prompt (~550 tokens total) with the last 2 turn-pairs to the already-loaded LLM. Assistant messages are truncated to 150 words to keep input small. 3. **Use**: The rewritten query goes to `performSimilaritySearch`. The **original** message is still used for the LLM chat completion and stored in chat history. 4. **Fallback**: On any error, the original query is used unchanged. Zero risk of breaking existing behavior. ### Selective rewriting with UNCHANGED signal Self-contained queries (ones that already contain their own subject/topic) are **not rewritten**. The LLM responds with a single token `UNCHANGED` instead of reproducing the full query, which provides two benefits: - **Latency**: 1 output token (~125ms) vs. reproducing the query (5-15 tokens, ~300ms) — roughly **40% faster** for queries that don't need rewriting - **Correctness**: No risk of the LLM subtly altering a self-contained query (e.g., changing "What yoga courses are there?" to "Which yoga courses are available?" — same meaning, different search vector) ### 3-layer output validation (model-agnostic) LLM output is non-deterministic. Different models (GPT-4, Mistral, Llama, Qwen, local models) may produce meta-responses like "no rewrite needed" or "The query is already self-contained" instead of clean output. If such text reaches vector search, it returns nothing — a silent failure. The validation guarantees that `rewriteQueryForSearch()` only ever returns the **original query** or a **valid rewritten query**, never meta-text: | Layer | Check | Catches | |-------|-------|---------| | 1 | `startsWith("UNCHANGED")` | Explicit signal from compliant LLMs (fast path, 1 token) | | 2 | Verbatim string comparison | LLM copies the query exactly (fallback for less capable models) | | 3 | Content word overlap | Valid rewrites share topic words (>3 chars) with the conversation context. Meta-responses like "no rewrite needed" do not. | Layer 3 is the core safety net — it works across all languages and models without keyword lists. A valid rewrite necessarily contains subject/topic words from the conversation ("German courses", "B1", "weekends"), while meta-responses contain only generic words ("rewrite", "needed", "query") that never appear in the conversation context. **Universal script support**: For non-space-separated scripts (CJK, Thai, etc.), word-level matching is supplemented with character-level overlap (characters with charCode > 127). This ensures the validation works for Chinese, Japanese, Korean, Thai, and other languages where spaces don't delimit words. ### Production results Tested in production across 65 RAG chatbot instances. Actual rewrite logs: ``` [QueryRewrite] "Yes, B1 please" → "German courses B1 integration course" (339ms) [QueryRewrite] "Is this course also available on weekends?" → "Integration courses German B1 on weekends" (258ms) [QueryRewrite] "What about regular German courses?" → "German courses B1 on weekends" (219ms) [QueryRewrite] "Let's search for Spanish beginner courses!" → "UNCHANGED" (125ms) ← self-contained, no rewrite [QueryRewrite] "Maybe I'll try B1 after all!" → "Spanish courses B1 advanced" (369ms) [QueryRewrite] "Is that the only one???" → "Spanish courses B1 available" (332ms) ``` **Overhead: 125–370ms per query** (~550 tokens in, 1-15 tokens out). Self-contained queries cost ~125ms (1 token). The rewrite uses the same LLM that's already loaded for chat — no additional model required. ## Configuration Three environment variables, no database changes: | Variable | Default | Description | |----------|---------|-------------| | `DISABLE_QUERY_REWRITING` | `false` | Kill switch to disable entirely | | `QUERY_REWRITE_WORD_THRESHOLD` | `12` | Queries above N words skip rewriting | | `QUERY_REWRITE_MAX_HISTORY` | `2` | Number of turn-pairs used for context | ## Changes | File | Change | |------|--------| | `server/utils/helpers/chat/queryRewriter.js` | **New** — rewrite logic, 3-layer validation, universal script support | | `server/utils/chats/embed.js` | Import + rewrite before vector search | | `server/utils/chats/stream.js` | Import + rewrite before vector search | | `server/utils/chats/apiChatHandler.js` | Import + rewrite before vector search (2 handlers) | | `server/utils/chats/openaiCompatible.js` | Import + rewrite before vector search (2 handlers) | All 6 `performSimilaritySearch` call sites across all chat interfaces (embed widget, workspace UI, API sync, API stream, OpenAI-compatible sync/stream) are covered. ## Why this approach | Alternative | Why not | |------------|---------| | String concatenation (keywords from history) | Fragile, unreliable for real conversations | | Embed history + query together | Dilutes the embedding signal | | Conditional on word count only | Misses context-dependent longer queries ("Is **this** course available on weekends?") | | Always rewrite | Wastes compute on first messages and long self-contained queries | | Trust LLM output directly | Different models produce unpredictable meta-text — needs validation | The LLM-based approach with a word-count threshold is the industry standard (LangChain's `create_history_aware_retriever`, Open WebUI's search query generation). This implementation goes further with selective rewriting (UNCHANGED signal) and model-agnostic output validation (3-layer check). One file, ~115 lines of logic, no dependencies. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
yindo added the pull-request label 2026-06-05 15:21:00 -04:00
yindo closed this issue 2026-06-05 15:21:01 -04:00
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: Mintplex-Labs/anything-llm#5335