Skip to content

feat: add LLM-based query router for RAG, MCP search, and MCP analysis - #142

Closed
GovindhKishore wants to merge 4 commits into
reactome:mainfrom
GovindhKishore:feat/mcp-query-routing
Closed

GovindhKishore wants to merge 4 commits into
reactome:mainfrom
GovindhKishore:feat/mcp-query-routing

Conversation

@GovindhKishore

Copy link
Copy Markdown

Summary

Adds LLM-based query routing to ReactToMeGraphBuilder so questions are directed to the correct retrieval path before any tool calls are made.

Closes #141

Context

PR #127 and #137 introduced MCP tools into the agent. Without routing, every question goes through the tool calling loop even when the vector database already has the answer. This PR adds a lightweight classification step that runs before retrieval and directs each question to the right path.

What Changed

src/mcp/query_router.py (new)

  • create_query_router(llm) returns an async route() function
  • Classifies questions into rag, mcp_search, or mcp_analysis
  • Uses a structured prompt with explicit criteria for each category
  • Falls back to rag on unexpected output
  • Intended to be called with a lightweight model (gpt-4o-mini)

src/agent/profiles/react_to_me.py (modified)

  • Two route-specific LLM instances built at init time - llm_with_search_tools
    and llm_with_analysis_tools - each bound to the relevant tool subset only
  • call_model now calls the router first and branches accordingly:
    • rag - existing RAG path, no MCP involved
    • mcp_search - LLM with search/lookup tools only
    • mcp_analysis - LLM with analyze_identifiers only
  • When MCP_SERVER_PATH is not set, routing is skipped entirely and existing
    RAG behaviour is unchanged

Routing Logic

rag - general knowledge questions answerable from static embeddings
mcp_search - live lookup by identifier, name, species list, or db metadata
mcp_analysis - pathway enrichment on an explicit list of gene/protein identifiers
               (requires both analysis intent AND a list - a list alone routes to rag)

Testing

Manual testing against a running MCP server is pending.

The routing logic has been reviewed against the prompt criteria for the following question types:

  • General knowledge questions → rag
  • Lookup by identifier or name → mcp_search
  • Enrichment analysis on a gene list → mcp_analysis
  • Ambiguous or unclear questions → rag fallback

Notes

  • Router uses gpt-4o-mini separately from the main generation model to keep
    classification cost low
  • Tool subset filtering happens at init time via bind_tools - no rebinding
    per message
  • Routing is only active when MCP_SERVER_PATH is set - zero impact on
    deployments without MCP

AI assistance was used in drafting and implementation. All changes reviewed and verified
by me.

@adamjohnwright

Copy link
Copy Markdown
Contributor

Closing this one, and I want to be clear about which part is the problem — because the idea is right and it is going into the design.

What is right: questions genuinely do split three ways. A general pathway question, a lookup of a specific identifier, and "analyse this gene list" want different machinery, and your three-way split with the rule that mcp_analysis needs both analysis intent and a list of identifiers is a sharper distinction than we had written down anywhere. That framing is now in specs/007-answer-cascade.

What is the problem: query_router.py is a second LLM classifier standing beside src/agent/tasks/intent_classifier.py, which already classifies every question. I put the two prompts side by side — they do the same job, they just have different destination sets.

That is a second model call on every question. Spec 001's retriever rewrite cut this pipeline from 21 model calls to 1, and adding one back per question gives up a meaningful part of that for a decision the existing classifier is already positioned to make. specs/007-answer-cascade FR-001 now says routing MUST extend the existing classifier rather than acquire a sibling, and that requirement exists because of this PR.

So the destination set you designed here is what we want; it should arrive as extra options on intent_classifier, not as a second classifier.

The other files in this PR are #137's, which is the one still open for the MCP client and tool wrappers.

One thing worth carrying over whichever way this lands: the package name src/mcp/ will collide with the official MCP Python SDK, which is published as mcp. Nothing breaks today because that dependency is not installed, but it will the day someone adds it.

Thank you — the classification design here is being used, with credit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: route questions to RAG or MCP tools based on query intent

2 participants