@convex-dev/agent 0.5.0-alpha.1 → 0.6.0-alpha.0

package/src/client/messages.ts

@@ -0,0 +1,237 @@
+import type { ModelMessage } from "ai";
+import type { PaginationOptions, PaginationResult } from "convex/server";
+import type { MessageDoc } from "../validators.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 "../mapping.js";
+import { toUIMessages, type UIMessage } from "../UIMessages.js";
+import type {
+  AgentComponent,
+  MutationCtx,
+  QueryCtx,
+  ActionCtx,
+} 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: QueryCtx | MutationCtx | ActionCtx,
+  component: AgentComponent,
+  {
+    threadId,
+    paginationOpts,
+    excludeToolMessages,
+    statuses,
+  }: {
+    threadId: string;
+    paginationOpts: PaginationOptions;
+    excludeToolMessages?: boolean;
+    statuses?: MessageStatus[];
+  },
+): Promise<PaginationResult<MessageDoc>> {
+  if (paginationOpts.numItems === 0) {
+    return {
+      page: [],
+      isDone: true,
+      continueCursor: paginationOpts.cursor ?? "",
+    };
+  }
+  return ctx.runQuery(component.messages.listMessagesByThreadId, {
+    order: "desc",
+    threadId,
+    paginationOpts,
+    excludeToolMessages,
+    statuses,
+  });
+}
+
+export async function listUIMessages(
+  ctx: QueryCtx | MutationCtx | ActionCtx,
+  component: AgentComponent,
+  args: {
+    threadId: string;
+    paginationOpts: PaginationOptions;
+  },
+): Promise<PaginationResult<UIMessage>> {
+  const result = await listMessages(ctx, component, args);
+  return { ...result, page: toUIMessages(result.page) };
+}
+
+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: MutationCtx | ActionCtx,
+  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 base = args.metadata?.[i];
+        const allFileIds = [...(base?.fileIds ?? [])];
+        if (fileIds) allFileIds.push(...fileIds);
+
+        return parse(vMessageWithMetadata, {
+          ...base,
+          message,
+          ...(allFileIds.length > 0 ? { fileIds: allFileIds } : {}),
+        });
+      }),
+    ),
+    failPendingSteps: args.failPendingSteps ?? false,
+  });
+  return { messages: result.messages };
+}
+
+export type SaveMessageArgs = {
+  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;
+  /**
+   * 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: MutationCtx | ActionCtx,
+  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,
+    promptMessageId: args.promptMessageId,
+    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/mockModel.ts

@@ -0,0 +1,245 @@
+import type {
+  LanguageModelV3,
+  LanguageModelV3Content,
+  LanguageModelV3StreamPart,
+} from "@ai-sdk/provider";
+import { simulateReadableStream, type ProviderMetadata } from "ai";
+import { assert, pick } from "convex-helpers";
+
+export const DEFAULT_TEXT = `
+A A A A A A A A A A A A A A A
+B B B B B B B B B B B B B B B
+C C C C C C C C C C C C C C C
+D D D D D D D D D D D D D D D
+`;
+const DEFAULT_USAGE = {
+  outputTokens: 10,
+  inputTokens: 3,
+  totalTokens: 13,
+  inputTokenDetails: undefined,
+  outputTokenDetails: undefined,
+};
+
+export type MockModelArgs = {
+  provider?: LanguageModelV3["provider"];
+  modelId?: LanguageModelV3["modelId"];
+  supportedUrls?:
+    | LanguageModelV3["supportedUrls"]
+    | (() => LanguageModelV3["supportedUrls"]);
+  chunkDelayInMs?: number;
+  initialDelayInMs?: number;
+  /** A list of the responses for multiple steps.
+   * For tool calls, the first list would include a tool call part,
+   * then the next list would be after the tool response or another tool call.
+   * Tool responses come from actual tool calls!
+   */
+  contentSteps?: LanguageModelV3Content[][];
+  /** A single list of content responded from each step.
+   * Provide contentSteps instead if you want to do multi-step responses with
+   * tool calls.
+   */
+  content?: LanguageModelV3Content[];
+  // provide either content, contentResponses or doGenerate & doStream
+  doGenerate?: LanguageModelV3["doGenerate"];
+  doStream?: LanguageModelV3["doStream"];
+  providerMetadata?: ProviderMetadata;
+  fail?:
+    | boolean
+    | {
+        probability?: number;
+        error?: string;
+      };
+};
+
+function atMostOneOf(...args: unknown[]) {
+  return args.filter(Boolean).length <= 1;
+}
+
+export function mockModel(args?: MockModelArgs): LanguageModelV3 {
+  return new MockLanguageModel(args ?? {});
+}
+
+export class MockLanguageModel implements LanguageModelV3 {
+  readonly specificationVersion = "v3";
+
+  private _supportedUrls: () => LanguageModelV3["supportedUrls"];
+
+  readonly provider: LanguageModelV3["provider"];
+  readonly modelId: LanguageModelV3["modelId"];
+
+  doGenerate: LanguageModelV3["doGenerate"];
+  doStream: LanguageModelV3["doStream"];
+
+  doGenerateCalls: Parameters<LanguageModelV3["doGenerate"]>[0][] = [];
+  doStreamCalls: Parameters<LanguageModelV3["doStream"]>[0][] = [];
+
+  constructor(args: MockModelArgs) {
+    assert(
+      atMostOneOf(
+        args.content,
+        args.contentSteps,
+        args.doGenerate && args.doStream,
+      ),
+      "Expected only one of content, contentSteps, or doGenerate and doStream",
+    );
+    this.provider = args.provider || "mock-provider";
+    this.modelId = args.modelId || "mock-model-id";
+    const {
+      content = [{ type: "text", text: DEFAULT_TEXT }],
+      contentSteps = [content],
+      chunkDelayInMs = 0,
+      initialDelayInMs = 0,
+      supportedUrls = {},
+    } = args;
+    const fail =
+      args.fail &&
+      (args.fail === true ||
+        !args.fail.probability ||
+        Math.random() < args.fail.probability);
+    const error =
+      (typeof args.fail === "object" && args.fail.error) ||
+      "Mock error message";
+    const metadata = pick(args, ["providerMetadata"]);
+
+    const chunkResponses: LanguageModelV3StreamPart[][] = contentSteps.map(
+      (content) => {
+        const chunks: LanguageModelV3StreamPart[] = [
+          { type: "stream-start", warnings: [] },
+        ];
+        chunks.push(
+          ...content.flatMap((c, ci): LanguageModelV3StreamPart[] => {
+            if (c.type !== "text" && c.type !== "reasoning") {
+              return [c];
+            }
+            const metadata = pick(c, ["providerMetadata"]);
+            const deltas = c.text.split(" ");
+            const parts: LanguageModelV3StreamPart[] = [];
+            if (c.type === "reasoning") {
+              parts.push({
+                type: "reasoning-start",
+                id: `reasoning-${ci}`,
+                ...metadata,
+              });
+              parts.push(
+                ...deltas.map(
+                  (delta, di) =>
+                    ({
+                      type: "reasoning-delta",
+                      delta: (di ? " " : "") + delta,
+                      id: `reasoning-${ci}`,
+                      ...metadata,
+                    }) satisfies LanguageModelV3StreamPart,
+                ),
+              );
+              parts.push({
+                type: "reasoning-end",
+                id: `reasoning-${ci}`,
+                ...metadata,
+              });
+            } else if (c.type === "text") {
+              parts.push({
+                type: "text-start",
+                id: `txt-${ci}`,
+                ...metadata,
+              });
+              parts.push(
+                ...deltas.map(
+                  (delta, di) =>
+                    ({
+                      type: "text-delta",
+                      delta: (di ? " " : "") + delta,
+                      id: `txt-${ci}`,
+                      ...metadata,
+                    }) satisfies LanguageModelV3StreamPart,
+                ),
+              );
+              parts.push({
+                type: "text-end",
+                id: `txt-${ci}`,
+                ...metadata,
+              });
+            }
+            return parts;
+          }),
+        );
+        if (fail) {
+          chunks.push({
+            type: "error",
+            error,
+          });
+        }
+        chunks.push({
+          type: "finish",
+          finishReason: fail ? "error" : "stop",
+          usage: DEFAULT_USAGE,
+          ...(metadata as any),
+        });
+        return chunks;
+      },
+    );
+    let callIndex = 0;
+    this.doGenerate = async (options) => {
+      this.doGenerateCalls.push(options);
+
+      if (fail) {
+        throw new Error(error);
+      }
+      if (typeof args.doGenerate === "function") {
+        return args.doGenerate(options);
+      } else if (Array.isArray(args.doGenerate)) {
+        return args.doGenerate[this.doGenerateCalls.length];
+      } else if (contentSteps.length) {
+        const result = {
+          content: contentSteps[callIndex % contentSteps.length],
+          finishReason: "stop" as const,
+          usage: DEFAULT_USAGE,
+          ...(metadata as any),
+          warnings: [],
+        };
+        callIndex++;
+        return result;
+      } else {
+        throw new Error("Unexpected: no content or doGenerate");
+      }
+    };
+    this.doStream = async (options) => {
+      this.doStreamCalls.push(options);
+
+      if (typeof args.doStream === "function") {
+        return args.doStream(options);
+      } else if (Array.isArray(args.doStream)) {
+        return args.doStream[this.doStreamCalls.length];
+      } else if (contentSteps) {
+        const stream = simulateReadableStream({
+          chunks: chunkResponses[callIndex % chunkResponses.length],
+          initialDelayInMs,
+          chunkDelayInMs,
+        });
+        callIndex++;
+
+        if (options.abortSignal) {
+          options.abortSignal.addEventListener("abort", () => {
+            console.warn("abortSignal in mock model not supported");
+          });
+        }
+        return {
+          stream,
+          request: { body: {} },
+          response: { headers: {} },
+        };
+      } else if (args.doStream) {
+        return args.doStream;
+      } else {
+        throw new Error("Provide either content or doStream");
+      }
+    };
+    this._supportedUrls =
+      typeof supportedUrls === "function"
+        ? supportedUrls
+        : async () => supportedUrls;
+  }
+
+  get supportedUrls() {
+    return this._supportedUrls();
+  }
+}