feat: implement query decomposition for search and enhance reranking … - #266
Merged
Conversation
📝 WalkthroughWalkthroughIntroduces multi-query support via a new SearchQueries model, replaces AsyncOpenAI with ChatOpenAI plus structured function output for query generation, adds batch retrieval (get_relevant_docs), adds RRF reranking and tests, updates prompt template for query decomposition, and extends mock vLLM to simulate tool calls. Error message formatting in on_message adjusted. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant App as App Frontend
participant QueryGen as QueryGenerator<br/>(ChatOpenAI)
participant Retriever as RetrieverPipeline
participant Reranker as Reranker (RRF)
participant Storage as Document Store
User->>App: Send message
App->>QueryGen: call generate_query(messages)
QueryGen->>QueryGen: detect_language + inject current date
QueryGen-->>App: return SearchQueries (multiple queries)
App->>Retriever: get_relevant_docs(partition, SearchQueries, top_k)
loop per sub-query
Retriever->>Storage: retrieve documents for sub-query
Storage-->>Retriever: documents
end
Retriever->>Reranker: rrf_reranking(doc_lists)
Reranker->>Reranker: fuse rankings, deduplicate
Reranker-->>Retriever: reranked documents
Retriever-->>App: aggregated ranked documents
App-->>User: return results
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
📝 Coding Plan for PR comments
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Ahmath-Gadji
force-pushed
the
feat/query_decomposition_for_search
branch
2 times, most recently
from
March 5, 2026 08:35
6b6fc2e to
5de9a99
Compare
Ahmath-Gadji
marked this pull request as ready for review
March 5, 2026 08:42
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
openrag/components/pipeline.py (1)
84-85: Alignpartitiontype annotation with actual usage.
get_relevant_docsdeclarespartition: str, but it is used/called aslist[str]in this pipeline path.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/pipeline.py` around lines 84 - 85, The partition parameter on get_relevant_docs is annotated as str but callers pass list[str]; update the function signature for get_relevant_docs to accept a sequence of strings (e.g., partition: list[str] or Sequence[str]) and adjust any internal handling to iterate over partitions instead of treating it as a single string; ensure related references to partition inside the function (and any callers expecting a str) are updated to use the new type and import typing.Sequence if chosen.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@openrag/components/pipeline.py`:
- Around line 86-90: The RRF fusion currently returns a combined list that can
exceed the intended top_k budget; after calling
self.reranker.rrf_reranking(doc_lists=results) you must enforce the top_k cap by
deduplicating results (by document id or unique key) and trimming to the top_k
highest-scoring entries before returning. Update the code around retrieve_docs,
self.reranker.rrf_reranking, and the return path to collapse duplicates, sort by
the reranker score, and slice to top_k so the final returned list respects the
top_k parameter.
- Line 21: The import uses a relative import for utils (importing
detect_language and format_context); replace that relative import with an
absolute import from the package root (openrag.utils) so the module imports
detect_language and format_context via the absolute package path rather than
using a leading dot; update any similar relative imports in the same module if
present to follow the same absolute-import pattern.
- Around line 131-133: The code is incorrectly putting model-specific parameters
into a RunnableConfig-like dict named params (setting
params["max_completion_tokens"] and params["extra_body"]) which won't be
forwarded to the OpenAI model; update the call site that constructs params in
the Pipeline (and the similar block around the second occurrence) to pass model
kwargs via the model's .bind() or by setting them directly on the model instance
(use self.model.bind(max_completion_tokens=self.max_contextualized_query_len,
extra_body={"chat_template_kwargs": {"enable_thinking": False}}) or set
equivalent attributes on the model before execution) instead of placing them
into params/config so the OpenAI API receives them.
In `@prompts/example1/query_contextualizer_tmpl.txt`:
- Line 15: Fix the typo in the template string inside
query_contextualizer_tmpl.txt by replacing "carbone footprint" with "carbon
footprint" and update the phrase "in the last 2" to "in the last 2 years" so the
example query reads "Evolution of carbon footprint in the last 2 years" (this
will produce two sub-queries as intended); locate and edit the example line
containing the quoted phrase to apply this change.
---
Nitpick comments:
In `@openrag/components/pipeline.py`:
- Around line 84-85: The partition parameter on get_relevant_docs is annotated
as str but callers pass list[str]; update the function signature for
get_relevant_docs to accept a sequence of strings (e.g., partition: list[str] or
Sequence[str]) and adjust any internal handling to iterate over partitions
instead of treating it as a single string; ensure related references to
partition inside the function (and any callers expecting a str) are updated to
use the new type and import typing.Sequence if chosen.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9e6b1640-0673-4c9c-aa1e-bf8cad421c65
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (6)
openrag/app_front.pyopenrag/components/pipeline.pyopenrag/components/reranker.pyopenrag/components/test_rrf_reranking.pyprompts/example1/query_contextualizer_tmpl.txttests/api_tests/api_run/mock_vllm.py
…functionality # Conflicts: # openrag/components/pipeline.py
…le, add trailing newline
EnjoyBacon7
force-pushed
the
feat/query_decomposition_for_search
branch
from
March 12, 2026 13:41
5de9a99 to
7e55e4f
Compare
Collaborator
|
Bug:
Fixed in rebase commit 55af5af — signature updated to |
Collaborator
|
Bug: The Fixed in rebase commit 55af5af — |
Collaborator
|
Web search multi-query strategy: Option C (one search per sub-query, concurrent) For requests combining RAG + web search, the implementation runs one Alternatives considered:
Implemented in rebase commit 55af5af. |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/api_tests/api_run/mock_vllm.py (1)
240-242: Emit more than one query forSearchQueriesarrays.Because every array field is hard-coded to
[user_text], this path still drives API tests through a single retrieval request. For the newSearchQueries.queriesschema, that means the API path never exercises the multi-query branch or any cross-query fusion behavior. A deterministic 2–3 item payload here would give the new flow real coverage.Also applies to: 283-285
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/api_run/mock_vllm.py` around lines 240 - 242, The current mock generator sets every array property to a single-item list ([user_text]), which prevents exercising multi-query behavior (e.g., SearchQueries.queries); update the logic in the loop that builds mock_args (the for prop_name, prop_schema in properties.items() block) to produce a deterministic multi-item array (2–3 entries) for array-typed schemas—use repeated or slightly varied values derived from user_text so tests deterministically hit the multi-query and cross-query fusion branches; apply the same change to the similar block around lines handling properties at the other occurrence (the block referenced as also applies to: 283-285).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@openrag/components/pipeline.py`:
- Around line 183-184: The current debug log prints raw generated queries
(queries from generate_query) which may contain PII; change the logging in the
code that calls generate_query so it does NOT include the full queries
string—use a safe summary instead (e.g., log the count of queries and a
redacted/hashed summary or the length of each query). Specifically, update the
logger.debug call that references queries (the variable returned by
generate_query) to emit only non-sensitive metadata (counts, lengths, or a
redacted preview) or call a helper redact_summary(queries) before logging.
- Around line 152-176: The query generator's returned SearchQueries object must
have its query_list normalized before return: after calling
self.query_generator.bind(...).ainvoke(messages), take output.query_list, strip
whitespace from each entry, remove empty strings, deduplicate while preserving
order, and if the resulting list is empty, replace it with a single-item list
containing the original user query (capture the original user message content
before you overwrite messages—e.g., save the pre-format user query from
messages[-1]["content"] or chat_history). Update output.query_list with the
cleaned list and then return output.
---
Nitpick comments:
In `@tests/api_tests/api_run/mock_vllm.py`:
- Around line 240-242: The current mock generator sets every array property to a
single-item list ([user_text]), which prevents exercising multi-query behavior
(e.g., SearchQueries.queries); update the logic in the loop that builds
mock_args (the for prop_name, prop_schema in properties.items() block) to
produce a deterministic multi-item array (2–3 entries) for array-typed
schemas—use repeated or slightly varied values derived from user_text so tests
deterministically hit the multi-query and cross-query fusion branches; apply the
same change to the similar block around lines handling properties at the other
occurrence (the block referenced as also applies to: 283-285).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: f9eddf5f-1ce2-434e-b3dc-887efdede102
📒 Files selected for processing (6)
openrag/app_front.pyopenrag/components/pipeline.pyopenrag/components/reranker.pyopenrag/components/test_rrf_reranking.pyprompts/example1/query_contextualizer_tmpl.txttests/api_tests/api_run/mock_vllm.py
🚧 Files skipped from review as they are similar to previous changes (2)
- prompts/example1/query_contextualizer_tmpl.txt
- openrag/components/test_rrf_reranking.py
Comment on lines
+152
to
+176
| query_language = detect_language(messages[-1]["content"]) | ||
|
|
||
| model_kwargs = { | ||
| "max_completion_tokens": self.max_contextualized_query_len, | ||
| "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, | ||
| } | ||
| prompt = QUERY_CONTEXTUALIZER_PROMPT.format( | ||
| query_language=query_language, | ||
| current_date=datetime.now().strftime("%Y-%m-%d"), | ||
| ) | ||
| contextualized_query = response.choices[0].message.content | ||
| return contextualized_query | ||
|
|
||
| messages = [ | ||
| { | ||
| "role": "system", | ||
| "content": prompt, | ||
| }, | ||
| { | ||
| "role": "user", | ||
| "content": f"Here is the chat history: \n{chat_history}\n", | ||
| }, | ||
| ] | ||
|
|
||
| # generate queries based on the chat history | ||
| output: SearchQueries = await self.query_generator.bind(**model_kwargs).ainvoke(messages) | ||
| return output |
There was a problem hiding this comment.
Normalize the generated query_list before returning it.
The structured output is used verbatim downstream. Blank or duplicate items trigger redundant retrieval/web-search calls and repeated queries get extra weight in RRF; an empty list skips retrieval entirely. Strip, dedupe, and fall back to the original user query here.
Suggested fix
- query_language = detect_language(messages[-1]["content"])
+ original_query = messages[-1]["content"]
+ query_language = detect_language(original_query)
@@
- messages = [
+ query_messages = [
{
"role": "system",
"content": prompt,
},
{
"role": "user",
"content": f"Here is the chat history: \n{chat_history}\n",
},
]
@@
- output: SearchQueries = await self.query_generator.bind(**model_kwargs).ainvoke(messages)
- return output
+ output: SearchQueries = await self.query_generator.bind(**model_kwargs).ainvoke(query_messages)
+
+ normalized_queries = []
+ seen = set()
+ for query in output.query_list:
+ normalized_query = query.strip()
+ if normalized_query and normalized_query not in seen:
+ seen.add(normalized_query)
+ normalized_queries.append(normalized_query)
+
+ return SearchQueries(query_list=normalized_queries or [original_query])🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@openrag/components/pipeline.py` around lines 152 - 176, The query generator's
returned SearchQueries object must have its query_list normalized before return:
after calling self.query_generator.bind(...).ainvoke(messages), take
output.query_list, strip whitespace from each entry, remove empty strings,
deduplicate while preserving order, and if the resulting list is empty, replace
it with a single-item list containing the original user query (capture the
original user message content before you overwrite messages—e.g., save the
pre-format user query from messages[-1]["content"] or chat_history). Update
output.query_list with the cleaned list and then return output.
Comment on lines
+183
to
+184
| queries: SearchQueries = await self.generate_query(messages) | ||
| logger.debug("Prepared query for chat completion", queries=str(queries)) |
There was a problem hiding this comment.
Don’t log raw generated queries.
These strings are derived from user chat history and can contain PII or secrets. Log counts or a redacted summary instead of the full query text.
Suggested fix
- logger.debug("Prepared query for chat completion", queries=str(queries))
+ logger.debug(
+ "Prepared query for chat completion",
+ query_count=len(queries.query_list),
+ decomposed=len(queries.query_list) > 1,
+ )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@openrag/components/pipeline.py` around lines 183 - 184, The current debug log
prints raw generated queries (queries from generate_query) which may contain
PII; change the logging in the code that calls generate_query so it does NOT
include the full queries string—use a safe summary instead (e.g., log the count
of queries and a redacted/hashed summary or the length of each query).
Specifically, update the logger.debug call that references queries (the variable
returned by generate_query) to emit only non-sensitive metadata (counts,
lengths, or a redacted preview) or call a helper redact_summary(queries) before
logging.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR Improves retrieval quality and reranking by adding multi-query generation and RRF-based result fusion.
Added
SearchQueriesmodel to generate multiple sub-queries from chat history when needed to handle advances queries that requires decomposition.Updated pipeline to retrieve with multiple queries and merge results using Reciprocal Rank Fusion (RRF).
Added tests for RRF logic.
Summary by CodeRabbit
New Features
Documentation
Tests
Bug Fixes