@convex-dev/agent 0.2.6 → 0.2.8-alpha.0

package/src/client/start.ts

@@ -21,7 +21,6 @@ import type {
   Options,
   UserActionCtx,
 } from "./types.js";
-import { saveMessages } from "./messages.js";
 import type { Message, MessageDoc } from "../validators.js";
 import {
   getModelName,
@@ -105,6 +104,7 @@ export async function start<
     system?: string;
     model: LanguageModel;
     messages: ModelMessage[];
+    prompt?: never;
     tools?: Tools;
   } & CallSettings;
   order: number;
@@ -200,6 +200,7 @@ export async function start<
   } as T & {
     model: LanguageModel;
     messages: ModelMessage[];
+    prompt?: never;
     tools?: Tools;
     _internal?: { generateId?: IdGenerator };
   } & CallSettings;

package/src/client/streaming.test.ts

@@ -0,0 +1,186 @@
+import { beforeEach, describe, expect, test } from "vitest";
+import { createThread } from "./index.js";
+import type { GenericSchema, SchemaDefinition } from "convex/server";
+import { streamText } from "ai";
+import { components, initConvexTest } from "./setup.test.js";
+import { mockModel } from "./mockModel.js";
+import { DeltaStreamer } from "./streaming.js";
+import { compressUIMessageChunks, getParts } from "../deltas.js";
+import type { TestConvex } from "convex-test";
+
+const defaultTestOptions = {
+  throttleMs: 0,
+  abortSignal: undefined,
+  compress: null,
+  onAsyncAbort: async () => {
+    throw new Error("unexpected");
+  },
+};
+
+const testMetadata = {
+  order: 0,
+  stepOrder: 0,
+  agentName: "test agent",
+  model: "test model",
+  provider: "test provider",
+  providerOptions: {},
+  format: "UIMessageChunk" as const,
+};
+
+describe("DeltaStreamer", () => {
+  let t: TestConvex<SchemaDefinition<GenericSchema, boolean>>;
+  let threadId: string;
+  beforeEach(async () => {
+    t = initConvexTest();
+    await t.run(async (ctx) => {
+      threadId = await createThread(ctx, components.agent, {});
+    });
+  });
+  test("should save chunks via DeltaStreamer", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer(
+        components.agent,
+        ctx,
+        { ...defaultTestOptions },
+        { ...testMetadata, threadId },
+      );
+      const result = streamText({
+        model: mockModel(),
+        prompt: "Test prompt",
+      });
+      await streamer.consumeStream(result.toUIMessageStream());
+      const streamId = streamer.streamId!;
+      expect(streamId).toBeDefined();
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
+      const { parts } = getParts(deltas);
+      const stream = result.toUIMessageStream();
+      for await (const part of stream) {
+        const expected = parts.shift();
+        expect(part).toEqual(expected);
+      }
+    });
+  });
+  test("should save all parts when throttleMs is 0", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer(
+        components.agent,
+        ctx,
+        { ...defaultTestOptions, throttleMs: 0 },
+        { ...testMetadata, threadId },
+      );
+      const result = streamText({
+        model: mockModel({
+          content: [
+            // The mockModel splits these into deltas based on spaces
+            { type: "text", text: "A B C" },
+            { type: "reasoning", text: "D E F" },
+          ],
+        }),
+        prompt: "Test prompt",
+      });
+      await streamer.consumeStream(result.toUIMessageStream());
+      const streamId = streamer.streamId!;
+      expect(streamId).toBeDefined();
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
+      const { parts } = getParts(deltas);
+      const expected = [
+        { type: "start" },
+        { type: "start-step" },
+        { type: "text-start" },
+        { type: "text-delta", delta: "A" },
+        { type: "text-delta", delta: " B" },
+        { type: "text-delta", delta: " C" },
+        { type: "text-end" },
+        { type: "reasoning-start" },
+        { type: "reasoning-delta", delta: "D" },
+        { type: "reasoning-delta", delta: " E" },
+        { type: "reasoning-delta", delta: " F" },
+        { type: "reasoning-end" },
+        { type: "finish-step" },
+        { type: "finish" },
+      ];
+      for (const expectedPart of expected) {
+        const part = parts.shift();
+        expect(part).toBeDefined();
+        expect(part).toMatchObject(expectedPart);
+      }
+    });
+  });
+
+  test("should save compressed parts via DeltaStreamer", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer(
+        components.agent,
+        ctx,
+        {
+          throttleMs: 1000,
+          abortSignal: undefined,
+          compress: compressUIMessageChunks,
+          onAsyncAbort: async () => {
+            throw new Error("async abort");
+          },
+        },
+        {
+          ...testMetadata,
+          threadId,
+        },
+      );
+      const result = streamText({
+        model: mockModel({
+          content: [
+            // The mockModel splits these into deltas based on spaces
+            { type: "text", text: "A B C" },
+            { type: "text", text: "D E F" },
+            { type: "reasoning", text: "J K L" },
+            { type: "text", text: "M N O" },
+          ],
+        }),
+        prompt: "Test prompt",
+        // experimental_transform: smoothStream({ chunking: "line" }),
+        onError: (error) => {
+          console.error(error);
+        },
+      });
+      await streamer.consumeStream(result.toUIMessageStream());
+      const streamId = streamer.streamId!;
+      expect(streamId).toBeDefined();
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
+      const { parts } = getParts(deltas);
+      const expected = [
+        { type: "start" },
+        { type: "start-step" },
+        { type: "text-start" },
+        // These are collapsed into a single delta
+        { type: "text-delta", delta: "A B C" },
+        { type: "text-end" },
+        { type: "text-start" },
+        { type: "text-delta", delta: "D E F" },
+        { type: "text-end" },
+        { type: "reasoning-start" },
+        { type: "reasoning-delta", delta: "J K L" },
+        { type: "reasoning-end" },
+        { type: "text-start" },
+        { type: "text-delta", delta: "M N O" },
+        { type: "text-end" },
+        { type: "finish-step" },
+        { type: "finish" },
+      ];
+      for (const expectedPart of expected) {
+        const part = parts.shift();
+        expect(part).toBeDefined();
+        expect(part).toMatchObject(expectedPart);
+      }
+    });
+  });
+  // TODO: test errors & aborted states
+  // TODO: test fetching partial stream data - syncStreams w/ cursors
+});

package/src/client/streaming.ts

@@ -1,4 +1,10 @@
-import { type ChunkDetector } from "ai";
+import {
+  type ChunkDetector,
+  type AsyncIterableStream,
+  type StreamTextTransform,
+  type ToolSet,
+  smoothStream,
+} from "ai";
 import {
   vStreamDelta,
   vStreamMessage,
@@ -152,6 +158,36 @@ export const DEFAULT_STREAMING_OPTIONS = {
   returnImmediately: false,
 } satisfies StreamingOptions;
 
+/**
+ *
+ * @param options The options passed to `agent.streamText` to decide whether to
+ * save deltas while streaming.
+ * @param existing The transforms passed to `agent.streamText` to merge with.
+ * @returns The merged transforms to pass to the underlying `streamText` call.
+ */
+export function mergeTransforms<TOOLS extends ToolSet>(
+  options: { chunking?: StreamingOptions["chunking"] } | boolean | undefined,
+  existing:
+    | StreamTextTransform<TOOLS>
+    | Array<StreamTextTransform<TOOLS>>
+    | undefined,
+) {
+  if (!options) {
+    return existing;
+  }
+  const chunking =
+    typeof options === "boolean"
+      ? DEFAULT_STREAMING_OPTIONS.chunking
+      : options.chunking;
+  const transforms = Array.isArray(existing)
+    ? existing
+    : existing
+      ? [existing]
+      : [];
+  transforms.push(smoothStream({ delayInMs: null, chunking }));
+  return transforms;
+}
+
 /**
  * DeltaStreamer can be used to save a stream of "parts" by writing
  * batches of them in "deltas" to the database so clients can subscribe
@@ -162,7 +198,7 @@ export const DEFAULT_STREAMING_OPTIONS = {
 export class DeltaStreamer<T> {
   public streamId: string | undefined;
   public readonly config: {
-    stream: Required<StreamingOptions>;
+    throttleMs: number;
     onAsyncAbort: (reason: string) => Promise<void>;
     compress: ((parts: T[]) => T[]) | null;
   };
@@ -176,7 +212,7 @@ export class DeltaStreamer<T> {
     public readonly component: AgentComponent,
     public readonly ctx: RunMutationCtx,
     config: {
-      stream: true | StreamingOptions;
+      throttleMs: number | undefined;
       onAsyncAbort: (reason: string) => Promise<void>;
       abortSignal: AbortSignal | undefined;
       compress: ((parts: T[]) => T[]) | null;
@@ -190,13 +226,11 @@ export class DeltaStreamer<T> {
       model?: string;
       provider?: string;
       providerOptions?: ProviderOptions;
+      format: "UIMessageChunk" | "TextStreamPart" | undefined;
     },
   ) {
     this.config = {
-      stream:
-        config.stream === true
-          ? DEFAULT_STREAMING_OPTIONS
-          : { ...DEFAULT_STREAMING_OPTIONS, ...config },
+      throttleMs: config.throttleMs ?? DEFAULT_STREAMING_OPTIONS.throttleMs,
       onAsyncAbort: config.onAsyncAbort,
       compress: config.compress,
     };
@@ -232,12 +266,19 @@ export class DeltaStreamer<T> {
     this.#nextParts.push(...parts);
     if (
       !this.#ongoingWrite &&
-      Date.now() - this.#latestWrite >= this.config.stream.throttleMs
+      Date.now() - this.#latestWrite >= this.config.throttleMs
     ) {
       this.#ongoingWrite = this.#sendDelta();
     }
   }
 
+  public async consumeStream(stream: AsyncIterableStream<T>) {
+    for await (const chunk of stream) {
+      await this.addParts([chunk]);
+    }
+    await this.finish();
+  }
+
   async #sendDelta() {
     if (this.abortController.signal.aborted) {
       return;
@@ -267,7 +308,7 @@ export class DeltaStreamer<T> {
     // Now that we've sent the delta, check if we need to send another one.
     if (
       this.#nextParts.length > 0 &&
-      Date.now() - this.#latestWrite >= this.config.stream.throttleMs
+      Date.now() - this.#latestWrite >= this.config.throttleMs
     ) {
       // We send again immediately with the accumulated deltas.
       this.#ongoingWrite = this.#sendDelta();

package/src/client/types.ts

@@ -388,7 +388,7 @@ export type GenerateObjectArgs<
     ? Array<InferSchema<SCHEMA>>
     : InferSchema<SCHEMA>,
 > = Omit<CallSettings, "stopSequences"> &
-  Prompt & {
+  Partial<Prompt> & {
     /**
      * If provided, this message will be used as the "prompt" for the LLM call,
      * instead of the prompt or messages.
@@ -555,7 +555,15 @@ export interface Thread<DefaultTools extends ToolSet> {
       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;
+    },
     options?: Options,
   ): Promise<
     GenerateTextResult<TOOLS extends undefined ? DefaultTools : TOOLS, OUTPUT> &
@@ -582,7 +590,15 @@ export interface Thread<DefaultTools extends ToolSet> {
       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;
+    },
     options?: Options & {
       /**
        * Whether to save incremental data (deltas) from streaming responses.
@@ -630,15 +646,6 @@ export interface Thread<DefaultTools extends ToolSet> {
        * 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.
-       */
     },
     options?: Options,
   ): Promise<GenerateObjectResult<RESULT> & ThreadOutputMetadata>;
@@ -672,15 +679,6 @@ export interface Thread<DefaultTools extends ToolSet> {
        * 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.
-       */
     },
     options?: Options,
   ): Promise<

package/src/component/_generated/api.d.ts

@@ -2189,6 +2189,7 @@ export type Mounts = {
       "public",
       {
         agentName?: string;
+        format?: "UIMessageChunk" | "TextStreamPart";
         model?: string;
         order: number;
         provider?: string;
@@ -2253,6 +2254,7 @@ export type Mounts = {
       },
       Array<{
         agentName?: string;
+        format?: "UIMessageChunk" | "TextStreamPart";
         model?: string;
         order: number;
         provider?: string;

package/src/component/messages.ts

@@ -344,7 +344,7 @@ export const finalizeMessage = mutation({
     const message = await ctx.db.get(messageId);
     assert(message, `Message ${messageId} not found`);
     if (message.status !== "pending") {
-      console.log(
+      console.debug(
         "Trying to finalize a message that's already",
         message.status,
       );

package/src/component/schema.ts

@@ -91,6 +91,12 @@ export const schema = defineSchema({
     model: v.optional(v.string()),
     provider: v.optional(v.string()),
     providerOptions: v.optional(vProviderOptions), // Sent to model
+    // The data format for the deltas. By default, we use UIMessageChunks.
+    // This format dictates how the messages are materialized for capturing
+    // partial messages during failure, as well as on the client side.
+    format: v.optional(
+      v.union(v.literal("UIMessageChunk"), v.literal("TextStreamPart")),
+    ),
 
     threadId: v.id("threads"),
     order: v.number(),

package/src/component/streams.ts

@@ -20,8 +20,8 @@ import { stream } from "convex-helpers/server/stream";
 import { mergedStream } from "convex-helpers/server/stream";
 import { paginator } from "convex-helpers/server/pagination";
 import type { WithoutSystemFields } from "convex/server";
-import { mergeDeltas } from "../react/deltas.js";
-import { serializeOrThrow } from "../mapping.js";
+import { deriveUIMessagesFromDeltas } from "../deltas.js";
+import { fromUIMessages } from "../UIMessages.js";
 
 const SECOND = 1000;
 const MINUTE = 60 * SECOND;
@@ -140,6 +140,7 @@ export const list = query({
       streamId: m._id,
       status: m.state.kind,
       ...pick(m, [
+        "format",
         "order",
         "stepOrder",
         "userId",
@@ -516,72 +517,54 @@ export async function getStreamingMessagesWithMetadata(
   );
   const messages = (
     await Promise.all(
-      streamingMessages.map((m) =>
-        getMessagesWithMetadataForStreamingMessage(
-          ctx,
+      streamingMessages.map(async (streamingMessage) => {
+        const deltas = await ctx.db
+          .query("streamDeltas")
+          .withIndex("streamId_start_end", (q) =>
+            q.eq("streamId", streamingMessage._id),
+          )
+          .take(1000);
+        const uiMessages = await deriveUIMessagesFromDeltas(
           threadId,
-          stepOrder,
-          m,
-          metadata,
-        ),
-      ),
+          [
+            {
+              ...streamingMessage,
+              status: "streaming",
+              streamId: streamingMessage._id,
+            },
+          ],
+          deltas,
+        );
+        // We don't save messages that have already been saved
+        const numToSkip = stepOrder - streamingMessage.stepOrder;
+        const messages = await Promise.all(
+          fromUIMessages(uiMessages, streamingMessage)
+            .slice(numToSkip)
+            .filter((m) => m.message !== undefined)
+            .map(async (msg) => {
+              return {
+                ...pick(msg, [
+                  "message",
+                  "fileIds",
+                  "status",
+                  "finishReason",
+                  "model",
+                  "provider",
+                  "providerMetadata",
+                  "sources",
+                  "reasoning",
+                  "reasoningDetails",
+                  "usage",
+                  "warnings",
+                  "error",
+                ]),
+                ...metadata,
+              } as MessageWithMetadataInternal;
+            }),
+        );
+        return messages;
+      }),
     )
   ).flat();
   return messages;
 }
-
-export async function getMessagesWithMetadataForStreamingMessage(
-  ctx: MutationCtx,
-  threadId: Id<"threads">,
-  stepOrder: number,
-  streamingMessage: Doc<"streamingMessages">,
-  metadata: { status: "success" | "failed"; error?: string },
-): Promise<MessageWithMetadataInternal[]> {
-  const deltas = await ctx.db
-    .query("streamDeltas")
-    .withIndex("streamId_start_end", (q) =>
-      q.eq("streamId", streamingMessage._id),
-    )
-    .take(1000);
-  const [messageDocs] = mergeDeltas(
-    threadId,
-    [
-      {
-        ...streamingMessage,
-        status: "streaming",
-        streamId: streamingMessage._id,
-      },
-    ],
-    [],
-    deltas,
-  );
-  // We don't save messages that have already been saved
-  const numToSkip = stepOrder - streamingMessage.stepOrder;
-  const messages = await Promise.all(
-    messageDocs
-      .slice(numToSkip)
-      .filter((m) => m.message !== undefined)
-      .map(async (msg) => {
-        const message = await serializeOrThrow(msg.message!);
-        return {
-          message,
-          ...pick(msg, [
-            "fileIds",
-            "status",
-            "finishReason",
-            "model",
-            "provider",
-            "providerMetadata",
-            "sources",
-            "reasoning",
-            "reasoningDetails",
-            "usage",
-            "warnings",
-            "error",
-          ]),
-          ...metadata,
-        } as MessageWithMetadataInternal;
-      }),
-  );
-  return messages;
-}