|
1 | 1 | import { MCPClientManager } from 'agents/mcp/client' |
2 | | -import { jsonSchema, streamText, tool } from 'ai' |
3 | | -import { z } from 'zod' |
4 | | - |
5 | | -import type { LanguageModelV1, StreamTextResult, ToolCallPart, ToolSet } from 'ai' |
6 | 2 |
|
7 | 3 | export async function initializeClient(): Promise<MCPClientManager> { |
8 | 4 | const clientManager = new MCPClientManager('test-client', '0.0.0') |
9 | 5 | await clientManager.connect('http://localhost:8977/sse') |
10 | 6 | return clientManager |
11 | 7 | } |
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 | | - const res = streamText({ |
53 | | - model, |
54 | | - system: |
55 | | - "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.", |
56 | | - tools: toolSet, |
57 | | - prompt: input, |
58 | | - maxRetries: 1, |
59 | | - maxSteps: 10, |
60 | | - }) |
61 | | - |
62 | | - // we need to consume the fill stream, so this is empty |
63 | | - // eslint-disable-next-line no-empty |
64 | | - for await (const _ of res.fullStream) { |
65 | | - } |
66 | | - |
67 | | - // convert into an LLM readable result so our factuality checker can validate tool calls |
68 | | - let messagesWithTools = '' |
69 | | - const toolCalls: ToolCallPart[] = [] |
70 | | - const response = await res.response |
71 | | - const messages = response.messages |
72 | | - |
73 | | - for (const message of messages) { |
74 | | - for (const messagePart of message.content) { |
75 | | - if (typeof messagePart === 'string') { |
76 | | - messagesWithTools += `<message_content type="text">${messagePart}</message_content>` |
77 | | - } else if (messagePart.type === 'tool-call') { |
78 | | - messagesWithTools += `<message_content type=${messagePart.type}> |
79 | | - <tool_name>${messagePart.toolName}</tool_name> |
80 | | - <tool_arguments>${JSON.stringify(messagePart.args)}</tool_arguments> |
81 | | -</message_content>` |
82 | | - toolCalls.push(messagePart) |
83 | | - } else if (messagePart.type === 'text') { |
84 | | - messagesWithTools += `<message_content type=${messagePart.type}>${messagePart.text}</message_content>` |
85 | | - } |
86 | | - } |
87 | | - } |
88 | | - |
89 | | - return { promptOutput: messagesWithTools, fullResult: res, toolCalls } |
90 | | -} |
0 commit comments