@convex-dev/agent 0.2.11-alpha.2 → 0.2.11-alpha.4

package/src/UIMessages.test.ts

@@ -0,0 +1,273 @@
+import { describe, it, expect } from "vitest";
+import { combineUIMessages } from "./UIMessages.js";
+
+describe("combineUIMessages", () => {
+  it("combines messages spanning two pages correctly", () => {
+    const message1 = {
+      id: "msg1",
+      key: "msg1-key",
+      _creationTime: Date.now(),
+      order: 1,
+      stepOrder: 1,
+      status: "success" as const,
+      role: "assistant" as const,
+      text: "",
+      parts: [
+        {
+          type: "dynamic-tool" as const,
+          state: "input-available" as const,
+          toolCallId: "call_123",
+          toolName: "calculator",
+          input: { operation: "add", a: 2, b: 3 },
+        },
+      ],
+    };
+
+    const message2 = {
+      id: "msg2",
+      key: "msg2-key",
+      _creationTime: Date.now() + 1,
+      order: 1,
+      stepOrder: 2,
+      status: "success" as const,
+      role: "assistant" as const,
+      text: "The result is 5.",
+      parts: [
+        {
+          type: "tool-calculator" as const,
+          state: "output-available" as const,
+          toolCallId: "call_123",
+          input: { operation: "add", a: 2, b: 3 },
+          output: { result: 5 },
+        },
+        {
+          type: "text" as const,
+          text: "The result is 5.",
+          state: "done" as const,
+        },
+      ],
+    };
+
+    const combined = combineUIMessages([message1, message2]);
+
+    expect(combined).toHaveLength(1);
+    expect(combined[0].role).toBe("assistant");
+    expect(combined[0].text).toBe("The result is 5.");
+    expect(combined[0].parts).toHaveLength(2);
+
+    const toolPart = combined[0].parts.find(
+      (p) => p.type === "tool-calculator",
+    );
+    expect(toolPart).toMatchObject({
+      type: "tool-calculator",
+      state: "output-available",
+      toolCallId: "call_123",
+      input: { operation: "add", a: 2, b: 3 },
+      output: { result: 5 },
+    });
+
+    const textPart = combined[0].parts.find((p) => p.type === "text");
+    expect(textPart).toMatchObject({
+      type: "text",
+      text: "The result is 5.",
+      state: "done",
+    });
+  });
+
+  it("preserves separate messages with different roles", () => {
+    const userMessage = {
+      id: "user1",
+      key: "user1-key",
+      _creationTime: Date.now(),
+      order: 1,
+      stepOrder: 0,
+      status: "success" as const,
+      role: "user" as const,
+      text: "Calculate 2 + 3",
+      parts: [{ type: "text" as const, text: "Calculate 2 + 3" }],
+    };
+
+    const assistantMessage = {
+      id: "assistant1",
+      key: "assistant1-key",
+      _creationTime: Date.now() + 1,
+      order: 2,
+      stepOrder: 0,
+      status: "success" as const,
+      role: "assistant" as const,
+      text: "The result is 5.",
+      parts: [
+        {
+          type: "text" as const,
+          text: "The result is 5.",
+          state: "done" as const,
+        },
+      ],
+    };
+
+    const combined = combineUIMessages([userMessage, assistantMessage]);
+
+    expect(combined).toHaveLength(2);
+    expect(combined[0]).toEqual(userMessage);
+    expect(combined[1]).toEqual(assistantMessage);
+  });
+
+  it("combines multiple tool calls across pages", () => {
+    const message1 = {
+      id: "msg1",
+      key: "msg1-key",
+      _creationTime: Date.now(),
+      order: 1,
+      stepOrder: 1,
+      status: "success" as const,
+      role: "assistant" as const,
+      text: "",
+      parts: [
+        {
+          type: "dynamic-tool" as const,
+          state: "input-available" as const,
+          toolCallId: "call_1",
+          toolName: "calculator",
+          input: { operation: "add", a: 2, b: 3 },
+        },
+        {
+          type: "dynamic-tool" as const,
+          state: "input-available" as const,
+          toolCallId: "call_2",
+          toolName: "formatter",
+          input: { text: "result" },
+        },
+      ],
+    };
+
+    const message2 = {
+      id: "msg2",
+      key: "msg2-key",
+      _creationTime: Date.now() + 1,
+      order: 1,
+      stepOrder: 2,
+      status: "success" as const,
+      role: "assistant" as const,
+      text: "The formatted result is: 5",
+      parts: [
+        {
+          type: "tool-calculator" as const,
+          state: "output-available" as const,
+          toolCallId: "call_1",
+          input: { operation: "add", a: 2, b: 3 },
+          output: { result: 5 },
+        },
+        {
+          type: "tool-formatter" as const,
+          state: "output-available" as const,
+          toolCallId: "call_2",
+          input: { text: "result" },
+          output: { formatted: "The formatted result is: 5" },
+        },
+        {
+          type: "text" as const,
+          text: "The formatted result is: 5",
+          state: "done" as const,
+        },
+      ],
+    };
+
+    const combined = combineUIMessages([message1, message2]);
+
+    expect(combined).toHaveLength(1);
+    expect(combined[0].role).toBe("assistant");
+    expect(combined[0].text).toBe("The formatted result is: 5");
+    expect(combined[0].parts).toHaveLength(3);
+
+    const calculatorPart = combined[0].parts.find(
+      (p) =>
+        p.type === "tool-calculator" &&
+        "toolCallId" in p &&
+        p.toolCallId === "call_1",
+    );
+    expect(calculatorPart).toMatchObject({
+      type: "tool-calculator",
+      state: "output-available",
+      toolCallId: "call_1",
+      input: { operation: "add", a: 2, b: 3 },
+      output: { result: 5 },
+    });
+
+    const formatterPart = combined[0].parts.find(
+      (p) =>
+        p.type === "tool-formatter" &&
+        "toolCallId" in p &&
+        p.toolCallId === "call_2",
+    );
+    expect(formatterPart).toMatchObject({
+      type: "tool-formatter",
+      state: "output-available",
+      toolCallId: "call_2",
+      input: { text: "result" },
+      output: { formatted: "The formatted result is: 5" },
+    });
+  });
+
+  it("handles tool call without corresponding output", () => {
+    const message1 = {
+      id: "msg1",
+      key: "msg1-key",
+      _creationTime: Date.now(),
+      order: 1,
+      stepOrder: 1,
+      status: "success" as const,
+      role: "assistant" as const,
+      text: "",
+      parts: [
+        {
+          type: "dynamic-tool" as const,
+          state: "input-available" as const,
+          toolCallId: "call_orphan",
+          toolName: "calculator",
+          input: { operation: "add", a: 2, b: 3 },
+        },
+      ],
+    };
+
+    const message2 = {
+      id: "msg2",
+      key: "msg2-key",
+      _creationTime: Date.now() + 1,
+      order: 1,
+      stepOrder: 2,
+      status: "success" as const,
+      role: "assistant" as const,
+      text: "Still processing...",
+      parts: [
+        {
+          type: "text" as const,
+          text: "Still processing...",
+          state: "done" as const,
+        },
+      ],
+    };
+
+    const combined = combineUIMessages([message1, message2]);
+
+    expect(combined).toHaveLength(1);
+    expect(combined[0].role).toBe("assistant");
+    expect(combined[0].text).toBe("Still processing...");
+    expect(combined[0].parts).toHaveLength(2);
+
+    const toolPart = combined[0].parts.find((p) => p.type === "dynamic-tool");
+    expect(toolPart).toMatchObject({
+      type: "dynamic-tool",
+      state: "input-available",
+      toolCallId: "call_orphan",
+      toolName: "calculator",
+      input: { operation: "add", a: 2, b: 3 },
+    });
+
+    const textPart = combined[0].parts.find((p) => p.type === "text");
+    expect(textPart).toMatchObject({
+      type: "text",
+      text: "Still processing...",
+      state: "done",
+    });
+  });
+});

package/src/UIMessages.ts

@@ -2,6 +2,7 @@ import {
   convertToModelMessages,
   type UIMessage as AIUIMessage,
   type DeepPartial,
+  type DynamicToolUIPart,
   type ReasoningUIPart,
   type SourceDocumentUIPart,
   type SourceUrlUIPart,
@@ -555,3 +556,66 @@ function toSourcePart(
     providerMetadata: part.providerMetadata,
   } satisfies SourceDocumentUIPart;
 }
+
+export function combineUIMessages(messages: UIMessage[]): UIMessage[] {
+  const combined = messages.reduce((acc, message) => {
+    if (!acc.length) {
+      return [message];
+    }
+    const previous = acc.at(-1)!;
+    if (
+      message.order !== previous.order ||
+      previous.role !== message.role ||
+      message.role !== "assistant"
+    ) {
+      acc.push(message);
+      return acc;
+    }
+    // We will replace it with a combined message
+    acc.pop();
+    const newParts = [...previous.parts];
+    for (const part of message.parts) {
+      const toolCallId = getToolCallId(part);
+      if (!toolCallId) {
+        newParts.push(part);
+        continue;
+      }
+      const previousPartIndex = newParts.findIndex(
+        (p) => getToolCallId(p) === toolCallId,
+      );
+      const previousPart = newParts.splice(previousPartIndex, 1)[0];
+      if (!previousPart) {
+        newParts.push(part);
+        continue;
+      }
+      newParts.push(mergeParts(previousPart, part));
+    }
+    acc.push({
+      ...previous,
+      ...pick(message, ["status", "metadata", "agentName"]),
+      parts: newParts,
+      text: joinText(newParts),
+    });
+    return acc;
+  }, [] as UIMessage[]);
+  return combined;
+}
+
+function getToolCallId(
+  part: UIMessage["parts"][number] & { toolCallId?: string },
+) {
+  return part.toolCallId;
+}
+
+function mergeParts(
+  previousPart: UIMessage["parts"][number],
+  part: UIMessage["parts"][number],
+): UIMessage["parts"][number] {
+  const merged: Record<string, unknown> = { ...previousPart };
+  for (const [key, value] of Object.entries(part)) {
+    if (value !== undefined) {
+      merged[key] = value;
+    }
+  }
+  return merged as ToolUIPart | DynamicToolUIPart;
+}

package/src/client/index.ts

@@ -1,3 +1,4 @@
+import type { JSONValue } from "@ai-sdk/provider";
 import type {
   FlexibleSchema,
   IdGenerator,
@@ -81,12 +82,10 @@ import type {
   AgentComponent,
   Config,
   ContextOptions,
-  DefaultObjectSchema,
   GenerateObjectArgs,
   GenerationOutputMetadata,
   MaybeCustomCtx,
   ObjectMode,
-  ObjectSchema,
   Options,
   RawRequestResponseHandler,
   MutationCtx,
@@ -99,6 +98,7 @@ import type {
   UsageHandler,
   UserActionCtx,
   QueryCtx,
+  AgentPrompt,
 } from "./types.js";
 
 export { stepCountIs } from "ai";
@@ -249,11 +249,9 @@ export class Agent<
        * When generating or streaming text with tools available, this
        * determines when to stop. Defaults to the AI SDK default.
        */
-      stopWhen?: StopCondition<AgentTools> | Array<StopCondition<AgentTools>>;
-      /**
-       * @deprecated Use `languageEmbeddingModel` instead.
-       */
-      chat?: LanguageModel;
+      stopWhen?:
+        | StopCondition<NoInfer<AgentTools>>
+        | Array<StopCondition<NoInfer<AgentTools>>>;
     },
   ) {}
 
@@ -382,54 +380,23 @@ export class Agent<
      * The type of the arguments returned infers from the type of the arguments
      * you pass here.
      */
-    args: T & {
-      /**
-       * If provided, this message will be used as the "prompt" for the LLM call,
-       * instead of the prompt or messages.
-       * This is useful if you want to first save a user message, then use it as
-       * the prompt for the LLM call in another call.
-       */
-      promptMessageId?: string;
-      /**
-       * The model to use for the LLM calls. This will override the model specified
-       * in the Agent constructor.
-       */
-      model?: LanguageModel;
-      /**
-       * The tools to use for the tool calls. This will override tools specified
-       * in the Agent constructor or createThread / continueThread.
-       */
-      tools?: TOOLS;
-      /**
-       * The single prompt message to use for the LLM call. This will be the
-       * last message in the context. If it's a string, it will be a user role.
-       */
-      prompt?: string | (ModelMessage | Message)[];
-      /**
-       * If provided alongside prompt, the ordering will be:
-       * 1. system prompt
-       * 2. search context
-       * 3. recent messages
-       * 4. these messages
-       * 5. prompt messages, including those already on the same `order` as
-       *   the promptMessageId message, if provided.
-       */
-      messages?: (ModelMessage | Message)[];
-      /**
-       * This will be the first message in the context, and overrides the
-       * agent's instructions.
-       */
-      system?: string;
-      /**
-       * The abort signal to be passed to the LLM call. If triggered, it will
-       * mark the pending message as failed. If the generation is asynchronously
-       * aborted, it will trigger this signal when detected.
-       */
-      abortSignal?: AbortSignal;
-      stopWhen?:
-        | StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>
-        | Array<StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>>;
-    },
+    args: T &
+      AgentPrompt & {
+        /**
+         * The tools to use for the tool calls. This will override tools specified
+         * in the Agent constructor or createThread / continueThread.
+         */
+        tools?: TOOLS;
+        /**
+         * The abort signal to be passed to the LLM call. If triggered, it will
+         * mark the pending message as failed. If the generation is asynchronously
+         * aborted, it will trigger this signal when detected.
+         */
+        abortSignal?: AbortSignal;
+        stopWhen?:
+          | StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>
+          | Array<StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>>;
+      },
     options?: Options & { userId?: string | null; threadId?: string },
   ): Promise<{
     args: T & {
@@ -481,9 +448,10 @@ export class Agent<
    * Use {@link continueThread} to get a version of this function already scoped
    * to a thread (and optionally userId).
    * @param ctx The context passed from the action function calling this.
-   * @param { userId, threadId }: The user and thread to associate the message with
-   * @param generateTextArgs The arguments to the generateText function, along with extra controls
-   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @param scope: The user and thread to associate the message with
+   * @param generateTextArgs The arguments to the generateText function, along
+   * with {@link AgentPrompt} options, such as promptMessageId.
+   * @param options Extra controls for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the generateText function.
    */
   async generateText<
@@ -493,15 +461,12 @@ export class Agent<
   >(
     ctx: ActionCtx & CustomCtx,
     threadOpts: { userId?: string | null; threadId?: string },
-    generateTextArgs: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL> & {
-      /**
-       * If provided, this message will be used as the "prompt" for the LLM call,
-       * instead of the prompt or messages.
-       * This is useful if you want to first save a user message, then use it as
-       * the prompt for the LLM call in another call.
-       */
-      promptMessageId?: string;
-    },
+    /**
+     * The arguments to the generateText function, similar to the ai sdk's
+     * {@link generateText} function, along with Agent prompt options.
+     */
+    generateTextArgs: AgentPrompt &
+      TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>,
     options?: Options,
   ): Promise<
     GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
@@ -557,22 +522,11 @@ export class Agent<
     ctx: ActionCtx & CustomCtx,
     threadOpts: { userId?: string | null; threadId?: string },
     /**
-     * The arguments to the streamText function, similar to the ai `streamText` function.
+     * The arguments to the streamText function, similar to the ai sdk's
+     * {@link streamText} function, along with Agent prompt options.
      */
-    streamTextArgs: StreamingTextArgs<
-      AgentTools,
-      TOOLS,
-      OUTPUT,
-      PARTIAL_OUTPUT
-    > & {
-      /**
-       * If provided, this message will be used as the "prompt" for the LLM call,
-       * instead of the prompt or messages.
-       * This is useful if you want to first save a user message, then use it as
-       * the prompt for the LLM call in another call.
-       */
-      promptMessageId?: string;
-    },
+    streamTextArgs: AgentPrompt &
+      StreamingTextArgs<AgentTools, TOOLS, OUTPUT, PARTIAL_OUTPUT>,
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
@@ -698,7 +652,7 @@ export class Agent<
    * to a thread (and optionally userId).
    */
   async generateObject<
-    SCHEMA extends ObjectSchema = DefaultObjectSchema,
+    SCHEMA extends FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
     OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
       ? "enum"
       : "object",
@@ -709,17 +663,11 @@ export class Agent<
     ctx: ActionCtx & CustomCtx,
     threadOpts: { userId?: string | null; threadId?: string },
     /**
-     * The arguments to the generateObject function, similar to the ai.generateObject function.
+     * The arguments to the generateObject function, similar to the ai sdk's
+     * {@link generateObject} function, along with Agent prompt options.
      */
-    generateObjectArgs: GenerateObjectArgs<SCHEMA, OUTPUT, RESULT> & {
-      /**
-       * If provided, this message will be used as the "prompt" for the LLM call,
-       * instead of the prompt or messages.
-       * This is useful if you want to first save a user message, then use it as
-       * the prompt for the LLM call in another call.
-       */
-      promptMessageId?: string;
-    },
+    generateObjectArgs: AgentPrompt &
+      GenerateObjectArgs<SCHEMA, OUTPUT, RESULT>,
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
@@ -756,7 +704,7 @@ export class Agent<
    * to a thread (and optionally userId).
    */
   async streamObject<
-    SCHEMA extends ObjectSchema = DefaultObjectSchema,
+    SCHEMA extends FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
     OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
       ? "enum"
       : "object",
@@ -767,17 +715,10 @@ export class Agent<
     ctx: ActionCtx & CustomCtx,
     threadOpts: { userId?: string | null; threadId?: string },
     /**
-     * The arguments to the streamObject function, similar to the ai `streamObject` function.
+     * The arguments to the streamObject function, similar to the ai sdk's
+     * {@link streamObject} function, along with Agent prompt options.
      */
-    streamObjectArgs: StreamObjectArgs<SCHEMA, OUTPUT, RESULT> & {
-      /**
-       * If provided, this message will be used as the "prompt" for the LLM call,
-       * instead of the prompt or messages.
-       * This is useful if you want to first save a user message, then use it as
-       * the prompt for the LLM call in another call.
-       */
-      promptMessageId?: string;
-    },
+    streamObjectArgs: AgentPrompt & StreamObjectArgs<SCHEMA, OUTPUT, RESULT>,
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
@@ -1563,7 +1504,7 @@ export class Agent<
    * and stopWhen.
    */
   asObjectAction<T, DataModel extends GenericDataModel>(
-    objectArgs: GenerateObjectArgs<FlexibleSchema<T>>,
+    objectArgs: GenerateObjectArgs<FlexibleSchema<T>> & Partial<AgentPrompt>,
     options?: Options & MaybeCustomCtx<CustomCtx, DataModel, AgentTools>,
   ) {
     return internalActionGeneric({
@@ -1604,24 +1545,7 @@ export class Agent<
   }
 
   /**
-   * Save messages to the thread.
-   * Useful as a step in Workflows, e.g.
-   * ```ts
-   * const saveMessages = agent.asSaveMessagesMutation();
-   *
-   * const myWorkflow = workflow.define({
-   *   args: {...},
-   *   handler: async (step, args) => {
-   *     // do things to create (but not save)messages
-   *     const { messageIds } = await step.runMutation(internal.foo.saveMessages, {
-   *       threadId: args.threadId,
-   *       messages: args.messages,
-   *     });
-   *     // ...
-   *   },
-   * })
-   * ```
-   * @returns A mutation that can be used to save messages to the thread.
+   * @deprecated Use {@link saveMessages} directly instead.
    */
   asSaveMessagesMutation() {
     return internalMutationGeneric({