@convex-dev/agent 0.1.18 → 0.2.0

package/src/client/messages.ts

@@ -0,0 +1,205 @@
+import type { ModelMessage } from "ai";
+import type { PaginationOptions, PaginationResult } from "convex/server";
+import type { MessageDoc } from "../component/schema.js";
+import { validateVectorDimension } from "../component/vector/tables.js";
+import {
+  vMessageWithMetadata,
+  type Message,
+  type MessageEmbeddings,
+  type MessageEmbeddingsWithDimension,
+  type MessageStatus,
+  type MessageWithMetadata,
+} from "../validators.js";
+import { serializeMessage } from "./index.js";
+import type { AgentComponent, RunMutationCtx, RunQueryCtx } from "./types.js";
+import { parse } from "convex-helpers/validators";
+
+/**
+ * List messages from a thread.
+ * @param ctx A ctx object from a query, mutation, or action.
+ * @param component The agent component, usually `components.agent`.
+ * @param args.threadId The thread to list messages from.
+ * @param args.paginationOpts Pagination options (e.g. via usePaginatedQuery).
+ * @param args.excludeToolMessages Whether to exclude tool messages.
+ *   False by default.
+ * @param args.statuses What statuses to include. All by default.
+ * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
+ */
+export async function listMessages(
+  ctx: RunQueryCtx,
+  component: AgentComponent,
+  args: {
+    threadId: string;
+    paginationOpts: PaginationOptions;
+    excludeToolMessages?: boolean;
+    statuses?: MessageStatus[];
+  },
+): Promise<PaginationResult<MessageDoc>> {
+  if (args.paginationOpts.numItems === 0) {
+    return {
+      page: [],
+      isDone: true,
+      continueCursor: args.paginationOpts.cursor ?? "",
+    };
+  }
+  return ctx.runQuery(component.messages.listMessagesByThreadId, {
+    order: "desc",
+    ...args,
+  });
+}
+
+export type SaveMessagesArgs = {
+  threadId: string;
+  userId?: string | null;
+  /**
+   * The message that these messages are in response to. They will be
+   * the same "order" as this message, at increasing stepOrder(s).
+   */
+  promptMessageId?: string;
+  /**
+   * The messages to save.
+   */
+  messages: (ModelMessage | Message)[];
+  /**
+   * Metadata to save with the messages. Each element corresponds to the
+   * message at the same index.
+   */
+  metadata?: Omit<MessageWithMetadata, "message">[];
+  /**
+   * If true, it will fail any pending steps.
+   * Defaults to false.
+   */
+  failPendingSteps?: boolean;
+  /**
+   * The embeddings to save with the messages.
+   */
+  embeddings?: MessageEmbeddings;
+  /**
+   * A pending message ID to replace when adding messages.
+   */
+  pendingMessageId?: string;
+};
+
+/**
+ * Explicitly save messages associated with the thread (& user if provided)
+ */
+export async function saveMessages(
+  ctx: RunMutationCtx,
+  component: AgentComponent,
+  args: SaveMessagesArgs & {
+    /**
+     * The agent name to associate with the messages.
+     */
+    agentName?: string;
+  },
+): Promise<{ messages: MessageDoc[] }> {
+  let embeddings: MessageEmbeddingsWithDimension | undefined;
+  if (args.embeddings) {
+    const dimension = args.embeddings.vectors.find((v) => v !== null)?.length;
+    if (dimension) {
+      validateVectorDimension(dimension);
+      embeddings = {
+        model: args.embeddings.model,
+        dimension,
+        vectors: args.embeddings.vectors,
+      };
+    }
+  }
+  const result = await ctx.runMutation(component.messages.addMessages, {
+    threadId: args.threadId,
+    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: allFileIds.length > 0 ? allFileIds : undefined,
+        });
+      }),
+    ),
+    failPendingSteps: args.failPendingSteps ?? false,
+  });
+  return { messages: result.messages };
+}
+
+export type SaveMessageArgs = {
+  threadId: string;
+  userId?: string | null;
+  /**
+   * Metadata to save with the messages. Each element corresponds to the
+   * message at the same index.
+   */
+  metadata?: Omit<MessageWithMetadata, "message">;
+  /**
+   * 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;
+      /**
+       * The message to save.
+       */
+      message: ModelMessage | Message;
+    }
+  | {
+      /*
+       * The prompt to save with the message.
+       */
+      prompt: string;
+      message?: undefined;
+    }
+);
+
+/**
+ * Save a message to the thread.
+ * @param ctx A ctx object from a mutation or action.
+ * @param args The message and what to associate it with (user / thread)
+ * You can pass extra metadata alongside the message, e.g. associated fileIds.
+ * @returns The messageId of the saved message.
+ */
+export async function saveMessage(
+  ctx: RunMutationCtx,
+  component: AgentComponent,
+  args: SaveMessageArgs & {
+    /**
+     * The agent name to associate with the message.
+     */
+    agentName?: string;
+  },
+) {
+  let embeddings: { vectors: number[][]; model: string } | undefined;
+  if (args.embedding && args.embedding.vector) {
+    embeddings = {
+      model: args.embedding.model,
+      vectors: [args.embedding.vector],
+    };
+  }
+  const { messages } = await saveMessages(ctx, component, {
+    threadId: args.threadId,
+    userId: args.userId ?? undefined,
+    agentName: args.agentName,
+    pendingMessageId: args.pendingMessageId,
+    messages:
+      args.prompt !== undefined
+        ? [{ role: "user", content: args.prompt }]
+        : [args.message],
+    metadata: args.metadata ? [args.metadata] : undefined,
+    embeddings,
+  });
+  const message = messages.at(-1)!;
+  return { messageId: message._id, message };
+}

package/src/client/search.ts

@@ -5,19 +5,22 @@ import type {
   RunQueryCtx,
 } from "./types.js";
 import type { MessageDoc } from "../component/schema.js";
-import type { CoreMessage } from "ai";
+import type { EmbeddingModel, LanguageModel, ModelMessage } from "ai";
 import { assert } from "convex-helpers";
 import {
   DEFAULT_MESSAGE_RANGE,
   DEFAULT_RECENT_MESSAGES,
   extractText,
 } from "../shared.js";
+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;
+  textEmbeddingModel: string | EmbeddingModel<string>;
 }>;
 
 /**
@@ -34,7 +37,7 @@ export async function fetchContextMessages(
   args: {
     userId: string | undefined;
     threadId: string | undefined;
-    messages: CoreMessage[];
+    messages: (ModelMessage | Message)[];
     /**
      * If provided, it will search for messages up to and including this message.
      * Note: if this is far in the past, text and vector search results may be more
@@ -95,20 +98,21 @@ 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,
       {
         searchAllMessagesForUserId: opts?.searchOtherThreads
-          ? args.userId ??
+          ? (args.userId ??
             (args.threadId &&
               (
                 await ctx.runQuery(component.threads.getThread, {
                   threadId: args.threadId,
                 })
-              )?.userId)
+              )?.userId))
           : undefined,
         threadId: args.threadId,
         beforeMessageId: args.upToAndIncludingMessageId,
@@ -122,7 +126,9 @@ export async function fetchContextMessages(
           opts.searchOptions?.vectorScoreThreshold ??
           DEFAULT_VECTOR_SCORE_THRESHOLD,
         embedding: embeddingFields?.embedding,
-        embeddingModel: embeddingFields?.embeddingModel,
+        embeddingModel: embeddingFields?.textEmbeddingModel
+          ? getModelName(embeddingFields.textEmbeddingModel)
+          : undefined,
       },
     );
     // TODO: track what messages we used for context
@@ -170,3 +176,24 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   }
   return result;
 }
+
+export function getModelName(
+  embeddingModel: string | EmbeddingModel<string> | LanguageModel,
+): string {
+  if (typeof embeddingModel === "string") {
+    if (embeddingModel.includes("/")) {
+      return embeddingModel.split("/").slice(1).join("/");
+    }
+    return embeddingModel;
+  }
+  return embeddingModel.modelId;
+}
+
+export function getProviderName(
+  embeddingModel: string | EmbeddingModel<string> | LanguageModel,
+): string {
+  if (typeof embeddingModel === "string") {
+    return embeddingModel.split("/").at(0)!;
+  }
+  return embeddingModel.provider;
+}

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,
-  TextStreamPart,
 } from "../validators.js";
-import type { MessageDoc } from "../component/schema.js";
 import type {
   AgentComponent,
   RunActionCtx,
@@ -20,6 +19,7 @@ import type {
   SyncStreamsReturnValue,
 } from "./types.js";
 import { omit } from "convex-helpers";
+import { serializeTextStreamingPartsV5 } from "../parts.js";
 
 /**
  * A function that handles fetching stream deltas, used with the React hooks
@@ -64,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, {
@@ -158,9 +159,7 @@ export function mergeTransforms<TOOLS extends ToolSet>(
 export class DeltaStreamer {
   public streamId: string | undefined;
   public readonly options: Required<StreamingOptions>;
-  #nextParts: TextStreamPart[] = [];
-  #nextOrder: number;
-  #nextStepOrder: number;
+  #nextParts: TextStreamPart<ToolSet>[] = [];
   #latestWrite: number = 0;
   #ongoingWrite: Promise<void> | undefined;
   #cursor: number = 0;
@@ -172,51 +171,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: TextStreamPart[]) {
+
+  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);
@@ -233,6 +227,9 @@ export class DeltaStreamer {
       return;
     }
     const delta = this.#createDelta();
+    if (!delta) {
+      return;
+    }
     this.#latestWrite = Date.now();
     try {
       const success = await this.ctx.runMutation(
@@ -258,46 +255,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,
     });
   }

package/src/client/threads.ts

@@ -0,0 +1,78 @@
+import type { WithoutSystemFields } from "convex/server";
+import type { ThreadDoc } from "../component/schema.js";
+import type { AgentComponent, RunMutationCtx, RunQueryCtx } from "./types.js";
+
+/**
+ * Create a thread to store messages with an Agent.
+ * @param ctx The context from a mutation or action.
+ * @param component The Agent component, usually `components.agent`.
+ * @param args The associated thread metadata.
+ * @returns The id of the created thread.
+ */
+export async function createThread(
+  ctx: RunMutationCtx,
+  component: AgentComponent,
+  args?: { userId?: string | null; title?: string; summary?: string },
+) {
+  const { _id: threadId } = await ctx.runMutation(
+    component.threads.createThread,
+    {
+      userId: args?.userId ?? undefined,
+      title: args?.title,
+      summary: args?.summary,
+    },
+  );
+  return threadId;
+}
+
+/**
+ * Get the metadata for a thread.
+ * @param ctx A ctx object from a query, mutation, or action.
+ * @param args.threadId The thread to get the metadata for.
+ * @returns The metadata for the thread.
+ */
+export async function getThreadMetadata(
+  ctx: RunQueryCtx,
+  component: AgentComponent,
+  args: { threadId: string },
+): Promise<ThreadDoc> {
+  const thread = await ctx.runQuery(component.threads.getThread, {
+    threadId: args.threadId,
+  });
+  if (!thread) {
+    throw new Error("Thread not found");
+  }
+  return thread;
+}
+
+export async function updateThreadMetadata(
+  ctx: RunMutationCtx,
+  component: AgentComponent,
+  args: { threadId: string; patch: Partial<WithoutSystemFields<ThreadDoc>> },
+) {
+  return ctx.runMutation(component.threads.updateThread, {
+    threadId: args.threadId,
+    patch: args.patch,
+  });
+}
+
+/**
+ * Search for threads by title, paginated.
+ * @param ctx The context passed from the query/mutation/action.
+ * @returns The threads matching the search, paginated.
+ */
+export async function searchThreadTitles(
+  ctx: RunQueryCtx,
+  component: AgentComponent,
+  {
+    userId,
+    query,
+    limit,
+  }: { userId?: string | undefined; query: string; limit?: number },
+): Promise<ThreadDoc[]> {
+  return ctx.runQuery(component.threads.searchThreadTitles, {
+    userId,
+    query,
+    limit: limit ?? 10,
+  });
+}