@convex-dev/agent 0.0.16-alpha.0 → 0.0.17-alpha.0

package/src/component/messages.test.ts

@@ -5,9 +5,32 @@ import { convexTest } from "convex-test";
 import schema from "./schema.js";
 import { api } from "./_generated/api.js";
 import { modules } from "./setup.test.js";
+import { Id } from "./_generated/dataModel.js";
+import { getMaxMessage } from "./messages.js";
 
 describe("agent", () => {
   test("add and subtract", async () => {
     const t = convexTest(schema, modules);
   });
+  test("getMaxMessage works for threads", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [
+        { message: { role: "user", content: "hello" } },
+        { message: { role: "assistant", content: "world" } },
+      ],
+    });
+    const maxMessage = await t.run(async (ctx) => {
+      return await getMaxMessage(ctx, thread._id as Id<"threads">, "test");
+    });
+    expect(maxMessage).toMatchObject({
+      _id: messages.at(-1)!._id,
+      order: 1,
+      stepOrder: 0,
+    });
+  });
 });

package/src/component/messages.ts

@@ -57,10 +57,11 @@ export async function deleteMessage(
   if (messageDoc.embeddingId) {
     await ctx.db.delete(messageDoc.embeddingId);
   }
-  if (messageDoc.fileId) {
-    const file = await ctx.db.get(messageDoc.fileId);
+  for (const { fileId } of messageDoc.files ?? []) {
+    if (!fileId) continue;
+    const file = await ctx.db.get(fileId);
     if (file) {
-      await ctx.db.patch(messageDoc.fileId, { refcount: file.refcount - 1 });
+      await ctx.db.patch(fileId, { refcount: file.refcount - 1 });
     }
   }
 }
@@ -117,9 +118,11 @@ async function addMessagesHandler(
   }
   const maxMessage = await getMaxMessage(ctx, threadId, userId);
   let order = maxMessage?.order ?? -1;
+  let stepOrder = maxMessage?.stepOrder ?? 0;
+  let lastMessageIsTool = maxMessage?.tool ?? false;
   const toReturn: Doc<"messages">[] = [];
   if (messages.length > 0) {
-    for (const { message, fileId, embedding, ...fields } of messages) {
+    for (const { message, files, embedding, ...fields } of messages) {
       let embeddingId: VectorTableId | undefined;
       if (embedding) {
         embeddingId = await insertVector(ctx, embedding.dimension, {
@@ -131,9 +134,13 @@ async function addMessagesHandler(
         });
       }
       const tool = isTool(message);
-      if (!tool) {
+      if (lastMessageIsTool) {
+        stepOrder++;
+      } else {
         order++;
+        stepOrder = 0;
       }
+      lastMessageIsTool = tool;
       const text = extractText(message);
       const messageId = await ctx.db.insert("messages", {
         ...rest,
@@ -145,11 +152,17 @@ async function addMessagesHandler(
         order,
         tool,
         text,
-        fileId,
+        files,
         status: pending ? "pending" : "success",
-        stepOrder: 0,
+        stepOrder,
       });
-      if (fileId) {
+      if (!fields.id) {
+        await ctx.db.patch(messageId, {
+          id: messageId,
+        });
+      }
+      for (const { fileId } of files ?? []) {
+        if (!fileId) continue;
         await ctx.db.patch(fileId, {
           refcount: (await ctx.db.get(fileId))!.refcount + 1,
         });
@@ -160,7 +173,8 @@ async function addMessagesHandler(
   return { messages: toReturn };
 }
 
-async function getMaxMessage(
+// exported for tests
+export async function getMaxMessage(
   ctx: QueryCtx,
   threadId: Id<"threads"> | undefined,
   userId: string | undefined
@@ -168,39 +182,32 @@ async function getMaxMessage(
   assert(threadId || userId, "One of threadId or userId is required");
   if (threadId) {
     return mergedStream(
-      ["success" as const, "pending" as const].map((status) =>
-        stream(ctx.db, schema)
-          .query("messages")
-          .withIndex("threadId_status_tool_order_stepOrder", (q) =>
-            q.eq("threadId", threadId).eq("status", status).eq("tool", false)
-          )
-          .order("desc")
+      [true, false].flatMap((tool) =>
+        ["success" as const, "pending" as const].map((status) =>
+          stream(ctx.db, schema)
+            .query("messages")
+            .withIndex("threadId_status_tool_order_stepOrder", (q) =>
+              q.eq("threadId", threadId).eq("status", status).eq("tool", tool)
+            )
+            .order("desc")
+        )
       ),
       ["order", "stepOrder"]
     ).first();
   } else {
-    // DO explicitly
-    const maxPending = await ctx.db
-      .query("messages")
-      .withIndex("userId_status_tool_order_stepOrder", (q) =>
-        q.eq("userId", userId).eq("status", "pending").eq("tool", false)
-      )
-      .order("desc")
-      .first();
-    const maxSuccess = await ctx.db
-      .query("messages")
-      .withIndex("userId_status_tool_order_stepOrder", (q) =>
-        q.eq("userId", userId).eq("status", "success").eq("tool", false)
-      )
-      .order("desc")
-      .first();
-    return maxPending
-      ? maxSuccess
-        ? maxPending.order > maxSuccess.order
-          ? maxPending
-          : maxSuccess
-        : maxPending
-      : maxSuccess ?? null;
+    return mergedStream(
+      [true, false].flatMap((tool) =>
+        ["success" as const, "pending" as const].map((status) =>
+          stream(ctx.db, schema)
+            .query("messages")
+            .withIndex("userId_status_tool_order_stepOrder", (q) =>
+              q.eq("userId", userId).eq("status", status).eq("tool", tool)
+            )
+            .order("desc")
+        )
+      ),
+      ["order", "stepOrder"]
+    ).first();
   }
 }
 
@@ -225,6 +232,7 @@ async function addStepHandler(
   assert(parentMessage, `Message ${args.parentMessageId} not found`);
   const order = parentMessage.order;
   assert(order !== undefined, `${args.parentMessageId} has no order`);
+  // TODO: only fetch the last one if we aren't failing pending steps
   let steps = await ctx.db
     .query("steps")
     .withIndex("parentMessageId_order_stepOrder", (q) =>
@@ -376,6 +384,12 @@ export const listMessagesByThreadId = query({
             return qq;
           })
           .order(order)
+          .filterWith(
+            async (m) =>
+              !before ||
+              m.order < before.order ||
+              (m.order === before.order && m.stepOrder < before.stepOrder)
+          )
       )
     );
     const messages = await mergedStream(streams, [
@@ -401,7 +415,6 @@ export const getThreadMessages = query({
   returns: paginationResultValidator(v.doc("messages")),
 });
 
-
 export const searchMessages = action({
   args: {
     userId: v.optional(v.string()),
@@ -420,6 +433,7 @@ export const searchMessages = action({
         threadId: args.threadId,
         text: args.text,
         limit,
+        beforeMessageId: args.beforeMessageId,
       });
     }
     if (args.vector) {
@@ -506,7 +520,10 @@ export const _fetchSearchMessages = internalQuery({
         m !== undefined &&
         m !== null &&
         !m.tool &&
-        (!beforeMessage || m.order <= beforeMessage.order)
+        (!beforeMessage ||
+          m.order < beforeMessage.order ||
+          (m.order === beforeMessage.order &&
+            m.stepOrder < beforeMessage.stepOrder))
     );
     messages.push(...(args.textSearchMessages ?? []));
     // TODO: prioritize more recent messages
@@ -597,9 +614,13 @@ export const textSearch = query({
     userId: v.optional(v.string()),
     text: v.string(),
     limit: v.number(),
+    beforeMessageId: v.optional(v.id("messages")),
   },
   handler: async (ctx, args) => {
     assert(args.userId || args.threadId, "Specify userId or threadId");
+    const beforeMessage =
+      args.beforeMessageId && (await ctx.db.get(args.beforeMessageId));
+    const order = beforeMessage?.order;
     const messages = await ctx.db
       .query("messages")
       .withSearchIndex("text_search", (q) =>
@@ -608,9 +629,21 @@ export const textSearch = query({
           : q.search("text", args.text).eq("threadId", args.threadId!)
       )
       // Just in case tool messages slip through
-      .filter((q) => q.eq(q.field("tool"), false))
+      .filter((q) => {
+        const qq = q.eq(q.field("tool"), false);
+        if (order) {
+          return q.and(qq, q.lte(q.field("order"), order));
+        }
+        return qq;
+      })
       .take(args.limit);
-    return messages;
+    return messages.filter(
+      (m) =>
+        !beforeMessage ||
+        m.order < beforeMessage.order ||
+        (m.order === beforeMessage.order &&
+          m.stepOrder < beforeMessage.stepOrder)
+    );
   },
   returns: v.array(v.doc("messages")),
 });

package/src/component/schema.ts

@@ -9,6 +9,10 @@ import {
   vSource,
   vLanguageModelV1CallWarning,
   vFinishReason,
+  vProviderOptions,
+  vProviderMetadata,
+  vReasoningDetails,
+  vFile,
 } from "../validators.js";
 import { typedV } from "convex-helpers/validators";
 import vectorTables, { vVectorId } from "./vector/tables.js";
@@ -29,35 +33,44 @@ export const schema = defineSchema({
     id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
     userId: v.optional(v.string()), // useful for future indexes (text search)
     threadId: v.id("threads"),
+    // TODO: is this redunant with message at last step @ order - 1?
     parentMessageId: v.optional(v.id("messages")),
     stepId: v.optional(v.id("steps")),
-    agentName: v.optional(v.string()),
-    message: v.optional(vMessage),
+    // Repeats until a non-tool message.
+    order: v.number(),
+    stepOrder: v.number(),
+    embeddingId: v.optional(vVectorId),
     error: v.optional(v.string()),
+    status: vMessageStatus,
+
+    // Context on how it was generated
+    agentName: v.optional(v.string()),
     model: v.optional(v.string()),
     provider: v.optional(v.string()),
+    providerOptions: v.optional(vProviderOptions), // Sent to model
+
+    // The result
+    message: v.optional(vMessage),
+    // Convenience fields extracted from the message
+    tool: v.boolean(), // either tool call (assistant) or tool result (tool)
     text: v.optional(v.string()),
-    embeddingId: v.optional(vVectorId),
-    // TODO: add sub-messages back in? or be able to skip them?
-    tool: v.boolean(),
-    // Repeats until a non-tool message.
-    // Unset if it's not in a thread.
-    order: v.number(),
-    stepOrder: v.number(),
+    files: v.optional(v.array(vFile)),
+
+    // Result metadata
     usage: v.optional(vUsage),
-    finishReason: v.optional(vFinishReason),
-    providerMetadata: v.optional(v.record(v.string(), v.any())),
+    providerMetadata: v.optional(vProviderMetadata), // Received from model
     sources: v.optional(v.array(vSource)),
     reasoning: v.optional(v.string()),
+    reasoningDetails: v.optional(vReasoningDetails),
     warnings: v.optional(v.array(vLanguageModelV1CallWarning)),
-    fileId: v.optional(v.id("files")),
-    status: vMessageStatus,
+    finishReason: v.optional(vFinishReason),
   })
     // Allows finding successful visible messages in order
     // Also surface pending messages separately to e.g. stream
     .index("threadId_status_tool_order_stepOrder", [
       "threadId",
       "status",
+      // TODO: we might not need this to be in the index..
       "tool",
       "order",
       "stepOrder",

package/src/component/vector/index.ts

@@ -65,14 +65,13 @@ export const deleteBatchForThread = mutation({
   handler: async (ctx, args) => {
     const tableName = getVectorTableName(args.vectorDimension);
     const vectors = await mergedStream(
-      ["thread", "memory"].map((kind) =>
+      ["thread", "memory"].map((table) =>
         stream(ctx.db, schema)
           .query(tableName)
-          // eslint-disable-next-line @typescript-eslint/no-explicit-any
-          .withIndex("model_kind_threadId" as any, (q) =>
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            (q.eq("model", args.model) as any)
-              .eq("kind", kind)
+          .withIndex("model_table_threadId", (q) =>
+            q
+              .eq("model", args.model)
+              .eq("table", table)
               .eq("threadId", args.threadId)
           )
       ),

package/src/component/vector/tables.ts

@@ -47,6 +47,21 @@ function table<D extends number>(dimensions: D): Table<D> {
     .index("model_table_threadId", ["model", "table", "threadId"]);
 }
 
+type Table<D extends number> = TableDefinition<
+  VObject<ObjectType<typeof embeddings>, typeof embeddings>,
+  { model_table_threadId: ["model", "table", "threadId", "_creationTime"] },
+  GenericTableSearchIndexes,
+  VectorIndex<D>
+>;
+
+type VectorIndex<D extends number> = {
+  vector: {
+    vectorField: "vector";
+    dimensions: D;
+    filterFields: "model_table_userId" | "model_table_threadId";
+  };
+};
+
 export type VectorSchema = SchemaDefinition<
   { [key in VectorTableName]: Table<128> },
   true
@@ -80,21 +95,6 @@ export const vVectorId = v.union(
   VId<(typeof VectorTableNames)[number]>[]
 >;
 
-type Table<D extends number> = TableDefinition<
-  VObject<ObjectType<typeof embeddings>, typeof embeddings>,
-  { id: ["id"] },
-  GenericTableSearchIndexes,
-  VectorIndex<D>
->;
-
-type VectorIndex<D extends number> = {
-  vector: {
-    vectorField: "vector";
-    dimensions: D;
-    filterFields: string;
-  };
-};
-
 export function getVectorTableName(dimension: VectorDimension) {
   return `embeddings_${dimension}` as VectorTableName;
 }

package/src/mapping.ts

@@ -11,6 +11,7 @@ import {
   type Message as AIMessage,
   type UserContent,
 } from "ai";
+import type { FileUIPart } from "@ai-sdk/ui-utils";
 import { assert } from "convex-helpers";
 import {
   MessageWithMetadata,
@@ -39,9 +40,11 @@ export type SerializedMessage = SerializeUrlsAndUint8Arrays<CoreMessage>;
 export function serializeMessage(
   messageWithId: CoreMessage & { id?: string }
 ): SerializedMessage {
-  const { id: _, ...message } = messageWithId;
+  const { id: _, experimental_providerMetadata, ...message } = messageWithId;
   const content = message.content;
   return {
+    // for backwards compatibility
+    providerOptions: experimental_providerMetadata,
     ...message,
     content: serializeContent(content),
   } as SerializedMessage;
@@ -90,6 +93,7 @@ export function serializeNewMessagesInStep<TOOLS extends ToolSet>(
     provider: metadata.provider,
     providerMetadata: step.providerMetadata,
     reasoning: step.reasoning,
+    reasoningDetails: step.reasoningDetails,
     usage: step.usage,
     warnings: step.warnings,
     finishReason: step.finishReason,
@@ -101,33 +105,41 @@ export function serializeNewMessagesInStep<TOOLS extends ToolSet>(
     step.toolResults.length > 0
       ? step.response.messages.slice(-2)
       : step.response.messages.slice(-1)
-  ).map((message) => ({
-    message: serializeMessage(message),
-    id: message.id,
-    ...(message.role === "tool" ? toolFields : assistantFields),
-    // fileId: message.fileId,
-  }));
+  ).map(
+    (message): MessageWithMetadata => ({
+      message: serializeMessage(message),
+      id: message.id,
+      ...(message.role === "tool" ? toolFields : assistantFields),
+      text: step.text,
+      // fileId: message.fileId,
+      files: step.files.map((file) => ({
+        mimeType: file.mimeType,
+        data: serializeDataOrUrl(file.uint8Array ?? file.base64),
+        // TODO: if the file is big, store it and populate url, fileId
+      })),
+    })
+  );
   return messages;
 }
 
 export function serializeObjectResult(
-  step: GenerateObjectResult<unknown>,
+  result: GenerateObjectResult<unknown>,
   metadata: { model: string; provider: string }
 ): StepWithMessagesWithMetadata {
-  const text = JSON.stringify(step.object);
+  const text = JSON.stringify(result.object);
 
   return {
     messages: [
       {
         message: { role: "assistant" as const, content: text },
-        id: step.response.id,
+        id: result.response.id,
         model: metadata.model,
         provider: metadata.provider,
-        providerMetadata: step.providerMetadata,
-        finishReason: step.finishReason,
+        providerMetadata: result.providerMetadata,
+        finishReason: result.finishReason,
         text,
-        usage: step.usage,
-        warnings: step.warnings,
+        usage: result.usage,
+        warnings: result.warnings,
       },
     ],
     step: {
@@ -136,23 +148,22 @@ export function serializeObjectResult(
       stepType: "initial",
       toolCalls: [],
       toolResults: [],
-      usage: step.usage,
-      warnings: step.warnings,
-      finishReason: step.finishReason,
-      request: step.request,
+      usage: result.usage,
+      warnings: result.warnings,
+      finishReason: result.finishReason,
+      providerMetadata: result.providerMetadata,
+      request: result.request,
       response: {
-        ...step.response,
-        timestamp: step.response.timestamp.getTime(),
+        ...result.response,
+        timestamp: result.response.timestamp.getTime(),
         messages: [
           serializeMessageWithId({
             role: "assistant" as const,
             content: text,
-            id: step.response.id,
+            id: result.response.id,
           }),
         ],
       },
-      providerMetadata: step.providerMetadata,
-      experimental_providerMetadata: step.experimental_providerMetadata,
     },
   };
 }
@@ -161,16 +172,19 @@ export function serializeContent(content: Content): SerializedContent {
   if (typeof content === "string") {
     return content;
   }
-  const serialized = content.map((part) => {
-    switch (part.type) {
-      case "image":
-        return { ...part, image: serializeDataOrUrl(part.image) };
-      case "file":
-        return { ...part, file: serializeDataOrUrl(part.data) };
-      default:
-        return part;
+  const serialized = content.map(
+    ({ experimental_providerMetadata, ...rest }) => {
+      const part = { providerOptions: experimental_providerMetadata, ...rest };
+      switch (part.type) {
+        case "image":
+          return { ...part, image: serializeDataOrUrl(part.image) };
+        case "file":
+          return { ...part, file: serializeDataOrUrl(part.data) };
+        default:
+          return part;
+      }
     }
-  });
+  );
   return serialized as SerializedContent;
 }
 
@@ -222,6 +236,25 @@ function deserializeUrl(urlOrString: string | ArrayBuffer): URL | DataContent {
   return urlOrString;
 }
 
+export function toUIFilePart(file: {
+  data?: ArrayBuffer | string;
+  url?: string;
+  mimeType: string;
+}): FileUIPart {
+  return {
+    type: "file",
+    data:
+      file.data instanceof ArrayBuffer
+        ? encodeBase64(file.data)
+        : file.url ?? file.data ?? "",
+    mimeType: file.mimeType,
+  };
+}
+
+function encodeBase64(data: ArrayBuffer): string {
+  return Buffer.from(data).toString("base64");
+}
+
 export function promptOrMessagesToCoreMessages(args: {
   prompt?: string;
   messages?: CoreMessage[] | AIMessageWithoutId[];