Router LLM should extend AugmentedLLM to enable composability #369
Conversation
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
WalkthroughThe changes introduce new asynchronous proxy methods to the Changes
Sequence Diagram(s)LLMRouter as AugmentedLLM: Routing Request FlowsequenceDiagram
participant User
participant LLMRouter
participant classifier_llm
participant Router
User->>LLMRouter: generate / generate_str / generate_structured(message)
LLMRouter->>Router: _route_with_llm(message)
Router->>classifier_llm: generate_structured(message)
classifier_llm-->>Router: StructuredResponse
Router-->>LLMRouter: Routing result
LLMRouter-->>User: Result (dict, str, or StructuredResponse)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~18 minutes Poem
Note 🔌 MCP (Model Context Protocol) integration is now available in Early Access!Pro users can now connect to remote MCP servers under the Integrations page to get reviews and chat conversations that understand additional development context. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
src/mcp_agent/workflows/router/router_llm.py (2)
398-425: Exposetop_kinstead of hard-coding 5
generate()always calls_route_with_llm(..., top_k=5), but callers have no way to tune the breadth of routing. Makingtop_ka parameter (or deriving it fromrequest_params) would give the API necessary flexibility.- async def generate( - self, - message: str | Any | List[Any], - request_params: RequestParams | None = None, - ) -> List[Any]: + async def generate( + self, + message: str | Any | List[Any], + *, + top_k: int = 5, + request_params: RequestParams | None = None, + ) -> List[Any]: … - results = await self._route_with_llm(str(message), top_k=5) + results = await self._route_with_llm(str(message), top_k=top_k)
336-337: Span attribute key"llm"may collideUsing the generic key
"llm"for the classifier name could overwrite or conflict with other tracing attributes. A more specific key (e.g.,"router.classifier_llm") avoids ambiguity.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/mcp_agent/workflows/llm/augmented_llm.py(2 hunks)src/mcp_agent/workflows/router/router_llm.py(7 hunks)src/mcp_agent/workflows/router/router_llm_anthropic.py(2 hunks)src/mcp_agent/workflows/router/router_llm_openai.py(2 hunks)
🧰 Additional context used
🧠 Learnings (5)
📓 Common learnings
Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/**/*.py : Use mcp-agent's Agent abstraction for ALL LLM interactions, including quality evaluation, to ensure consistent tool access, logging, and error handling.
📚 Learning: 2025-07-22T18:59:49.368Z
Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/**/*.py : Use mcp-agent's Agent abstraction for ALL LLM interactions, including quality evaluation, to ensure consistent tool access, logging, and error handling.
Applied to files:
src/mcp_agent/workflows/router/router_llm_anthropic.pysrc/mcp_agent/workflows/router/router_llm_openai.pysrc/mcp_agent/workflows/llm/augmented_llm.pysrc/mcp_agent/workflows/router/router_llm.py
📚 Learning: 2025-07-22T18:59:49.368Z
Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/utils/config.py : Configuration values such as quality_threshold, max_refinement_attempts, consolidation_interval, and evaluator_model_provider must be loaded from mcp_agent.config.yaml.
Applied to files:
src/mcp_agent/workflows/router/router_llm_anthropic.pysrc/mcp_agent/workflows/router/router_llm_openai.pysrc/mcp_agent/workflows/router/router_llm.py
📚 Learning: 2025-05-27T19:10:22.911Z
Learnt from: rholinshead
PR: lastmile-ai/mcp-agent#241
File: src/mcp_agent/mcp/mcp_agent_client_session.py:174-187
Timestamp: 2025-05-27T19:10:22.911Z
Learning: In the MCP Agent codebase, all request parameter types inherit from RequestParams or NotificationParams with consistent Meta class structures, so trace context injection using the base Meta classes doesn't cause type-mismatch issues.
Applied to files:
src/mcp_agent/workflows/router/router_llm.py
📚 Learning: 2025-05-27T19:10:22.911Z
Learnt from: rholinshead
PR: lastmile-ai/mcp-agent#241
File: src/mcp_agent/mcp/mcp_agent_client_session.py:174-187
Timestamp: 2025-05-27T19:10:22.911Z
Learning: In MCP Agent trace context injection, creating params objects when trace context exists is justified because param types either require other non-meta parameters (so params must exist) or are optional (so creating them for trace context is appropriate).
Applied to files:
src/mcp_agent/workflows/router/router_llm.py
🧬 Code Graph Analysis (3)
src/mcp_agent/workflows/router/router_llm_anthropic.py (1)
src/mcp_agent/workflows/llm/augmented_llm.py (1)
AugmentedLLM(220-698)
src/mcp_agent/workflows/router/router_llm_openai.py (1)
src/mcp_agent/workflows/llm/augmented_llm.py (1)
AugmentedLLM(220-698)
src/mcp_agent/workflows/llm/augmented_llm.py (4)
src/mcp_agent/agents/agent.py (2)
list_tools(394-485)close(302-307)src/mcp_agent/mcp/mcp_aggregator.py (2)
list_tools(611-654)close(198-279)tests/mcp/test_mcp_aggregator.py (1)
list_tools(786-793)src/mcp_agent/logging/transport.py (1)
close(186-188)
🔇 Additional comments (1)
src/mcp_agent/workflows/router/router_llm.py (1)
114-122: Double-passing**kwargsrisks collision
__init__forwards the same**kwargsto bothRouter.__init__andAugmentedLLM.__init__.
If either base mutates or consumes a key the second call may receive unexpected/duplicate parameters, raisingTypeError. Consider splitting or explicitly listing shared kwargs.
Comment on lines
+697
to
+698
| async def __aexit__(self, exc_type, exc_val, exc_tb): | ||
| await self.agent.__aexit__(exc_type, exc_val, exc_tb) |
There was a problem hiding this comment.
Propagate the underlying __aexit__ return value
Agent.__aexit__() may return True to suppress an exception; dropping that value changes behaviour to “always re-raise”. Forward the result explicitly:
- async def __aexit__(self, exc_type, exc_val, exc_tb):
- await self.agent.__aexit__(exc_type, exc_val, exc_tb)
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
+ return await self.agent.__aexit__(exc_type, exc_val, exc_tb)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def __aexit__(self, exc_type, exc_val, exc_tb): | |
| await self.agent.__aexit__(exc_type, exc_val, exc_tb) | |
| async def __aexit__(self, exc_type, exc_val, exc_tb): | |
| return await self.agent.__aexit__(exc_type, exc_val, exc_tb) |
🤖 Prompt for AI Agents
In src/mcp_agent/workflows/llm/augmented_llm.py around lines 697 to 698, the
async __aexit__ method calls self.agent.__aexit__ but does not return its
result. Modify the method to return the awaited result of self.agent.__aexit__
so that any True value indicating exception suppression is properly propagated.
Comment on lines
+53
to
54
| llm: AugmentedLLM | None = None, | ||
| server_names: List[str] | None = None, |
There was a problem hiding this comment.
🛠️ Refactor suggestion
llm parameter is unused
create() now accepts llm: AugmentedLLM | None but doesn’t pass it to cls(...) or store it, so callers cannot inject a custom classifier LLM. Either:
- Wire it through (
cls(llm=llm, …)) and addllmto__init__, or - Drop the parameter until it’s needed.
Leaving it unused will trigger linters and confuse API consumers.
🤖 Prompt for AI Agents
In src/mcp_agent/workflows/router/router_llm_anthropic.py around lines 53 to 54,
the create() method has an llm parameter that is not used or passed to the class
constructor, causing confusion and linter warnings. To fix this, either pass llm
to cls(...) in create() and add llm as an __init__ parameter to store it
properly, or remove the llm parameter from create() until it is actually needed.
Comment on lines
+53
to
54
| llm: AugmentedLLM | None = None, | ||
| server_names: List[str] | None = None, |
There was a problem hiding this comment.
🛠️ Refactor suggestion
llm parameter currently serves no purpose
The optional llm argument is accepted but ignored. Align implementation with intent (thread it into the instance or remove it) to avoid dead parameters and static-analysis warnings.
🤖 Prompt for AI Agents
In src/mcp_agent/workflows/router/router_llm_openai.py around lines 53 to 54,
the llm parameter is accepted but not used anywhere in the code, causing dead
code and static-analysis warnings. To fix this, either remove the llm parameter
entirely if it is not needed, or if it is intended to be used, assign it to an
instance attribute or otherwise integrate it into the class or function logic so
it serves a purpose.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Leading up to a wider revamp that removes the need to use both
AgentandAugmentedLLMclasses (stay tuned)Summary by CodeRabbit
New Features
Chores