Skip to content

Commit e1daa91

Browse files
cmsparkscmsparks
committed
Add ModelHeuristicRouter
1 parent e3cd235 commit e1daa91

File tree

3 files changed

+252
-33
lines changed

3 files changed

+252
-33
lines changed

‎.changeset/sharp-bears-appear.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"agents": patch
3+
---
4+
5+
Add ContextRouter to allow users to provide an LLM with a list of filtered tools and resources

‎packages/agents/src/mcp/client.ts‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,10 @@ export class MCPClientManager {
3434
constructor(
3535
private _name: string,
3636
private _version: string,
37-
public unstable_contextRouter = new BaseContextRouter(),
38-
) {}
37+
public unstable_contextRouter: ContextRouter = new BaseContextRouter()
38+
) {
39+
this.unstable_contextRouter.clientManager = this;
40+
}
3941

4042
/**
4143
* Connect to and register an MCP server
@@ -204,21 +206,21 @@ export class MCPClientManager {
204206
* @returns namespaced list of tools
205207
*/
206208
listTools(): NamespacedData["tools"] {
207-
return this.unstable_contextRouter.listTools(this)
209+
return this.unstable_contextRouter.listTools();
208210
}
209211

210212
/**
211213
* @returns a set of tools that you can use with the AI SDK
212214
*/
213215
unstable_getAITools(): ToolSet {
214-
return this.unstable_contextRouter.getAITools(this)
216+
return this.unstable_contextRouter.getAITools();
215217
}
216218

217219
/**
218220
* @returns namespaced list of tools
219221
*/
220222
listResources(): NamespacedData["resources"] {
221-
return this.unstable_contextRouter.listResources(this);
223+
return this.unstable_contextRouter.listResources();
222224
}
223225

224226
/**
@@ -288,7 +290,7 @@ export class MCPClientManager {
288290
* @param includeResources Whether to include resources in the prompt. This may be useful if you include a tool to fetch resources OR if the servers you connect to interact with resources.
289291
*/
290292
unstable_systemPrompt(): string {
291-
return this.unstable_contextRouter.systemPrompt(this)
293+
return this.unstable_contextRouter.systemPrompt();
292294
}
293295
}
294296

‎packages/agents/src/mcp/context-router.ts‎

Lines changed: 239 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,43 +1,106 @@
1-
import { jsonSchema, type CoreMessage, type ToolSet, type LanguageModelV1, type UserContent } from "ai";
2-
import { getNamespacedData, type MCPClientManager, type NamespacedData } from "./client";
1+
import {
2+
jsonSchema,
3+
type CoreMessage,
4+
type ToolSet,
5+
generateText,
6+
type LanguageModelV1,
7+
} from "ai";
8+
import {
9+
getNamespacedData,
10+
type MCPClientManager,
11+
type NamespacedData,
12+
} from "./client";
313
import type { MCPClientConnection } from "./client-connection";
4-
import type { Resource } from "@modelcontextprotocol/sdk/types.js"
14+
import type { Resource } from "@modelcontextprotocol/sdk/types.js";
15+
import { z } from "zod";
516

617
/**
718
* A context router provides an interface for:
819
* - Filtering tools and resources
9-
* - Exposes a system prompt based on the MCP state
20+
* - Exposes a system prompt based on MCP state
21+
*
22+
* This filtering and system prompt is intended to be determined by the internal messages state, set via
23+
* setMessages()
1024
*/
11-
export interface ContextRouter {
12-
systemPrompt(connManager: MCPClientManager): string
13-
14-
listTools(connManager: MCPClientManager): NamespacedData["tools"]
15-
getAITools(connManager: MCPClientManager): ToolSet
16-
listResources(connManager: MCPClientManager): NamespacedData["resources"]
25+
export abstract class ContextRouter {
26+
private _clientManager: MCPClientManager | undefined;
27+
28+
get clientManager() {
29+
if (!this._clientManager) {
30+
throw new Error(
31+
"Tried to get the client manager before it was set. Is the ContextRouter being correctly used in an MCPClientManager?"
32+
);
33+
}
34+
return this.clientManager;
35+
}
36+
37+
set clientManager(clientManager: MCPClientManager) {
38+
this._clientManager = clientManager;
39+
}
40+
41+
/**
42+
* Return the system prompt
43+
* @param clientManager
44+
*/
45+
abstract systemPrompt(): string;
46+
47+
/**
48+
* Set internal context router state. This state is intended to allow
49+
* the ContextRouter to filter tools, resources, or construct the
50+
* system prompt.
51+
*
52+
* @param messages
53+
*/
54+
abstract setMessages(messages: CoreMessage[]): void | Promise<void>;
55+
56+
/**
57+
* List tools from the client manager based on internal context router state
58+
* @param clientManager
59+
*/
60+
abstract listTools(): NamespacedData["tools"];
61+
62+
/**
63+
* List AI tools from the client manager based on internal context router state. This could include
64+
* "synthetic" tools not sourced from an MCP server, such as a `read_resource` or `list_resources` tool.
65+
*
66+
* @param clientManager
67+
*/
68+
abstract getAITools(): ToolSet;
69+
70+
/**
71+
* List tools from the client manager based on internal context router state
72+
* @param clientManager
73+
*/
74+
abstract listResources(): NamespacedData["resources"];
1775
}
1876

1977
/**
20-
* The base ContextRouter:
78+
* The BaseContextRouter:
2179
* - Does not filter tools or resources
22-
* - Provides a system prompt
80+
* - Exposes the setMessages interface, but setting it is a no-op
81+
* - Provides a system prompt based on ALL tools & resources
2382
*/
24-
export class BaseContextRouter implements ContextRouter {
25-
constructor(public includeResources = true) {}
83+
export class BaseContextRouter extends ContextRouter {
84+
constructor(private includeResources = true) {
85+
super();
86+
}
87+
88+
setMessages(_messages: CoreMessage[]): void {}
2689

27-
listTools(connManager: MCPClientManager) {
28-
return getNamespacedData(connManager.mcpConnections, "tools");
90+
listTools() {
91+
return getNamespacedData(this.clientManager.mcpConnections, "tools");
2992
}
3093

31-
getAITools(connManager: MCPClientManager): ToolSet {
94+
getAITools(): ToolSet {
3295
return Object.fromEntries(
33-
this.listTools(connManager).map((tool) => {
96+
this.listTools().map((tool) => {
3497
return [
3598
`${tool.serverId}_${tool.name}`,
3699
{
37100
parameters: jsonSchema(tool.inputSchema),
38101
description: tool.description,
39102
execute: async (args) => {
40-
const result = await connManager.callTool({
103+
const result = await this.clientManager.callTool({
41104
name: tool.name,
42105
arguments: args,
43106
serverId: tool.serverId,
@@ -51,23 +114,23 @@ export class BaseContextRouter implements ContextRouter {
51114
},
52115
];
53116
})
54-
)
117+
);
55118
}
56119

57-
listResources(connManager: MCPClientManager) {
58-
return getNamespacedData(connManager.mcpConnections, "resources");
120+
listResources() {
121+
return getNamespacedData(this.clientManager.mcpConnections, "resources");
59122
}
60123

61-
systemPrompt(connManager: MCPClientManager): string {
124+
systemPrompt(): string {
62125
return `<integrations_list>
63126
You have access to multiple integrations via Model Context Protocol (MCP). These integrations provide you with tools which you can use to execute to complete tasks or retrieive information.
64127
65128
${this.includeResources && "Each integration, provides a list of resources, which are included in the list of integrations below."}
66129
67130
Here is a list of all of the integrations you have access to, with instructions if necessary:
68131
69-
${Object.entries(connManager.mcpConnections).map(([_id, conn]) => BaseContextRouter.serverContext(conn, this.includeResources))}
70-
<integrations_list>`
132+
${Object.entries(this.clientManager.mcpConnections).map(([_id, conn]) => BaseContextRouter.serverContext(conn, this.includeResources))}
133+
<integrations_list>`;
71134
}
72135

73136
static serverContext(conn: MCPClientConnection, includeResources: boolean) {
@@ -77,7 +140,7 @@ export class BaseContextRouter implements ContextRouter {
77140
${includeResources && `<resources_list>${conn.resources.map((resource) => BaseContextRouter.resourceContext(resource))}</resources_list>`}
78141
<integration>`;
79142
}
80-
143+
81144
static resourceContext(resource: Resource) {
82145
return `<resource>
83146
<name>${resource.name}</name>
@@ -86,4 +149,153 @@ export class BaseContextRouter implements ContextRouter {
86149
<mimeType>${resource.mimeType}</mimeType>
87150
</resource>`;
88151
}
89-
}
152+
}
153+
154+
/**
155+
* The ModelHeuristicRouter:
156+
* - Filters tools and resources using an LLM
157+
*/
158+
export class ModelHeuristicRouter extends BaseContextRouter {
159+
private tools: NamespacedData["tools"];
160+
private resources: NamespacedData["resources"];
161+
162+
constructor(
163+
private model: LanguageModelV1,
164+
private options: {
165+
toolLimit: number;
166+
resourceLimit: number;
167+
} = { toolLimit: 10, resourceLimit: 5 }
168+
) {
169+
super(true);
170+
this.tools = getNamespacedData(this.clientManager.mcpConnections, "tools");
171+
this.resources = getNamespacedData(
172+
this.clientManager.mcpConnections,
173+
"resources"
174+
);
175+
}
176+
177+
async setMessages(messages: CoreMessage[]): Promise<void> {
178+
this.tools = [];
179+
this.resources = [];
180+
181+
const tools = super.listTools();
182+
const resources = super.listResources();
183+
184+
await Promise.all([
185+
generateText({
186+
model: this.model,
187+
system: this.toolSelectionPrompt(tools),
188+
messages,
189+
tools: {
190+
add_tool: {
191+
parameters: z.object({
192+
tool_name: z.string().describe("The exact tool name to add"),
193+
}),
194+
description: "Add a tool to context",
195+
execute: async ({ tool_name }) => {
196+
const toolToAdd = tools.find((tool) =>
197+
tool.name.includes(tool_name)
198+
);
199+
if (!toolToAdd) {
200+
return "Failed to find tool by the name `tool_name`";
201+
}
202+
if (this.tools.length >= this.options.toolLimit) {
203+
return "Failed to add tool. There are already too many active tools.";
204+
}
205+
this.tools.push(toolToAdd);
206+
return "Successfully added tool";
207+
},
208+
},
209+
},
210+
}),
211+
generateText({
212+
model: this.model,
213+
system: this.resourceSelectionPrompt(resources),
214+
messages,
215+
tools: {
216+
add_resource: {
217+
parameters: z.object({
218+
resource_uri: z.string(),
219+
}),
220+
description: "Add a tool to context",
221+
execute: async ({ resource_uri }, options) => {
222+
const resourceToAdd = resources.find(
223+
(resource) => resource.uri === resource_uri
224+
);
225+
if (!resourceToAdd) {
226+
return `Failed to find resource by the uri "${resource_uri}"`;
227+
}
228+
if (this.resources.length >= this.options.resourceLimit) {
229+
return "Failed to add resource. There are already too many active resources.";
230+
}
231+
this.resources.push(resourceToAdd);
232+
return "Successfully added resource";
233+
},
234+
},
235+
},
236+
}),
237+
]);
238+
}
239+
240+
listTools() {
241+
return this.tools;
242+
}
243+
244+
listResources() {
245+
return this.resources;
246+
}
247+
248+
toolSelectionPrompt(tools: NamespacedData["tools"]) {
249+
return `
250+
You are an expert at selecting relevant tools for a conversation with a large language model. You are given access to only one tool: "add_tool".
251+
252+
You will be provided a list of tools, their descriptions, and a series of messages below.
253+
Based on the messages, call the tool "add_tool" with the exact tool_name of the tool you would like to add.
254+
255+
* You MUST use "add_tool" with only the tool_name argument
256+
* You SHOULD NOT ask for permission to use the "add_tool" call.
257+
* You MUST NOT respond to the attached messages below, ONLY use the system prompt
258+
* You MUST NOT attempt to call any other tools
259+
* Limit the amount of tools selected to ${this.options.toolLimit}
260+
261+
<tool_options>
262+
${tools.map((tool) => {
263+
return `
264+
<tool>
265+
<name>${tool.name}</name>
266+
<description>${tool.description}</description>
267+
</tool>
268+
`;
269+
})}
270+
</tool_options>
271+
`;
272+
}
273+
274+
resourceSelectionPrompt(resources: NamespacedData["resources"]) {
275+
return `
276+
You are an expert at selecting relevant resources for a conversation with a large language model. You are given access to only one tool: "add_resource".
277+
278+
You will be provided a list of resources, their descriptions, and a series of messages below.
279+
Based on the messages, call the tool "add_resource" with the exact resource_uri of the resource you would like to add.
280+
281+
* You MUST use "add_resource" with only the resource_uri argument
282+
* You SHOULD NOT ask for permission to use the "add_resource" call.
283+
* You MUST NOT respond to the attached messages below, ONLY use the system prompt
284+
* You MUST NOT attempt to call any other tools
285+
* Limit the amount of resources selected to ${this.options.resourceLimit}
286+
287+
<resource_options>
288+
${resources.map((resource) => {
289+
return `
290+
<resource>
291+
<uri>${resource.uri}</uri>
292+
<name>${resource.name}</name>
293+
<mimeType>${resource.mimeType}</mimeType>
294+
<description>${resource.description}</description>
295+
</resource>
296+
`;
297+
})}
298+
</resource_options>
299+
`;
300+
}
301+
}

0 commit comments

Comments
 (0)