[Feat] add RaragQueryAgent, an optimized ChainOfQueryAgent for multi-hop retrieval#1477
[Feat] add RaragQueryAgent, an optimized ChainOfQueryAgent for multi-hop retrieval#1477junttang wants to merge 6 commits into
Conversation
| """ | ||
|
|
||
|
|
||
| class RaragQueryAgent(AgentToolBase): |
There was a problem hiding this comment.
I think we need to plug this agent into the tool select agent.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Since in your test, the rragqueryagent can achieve the same result with lower cost, I think option A should be good.
|
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. |
b0b6993 to
834ad6f
Compare
|
@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). |
f923dba to
e79488d
Compare
Add a spaCy-based rule decomposer as an embedded helper under retrieval_agent/agents/decomposer/. It splits compositional multi-hop questions into hop-level sub-queries without an LLM call, to be used by the upcoming RaragQueryAgent. Signed-off-by: Junhyeok Park <junhyeok0.park@sk.com>
Add RaragQueryAgent, an optimized variant of ChainOfQueryAgent for multi-hop retrieval. It decomposes multi-hop queries (non-LLM decomposer or LLM fallback), searches A and A->C in parallel, finds UID-based overlaps, builds combined queries from the top overlaps, deduplicates by UID, and returns up to 200 episodes. Reports agent_name as "ChainOfQueryAgent" so it can drop into the ChainOfQueryAgent slot of ToolSelectAgent when configured. Registered in the agents package __init__. - Fall back to LLM-based splitting when the decomposer cannot handle the query (not multi-hop per its rules, or it raises). - Fall back to the A->C episodes when A and A->C have no overlap, so a zero-overlap multi-hop question still yields memories. - Demote the per-query "CALLING" log to debug. Signed-off-by: Junhyeok Park <junhyeok0.park@sk.com>
In init_agent, add multi_hop_decomposer / multi_hop_sub_limit / use_optimized_coq params. When use_optimized_coq is true, fill the ChainOfQueryAgent slot with RaragQueryAgent (it reports agent_name as "ChainOfQueryAgent", keeping the slot consistent); otherwise use the original ChainOfQueryAgent. Pass the decomposer/sub-limit params through extra_params. In init_memmachine_params, read use_optimized_coq and the nested optimized_coq config (only when enabled) and forward them to init_agent, logging the RaragQueryAgent setup only when it is active. Signed-off-by: Junhyeok Park <junhyeok0.park@sk.com>
Add retrieval_agent.use_optimized_coq (bool, default false) to opt into the RaragQueryAgent optimized variant, and a nested retrieval_agent.optimized_coq config (OptimizedCoqConf) holding multi_hop_decomposer (bool) and multi_hop_sub_limit (int, default 20) which control RaragQueryAgent's hop splitting and per-sub-search limit. Signed-off-by: Junhyeok Park <junhyeok0.park@sk.com>
d6c5408 to
57b6bbb
Compare
Add spacy>=3.8.0,<3.9.0 and the en_core_web_sm pipeline model so the embedded non-LLM decomposer works out of the box, but keep them OUT of the server package's hard runtime dependencies: - spaCy model wheels aren't on PyPI, so pip-based install jobs (test-server-package, installation-test) cannot resolve en-core-web-sm when it is declared as a server dep. - spaCy 3.8.x only ships cp312/cp313 wheels (no cp314), so a hard dep breaks the Python 3.14 test matrix. Instead, expose them as an optional multihop dependency-group (python_version < '3.14' so 3.14 is skipped), with the en_core_web_sm wheel still resolved via tool.uv.sources (spaCy models aren't on PyPI). Install on demand with: uv sync --group multihop. The decomposer package import is lazy-guarded and RaragQueryAgent already falls back to LLM-based query splitting when spaCy is absent (DECOMPOSER_AVAILABLE), so base installs and the 3.14 matrix run without spaCy. Signed-off-by: Junhyeok Park <junhyeok0.park@sk.com>
- Auto-fix quote style (Q000), import sorting (I001), unused imports (F401), and f-string / unused-variable issues across the decomposer and RaragQueryAgent. - decomposer: annotate spaCy args with Doc/Token instead of Any (ANN401), make the WH/PROPERTY/VERB_PREP pattern class attrs ClassVar tuples (RUF012), collapse nested ifs (SIM102), use list.extend (PERF401), fix docstring formatting (D205/D101/D107/D103), and split the large hop-splitting methods into helpers to bring cyclomatic complexity under the limit (C901). Behavior is unchanged for the supported English two-hop patterns. Signed-off-by: Junhyeok Park <junhyeok0.park@sk.com>
57b6bbb to
3e8f215
Compare
|
Note on the failing check types: |
|
we need address the static type check. And use the RaRag as an optimized ChainOfQueryAgent. |
|
@malatewang Thanks. here're two things. Looks mergeable from my side, please let me know if anything else is needed. |
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.
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)
Evaluation
Benchmark: 2WikiMultiHopQA, MemMachine with top-20 search per question.
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
use_optimized_coq: false(default) keeps the originalChainOfQueryAgent.RaragQueryAgentreportsagent_name == "ChainOfQueryAgent", so it drops into the existing CoQ slot andToolSelectAgentneeds no changes.Fixes/Closes
No
Type of change
[Please delete options that are not relevant.]
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.]
Test Results: see above results.
Checklist
[Please delete options that are not relevant.]
Maintainer Checklist
Further comments
Lint errors will be addressed in a follow-up.