- "content": "import {\n AgentConfig,\n AbstractAgent,\n EventType,\n BaseEvent,\n Message,\n AssistantMessage,\n RunAgentInput,\n MessagesSnapshotEvent,\n RunFinishedEvent,\n RunStartedEvent,\n TextMessageChunkEvent,\n ToolCallArgsEvent,\n ToolCallEndEvent,\n ToolCallStartEvent,\n ToolCall,\n ToolMessage,\n} from \"@ag-ui/client\";\nimport { Observable } from \"rxjs\";\nimport {\n CoreMessage,\n LanguageModelV1,\n processDataStream,\n streamText,\n tool as createVercelAISDKTool,\n ToolChoice,\n ToolSet,\n} from \"ai\";\nimport { randomUUID } from \"crypto\";\nimport { z } from \"zod\";\n\ntype ProcessedEvent =\n | MessagesSnapshotEvent\n | RunFinishedEvent\n | RunStartedEvent\n | TextMessageChunkEvent\n | ToolCallArgsEvent\n | ToolCallEndEvent\n | ToolCallStartEvent;\n\ninterface VercelAISDKAgentConfig extends AgentConfig {\n model: LanguageModelV1;\n maxSteps?: number;\n toolChoice?: ToolChoice<Record<string, unknown>>;\n}\n\nexport class VercelAISDKAgent extends AbstractAgent {\n model: LanguageModelV1;\n maxSteps: number;\n toolChoice: ToolChoice<Record<string, unknown>>;\n constructor({ model, maxSteps, toolChoice, ...rest }: VercelAISDKAgentConfig) {\n super({ ...rest });\n this.model = model;\n this.maxSteps = maxSteps ?? 1;\n this.toolChoice = toolChoice ?? \"auto\";\n }\n\n protected run(input: RunAgentInput): Observable<BaseEvent> {\n const finalMessages: Message[] = input.messages;\n\n return new Observable<ProcessedEvent>((subscriber) => {\n subscriber.next({\n type: EventType.RUN_STARTED,\n threadId: input.threadId,\n runId: input.runId,\n } as RunStartedEvent);\n\n const response = streamText({\n model: this.model,\n messages: convertMessagesToVercelAISDKMessages(input.messages),\n tools: convertToolToVerlAISDKTools(input.tools),\n maxSteps: this.maxSteps,\n toolChoice: this.toolChoice,\n });\n\n let messageId = randomUUID();\n let assistantMessage: AssistantMessage = {\n id: messageId,\n role: \"assistant\",\n content: \"\",\n toolCalls: [],\n };\n finalMessages.push(assistantMessage);\n\n processDataStream({\n stream: response.toDataStreamResponse().body!,\n onTextPart: (text) => {\n assistantMessage.content += text;\n const event: TextMessageChunkEvent = {\n type: EventType.TEXT_MESSAGE_CHUNK,\n role: \"assistant\",\n messageId,\n delta: text,\n };\n subscriber.next(event);\n },\n onFinishMessagePart: () => {\n // Emit message snapshot\n const event: MessagesSnapshotEvent = {\n type: EventType.MESSAGES_SNAPSHOT,\n messages: finalMessages,\n };\n subscriber.next(event);\n\n // Emit run finished event\n subscriber.next({\n type: EventType.RUN_FINISHED,\n threadId: input.threadId,\n runId: input.runId,\n } as RunFinishedEvent);\n\n // Complete the observable\n subscriber.complete();\n },\n onToolCallPart(streamPart) {\n let toolCall: ToolCall = {\n id: streamPart.toolCallId,\n type: \"function\",\n function: {\n name: streamPart.toolName,\n arguments: JSON.stringify(streamPart.args),\n },\n };\n assistantMessage.toolCalls!.push(toolCall);\n\n const startEvent: ToolCallStartEvent = {\n type: EventType.TOOL_CALL_START,\n parentMessageId: messageId,\n toolCallId: streamPart.toolCallId,\n toolCallName: streamPart.toolName,\n };\n subscriber.next(startEvent);\n\n const argsEvent: ToolCallArgsEvent = {\n type: EventType.TOOL_CALL_ARGS,\n toolCallId: streamPart.toolCallId,\n delta: JSON.stringify(streamPart.args),\n };\n subscriber.next(argsEvent);\n\n const endEvent: ToolCallEndEvent = {\n type: EventType.TOOL_CALL_END,\n toolCallId: streamPart.toolCallId,\n };\n subscriber.next(endEvent);\n },\n onToolResultPart(streamPart) {\n const toolMessage: ToolMessage = {\n role: \"tool\",\n id: randomUUID(),\n toolCallId: streamPart.toolCallId,\n content: JSON.stringify(streamPart.result),\n };\n finalMessages.push(toolMessage);\n },\n onErrorPart(streamPart) {\n subscriber.error(streamPart);\n },\n }).catch((error) => {\n console.error(\"catch error\", error);\n // Handle error\n subscriber.error(error);\n });\n\n return () => {};\n });\n }\n}\n\nexport function convertMessagesToVercelAISDKMessages(messages: Message[]): CoreMessage[] {\n const result: CoreMessage[] = [];\n\n for (const message of messages) {\n if (message.role === \"assistant\") {\n const parts: any[] = message.content ? [{ type: \"text\", text: message.content }] : [];\n for (const toolCall of message.toolCalls ?? []) {\n parts.push({\n type: \"tool-call\",\n toolCallId: toolCall.id,\n toolName: toolCall.function.name,\n args: JSON.parse(toolCall.function.arguments),\n });\n }\n result.push({\n role: \"assistant\",\n content: parts,\n });\n } else if (message.role === \"user\") {\n result.push({\n role: \"user\",\n content: message.content || \"\",\n });\n } else if (message.role === \"tool\") {\n let toolName = \"unknown\";\n for (const msg of messages) {\n if (msg.role === \"assistant\") {\n for (const toolCall of msg.toolCalls ?? []) {\n if (toolCall.id === message.toolCallId) {\n toolName = toolCall.function.name;\n break;\n }\n }\n }\n }\n result.push({\n role: \"tool\",\n content: [\n {\n type: \"tool-result\",\n toolCallId: message.toolCallId,\n toolName: toolName,\n result: message.content,\n },\n ],\n });\n }\n }\n\n return result;\n}\n\nexport function convertJsonSchemaToZodSchema(jsonSchema: any, required: boolean): z.ZodSchema {\n if (jsonSchema.type === \"object\") {\n const spec: { [key: string]: z.ZodSchema } = {};\n\n if (!jsonSchema.properties || !Object.keys(jsonSchema.properties).length) {\n return !required ? z.object(spec).optional() : z.object(spec);\n }\n\n for (const [key, value] of Object.entries(jsonSchema.properties)) {\n spec[key] = convertJsonSchemaToZodSchema(\n value,\n jsonSchema.required ? jsonSchema.required.includes(key) : false,\n );\n }\n let schema = z.object(spec).describe(jsonSchema.description);\n return required ? schema : schema.optional();\n } else if (jsonSchema.type === \"string\") {\n let schema = z.string().describe(jsonSchema.description);\n return required ? schema : schema.optional();\n } else if (jsonSchema.type === \"number\") {\n let schema = z.number().describe(jsonSchema.description);\n return required ? schema : schema.optional();\n } else if (jsonSchema.type === \"boolean\") {\n let schema = z.boolean().describe(jsonSchema.description);\n return required ? schema : schema.optional();\n } else if (jsonSchema.type === \"array\") {\n let itemSchema = convertJsonSchemaToZodSchema(jsonSchema.items, true);\n let schema = z.array(itemSchema).describe(jsonSchema.description);\n return required ? schema : schema.optional();\n }\n throw new Error(\"Invalid JSON schema\");\n}\n\nexport function convertToolToVerlAISDKTools(tools: RunAgentInput[\"tools\"]): ToolSet {\n return tools.reduce(\n (acc: ToolSet, tool: RunAgentInput[\"tools\"][number]) => ({\n ...acc,\n [tool.name]: createVercelAISDKTool({\n description: tool.description,\n parameters: convertJsonSchemaToZodSchema(tool.parameters, true),\n }),\n }),\n {},\n );\n}\n",
0 commit comments