[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.

[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).

[junttang]

Note on the failing check types:
uv.lock — generated locally (glibc 2.35); CI re-resolves on glibc 2.39, so the count differs. Running uv lock in CI at merge time will fix it.
ty — all errors in this PR's files are fixed; remaining ty failures are from pre-existing code outside this PR.

[malatewang]

we need address the static type check. And use the RaRag as an optimized ChainOfQueryAgent.

[junttang]

@malatewang Thanks. here're two things.
1. Static type check — the 4 remaining ty failures are all in files this PR doesn't touch (langgraph.py, database_manager.py, nebula_graph_vector_graph_store.py).
2. RaRag positioning — already implemented as an optimized variant of ChainOfQueryAgent swapped in via use_optimized_coq, no ToolSelectAgent/prompt changes.

Looks mergeable from my side, please let me know if anything else is needed.

[malatewang]

@malatewang Thanks. here're two things. 1. Static type check — the 4 remaining ty failures are all in files this PR doesn't touch (langgraph.py, database_manager.py, nebula_graph_vector_graph_store.py). 2. RaRag positioning — already implemented as an optimized variant of ChainOfQueryAgent swapped in via use_optimized_coq, no ToolSelectAgent/prompt changes.

Looks mergeable from my side, please let me know if anything else is needed.

You need to update this file: https://github.com/MemMachine/MemMachine/blob/main/packages/server/src/memmachine_server/retrieval_agent/service_locator.py to use the new agent.

[malatewang]

Need to honor the limit from the parameter.

[junttang]

The fixed 200 cap is intentional: the gold episode has low similarity to the question and its hops, so capping to a small top-K — even with reranking — collapses recall. Rarag's point is to be a faster, cheaper, deterministic alternative to CoT; returning a wider candidate set is what makes it work (equal quality, far fewer tokens than CoT). Capping to query.limit would silently turn Rarag back into plain search.

That said, the tradeoff is real if the answer LLM (which consumes the returned episodes) is the bottleneck. Three options:

• (1) Document and keep the fixed cap — note that Rarag returns up to ~200 episodes and ignores query.limit on the final return; users should raise top-k / context budget when using Rarag.
• (2) Make the cap a config knob — add multi_hop_final_limit (default 200) to optimized_coq and cap to that, still not bound to query.limit.
• (3) Honor query.limit but document the characteristic — cap to query.limit, and clearly warn that small limits degrade Rarag's recall advantage, so users should raise top-k when opting into Rarag.

I'd lean toward (2) or (3) for now. Which do you prefer?

[malatewang]

When the query is splitted by LLM, there could be more than two sub queries.

[junttang]

Thanks for the comment. Rarag is intentionally a 2-hop design — its A / A→C overlap + B'→C mechanism is built around exactly two hops. So 3+ hop questions shouldn't be handled by Rarag.

My plan: keep Rarag strictly 2-hop and route 3+ hop questions to CoT. The cleanest place for that routing is ToolSelectAgent — steer ≤2-hop questions to Rarag and ≥3-hop questions to ChainOfQueryAgent (when RaragQueryAgent is enabled by users), so the boundary is explicit and Rarag stays focused on what it's good at. I'll implement that routing if you agree. Extending Rarag to deeper chains can come later.

[junttang]

Re your service_locator.py comment: RaragQueryAgent currently drops into the ChainOfQueryAgent slot via agent_name matching, so service_locator.py wasn't touched.

I'll make it explicit — register both RaragQueryAgent and ChainOfQueryAgent there. But that change becomes relevant once the 2-hop vs 3+ hop routing in ToolSelectAgent I proposed above is in place, so this would hinge on whether you agree to that routing. If you agree with that direction, I'd bundle it with the routing changes into one update.

[malatewang]

Can you add a comment here so other people know that at most two sub queries is done intentionally.

[malatewang]

It is ok for me to use RaragQueryAgent to replace the ChainOfQueryAgent if configured. But in the agent_utils.py, you did replace the COT slot. But in the service, I am not sure if I misunderstood but I did not see how it worked.

[junttang]

Thanks for the review. Added a comment noting Rarag is a 2-hop design, so only the first two sub-queries are consumed intentionally and any extra LLM-split sub-queries are ignored.

For the server path that was missing in the previous push, I fixed it too: memmachine.py now reads retrieval_agent.use_optimized_coq (plus the sub-configs) and passes them to create_retrieval_agent, which fills the ChainOfQueryAgent slot with RaragQueryAgent when configured and forwards the sub-configs (multi_hop_decomposer / multi_hop_sub_limit) via extra_params, matching the eval path in agent_utils.init_agent.

Note) The failing checks are unrelated to this PR, except for the uv.lock, which will be resolved by running uv lock in CI at merge time.