|
| 1 | +import { MCPClientManager } from 'agents/mcp/client' |
| 2 | +import { streamText, tool, jsonSchema } from 'ai' |
| 3 | +import { z } from 'zod' |
| 4 | + |
| 5 | +import type { LanguageModelV1, StreamTextResult, ToolCallPart, ToolSet } from 'ai' |
| 6 | + |
| 7 | +export async function initializeClient(): Promise<MCPClientManager> { |
| 8 | + const clientManager = new MCPClientManager('test-client', '0.0.0') |
| 9 | + await clientManager.connect('http://localhost:8976/sse') |
| 10 | + return clientManager |
| 11 | +} |
| 12 | + |
| 13 | +export async function runTask( |
| 14 | + clientManager: MCPClientManager, |
| 15 | + model: LanguageModelV1, |
| 16 | + input: string |
| 17 | +): Promise<{ |
| 18 | + promptOutput: string |
| 19 | + fullResult: StreamTextResult<ToolSet, never> |
| 20 | + toolCalls: ToolCallPart[] |
| 21 | +}> { |
| 22 | + const tools = clientManager.listTools() |
| 23 | + const toolSet: ToolSet = tools.reduce((acc, v) => { |
| 24 | + if (!v.inputSchema.properties) { |
| 25 | + v.inputSchema.properties = {} |
| 26 | + } |
| 27 | + |
| 28 | + acc[v.name] = tool({ |
| 29 | + parameters: jsonSchema(v.inputSchema as any), |
| 30 | + description: v.description, |
| 31 | + execute: async (args: any, opts) => { |
| 32 | + try { |
| 33 | + const res = await clientManager.callTool( |
| 34 | + { |
| 35 | + ...v, |
| 36 | + arguments: { ...args }, |
| 37 | + }, |
| 38 | + z.any() as any, |
| 39 | + { signal: opts.abortSignal } |
| 40 | + ) |
| 41 | + return res.content |
| 42 | + } catch (e) { |
| 43 | + console.log('Error calling tool') |
| 44 | + console.log(e) |
| 45 | + return e |
| 46 | + } |
| 47 | + }, |
| 48 | + }) |
| 49 | + return acc |
| 50 | + }, {} as ToolSet) |
| 51 | + |
| 52 | + |
| 53 | + const res = streamText({ |
| 54 | + model, |
| 55 | + system: |
| 56 | + "You are an assistant responsible for evaluating the results of calling various tools. Given the user's query, use the tools available to you to answer the question.", |
| 57 | + tools: toolSet, |
| 58 | + prompt: input, |
| 59 | + maxRetries: 1, |
| 60 | + maxSteps: 10, |
| 61 | + }) |
| 62 | + |
| 63 | + for await (const part of res.fullStream) { |
| 64 | + } |
| 65 | + |
| 66 | + // convert into an LLM readable result so our factuality checker can validate tool calls |
| 67 | + let messagesWithTools = '' |
| 68 | + const toolCalls: ToolCallPart[] = [] |
| 69 | + const response = await res.response |
| 70 | + const messages = response.messages |
| 71 | + |
| 72 | + for (const message of messages) { |
| 73 | + for (const messagePart of message.content) { |
| 74 | + if (typeof messagePart === 'string') { |
| 75 | + messagesWithTools += `<message_content type="text">${messagePart}</message_content>` |
| 76 | + } else if (messagePart.type === 'tool-call') { |
| 77 | + messagesWithTools += `<message_content type=${messagePart.type}> |
| 78 | + <tool_name>${messagePart.toolName}</tool_name> |
| 79 | + <tool_arguments>${JSON.stringify(messagePart.args)}</tool_arguments> |
| 80 | +</message_content>` |
| 81 | + toolCalls.push(messagePart) |
| 82 | + } else if (messagePart.type === 'text') { |
| 83 | + messagesWithTools += `<message_content type=${messagePart.type}>${messagePart.text}</message_content>` |
| 84 | + } |
| 85 | + } |
| 86 | + } |
| 87 | + |
| 88 | + return { promptOutput: messagesWithTools, fullResult: res, toolCalls } |
| 89 | +} |
0 commit comments