ai 6.0.33 → 6.0.35

package/src/agent/tool-loop-agent.test.ts

@@ -0,0 +1,442 @@
+import { LanguageModelV3CallOptions } from '@ai-sdk/provider';
+import { convertArrayToReadableStream } from '@ai-sdk/provider-utils/test';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { MockLanguageModelV3 } from '../test/mock-language-model-v3';
+import { ToolLoopAgent } from './tool-loop-agent';
+
+describe('ToolLoopAgent', () => {
+  describe('generate', () => {
+    let doGenerateOptions: LanguageModelV3CallOptions | undefined;
+    let mockModel: MockLanguageModelV3;
+
+    beforeEach(() => {
+      doGenerateOptions = undefined;
+      mockModel = new MockLanguageModelV3({
+        doGenerate: async options => {
+          doGenerateOptions = options;
+          return {
+            content: [{ type: 'text', text: 'reply' }],
+            finishReason: { unified: 'stop', raw: 'stop' },
+            usage: {
+              cachedInputTokens: undefined,
+              inputTokens: {
+                total: 3,
+                noCache: 3,
+                cacheRead: undefined,
+                cacheWrite: undefined,
+              },
+              outputTokens: {
+                total: 10,
+                text: 10,
+                reasoning: undefined,
+              },
+            },
+            warnings: [],
+          };
+        },
+      });
+    });
+
+    it('should use prepareCall', async () => {
+      const agent = new ToolLoopAgent<{ value: string }>({
+        model: mockModel,
+        prepareCall: ({ options, ...rest }) => {
+          return {
+            ...rest,
+            providerOptions: {
+              test: { value: options.value },
+            },
+          };
+        },
+      });
+
+      await agent.generate({
+        prompt: 'Hello, world!',
+        options: { value: 'test' },
+      });
+
+      expect(doGenerateOptions?.providerOptions).toMatchInlineSnapshot(`
+        {
+          "test": {
+            "value": "test",
+          },
+        }
+      `);
+    });
+
+    it('should pass abortSignal to generateText', async () => {
+      const abortController = new AbortController();
+
+      const agent = new ToolLoopAgent({ model: mockModel });
+
+      await agent.generate({
+        prompt: 'Hello, world!',
+        abortSignal: abortController.signal,
+      });
+
+      expect(doGenerateOptions?.abortSignal).toBe(abortController.signal);
+    });
+
+    it('should pass timeout to generateText', async () => {
+      const agent = new ToolLoopAgent({ model: mockModel });
+
+      await agent.generate({
+        prompt: 'Hello, world!',
+        timeout: 5000,
+      });
+
+      // timeout is merged into abortSignal, so we check that an abort signal was created
+      expect(doGenerateOptions?.abortSignal).toBeDefined();
+    });
+
+    it('should pass experimental_download to generateText', async () => {
+      const downloadFunction = vi
+        .fn()
+        .mockResolvedValue([
+          { data: new Uint8Array([1, 2, 3]), mediaType: 'image/png' },
+        ]);
+
+      const agent = new ToolLoopAgent({
+        model: mockModel,
+        experimental_download: downloadFunction,
+      });
+
+      await agent.generate({
+        prompt: [
+          {
+            role: 'user',
+            content: [
+              {
+                type: 'image',
+                image: new URL('https://example.com/image.png'),
+              },
+            ],
+          },
+        ],
+      });
+
+      expect(downloadFunction).toHaveBeenCalledWith([
+        {
+          url: new URL('https://example.com/image.png'),
+          isUrlSupportedByModel: false,
+        },
+      ]);
+    });
+
+    describe('instructions', () => {
+      it('should pass string instructions', async () => {
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+          instructions: 'INSTRUCTIONS',
+        });
+
+        await agent.generate({
+          prompt: 'Hello, world!',
+        });
+
+        expect(doGenerateOptions?.prompt).toMatchInlineSnapshot(`
+          [
+            {
+              "content": "INSTRUCTIONS",
+              "role": "system",
+            },
+            {
+              "content": [
+                {
+                  "text": "Hello, world!",
+                  "type": "text",
+                },
+              ],
+              "providerOptions": undefined,
+              "role": "user",
+            },
+          ]
+        `);
+      });
+
+      it('should pass system message instructions', async () => {
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+          instructions: {
+            role: 'system',
+            content: 'INSTRUCTIONS',
+            providerOptions: { test: { value: 'test' } },
+          },
+        });
+
+        await agent.generate({
+          prompt: 'Hello, world!',
+        });
+
+        expect(doGenerateOptions?.prompt).toMatchInlineSnapshot(`
+        [
+          {
+            "content": "INSTRUCTIONS",
+            "providerOptions": {
+              "test": {
+                "value": "test",
+              },
+            },
+            "role": "system",
+          },
+          {
+            "content": [
+              {
+                "text": "Hello, world!",
+                "type": "text",
+              },
+            ],
+            "providerOptions": undefined,
+            "role": "user",
+          },
+        ]
+      `);
+      });
+
+      it('should pass array of system message instructions', async () => {
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+          instructions: [
+            {
+              role: 'system',
+              content: 'INSTRUCTIONS',
+              providerOptions: { test: { value: 'test' } },
+            },
+            {
+              role: 'system',
+              content: 'INSTRUCTIONS 2',
+              providerOptions: { test: { value: 'test 2' } },
+            },
+          ],
+        });
+
+        await agent.generate({
+          prompt: 'Hello, world!',
+        });
+
+        expect(doGenerateOptions?.prompt).toMatchInlineSnapshot(`
+          [
+            {
+              "content": "INSTRUCTIONS",
+              "providerOptions": {
+                "test": {
+                  "value": "test",
+                },
+              },
+              "role": "system",
+            },
+            {
+              "content": "INSTRUCTIONS 2",
+              "providerOptions": {
+                "test": {
+                  "value": "test 2",
+                },
+              },
+              "role": "system",
+            },
+            {
+              "content": [
+                {
+                  "text": "Hello, world!",
+                  "type": "text",
+                },
+              ],
+              "providerOptions": undefined,
+              "role": "user",
+            },
+          ]
+        `);
+      });
+    });
+  });
+
+  describe('stream', () => {
+    let doStreamOptions: LanguageModelV3CallOptions | undefined;
+    let mockModel: MockLanguageModelV3;
+
+    beforeEach(() => {
+      doStreamOptions = undefined;
+      mockModel = new MockLanguageModelV3({
+        doStream: async options => {
+          doStreamOptions = options;
+          return {
+            stream: convertArrayToReadableStream([
+              {
+                type: 'stream-start',
+                warnings: [],
+              },
+              {
+                type: 'response-metadata',
+                id: 'id-0',
+                modelId: 'mock-model-id',
+                timestamp: new Date(0),
+              },
+              { type: 'text-start', id: '1' },
+              { type: 'text-delta', id: '1', delta: 'Hello' },
+              { type: 'text-delta', id: '1', delta: ', ' },
+              { type: 'text-delta', id: '1', delta: `world!` },
+              { type: 'text-end', id: '1' },
+              {
+                type: 'finish',
+                finishReason: { unified: 'stop', raw: 'stop' },
+                usage: {
+                  inputTokens: {
+                    total: 3,
+                    noCache: 3,
+                    cacheRead: undefined,
+                    cacheWrite: undefined,
+                  },
+                  outputTokens: {
+                    total: 10,
+                    text: 10,
+                    reasoning: undefined,
+                  },
+                },
+                providerMetadata: {
+                  testProvider: { testKey: 'testValue' },
+                },
+              },
+            ]),
+          };
+        },
+      });
+    });
+
+    it('should use prepareCall', async () => {
+      const agent = new ToolLoopAgent<{ value: string }>({
+        model: mockModel,
+        prepareCall: ({ options, ...rest }) => {
+          return {
+            ...rest,
+            providerOptions: {
+              test: { value: options.value },
+            },
+          };
+        },
+      });
+
+      const result = await agent.stream({
+        prompt: 'Hello, world!',
+        options: { value: 'test' },
+      });
+
+      await result.consumeStream();
+
+      expect(doStreamOptions?.providerOptions).toMatchInlineSnapshot(
+        `
+        {
+          "test": {
+            "value": "test",
+          },
+        }
+      `,
+      );
+    });
+
+    it('should pass abortSignal to streamText', async () => {
+      const abortController = new AbortController();
+
+      const agent = new ToolLoopAgent({
+        model: mockModel,
+      });
+
+      const result = await agent.stream({
+        prompt: 'Hello, world!',
+        abortSignal: abortController.signal,
+      });
+
+      await result.consumeStream();
+
+      expect(doStreamOptions?.abortSignal).toBe(abortController.signal);
+    });
+
+    it('should pass timeout to streamText', async () => {
+      const agent = new ToolLoopAgent({
+        model: mockModel,
+      });
+
+      const result = await agent.stream({
+        prompt: 'Hello, world!',
+        timeout: 5000,
+      });
+
+      await result.consumeStream();
+
+      // timeout is merged into abortSignal, so we check that an abort signal was created
+      expect(doStreamOptions?.abortSignal).toBeDefined();
+    });
+
+    it('should pass string instructions', async () => {
+      const agent = new ToolLoopAgent({
+        model: mockModel,
+        instructions: 'INSTRUCTIONS',
+      });
+
+      const result = await agent.stream({
+        prompt: 'Hello, world!',
+      });
+
+      await result.consumeStream();
+
+      expect(doStreamOptions?.prompt).toMatchInlineSnapshot(`
+        [
+          {
+            "content": "INSTRUCTIONS",
+            "role": "system",
+          },
+          {
+            "content": [
+              {
+                "text": "Hello, world!",
+                "type": "text",
+              },
+            ],
+            "providerOptions": undefined,
+            "role": "user",
+          },
+        ]
+      `);
+    });
+
+    it('should pass system message instructions', async () => {
+      const agent = new ToolLoopAgent({
+        model: mockModel,
+        instructions: {
+          role: 'system',
+          content: 'INSTRUCTIONS',
+          providerOptions: { test: { value: 'test' } },
+        },
+      });
+
+      const result = await agent.stream({
+        prompt: 'Hello, world!',
+      });
+
+      await result.consumeStream();
+
+      expect(doStreamOptions?.prompt).toMatchInlineSnapshot(`
+      [
+        {
+          "content": "INSTRUCTIONS",
+          "providerOptions": {
+            "test": {
+              "value": "test",
+            },
+          },
+          "role": "system",
+        },
+        {
+          "content": [
+            {
+              "text": "Hello, world!",
+              "type": "text",
+            },
+          ],
+          "providerOptions": undefined,
+          "role": "user",
+        },
+      ]
+    `);
+    });
+  });
+});

package/src/agent/tool-loop-agent.ts

@@ -0,0 +1,114 @@
+import { generateText } from '../generate-text/generate-text';
+import { GenerateTextResult } from '../generate-text/generate-text-result';
+import { Output } from '../generate-text/output';
+import { stepCountIs } from '../generate-text/stop-condition';
+import { streamText } from '../generate-text/stream-text';
+import { StreamTextResult } from '../generate-text/stream-text-result';
+import { ToolSet } from '../generate-text/tool-set';
+import { Prompt } from '../prompt';
+import { Agent, AgentCallParameters, AgentStreamParameters } from './agent';
+import { ToolLoopAgentSettings } from './tool-loop-agent-settings';
+
+/**
+ * A tool loop agent is an agent that runs tools in a loop. In each step,
+ * it calls the LLM, and if there are tool calls, it executes the tools
+ * and calls the LLM again in a new step with the tool results.
+ *
+ * The loop continues until:
+ * - A finish reasoning other than tool-calls is returned, or
+ * - A tool that is invoked does not have an execute function, or
+ * - A tool call needs approval, or
+ * - A stop condition is met (default stop condition is stepCountIs(20))
+ */
+export class ToolLoopAgent<
+  CALL_OPTIONS = never,
+  TOOLS extends ToolSet = {},
+  OUTPUT extends Output = never,
+> implements Agent<CALL_OPTIONS, TOOLS, OUTPUT>
+{
+  readonly version = 'agent-v1';
+
+  private readonly settings: ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>;
+
+  constructor(settings: ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>) {
+    this.settings = settings;
+  }
+
+  /**
+   * The id of the agent.
+   */
+  get id(): string | undefined {
+    return this.settings.id;
+  }
+
+  /**
+   * The tools that the agent can use.
+   */
+  get tools(): TOOLS {
+    return this.settings.tools as TOOLS;
+  }
+
+  private async prepareCall(
+    options: AgentCallParameters<CALL_OPTIONS>,
+  ): Promise<
+    Omit<
+      ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>,
+      'prepareCall' | 'instructions'
+    > &
+      Prompt
+  > {
+    const baseCallArgs = {
+      ...this.settings,
+      stopWhen: this.settings.stopWhen ?? stepCountIs(20),
+      ...options,
+    };
+
+    const preparedCallArgs =
+      (await this.settings.prepareCall?.(baseCallArgs)) ?? baseCallArgs;
+
+    const { instructions, messages, prompt, ...callArgs } = preparedCallArgs;
+
+    return {
+      ...callArgs,
+
+      // restore prompt types
+      ...({ system: instructions, messages, prompt } as Prompt),
+    };
+  }
+
+  /**
+   * Generates an output from the agent (non-streaming).
+   */
+  async generate({
+    abortSignal,
+    timeout,
+    ...options
+  }: AgentCallParameters<CALL_OPTIONS>): Promise<
+    GenerateTextResult<TOOLS, OUTPUT>
+  > {
+    return generateText({
+      ...(await this.prepareCall(options)),
+      abortSignal,
+      timeout,
+    });
+  }
+
+  /**
+   * Streams an output from the agent (streaming).
+   */
+  async stream({
+    abortSignal,
+    timeout,
+    experimental_transform,
+    ...options
+  }: AgentStreamParameters<CALL_OPTIONS, TOOLS>): Promise<
+    StreamTextResult<TOOLS, OUTPUT>
+  > {
+    return streamText({
+      ...(await this.prepareCall(options)),
+      abortSignal,
+      timeout,
+      experimental_transform,
+    });
+  }
+}