@convex-dev/agent 0.2.0-alpha.2 → 0.2.1

package/src/client/messages.ts

@@ -74,6 +74,10 @@ export type SaveMessagesArgs = {
    * The embeddings to save with the messages.
    */
   embeddings?: MessageEmbeddings;
+  /**
+   * A pending message ID to replace when adding messages.
+   */
+  pendingMessageId?: string;
 };
 
 /**
@@ -88,7 +92,7 @@ export async function saveMessages(
      */
     agentName?: string;
   },
-) {
+): Promise<{ messages: MessageDoc[] }> {
   let embeddings: MessageEmbeddingsWithDimension | undefined;
   if (args.embeddings) {
     const dimension = args.embeddings.vectors.find((v) => v !== null)?.length;
@@ -106,14 +110,19 @@ export async function saveMessages(
     userId: args.userId ?? undefined,
     agentName: args.agentName,
     promptMessageId: args.promptMessageId,
+    pendingMessageId: args.pendingMessageId,
     embeddings,
     messages: await Promise.all(
       args.messages.map(async (m, i) => {
         const { message, fileIds } = await serializeMessage(ctx, component, m);
+        const allFileIds = args.metadata?.[i]?.fileIds ?? [];
+        if (fileIds) {
+          allFileIds.push(...fileIds);
+        }
         return parse(vMessageWithMetadata, {
           ...args.metadata?.[i],
           message,
-          fileIds,
+          fileIds: allFileIds.length > 0 ? allFileIds : undefined,
         });
       }),
     ),
@@ -134,6 +143,10 @@ export type SaveMessageArgs = {
    * The embedding to save with the message.
    */
   embedding?: { vector: number[]; model: string };
+  /**
+   * A pending message ID to replace with this message.
+   */
+  pendingMessageId?: string;
 } & (
   | {
       prompt?: undefined;
@@ -179,6 +192,7 @@ export async function saveMessage(
     threadId: args.threadId,
     userId: args.userId ?? undefined,
     agentName: args.agentName,
+    pendingMessageId: args.pendingMessageId,
     messages:
       args.prompt !== undefined
         ? [{ role: "user", content: args.prompt }]

package/src/client/search.ts

@@ -16,9 +16,11 @@ import type { Message } from "../validators.js";
 
 const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0;
 
-export type GetEmbedding = (text: string) => Promise<{
+export type GetEmbedding = (
+  text: string,
+) => Promise<{
   embedding: number[];
-  embeddingModel: string | EmbeddingModel<string>;
+  textEmbeddingModel: string | EmbeddingModel<string>;
 }>;
 
 /**
@@ -96,9 +98,10 @@ export async function fetchContextMessages(
         "You must provide an embedding and embeddingModel to use vector search",
       );
     }
-    const embeddingFields = opts.searchOptions?.vectorSearch
-      ? await args.getEmbedding?.(text)
-      : undefined;
+    const embeddingFields =
+      opts.searchOptions?.vectorSearch && text
+        ? await args.getEmbedding?.(text)
+        : undefined;
     const searchMessages = await ctx.runAction(
       component.messages.searchMessages,
       {
@@ -123,8 +126,8 @@ export async function fetchContextMessages(
           opts.searchOptions?.vectorScoreThreshold ??
           DEFAULT_VECTOR_SCORE_THRESHOLD,
         embedding: embeddingFields?.embedding,
-        embeddingModel: embeddingFields?.embeddingModel
-          ? getModelName(embeddingFields.embeddingModel)
+        embeddingModel: embeddingFields?.textEmbeddingModel
+          ? getModelName(embeddingFields.textEmbeddingModel)
           : undefined,
       },
     );

package/src/client/streaming.ts

@@ -2,6 +2,7 @@ import {
   type ChunkDetector,
   smoothStream,
   type StreamTextTransform,
+  type TextStreamPart,
   type ToolSet,
 } from "ai";
 import type {
@@ -9,9 +10,7 @@ import type {
   StreamArgs,
   StreamDelta,
   StreamMessage,
-  vTextStreamPartV5,
 } from "../validators.js";
-import type { MessageDoc } from "../component/schema.js";
 import type {
   AgentComponent,
   RunActionCtx,
@@ -20,7 +19,7 @@ import type {
   SyncStreamsReturnValue,
 } from "./types.js";
 import { omit } from "convex-helpers";
-import type { Infer } from "convex/values";
+import { serializeTextStreamingPartsV5 } from "../parts.js";
 
 /**
  * A function that handles fetching stream deltas, used with the React hooks
@@ -65,9 +64,10 @@ export async function syncStreams(
 export async function abortStream(
   ctx: RunMutationCtx,
   component: AgentComponent,
-  args: {
-    reason: string;
-  } & ({ streamId: string } | { threadId: string; order: number }),
+  args: { reason: string } & (
+    | { streamId: string }
+    | { threadId: string; order: number }
+  ),
 ): Promise<boolean> {
   if ("streamId" in args) {
     return await ctx.runMutation(component.streams.abort, {
@@ -126,11 +126,20 @@ export type StreamingOptions = {
    * Defaults to 250.
    */
   throttleMs?: number;
+  /**
+   * If set to true, this will return immediately, as it would if you weren't
+   * saving the deltas. Otherwise, the call will "consume" the stream with
+   * .consumeStream(), which waits for the stream to finish before returning.
+   *
+   * When saving deltas, you're often not interactin with the stream otherwise.
+   */
+  returnImmediately?: boolean;
 };
 export const DEFAULT_STREAMING_OPTIONS = {
   // This chunks by sentences / clauses. Punctuation followed by whitespace.
   chunking: /[\p{P}\s]/u,
   throttleMs: 250,
+  returnImmediately: false,
 } satisfies StreamingOptions;
 
 export function mergeTransforms<TOOLS extends ToolSet>(
@@ -159,9 +168,7 @@ export function mergeTransforms<TOOLS extends ToolSet>(
 export class DeltaStreamer {
   public streamId: string | undefined;
   public readonly options: Required<StreamingOptions>;
-  #nextParts: Infer<typeof vTextStreamPartV5>[] = [];
-  #nextOrder: number;
-  #nextStepOrder: number;
+  #nextParts: TextStreamPart<ToolSet>[] = [];
   #latestWrite: number = 0;
   #ongoingWrite: Promise<void> | undefined;
   #cursor: number = 0;
@@ -173,51 +180,46 @@ export class DeltaStreamer {
     options: true | StreamingOptions,
     public readonly metadata: {
       threadId: string;
-      agentName: string | undefined;
-      model: string | undefined;
-      provider: string | undefined;
-      providerOptions: ProviderOptions | undefined;
-      userId: string | undefined;
-      order: number | undefined;
-      stepOrder: number | undefined;
-      abortSignal: AbortSignal | undefined;
+      userId?: string;
+      order: number;
+      stepOrder: number;
+      agentName?: string;
+      model?: string;
+      provider?: string;
+      providerOptions?: ProviderOptions;
+      abortSignal?: AbortSignal;
     },
   ) {
     this.options =
       typeof options === "boolean"
         ? DEFAULT_STREAMING_OPTIONS
-        : {
-            ...DEFAULT_STREAMING_OPTIONS,
-            ...options,
-          };
+        : { ...DEFAULT_STREAMING_OPTIONS, ...options };
     this.#nextParts = [];
-    this.#nextOrder = metadata.order ?? 0;
-    this.#nextStepOrder = (metadata.stepOrder ?? 0) + 1;
     this.abortController = new AbortController();
     if (metadata.abortSignal) {
       metadata.abortSignal.addEventListener("abort", async () => {
         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,
           });
         }
-        this.abortController.abort();
       });
     }
   }
-  public async addParts(parts: Infer<typeof vTextStreamPartV5>[]) {
+
+  public async addParts(parts: TextStreamPart<ToolSet>[]) {
     if (this.abortController.signal.aborted) {
       return;
     }
     if (!this.streamId) {
       this.streamId = await this.ctx.runMutation(
         this.component.streams.create,
-        {
-          ...omit(this.metadata, ["abortSignal"]),
-          order: this.#nextOrder,
-          stepOrder: this.#nextStepOrder,
-        },
+        omit(this.metadata, ["abortSignal"]),
       );
     }
     this.#nextParts.push(...parts);
@@ -234,6 +236,9 @@ export class DeltaStreamer {
       return;
     }
     const delta = this.#createDelta();
+    if (!delta) {
+      return;
+    }
     this.#latestWrite = Date.now();
     try {
       const success = await this.ctx.runMutation(
@@ -259,46 +264,46 @@ export class DeltaStreamer {
     }
   }
 
-  #createDelta(): StreamDelta {
+  #createDelta(): StreamDelta | undefined {
+    if (this.#nextParts.length === 0) {
+      return undefined;
+    }
     const start = this.#cursor;
     const end = start + this.#nextParts.length;
     this.#cursor = end;
-    const parts = this.#nextParts;
+    const parts = serializeTextStreamingPartsV5(this.#nextParts);
     this.#nextParts = [];
     if (!this.streamId) {
       throw new Error("Creating a delta before the stream is created");
     }
-    return {
+    return { streamId: this.streamId, start, end, parts };
+  }
+
+  public async finish() {
+    if (!this.streamId) {
+      return;
+    }
+    const finalDelta = this.#createDelta();
+    await this.#ongoingWrite;
+    await this.ctx.runMutation(this.component.streams.finish, {
       streamId: this.streamId,
-      start,
-      end,
-      parts,
-    };
+      finalDelta,
+    });
   }
 
-  public async finish(messages: MessageDoc[]) {
-    if (this.#ongoingWrite) {
-      await this.#ongoingWrite;
-      this.#ongoingWrite = undefined;
+  public async fail(reason: string) {
+    if (this.abortController.signal.aborted) {
+      return;
     }
+    this.abortController.abort();
     if (!this.streamId) {
-      throw new Error("Finish called before stream is created");
-    }
-    const lastMessage = messages.at(-1);
-    if (lastMessage) {
-      this.#nextOrder = lastMessage.order;
-      this.#nextStepOrder = lastMessage.stepOrder + 1;
-    } else {
-      console.warn("Step finished without generating a message");
+      return;
     }
-    const finalDelta =
-      this.#nextParts.length > 0 ? this.#createDelta() : undefined;
-    this.#nextParts = [];
-    const streamId = this.streamId;
-    this.streamId = undefined;
-    this.#cursor = 0;
-    await this.ctx.runMutation(this.component.streams.finish, {
-      streamId,
+    const finalDelta = this.#createDelta();
+    await this.#ongoingWrite;
+    await this.ctx.runMutation(this.component.streams.abort, {
+      streamId: this.streamId,
+      reason,
       finalDelta,
     });
   }