@convex-dev/agent 0.1.1-alpha.0 → 0.1.2-alpha.0

package/src/component/messages.ts

@@ -8,7 +8,7 @@ import {
   isTool,
 } from "../shared.js";
 import {
-  paginationResultValidator,
+  vPaginationResult,
   vMessageEmbeddings,
   vMessageStatus,
   vMessageWithMetadata,
@@ -180,11 +180,12 @@ async function addMessagesHandler(
         error: fail ? "Parent message failed" : undefined,
         stepOrder,
       });
-      if (!message.id) {
-        await ctx.db.patch(messageId, {
-          id: messageId,
-        });
-      }
+      // Let's just not set the id field and have it set only in explicit cases.
+      // if (!message.id) {
+      //   await ctx.db.patch(messageId, {
+      //     id: messageId,
+      //   });
+      // }
       for (const { fileId } of message.files ?? []) {
         if (!fileId) continue;
         await ctx.db.patch(fileId, {
@@ -214,7 +215,7 @@ function orderedMessagesStream(
 ) {
   return mergedStream(
     [true, false].flatMap((tool) =>
-      ["success" as const, "pending" as const].map((status) =>
+      messageStatuses.map((status) =>
         stream(ctx.db, schema)
           .query("messages")
           .withIndex("threadId_status_tool_order_stepOrder", (q) => {
@@ -244,7 +245,7 @@ const addStepArgs = {
 
 export const addStep = mutation({
   args: addStepArgs,
-  returns: v.array(v.doc("steps")),
+  returns: v.array(vMessageDoc),
   handler: addStepHandler,
 });
 async function addStepHandler(
@@ -279,7 +280,7 @@ async function addStepHandler(
     status: step.finishReason === "stop" ? "success" : "pending",
     step,
   });
-  await addMessagesHandler(ctx, {
+  const added = await addMessagesHandler(ctx, {
     userId: args.userId,
     threadId: args.threadId,
     stepId,
@@ -293,8 +294,7 @@ async function addStepHandler(
   if (step.finishReason === "stop") {
     await commitMessageHandler(ctx, { messageId: args.promptMessageId });
   }
-  steps.push((await ctx.db.get(stepId))!);
-  return steps;
+  return added.messages.map(publicMessage);
 }
 
 export const rollbackMessage = mutation({
@@ -388,7 +388,8 @@ export const listMessagesByThreadId = query({
     excludeToolMessages: v.optional(v.boolean()),
     /** @deprecated Use excludeToolMessages instead. */
     isTool: v.optional(v.literal("use excludeToolMessages instead of this")),
-    order: v.optional(v.union(v.literal("asc"), v.literal("desc"))),
+    /** What order to sort the messages in. To get the latest, use "desc". */
+    order: v.union(v.literal("asc"), v.literal("desc")),
     paginationOpts: v.optional(paginationOptsValidator),
     statuses: v.optional(v.array(vMessageStatus)),
     upToAndIncludingMessageId: v.optional(v.id("messages")),
@@ -439,7 +440,7 @@ export const listMessagesByThreadId = query({
     );
     return { ...messages, page: messages.page.map(publicMessage) };
   },
-  returns: paginationResultValidator(vMessageDoc),
+  returns: vPaginationResult(vMessageDoc),
 });
 
 /** @deprecated Use listMessagesByThreadId instead. */
@@ -448,7 +449,7 @@ export const getThreadMessages = query({
   handler: async () => {
     throw new Error("Use listMessagesByThreadId instead of getThreadMessages");
   },
-  returns: paginationResultValidator(vMessageDoc),
+  returns: vPaginationResult(vMessageDoc),
 });
 
 export const searchMessages = action({

package/src/component/schema.ts

@@ -14,6 +14,7 @@ import {
   vReasoningDetails,
   vFile,
   vFileWithStringId,
+  vTextStreamPart,
 } from "../validators.js";
 import { typedV } from "convex-helpers/validators";
 import vectorTables, { vVectorId } from "./vector/tables.js";
@@ -105,6 +106,51 @@ export const schema = defineSchema({
       "stepOrder",
     ]),
 
+  // Status: if it's done, it's deleted, then deltas are vacuumed
+  streamingMessages: defineTable({
+    // extra metadata?
+    userId: v.optional(v.string()),
+    agentName: v.optional(v.string()),
+    model: v.optional(v.string()),
+    provider: v.optional(v.string()),
+    providerOptions: v.optional(vProviderOptions), // Sent to model
+
+    threadId: v.id("threads"),
+    order: v.number(),
+    /**
+     * The step order of the first message in the stream.
+     * If the stream ends up with both a tool call and a tool result,
+     * the stepOrder of the result will be +1 of the tool call.
+     */
+    stepOrder: v.number(),
+    state: v.union(
+      v.object({
+        kind: v.literal("streaming"),
+        lastHeartbeat: v.number(),
+        timeoutFnId: v.optional(v.id("_scheduled_functions")),
+      }),
+      v.object({
+        kind: v.literal("finished"),
+        endedAt: v.number(),
+      }),
+      v.object({
+        kind: v.literal("error"),
+        error: v.string(),
+      })
+    ),
+  })
+    // There should only be one per "order" index
+    // If another exists, it's deleted and replaced
+    .index("threadId_order_stepOrder", ["threadId", "order", "stepOrder"]),
+
+  streamDeltas: defineTable({
+    streamId: v.id("streamingMessages"),
+    // the indexes work like: 0 <first> 1 <second> 2 <third> 3 ...
+    start: v.number(), // inclusive
+    end: v.number(), // exclusive
+    parts: v.array(vTextStreamPart),
+  }).index("streamId_start_end", ["streamId", "start", "end"]),
+
   memories: defineTable({
     threadId: v.optional(v.id("threads")),
     userId: v.optional(v.string()),

package/src/component/streams.ts

@@ -0,0 +1,206 @@
+import { v } from "convex/values";
+import {
+  internalMutation,
+  mutation,
+  MutationCtx,
+  query,
+} from "./_generated/server";
+import { omit, pick } from "convex-helpers";
+import { StreamDelta, vStreamDelta, vStreamMessage } from "../validators";
+import schema from "./schema";
+import { internal } from "./_generated/api";
+import { Id } from "./_generated/dataModel";
+
+const MAX_DELTAS_PER_REQUEST = 1000;
+const MAX_DELTAS_PER_STREAM = 100;
+const TIMEOUT_INTERVAL = 1000 * 60; // 1 minute
+
+const deltaValidator = schema.tables.streamDeltas.validator;
+
+export const addDelta = mutation({
+  args: deltaValidator,
+  returns: v.null(),
+  handler: async (ctx, args) => {
+    await ctx.db.insert("streamDeltas", args);
+    await heartbeatStream(ctx, { streamId: args.streamId });
+  },
+});
+
+export const listDeltas = query({
+  args: {
+    threadId: v.id("threads"),
+    cursors: v.array(
+      v.object({
+        streamId: v.id("streamingMessages"),
+        cursor: v.number(),
+      })
+    ),
+  },
+  returns: v.array(vStreamDelta),
+  handler: async (ctx, args): Promise<StreamDelta[]> => {
+    let totalDeltas = 0;
+    const deltas: StreamDelta[] = [];
+    for (const cursor of args.cursors) {
+      const streamDeltas = await ctx.db
+        .query("streamDeltas")
+        .withIndex("streamId_start_end", (q) =>
+          q.eq("streamId", cursor.streamId).gte("start", cursor.cursor)
+        )
+        .take(
+          Math.min(MAX_DELTAS_PER_STREAM, MAX_DELTAS_PER_REQUEST - totalDeltas)
+        );
+      totalDeltas += streamDeltas.length;
+      deltas.push(
+        ...streamDeltas.map((d) => ({
+          streamId: d.streamId.toString(),
+          ...pick(d, ["start", "end", "parts"]),
+        }))
+      );
+      if (totalDeltas >= MAX_DELTAS_PER_REQUEST) {
+        break;
+      }
+    }
+    return deltas;
+  },
+});
+
+export const create = mutation({
+  args: omit(schema.tables.streamingMessages.validator.fields, ["state"]),
+  returns: v.id("streamingMessages"),
+  handler: async (ctx, args) => {
+    const state = {
+      kind: "streaming" as const,
+      lastHeartbeat: Date.now(),
+    };
+    const streamId = await ctx.db.insert("streamingMessages", {
+      ...args,
+      state,
+    });
+    const timeoutFnId = await ctx.scheduler.runAfter(
+      TIMEOUT_INTERVAL,
+      internal.streams.timeoutStream,
+      { streamId }
+    );
+    await ctx.db.patch(streamId, { state: { ...state, timeoutFnId } });
+    return streamId;
+  },
+});
+
+export const list = query({
+  args: {
+    threadId: v.id("threads"),
+  },
+  returns: v.array(vStreamMessage),
+  handler: async (ctx, args) => {
+    return ctx.db
+      .query("streamingMessages")
+      .withIndex("threadId_order_stepOrder", (q) =>
+        q.eq("threadId", args.threadId)
+      )
+      .order("desc")
+      .take(100)
+      .then((msgs) =>
+        msgs.map((m) => ({
+          streamId: m._id,
+          ...pick(m, [
+            "order",
+            "stepOrder",
+            "userId",
+            "agentName",
+            "model",
+            "provider",
+            "providerOptions",
+          ]),
+        }))
+      );
+  },
+});
+
+export const finish = mutation({
+  args: {
+    streamId: v.id("streamingMessages"),
+    finalDelta: v.optional(deltaValidator),
+  },
+  returns: v.null(),
+  handler: async (ctx, args) => {
+    if (args.finalDelta) {
+      await ctx.db.insert("streamDeltas", args.finalDelta);
+    }
+    const stream = await ctx.db.get(args.streamId);
+    if (!stream) {
+      throw new Error(`Stream not found: ${args.streamId}`);
+    }
+    if (stream.state.kind !== "streaming") {
+      throw new Error(`Stream is not streaming: ${args.streamId}`);
+    }
+    if (stream.state.timeoutFnId) {
+      const timeoutFn = await ctx.db.system.get(stream.state.timeoutFnId);
+      if (timeoutFn?.state.kind === "pending") {
+        await ctx.scheduler.cancel(stream.state.timeoutFnId);
+      }
+    }
+    await ctx.db.patch(args.streamId, {
+      state: { kind: "finished", endedAt: Date.now() },
+    });
+  },
+});
+
+async function heartbeatStream(
+  ctx: MutationCtx,
+  args: { streamId: Id<"streamingMessages"> }
+) {
+  const stream = await ctx.db.get(args.streamId);
+  if (!stream) {
+    console.warn("Stream not found", args.streamId);
+    return;
+  }
+  if (stream.state.kind !== "streaming") {
+    console.warn("Stream is not streaming", args.streamId);
+    return;
+  }
+  if (Date.now() - stream.state.lastHeartbeat < TIMEOUT_INTERVAL / 4) {
+    // Debounce heartbeating.
+    return;
+  }
+  if (!stream.state.timeoutFnId) {
+    throw new Error("Stream has no timeout function");
+  }
+  const timeoutFn = await ctx.db.system.get(stream.state.timeoutFnId);
+  if (!timeoutFn) {
+    throw new Error("Timeout function not found");
+  }
+  if (timeoutFn.state.kind !== "pending") {
+    throw new Error("Timeout function is not pending");
+  }
+  await ctx.scheduler.cancel(stream.state.timeoutFnId);
+  const timeoutFnId = await ctx.scheduler.runAfter(
+    TIMEOUT_INTERVAL,
+    internal.streams.timeoutStream,
+    { streamId: args.streamId }
+  );
+  await ctx.db.patch(args.streamId, {
+    state: {
+      kind: "streaming",
+      lastHeartbeat: Date.now(),
+      timeoutFnId,
+    },
+  });
+}
+
+export const timeoutStream = internalMutation({
+  args: { streamId: v.id("streamingMessages") },
+  returns: v.null(),
+  handler: async (ctx, args) => {
+    const stream = await ctx.db.get(args.streamId);
+    if (!stream) {
+      console.warn("Stream not found", args.streamId);
+      return;
+    }
+    await ctx.db.patch(args.streamId, {
+      state: {
+        kind: "finished",
+        endedAt: Date.now(),
+      },
+    });
+  },
+});

package/src/component/threads.ts

@@ -2,7 +2,7 @@ import { assert, omit, pick } from "convex-helpers";
 import { paginator } from "convex-helpers/server/pagination";
 import { partial } from "convex-helpers/validators";
 import { ObjectType } from "convex/values";
-import { paginationResultValidator } from "../validators.js";
+import { vPaginationResult } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import { Doc } from "./_generated/dataModel.js";
 import {
@@ -38,7 +38,7 @@ export const getThread = query({
 
 export const listThreadsByUserId = query({
   args: {
-    userId: v.string(),
+    userId: v.optional(v.string()),
     order: v.optional(v.union(v.literal("asc"), v.literal("desc"))),
     paginationOpts: v.optional(paginationOptsValidator),
   },
@@ -53,7 +53,7 @@ export const listThreadsByUserId = query({
       page: threads.page.map(publicThread),
     };
   },
-  returns: paginationResultValidator(vThreadDoc),
+  returns: vPaginationResult(vThreadDoc),
 });
 
 const vThread = schema.tables.threads.validator;

package/src/component/users.ts

@@ -13,7 +13,7 @@ import { schema, v } from "./schema.js";
 import { deleteMessage } from "./messages.js";
 import { paginationOptsValidator } from "convex/server";
 import { stream } from "convex-helpers/server/stream";
-import { paginationResultValidator } from "../validators.js";
+import { vPaginationResult } from "../validators.js";
 import { Id } from "./_generated/dataModel.js";
 
 // Note: it only searches for users with threads
@@ -33,7 +33,7 @@ export const listUsersWithThreads = query({
       page: results.page.map((t) => t.userId).filter((t): t is string => !!t),
     };
   },
-  returns: paginationResultValidator(v.string()),
+  returns: vPaginationResult(v.string()),
 });
 
 export const deleteAllForUserId = action({

package/src/component/vector/tables.ts

@@ -68,7 +68,7 @@ export type VectorSchema = SchemaDefinition<
 >;
 
 export const VectorDimensions = [
-  128, 256, 512, 768, 1024, 1536, 2048, 3072, 4096,
+  128, 256, 512, 768, 1024, 1408, 1536, 2048, 3072, 4096,
 ] as const;
 export function validateVectorDimension(
   dimension: number

package/src/mapping.ts

@@ -108,7 +108,9 @@ export function serializeNewMessagesInStep<TOOLS extends ToolSet>(
   ).map(
     (message): MessageWithMetadata => ({
       message: serializeMessage(message),
-      id: message.id,
+      // Let's not store the ID by default here. It's being generated internally
+      // and not referenced elsewhere that we know of.
+      // id: message.id,
       ...(message.role === "tool" ? toolFields : assistantFields),
       text: step.text,
       // fileId: message.fileId,