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

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/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,