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

package/src/client/mockModel.ts

@@ -0,0 +1,195 @@
+import type {
+  LanguageModelV2,
+  LanguageModelV2StreamPart,
+} from "@ai-sdk/provider";
+import type { ReasoningPart, TextPart } from "@ai-sdk/provider-utils";
+import { simulateReadableStream } from "ai";
+
+const longDefaultText = `
+A A A A A A A A A A A A A A A
+B B B B B B B B B B B B B B B
+C C C C C C C C C C C C C C C
+D D D D D D D D D D D D D D D
+`;
+const defaultUsage = { outputTokens: 10, inputTokens: 3, totalTokens: 13 };
+
+export type MockModelArgs = {
+  provider?: LanguageModelV2["provider"];
+  modelId?: LanguageModelV2["modelId"];
+  supportedUrls?:
+    | LanguageModelV2["supportedUrls"]
+    | (() => LanguageModelV2["supportedUrls"]);
+  chunkDelayInMs?: number;
+  initialDelayInMs?: number;
+  // provide either content or doGenerate & doStream
+  content?: (TextPart | ReasoningPart)[];
+  doGenerate?: LanguageModelV2["doGenerate"];
+  doStream?: LanguageModelV2["doStream"];
+  fail?:
+    | boolean
+    | {
+        probability?: number;
+        error?: string;
+      };
+};
+
+export function mockModel(args?: MockModelArgs): LanguageModelV2 {
+  return new MockLanguageModel(args ?? {});
+}
+
+export class MockLanguageModel implements LanguageModelV2 {
+  readonly specificationVersion = "v2";
+
+  private _supportedUrls: () => LanguageModelV2["supportedUrls"];
+
+  readonly provider: LanguageModelV2["provider"];
+  readonly modelId: LanguageModelV2["modelId"];
+
+  doGenerate: LanguageModelV2["doGenerate"];
+  doStream: LanguageModelV2["doStream"];
+
+  doGenerateCalls: Parameters<LanguageModelV2["doGenerate"]>[0][] = [];
+  doStreamCalls: Parameters<LanguageModelV2["doStream"]>[0][] = [];
+
+  constructor(args: MockModelArgs) {
+    this.provider = args.provider || "mock-provider";
+    this.modelId = args.modelId || "mock-model-id";
+    const {
+      content = [{ type: "text", text: longDefaultText }],
+      chunkDelayInMs = 200,
+      initialDelayInMs = 1000,
+      supportedUrls = {},
+    } = args;
+    const fail =
+      args.fail &&
+      (args.fail === true ||
+        !args.fail.probability ||
+        Math.random() < args.fail.probability);
+    const error =
+      (typeof args.fail === "object" && args.fail.error) ||
+      "Mock error message";
+
+    const chunks: LanguageModelV2StreamPart[] = [
+      { type: "stream-start", warnings: [] },
+    ];
+    chunks.push(
+      ...content.flatMap((c, ci) => {
+        const deltas = c.text.split(" ");
+        let parts: LanguageModelV2StreamPart[] = [];
+        if (c.type === "reasoning") {
+          parts.push({
+            type: "reasoning-start",
+            id: `${ci}-reasoning-start`,
+          });
+          parts.push(
+            ...deltas.map(
+              (delta, di) =>
+                ({
+                  type: "reasoning-delta",
+                  delta,
+                  id: `${ci}-reasoning-${di}`,
+                  providerMetadata: {
+                    mockProvider: { mock: { reasoningDetails: null } },
+                  },
+                }) satisfies LanguageModelV2StreamPart,
+            ),
+          );
+          parts.push({
+            type: "reasoning-end",
+            id: `${ci}-reasoning-end`,
+          });
+        } else if (c.type === "text") {
+          parts.push({
+            type: "text-start",
+            id: `${ci}-text-start`,
+          });
+          parts = deltas.map((delta, di) => ({
+            type: "text-delta",
+            delta,
+            id: `${ci}-text-${di}`,
+          }));
+          parts.push({
+            type: "text-end",
+            id: `${ci}-text-end`,
+          });
+        }
+        return parts;
+      }),
+    );
+    if (fail) {
+      chunks.push({
+        type: "error",
+        error,
+      });
+    }
+    chunks.push({
+      type: "finish",
+      finishReason: fail ? "error" : "stop",
+      usage: defaultUsage,
+      providerMetadata: {
+        mockProvider: { mock: "mock metadata" },
+      },
+    });
+    this.doGenerate = async (options) => {
+      this.doGenerateCalls.push(options);
+
+      if (fail) {
+        throw new Error(error);
+      }
+      if (typeof args.doGenerate === "function") {
+        return args.doGenerate(options);
+      } else if (Array.isArray(args.doGenerate)) {
+        return args.doGenerate[this.doGenerateCalls.length];
+      } else if (content) {
+        return {
+          content,
+          finishReason: "stop",
+          usage: defaultUsage,
+          providerMetadata: { mockProvider: { mock: "mock metadata" } },
+          warnings: [],
+        };
+      } else {
+        throw new Error("Unexpected: no content or doGenerate");
+      }
+    };
+    this._supportedUrls =
+      typeof supportedUrls === "function"
+        ? supportedUrls
+        : async () => supportedUrls;
+    this.doStream = async (options) => {
+      this.doStreamCalls.push(options);
+
+      if (typeof args.doStream === "function") {
+        return args.doStream(options);
+      } else if (Array.isArray(args.doStream)) {
+        return args.doStream[this.doStreamCalls.length];
+      } else if (content) {
+        const stream = simulateReadableStream({
+          chunks,
+          initialDelayInMs,
+          chunkDelayInMs,
+        });
+        if (options.abortSignal) {
+          throw new Error("abortSignal in mock model");
+        }
+        return {
+          stream,
+          request: { body: {} },
+          response: { headers: {} },
+        };
+      } else if (args.doStream) {
+        return args.doStream;
+      } else {
+        throw new Error("Provide either content or doStream");
+      }
+    };
+    this._supportedUrls =
+      typeof supportedUrls === "function"
+        ? supportedUrls
+        : async () => supportedUrls;
+  }
+
+  get supportedUrls() {
+    return this._supportedUrls();
+  }
+}

package/src/client/streaming.ts

@@ -1,10 +1,4 @@
-import {
-  type ChunkDetector,
-  smoothStream,
-  type StreamTextTransform,
-  type TextStreamPart,
-  type ToolSet,
-} from "ai";
+import { type ChunkDetector } from "ai";
 import {
   vStreamDelta,
   vStreamMessage,
@@ -21,8 +15,6 @@ import type {
   RunQueryCtx,
   SyncStreamsReturnValue,
 } from "./types.js";
-import { omit } from "convex-helpers";
-import { serializeTextStreamingPartsV5 } from "../parts.js";
 import { v } from "convex/values";
 import { vMessageDoc } from "../component/schema.js";
 
@@ -161,33 +153,21 @@ export const DEFAULT_STREAMING_OPTIONS = {
   returnImmediately: false,
 } satisfies StreamingOptions;
 
-export function mergeTransforms<TOOLS extends ToolSet>(
-  options: StreamingOptions | 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;
-}
-
-export class DeltaStreamer {
+/**
+ * DeltaStreamer can be used to save a stream of "parts" by writing
+ * batches of them in "deltas" to the database so clients can subscribe
+ * (using the syncStreams utility and client hooks) and re-hydrate the stream.
+ * You can optionally compress the parts, e.g. concatenating text deltas, to
+ * optimize the data in transit.
+ */
+export class DeltaStreamer<T> {
   public streamId: string | undefined;
-  public readonly options: Required<StreamingOptions>;
-  #nextParts: TextStreamPart<ToolSet>[] = [];
+  public readonly config: {
+    stream: Required<StreamingOptions>;
+    onAsyncAbort: (reason: string) => Promise<void>;
+    compress: ((parts: T[]) => T[]) | null;
+  };
+  #nextParts: T[] = [];
   #latestWrite: number = 0;
   #ongoingWrite: Promise<void> | undefined;
   #cursor: number = 0;
@@ -196,7 +176,12 @@ export class DeltaStreamer {
   constructor(
     public readonly component: AgentComponent,
     public readonly ctx: RunActionCtx,
-    options: true | StreamingOptions,
+    config: {
+      stream: true | StreamingOptions;
+      onAsyncAbort: (reason: string) => Promise<void>;
+      abortSignal: AbortSignal | undefined;
+      compress: ((parts: T[]) => T[]) | null;
+    },
     public readonly metadata: {
       threadId: string;
       userId?: string;
@@ -206,45 +191,49 @@ export class DeltaStreamer {
       model?: string;
       provider?: string;
       providerOptions?: ProviderOptions;
-      abortSignal?: AbortSignal;
     },
   ) {
-    this.options =
-      typeof options === "boolean"
-        ? DEFAULT_STREAMING_OPTIONS
-        : { ...DEFAULT_STREAMING_OPTIONS, ...options };
+    this.config = {
+      stream:
+        config.stream === true
+          ? DEFAULT_STREAMING_OPTIONS
+          : { ...DEFAULT_STREAMING_OPTIONS, ...config },
+      onAsyncAbort: config.onAsyncAbort,
+      compress: config.compress,
+    };
     this.#nextParts = [];
     this.abortController = new AbortController();
-    if (metadata.abortSignal) {
-      metadata.abortSignal.addEventListener("abort", async () => {
+    if (config.abortSignal) {
+      config.abortSignal.addEventListener("abort", async () => {
+        if (this.abortController.signal.aborted) {
+          return;
+        }
         if (this.streamId) {
           this.abortController.abort();
-          const finalDelta = this.#createDelta();
           await this.#ongoingWrite;
           await this.ctx.runMutation(this.component.streams.abort, {
             streamId: this.streamId,
             reason: "abortSignal",
-            finalDelta,
           });
         }
       });
     }
   }
 
-  public async addParts(parts: TextStreamPart<ToolSet>[]) {
+  public async addParts(parts: T[]) {
     if (this.abortController.signal.aborted) {
       return;
     }
     if (!this.streamId) {
       this.streamId = await this.ctx.runMutation(
         this.component.streams.create,
-        omit(this.metadata, ["abortSignal"]),
+        this.metadata,
       );
     }
     this.#nextParts.push(...parts);
     if (
       !this.#ongoingWrite &&
-      Date.now() - this.#latestWrite >= this.options.throttleMs
+      Date.now() - this.#latestWrite >= this.config.stream.throttleMs
     ) {
       this.#ongoingWrite = this.#sendDelta();
     }
@@ -265,16 +254,21 @@ export class DeltaStreamer {
         delta,
       );
       if (!success) {
+        await this.config.onAsyncAbort("async abort");
         this.abortController.abort();
+        return;
       }
     } catch (e) {
+      await this.config.onAsyncAbort(
+        e instanceof Error ? e.message : "unknown error",
+      );
       this.abortController.abort();
       throw e;
     }
     // 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.options.throttleMs
+      Date.now() - this.#latestWrite >= this.config.stream.throttleMs
     ) {
       // We send again immediately with the accumulated deltas.
       this.#ongoingWrite = this.#sendDelta();
@@ -290,7 +284,9 @@ export class DeltaStreamer {
     const start = this.#cursor;
     const end = start + this.#nextParts.length;
     this.#cursor = end;
-    const parts = serializeTextStreamingPartsV5(this.#nextParts);
+    const parts = this.config.compress
+      ? this.config.compress(this.#nextParts)
+      : this.#nextParts;
     this.#nextParts = [];
     if (!this.streamId) {
       throw new Error("Creating a delta before the stream is created");
@@ -302,11 +298,9 @@ export class DeltaStreamer {
     if (!this.streamId) {
       return;
     }
-    const finalDelta = this.#createDelta();
     await this.#ongoingWrite;
     await this.ctx.runMutation(this.component.streams.finish, {
       streamId: this.streamId,
-      finalDelta,
     });
   }
 
@@ -318,12 +312,10 @@ export class DeltaStreamer {
     if (!this.streamId) {
       return;
     }
-    const finalDelta = this.#createDelta();
     await this.#ongoingWrite;
     await this.ctx.runMutation(this.component.streams.abort, {
       streamId: this.streamId,
       reason,
-      finalDelta,
     });
   }
 }

package/src/{parts.ts → client/textStreamParts.ts}

@@ -1,4 +1,13 @@
-import type { TextStreamPart, ToolSet } from "ai";
+import {
+  smoothStream,
+  type StreamTextTransform,
+  type ToolSet,
+  type TextStreamPart,
+} from "ai";
+import {
+  DEFAULT_STREAMING_OPTIONS,
+  type StreamingOptions,
+} from "./streaming.js";
 
 export function serializeTextStreamingPartsV5(
   parts: TextStreamPart<ToolSet>[],
@@ -37,3 +46,26 @@ export function serializeTextStreamingPartsV5(
   }
   return compressed;
 }
+
+export function mergeTransforms<TOOLS extends ToolSet>(
+  options: StreamingOptions | 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;
+}