[Feat] add RaragQueryAgent, an optimized ChainOfQueryAgent for multi-hop retrieval

[junttang]

Purpose of the change

Adds RaragQueryAgent, an optimized variant of ChainOfQueryAgent for multi-hop retrieval — Rarag = Retrieval-Augmented Retrieval-Augmented Generation, i.e. using retrieval results to drive further retrieval (search-then-search) before generation.
It is positioned as a cheaper drop-in for the ChainOfQueryAgent, and includes an embedded non-LLM (spaCy) decomposer, the config to opt into it, and the evaluation wiring.

Description

Motivation

Multi-hop reasoning targets the structure A -> B' -> C = D, where B is a missing fact not stated in the question.

• Example: "Where was the father of the founder of Versus born?"
  • A: "founder of Versus", B': "father of", C: "Where was ... born?", D: the birthplace
  • B (e.g. "Gianni Versace") must be recovered first, then the second hop B -> C — "Where was Gianni Versace's father born?" — is the query that actually surfaces D.

The gold evidence D is only weakly retrieved by a plain search on the original question A -> B' -> C, because the original question and the gold episode have low similarity. Retrieving D requires a query that explicitly contains the missing fact — i.e. the second hop B -> C.

ChainOfQueryAgent (CoT) reaches this with an LLM: search the original question, then iteratively rewrite the query toward B -> C using LLM-driven sufficiency checks until D is found. This works, but every iteration spends LLM tokens on long prompts.

Key observation

On the 2WikiMultiHopQA benchmark we observed that splitting the original question at B' into two hops and searching both independently surfaces B with high probability. In particular, the overlap between the original-question search and the first-hop A -> B' search consistently concentrates the episodes that contain the missing fact B — so combining each overlapping episode with B' -> C as a new query reliably raises the chance of retrieving D.

Approach (RaragQueryAgent)

1. Decompose the question into hops — LLM-based splitting by default; the non-LLM spaCy decomposer is an opt-in alternative that applies only when the question is grammar-correct English and a two-hop split is possible.
2. Search the first hop A and the original question A -> C in parallel (top-N each).
3. Find the UID-based overlap between the two result sets (these are the strongest candidates for containing B).
4. For each top overlap, build a combined query overlap_content + (B' -> C) and search it.
5. Deduplicate by UID and return up to 200 episodes.

Evaluation

Benchmark: 2WikiMultiHopQA, MemMachine with top-20 search per question.

Method	LLM Score (%)	Recall (%)	Avg. Episodes Cnt.	Avg. Answer LLM I/O Token	Total Runtime (s)
Baseline (Top-20)	~67	~58	20	~1200	~1000
Baseline (Top-100)	~75	~65	100	—	—
CoT (Top-20)	~90	~77	20	~13300	~3200
CoT (Top-100)	~90	~80	100	—	—
Rarag — LLM decomposer (Top-20)	~60	~55	~20	~5500	~2300
Rarag — LLM decomposer (Top-k X)	~90	~80	~94	~9950	~2300
Rarag — non-LLM decomposer (Top-k X)	~90	~80	~72	~2500	~1600

Limitations

Capping the final return to a small top-K (e.g. Top-20) collapses recall: even with reranking on the final set, the gold evidence does not reliably rank near the top, because the original question (and the hops) have low similarity to the gold turn. We tried reranking the final results and reranking the intermediate overlaps; both failed for the same root cause. The approach's clear win is over CoT — equal quality, far cheaper — so it is positioned as an optimization of CoT, not of plain MemMachine search.

The non-LLM decomposer currently supports English only and with only two-hops.

How to use

retrieval_agent:
  use_optimized_coq: true
  optimized_coq:
    multi_hop_decomposer: true   # non-LLM decomposer; false (default) -> LLM-based splitting
    multi_hop_sub_limit: 20       # per-sub-search limit; 20 by default (independent of user top-k)

use_optimized_coq: false (default) keeps the original ChainOfQueryAgent. RaragQueryAgent reports agent_name == "ChainOfQueryAgent", so it drops into the existing CoQ slot and ToolSelectAgent needs no changes.

Fixes/Closes

No

Type of change

[Please delete options that are not relevant.]

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Refactor (does not change functionality, e.g., code style improvements, linting)
• Documentation update
• Project Maintenance (updates to build scripts, CI, etc., that do not affect the main project)
• Security (improves security without changing functionality)

How Has This Been Tested?

The evaluation results above were obtained solely from evaluation/retrieval_agent against 2WikiMultiHopQA, using the first 213 compositional-type questions.

Set retrieval_agent.use_optimized_coq: true (with the optimized_coq sub-config) in configuration.yml to enable RaragQueryAgent; false falls back to the original ChainOfQueryAgent for baseline/CoT comparison.

[Please delete options that are not relevant.]

• Unit Test
• Integration Test
• End-to-end Test
• Test Script (please provide)
• Manual verification (list step-by-step instructions)

Test Results: see above results.

Checklist

[Please delete options that are not relevant.]

• I have signed the commit(s) within this pull request
• My code follows the style guidelines of this project (See STYLE_GUIDE.md)
• I have performed a self-review of my own code
• I have commented my code
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added unit tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• Any dependent changes have been merged and published in downstream modules
• I have checked my code and corrected any misspellings

Maintainer Checklist

• Confirmed all checks passed
• Contributor has signed the commit(s)
• Reviewed the code
• Run, Tested, and Verified the change(s) work as expected

Further comments

Lint errors will be addressed in a follow-up.

[malatewang]

Put this to debug

[junttang]

Done in commit chore(rarag): demote per-query CALLING log to debug.

[junttang]

I squashed the referenced commit to the original one.

[malatewang]

My understanding is that the decomposer only can handle certain predefined multihop patterns. If the decomposer can not handle the query, should it fall back to the LLM?

[junttang]

Done in fix(rarag): fall back to LLM when the decomposer cannot handle a query.
The decomposer now falls back to the LLM-based splitter in both cases — when it reports the query as not multi-hop, and when it raises.

[honggyukim]

@junttang You better not to introduce a new commit for this modification since the previous change was introduced in this PR. Please squash e8c026b into c02438a where the change was originally introduced.

[junttang]

@honggyukim I squashed the referenced commit to the original one.

[malatewang]

I think we need to plug this agent into the tool select agent.

[junttang]

I think here're two options.

(A) Keep the current design (drop-in variant): RaragQueryAgent reuses the existing ChainOfQueryAgent slot — it reports agent_name == "ChainOfQueryAgent", and use_optimized_coq swaps the slot in init_agent. ToolSelectAgent and its prompt stay unchanged. This was the original PR intent: position Rarag as an optimized variant of ChainOfQueryAgent rather than a separate tool.

(B) Promote it to a peer tool: add RaragQueryAgent as a 4th selectable tool in ToolSelectAgent. This requires updating the tool-selection prompt, e.g. route 2-hop multi-hop questions to Rarag and deeper chains (≥3 hops) to ChainOfQueryAgent.

I see either is fine; I'll defer to your preference. If you'd like (B), I'll update the prompt and wiring in a follow-up.

[honggyukim]

As mentioned #1477 (comment), b0b6993 is not needed as a separate commit.

If you reduce the number of patches, then it will be easier to be reviewed.

[junttang]

@honggyukim Thanks for the advice! I squashed those commits responding to @malatewang 's comments to the original commit feat(retrieval_agent): add RaragQueryAgent (optimized ChainOfQueryAgent).