@convex-dev/agent 0.0.1-alpha.1 → 0.0.1-alpha.2

package/src/client/types.ts

@@ -7,6 +7,10 @@ import {
   GenericActionCtx,
 } from "convex/server";
 import { GenericId } from "convex/values";
+import type { Doc } from "../component/_generated/dataModel";
+
+export type ChatDoc = OpaqueIds<Doc<"chats">>;
+export type MessageDoc = OpaqueIds<Doc<"messages">>;
 
 /* Type utils follow */
 export type RunQueryCtx = {

package/src/component/_generated/api.d.ts

@@ -40,7 +40,7 @@ export type Mounts = {
       "public",
       {
         agentName?: string;
-        chatId: string;
+        chatId?: string;
         failPendingSteps?: boolean;
         messages: Array<{
           fileId?: string;
@@ -147,13 +147,14 @@ export type Mounts = {
         parentMessageId?: string;
         pending?: boolean;
         stepId?: string;
+        userId?: string;
       },
       {
         messages: Array<{
           _creationTime: number;
           _id: string;
           agentName?: string;
-          chatId: string;
+          chatId?: string;
           embeddingId?:
             | string
             | string
@@ -277,7 +278,7 @@ export type Mounts = {
           _creationTime: number;
           _id: string;
           agentName?: string;
-          chatId: string;
+          chatId?: string;
           embeddingId?:
             | string
             | string
@@ -1021,7 +1022,7 @@ export type Mounts = {
           _creationTime: number;
           _id: string;
           agentName?: string;
-          chatId: string;
+          chatId?: string;
           embeddingId?:
             | string
             | string
@@ -1208,7 +1209,7 @@ export type Mounts = {
         _creationTime: number;
         _id: string;
         agentName?: string;
-        chatId: string;
+        chatId?: string;
         embeddingId?:
           | string
           | string
@@ -1337,7 +1338,7 @@ export type Mounts = {
         _creationTime: number;
         _id: string;
         agentName?: string;
-        chatId: string;
+        chatId?: string;
         embeddingId?:
           | string
           | string

package/src/component/messages.ts

@@ -2,19 +2,14 @@ import { assert, omit, pick } from "convex-helpers";
 import { paginator } from "convex-helpers/server/pagination";
 import { mergedStream, stream } from "convex-helpers/server/stream";
 import { nullable, partial } from "convex-helpers/validators";
-import { Infer, ObjectType } from "convex/values";
+import { ObjectType } from "convex/values";
 import { DEFAULT_MESSAGE_RANGE, extractText, isTool } from "../shared.js";
 import {
-  Message,
-  MessageWithFileAndId,
-  vAssistantMessage,
   vChatStatus,
   vMessageStatus,
   vMessageWithFileAndId,
   vSearchOptions,
-  vStep,
   vStepWithMessagesWithFileAndId,
-  vToolMessage,
 } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import { Doc, Id } from "./_generated/dataModel.js";
@@ -63,7 +58,7 @@ export const getChatsByUserId = query({
             .gte("order", args.offset ?? 0)
         )
     );
-    const chats = await mergedStream(streams, ["order", "stepOrder"]).paginate({
+    const chats = await mergedStream(streams, ["order"]).paginate({
       numItems: args.limit ?? 100,
       cursor: args.cursor ?? null,
     });
@@ -377,7 +372,8 @@ export const messageStatuses = vMessageDoc.fields.status.members.map(
 );
 
 const addMessagesArgs = {
-  chatId: v.id("chats"),
+  userId: v.optional(v.string()),
+  chatId: v.optional(v.id("chats")),
   stepId: v.optional(v.id("steps")),
   parentMessageId: v.optional(v.id("messages")),
   messages: v.array(vMessageWithFileAndId),
@@ -398,15 +394,22 @@ async function addMessagesHandler(
   ctx: MutationCtx,
   args: ObjectType<typeof addMessagesArgs>
 ) {
-  const chat = await ctx.db.get(args.chatId);
-  assert(chat, `Chat ${args.chatId} not found`);
+  let userId = args.userId;
+  const chatId = args.chatId;
+  if (!userId && args.chatId) {
+    const chat = await ctx.db.get(args.chatId);
+    assert(chat, `Chat ${args.chatId} not found`);
+    userId = chat._id;
+  }
   const { failPendingSteps, pending, messages, parentMessageId, ...rest } =
     args;
-  if (failPendingSteps) {
+  const parent = parentMessageId && (await ctx.db.get(parentMessageId));
+  if (failPendingSteps && parent?.status !== "pending") {
+    assert(args.chatId, "chatId is required to fail pending steps");
     const pendingMessages = await ctx.db
       .query("messages")
       .withIndex("chatId_status_tool_order_stepOrder", (q) =>
-        q.eq("chatId", args.chatId).eq("status", "pending")
+        q.eq("chatId", chatId).eq("status", "pending")
       )
       .collect();
     await Promise.all(
@@ -415,14 +418,14 @@ async function addMessagesHandler(
       )
     );
   }
-  let order: number | undefined;
-  const maxMessage = await getMaxMessage(ctx, args.chatId);
+  let threadId = parentMessageId;
+  const maxMessage = await getMaxMessage(ctx, chatId, userId);
   // If the previous message isn't our parent, we make a new thread.
-  const threadId =
+  threadId =
     parentMessageId && maxMessage?._id === parentMessageId
       ? maxMessage.threadId ?? parentMessageId
       : parentMessageId;
-  order = maxMessage?.order ?? -1;
+  let order = maxMessage?.order ?? -1;
   const toReturn: Doc<"messages">[] = [];
   if (messages.length > 0) {
     for (const { message, fileId, id } of messages) {
@@ -434,7 +437,7 @@ async function addMessagesHandler(
       const messageId = await ctx.db.insert("messages", {
         ...rest,
         threadId,
-        userId: chat.userId,
+        userId,
         message,
         id,
         order,
@@ -449,18 +452,37 @@ async function addMessagesHandler(
   return { messages: toReturn };
 }
 
-async function getMaxMessage(ctx: QueryCtx, chatId: Id<"chats">) {
-  return mergedStream(
-    ["success" as const, "pending" as const].map((status) =>
-      stream(ctx.db, schema)
-        .query("messages")
-        .withIndex("chatId_status_tool_order_stepOrder", (q) =>
-          q.eq("chatId", chatId).eq("status", status).eq("tool", false)
-        )
-        .order("desc")
-    ),
-    ["order", "stepOrder"]
-  ).first();
+async function getMaxMessage(
+  ctx: QueryCtx,
+  chatId: Id<"chats"> | undefined,
+  userId: string | undefined
+) {
+  assert(chatId || userId, "One of chatId or userId is required");
+  if (chatId) {
+    return mergedStream(
+      ["success" as const, "pending" as const].map((status) =>
+        stream(ctx.db, schema)
+          .query("messages")
+          .withIndex("chatId_status_tool_order_stepOrder", (q) =>
+            q.eq("chatId", chatId).eq("status", status).eq("tool", false)
+          )
+          .order("desc")
+      ),
+      ["order", "stepOrder"]
+    ).first();
+  } else {
+    return mergedStream(
+      ["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", false)
+          )
+          .order("desc")
+      ),
+      ["order", "stepOrder"]
+    ).first();
+  }
 }
 
 const addStepsArgs = {
@@ -518,7 +540,8 @@ async function addStepsHandler(
       pending: step.finishReason === "stop" ? false : true,
       failPendingSteps: false,
     });
-    if (step.finishReason === "stop") {
+    // We don't commit if the parent is still pending.
+    if (step.finishReason === "stop" && parentMessage.status === "success") {
       await commitMessageHandler(ctx, { messageId: args.messageId });
     }
     steps.push((await ctx.db.get(stepId))!);
@@ -585,6 +608,7 @@ async function commitMessageHandler(
   ).collect();
   for (const message of messages) {
     await ctx.db.patch(message._id, { status: "success" });
+    // TODO: recursively commit steps & messages that might depend on this one.
   }
 }
 
@@ -727,18 +751,19 @@ export const _fetchVectorMessages = internalQuery({
   handler: async (ctx, args): Promise<Doc<"messages">[]> => {
     const parent =
       args.parentMessageId && (await ctx.db.get(args.parentMessageId));
+    const { userId, chatId } = args;
+    assert(userId || chatId, "Specify userId or chatId to search");
     let messages = (
       await Promise.all(
         args.vectorIds.map((embeddingId) =>
           ctx.db
             .query("messages")
             .withIndex("embeddingId", (q) => q.eq("embeddingId", embeddingId))
-            .filter(
-              (q) =>
-                args.userId
-                  ? q.eq("userId", args.userId)
-                  : // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                    q.eq("chatId", args.chatId as any) // not sure why it's failing...
+            .filter((q) =>
+              userId
+                ? q.eq("userId", userId)
+                : // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                  q.eq("chatId", chatId as any)
             )
             .first()
         )
@@ -752,48 +777,71 @@ export const _fetchVectorMessages = internalQuery({
     messages.sort((a, b) => a.order! - b.order!);
     messages = messages.slice(0, args.limit);
     // Fetch the surrounding messages
-    const included: Record<Id<"chats">, Set<number>> = {};
+    if (!chatId) {
+      return messages.sort((a, b) => a.order - b.order);
+    }
+    const included: Record<string, Set<number>> = {};
     for (const m of messages) {
-      if (!included[m.chatId]) {
-        included[m.chatId] = new Set();
+      const searchId = m.chatId ?? m.userId!;
+      if (!included[searchId]) {
+        included[searchId] = new Set();
       }
-      included[m.chatId].add(m.order!);
+      included[searchId].add(m.order!);
     }
-    const ranges: Record<Id<"chats">, Doc<"messages">[]> = {};
+    const ranges: Record<string, Doc<"messages">[]> = {};
     const { before, after } = args.messageRange;
     for (const m of messages) {
+      const searchId = m.chatId ?? m.userId!;
       const order = m.order!;
       let earliest = order - before;
       let latest = order + after;
       for (; earliest <= latest; earliest++) {
-        if (!included[m.chatId].has(earliest)) {
+        if (!included[searchId].has(earliest)) {
           break;
         }
       }
       for (; latest >= earliest; latest--) {
-        if (!included[m.chatId].has(latest)) {
+        if (!included[searchId].has(latest)) {
           break;
         }
       }
       for (let i = earliest; i <= latest; i++) {
-        included[m.chatId].add(i);
+        included[searchId].add(i);
       }
       if (earliest !== latest) {
-        const surrounding = await ctx.db
-          .query("messages")
-          .withIndex("chatId_status_tool_order_stepOrder", (q) =>
-            q
-              .eq("chatId", m.chatId)
-              .eq("status", "success")
-              .eq("tool", false)
-              .gt("order", earliest)
-              .lt("order", latest)
-          )
-          .collect();
-        if (!ranges[m.chatId]) {
-          ranges[m.chatId] = [];
+        if (m.chatId) {
+          const surrounding = await ctx.db
+            .query("messages")
+            .withIndex("chatId_status_tool_order_stepOrder", (q) =>
+              q
+                .eq("chatId", m.chatId)
+                .eq("status", "success")
+                .eq("tool", false)
+                .gt("order", earliest)
+                .lt("order", latest)
+            )
+            .collect();
+          if (!ranges[searchId]) {
+            ranges[searchId] = [];
+          }
+          ranges[searchId].push(...surrounding);
+        } else {
+          const surrounding = await ctx.db
+            .query("messages")
+            .withIndex("userId_status_tool_order_stepOrder", (q) =>
+              q
+                .eq("userId", m.userId!)
+                .eq("status", "success")
+                .eq("tool", false)
+                .gt("order", earliest)
+                .lt("order", latest)
+            )
+            .collect();
+          if (!ranges[searchId]) {
+            ranges[searchId] = [];
+          }
+          ranges[searchId].push(...surrounding);
         }
-        ranges[m.chatId].push(...surrounding);
       }
     }
     return Object.values(ranges)

package/src/component/schema.ts

@@ -22,7 +22,7 @@ export const schema = defineSchema({
   messages: defineTable({
     id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
     userId: v.optional(v.string()), // useful for future indexes (text search)
-    chatId: v.id("chats"),
+    chatId: v.optional(v.id("chats")),
     threadId: v.optional(v.id("messages")),
     stepId: v.optional(v.id("steps")),
     agentName: v.optional(v.string()),
@@ -33,7 +33,7 @@ export const schema = defineSchema({
     // TODO: add sub-messages back in? or be able to skip them?
     tool: v.boolean(),
     // Repeats until a non-tool message.
-    // Set when the message is finished
+    // Unset if it's not in a chat.
     order: v.number(),
     stepOrder: v.optional(v.number()),
     fileId: v.optional(v.id("files")),
@@ -48,6 +48,13 @@ export const schema = defineSchema({
       "order",
       "stepOrder",
     ])
+    .index("userId_status_tool_order_stepOrder", [
+      "userId",
+      "status",
+      "tool",
+      "order",
+      "stepOrder",
+    ])
     // Allows finding all threaded messages in order
     // Allows finding all failed messages to evaluate
     // .index("status_threadId_order_stepOrder", [

package/src/mapping.ts

@@ -1,16 +1,20 @@
-import type {
-  AssistantContent,
-  CoreAssistantMessage,
-  CoreMessage,
-  CoreToolMessage,
-  DataContent,
-  GenerateTextResult,
-  StepResult,
-  ToolContent,
-  ToolSet,
-  UserContent,
+import {
+  convertToCoreMessages,
+  coreMessageSchema,
+  type AssistantContent,
+  type CoreAssistantMessage,
+  type CoreMessage,
+  type CoreToolMessage,
+  type DataContent,
+  type GenerateTextResult,
+  type StepResult,
+  type ToolContent,
+  type ToolSet,
+  type Message as UIMessage,
+  type UserContent,
 } from "ai";
 import { MessageWithFileAndId, Step } from "./validators";
+import { assert } from "convex-helpers";
 
 export type SerializeUrlsAndUint8Arrays<T> = T extends URL
   ? string
@@ -158,3 +162,34 @@ function deserializeUrl(urlOrString: string | ArrayBuffer): URL | DataContent {
   }
   return urlOrString;
 }
+
+export function promptOrMessagesToCoreMessages(args: {
+  system?: string;
+  prompt?: string;
+  messages?: CoreMessage[] | Omit<UIMessage, "id">[];
+}): CoreMessage[] {
+  const messages: CoreMessage[] = [];
+  if (args.system) {
+    messages.push({ role: "system", content: args.system });
+  }
+  if (!args.messages) {
+    assert(args.prompt, "messages or prompt is required");
+    messages.push({ role: "user", content: args.prompt });
+  } else if (
+    args.messages.some(
+      (m) =>
+        typeof m === "object" &&
+        m !== null &&
+        (m.role === "data" || // UI-only role
+          "toolInvocations" in m || // UI-specific field
+          "parts" in m || // UI-specific field
+          "experimental_attachments" in m)
+    )
+  ) {
+    messages.push(...convertToCoreMessages(args.messages as UIMessage[]));
+  } else {
+    messages.push(...coreMessageSchema.array().parse(args.messages));
+  }
+  assert(messages.length > 0, "Messages must contain at least one message");
+  return messages;
+}

package/src/validators.test.ts

@@ -3,11 +3,13 @@ import { expectTypeOf, test } from "vitest";
 import {
   vAssistantContent,
   vAssistantMessage,
+  vContextOptions,
   vFilePart,
   vImagePart,
   vMessage,
   vReasoningPart,
   vRedactedReasoningPart,
+  vStorageOptions,
   vSystemMessage,
   vTextPart,
   vToolCallPart,
@@ -31,6 +33,7 @@ import {
   UserContent,
 } from "ai";
 import { SerializeUrlsAndUint8Arrays } from "./mapping";
+import { ContextOptions, StorageOptions } from "./client";
 
 // type assertion
 type OurUserContent = SerializeUrlsAndUint8Arrays<UserContent>;
@@ -98,4 +101,10 @@ type OurMessage = SerializeUrlsAndUint8Arrays<CoreMessage>;
 expectTypeOf<OurMessage>().toExtend<Infer<typeof vMessage>>();
 expectTypeOf<Infer<typeof vMessage>>().toExtend<OurMessage>();
 
+expectTypeOf<Infer<typeof vContextOptions>>().toExtend<ContextOptions>();
+expectTypeOf<ContextOptions>().toExtend<Infer<typeof vContextOptions>>();
+
+expectTypeOf<Infer<typeof vStorageOptions>>().toExtend<StorageOptions>();
+expectTypeOf<StorageOptions>().toExtend<Infer<typeof vStorageOptions>>();
+
 test("noop", () => { });

package/src/validators.ts

@@ -256,3 +256,35 @@ export const vSearchOptions = v.object({
   messageRange: v.optional(v.object({ before: v.number(), after: v.number() })),
 });
 export type SearchOptions = Infer<typeof vSearchOptions>;
+
+export const vContextOptionsSearchOptions = v.object({
+  limit: v.number(),
+  textSearch: v.optional(v.boolean()),
+  vectorSearch: v.optional(v.boolean()),
+  messageRange: v.optional(v.object({ before: v.number(), after: v.number() })),
+});
+
+export const vContextOptions = v.object({
+  includeToolCalls: v.optional(v.boolean()),
+  recentMessages: v.optional(v.number()),
+  searchOptions: v.optional(vContextOptionsSearchOptions),
+  searchOtherChats: v.optional(v.boolean()),
+});
+
+export const vStorageOptions = v.object({
+  saveAllInputMessages: v.optional(v.boolean()),
+  saveAllOutputMessages: v.optional(v.boolean()),
+});
+
+export const vChatArgs = v.object({
+  maxSteps: v.optional(v.number()),
+  prompt: v.optional(v.string()),
+  messages: v.optional(v.array(vMessage)),
+});
+
+export const vObjectArgs = v.object({
+  output: v.optional(v.any()),
+  mode: v.optional(v.literal("json")),
+  prompt: v.optional(v.string()),
+  messages: v.optional(v.array(vMessage)),
+});