|
| 1 | +import { jsonSchemaToZod } from '@n8n/json-schema-to-zod' |
| 2 | +import { MCPClientManager } from 'agents/mcp/client' |
| 3 | +import { LanguageModelV1, streamText, tool, ToolSet } from 'ai' |
| 4 | + |
| 5 | +import type { JsonSchemaObject } from '@n8n/json-schema-to-zod' |
| 6 | + |
| 7 | +export async function runTask(model: LanguageModelV1, input: string) { |
| 8 | + const clientManager = new MCPClientManager('test-client', '0.0.0') |
| 9 | + await clientManager.connect('http://localhost:8787/sse') |
| 10 | + |
| 11 | + const tools = clientManager.listTools() |
| 12 | + const toolSet: ToolSet = tools.reduce((acc, v) => { |
| 13 | + acc[v.name] = tool({ |
| 14 | + parameters: jsonSchemaToZod(v.inputSchema as JsonSchemaObject), |
| 15 | + description: v.description, |
| 16 | + execute: async (args, opts) => { |
| 17 | + const res = await clientManager.callTool(v, args, { signal: opts.abortSignal }) |
| 18 | + console.log(res.toolResult) |
| 19 | + return res.content |
| 20 | + }, |
| 21 | + }) |
| 22 | + return acc |
| 23 | + }, {} as ToolSet) |
| 24 | + |
| 25 | + const res = streamText({ |
| 26 | + model, |
| 27 | + system: |
| 28 | + "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.", |
| 29 | + tools: toolSet, |
| 30 | + prompt: input, |
| 31 | + maxRetries: 1, |
| 32 | + maxSteps: 10, |
| 33 | + }) |
| 34 | + |
| 35 | + for await (const part of res.fullStream) { |
| 36 | + } |
| 37 | + |
| 38 | + // convert into an LLM readable result so our factuality checker can validate tool calls |
| 39 | + let messagesWithTools = '' |
| 40 | + const messages = (await res.response).messages |
| 41 | + for (const message of messages) { |
| 42 | + console.log(message.content) |
| 43 | + for (const messagePart of message.content) { |
| 44 | + if (typeof messagePart === 'string') { |
| 45 | + messagesWithTools += `<message_content type="text">${messagePart}</message_content>` |
| 46 | + } else if (messagePart.type === 'tool-call') { |
| 47 | + messagesWithTools += `<message_content type=${messagePart.type}> |
| 48 | + <tool_name>${messagePart.toolName}</tool_name> |
| 49 | + <tool_arguments>${JSON.stringify(messagePart.args)}</tool_arguments> |
| 50 | +</message_content>` |
| 51 | + } else if (messagePart.type === 'text') { |
| 52 | + messagesWithTools += `<message_content type=${messagePart.type}>${messagePart.text}</message_content>` |
| 53 | + } |
| 54 | + } |
| 55 | + } |
| 56 | + |
| 57 | + return messagesWithTools |
| 58 | +} |
0 commit comments