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

package/src/client/start.ts

@@ -0,0 +1,310 @@
+import {
+  stepCountIs,
+  type CallSettings,
+  type GenerateObjectResult,
+  type IdGenerator,
+  type LanguageModel,
+  type ModelMessage,
+  type StepResult,
+  type StopCondition,
+  type ToolSet,
+} from "ai";
+import {
+  serializeNewMessagesInStep,
+  serializeObjectResult,
+} from "../mapping.js";
+import { embedMessages, fetchContextWithPrompt } from "./search.js";
+import type { ActionCtx, AgentComponent, Config, Options } from "./types.js";
+import type { Message, MessageDoc } from "../validators.js";
+import {
+  getModelName,
+  getProviderName,
+  type ModelOrMetadata,
+} from "../shared.js";
+import { wrapTools, type ToolCtx } from "./createTool.js";
+import type { Agent } from "./index.js";
+import { assert, omit } from "convex-helpers";
+import { saveInputMessages } from "./saveInputMessages.js";
+import type { GenericActionCtx, GenericDataModel } from "convex/server";
+
+export async function startGeneration<
+  T,
+  Tools extends ToolSet = ToolSet,
+  CustomCtx extends object = object,
+>(
+  ctx: ActionCtx & CustomCtx,
+  component: AgentComponent,
+  /**
+   * These are the arguments you'll pass to the LLM call such as
+   * `generateText` or `streamText`. This function will look up the context
+   * and provide functions to save the steps, abort the generation, and more.
+   * The type of the arguments returned infers from the type of the arguments
+   * you pass here.
+   */
+  args: T & {
+    /**
+     * If provided, this message will be used as the "prompt" for the LLM call,
+     * instead of the prompt or messages.
+     * This is useful if you want to first save a user message, then use it as
+     * the prompt for the LLM call in another call.
+     */
+    promptMessageId?: string;
+    /**
+     * The model to use for the LLM calls. This will override the model specified
+     * in the Agent constructor.
+     */
+    model?: LanguageModel;
+    /**
+     * The tools to use for the tool calls. This will override tools specified
+     * in the Agent constructor or createThread / continueThread.
+     */
+    tools?: Tools;
+    /**
+     * The single prompt message to use for the LLM call. This will be the
+     * last message in the context. If it's a string, it will be a user role.
+     */
+    prompt?: string | (ModelMessage | Message)[];
+    /**
+     * If provided alongside prompt, the ordering will be:
+     * 1. system prompt
+     * 2. search context
+     * 3. recent messages
+     * 4. these messages
+     * 5. prompt messages, including those already on the same `order` as
+     *   the promptMessageId message, if provided.
+     */
+    messages?: (ModelMessage | Message)[];
+    /**
+     * The abort signal to be passed to the LLM call. If triggered, it will
+     * mark the pending message as failed. If the generation is asynchronously
+     * aborted, it will trigger this signal when detected.
+     */
+    abortSignal?: AbortSignal;
+    stopWhen?: StopCondition<Tools> | Array<StopCondition<Tools>>;
+    _internal?: { generateId?: IdGenerator };
+  },
+  {
+    threadId,
+    ...opts
+  }: Options &
+    Config & {
+      userId?: string | null;
+      threadId?: string;
+      languageModel?: LanguageModel;
+      agentName: string;
+      agentForToolCtx?: Agent;
+    },
+): Promise<{
+  args: T & {
+    system?: string;
+    model: LanguageModel;
+    messages: ModelMessage[];
+    prompt?: never;
+    tools?: Tools;
+  } & CallSettings;
+  order: number;
+  stepOrder: number;
+  userId: string | undefined;
+  promptMessageId: string | undefined;
+  updateModel: (model: ModelOrMetadata | undefined) => void;
+  save: <TOOLS extends ToolSet>(
+    toSave:
+      | { step: StepResult<TOOLS> }
+      | { object: GenerateObjectResult<unknown> },
+    createPendingMessage?: boolean,
+  ) => Promise<void>;
+  fail: (reason: string) => Promise<void>;
+  getSavedMessages: () => MessageDoc[];
+}> {
+  const userId =
+    opts.userId ??
+    (threadId &&
+      (await ctx.runQuery(component.threads.getThread, { threadId }))
+        ?.userId) ??
+    undefined;
+
+  const context = await fetchContextWithPrompt(ctx, component, {
+    ...opts,
+    userId,
+    threadId,
+    messages: args.messages,
+    prompt: args.prompt,
+    promptMessageId: args.promptMessageId,
+  });
+
+  const saveMessages = opts.storageOptions?.saveMessages ?? "promptAndOutput";
+  const { promptMessageId, pendingMessage, savedMessages } =
+    threadId && saveMessages !== "none"
+      ? await saveInputMessages(ctx, component, {
+          ...opts,
+          userId,
+          threadId,
+          prompt: args.prompt,
+          messages: args.messages,
+          promptMessageId: args.promptMessageId,
+          storageOptions: { saveMessages },
+        })
+      : {
+          promptMessageId: args.promptMessageId,
+          pendingMessage: undefined,
+          savedMessages: [] as MessageDoc[],
+        };
+
+  const order = pendingMessage?.order ?? context.order;
+  const stepOrder = pendingMessage?.stepOrder ?? context.stepOrder;
+  let pendingMessageId = pendingMessage?._id;
+
+  const model = args.model ?? opts.languageModel;
+  assert(model, "model is required");
+  let activeModel: ModelOrMetadata = model;
+
+  const fail = async (reason: string) => {
+    if (pendingMessageId) {
+      await ctx.runMutation(component.messages.finalizeMessage, {
+        messageId: pendingMessageId,
+        result: { status: "failed", error: reason },
+      });
+    }
+  };
+  if (args.abortSignal) {
+    const abortSignal = args.abortSignal;
+    abortSignal.addEventListener(
+      "abort",
+      async () => {
+        await fail(abortSignal.reason?.toString() ?? "abortSignal");
+      },
+      { once: true },
+    );
+  }
+  const toolCtx = {
+    ...(ctx as GenericActionCtx<GenericDataModel> & CustomCtx),
+    userId,
+    threadId,
+    promptMessageId,
+    agent: opts.agentForToolCtx,
+  } satisfies ToolCtx;
+  const tools = wrapTools(toolCtx, args.tools) as Tools;
+  const aiArgs = {
+    ...opts.callSettings,
+    providerOptions: opts.providerOptions,
+    ...omit(args, ["promptMessageId", "messages", "prompt"]),
+    model,
+    messages: context.messages,
+    stopWhen:
+      args.stopWhen ?? (opts.maxSteps ? stepCountIs(opts.maxSteps) : undefined),
+    tools,
+  } as T & {
+    model: LanguageModel;
+    messages: ModelMessage[];
+    prompt?: never;
+    tools?: Tools;
+    _internal?: { generateId?: IdGenerator };
+  } & CallSettings;
+  if (pendingMessageId) {
+    if (!aiArgs._internal?.generateId) {
+      aiArgs._internal = {
+        ...aiArgs._internal,
+        generateId: pendingMessageId
+          ? () => pendingMessageId ?? crypto.randomUUID()
+          : undefined,
+      };
+    }
+  }
+  return {
+    args: aiArgs,
+    order: order ?? 0,
+    stepOrder: stepOrder ?? 0,
+    userId,
+    promptMessageId,
+    getSavedMessages: () => savedMessages,
+    updateModel: (model: ModelOrMetadata | undefined) => {
+      if (model) {
+        activeModel = model;
+      }
+    },
+    fail,
+    save: async <TOOLS extends ToolSet>(
+      toSave:
+        | { step: StepResult<TOOLS> }
+        | { object: GenerateObjectResult<unknown> },
+      createPendingMessage?: boolean,
+    ) => {
+      if (threadId && saveMessages !== "none") {
+        const serialized =
+          "object" in toSave
+            ? await serializeObjectResult(
+                ctx,
+                component,
+                toSave.object,
+                activeModel,
+              )
+            : await serializeNewMessagesInStep(
+                ctx,
+                component,
+                toSave.step,
+                activeModel,
+              );
+        const embeddings = await embedMessages(
+          ctx,
+          { threadId, ...opts, userId },
+          serialized.messages.map((m) => m.message),
+        );
+        if (createPendingMessage) {
+          serialized.messages.push({
+            message: { role: "assistant", content: [] },
+            status: "pending",
+          });
+          embeddings?.vectors.push(null);
+        }
+        const saved = await ctx.runMutation(component.messages.addMessages, {
+          userId,
+          threadId,
+          agentName: opts.agentName,
+          promptMessageId,
+          pendingMessageId,
+          messages: serialized.messages,
+          embeddings,
+          failPendingSteps: false,
+        });
+        const lastMessage = saved.messages.at(-1)!;
+        if (createPendingMessage) {
+          if (lastMessage.status === "failed") {
+            pendingMessageId = undefined;
+            savedMessages.push(...saved.messages);
+            await fail(
+              lastMessage.error ??
+                "Aborting - the pending message was marked as failed",
+            );
+          } else {
+            pendingMessageId = lastMessage._id;
+            savedMessages.push(...saved.messages.slice(0, -1));
+          }
+        } else {
+          pendingMessageId = undefined;
+          savedMessages.push(...saved.messages);
+        }
+      }
+      const output = "object" in toSave ? toSave.object : toSave.step;
+      if (opts.rawRequestResponseHandler) {
+        await opts.rawRequestResponseHandler(ctx, {
+          userId,
+          threadId,
+          agentName: opts.agentName,
+          request: output.request,
+          response: output.response,
+        });
+      }
+      if (opts.usageHandler && output.usage) {
+        await opts.usageHandler(ctx, {
+          userId,
+          threadId,
+          agentName: opts.agentName,
+          model: getModelName(activeModel),
+          provider: getProviderName(activeModel),
+          usage: output.usage,
+          providerMetadata: output.providerMetadata,
+        });
+      }
+    },
+  };
+}

package/src/client/streamText.ts

@@ -0,0 +1,163 @@
+import type {
+  StepResult,
+  StreamTextResult,
+  ToolSet,
+  UIMessage as AIUIMessage,
+} from "ai";
+import { streamText as streamTextAi } from "ai";
+import {
+  compressUIMessageChunks,
+  DeltaStreamer,
+  mergeTransforms,
+  type StreamingOptions,
+} from "./streaming.js";
+import type {
+  ActionCtx,
+  AgentComponent,
+  AgentPrompt,
+  GenerationOutputMetadata,
+  Options,
+  Output,
+} from "./types.js";
+import { startGeneration } from "./start.js";
+import type { Agent } from "./index.js";
+import { getModelName, getProviderName } from "../shared.js";
+import { errorToString, willContinue } from "./utils.js";
+
+/**
+ * This behaves like {@link streamText} from the "ai" package except that
+ * it add context based on the userId and threadId and saves the input and
+ * resulting messages to the thread, if specified.
+ * Use {@link continueThread} to get a version of this function already scoped
+ * to a thread (and optionally userId).
+ */
+export async function streamText<
+  TOOLS extends ToolSet,
+  OUTPUT extends Output<any, any, any> = never,
+>(
+  ctx: ActionCtx,
+  component: AgentComponent,
+  /**
+   * The arguments to the streamText function, similar to the ai sdk's
+   * {@link streamText} function, along with Agent prompt options.
+   */
+  streamTextArgs: AgentPrompt &
+    Omit<
+      Parameters<typeof streamTextAi<TOOLS, OUTPUT>>[0],
+      "model" | "prompt" | "messages"
+    > & {
+      /**
+       * The tools to use for the tool calls. This will override tools specified
+       * in the Agent constructor or createThread / continueThread.
+       */
+      tools?: TOOLS;
+    },
+  /**
+   * The {@link ContextOptions} and {@link StorageOptions}
+   * options to use for fetching contextual messages and saving input/output messages.
+   */
+  options: Options & {
+    agentName: string;
+    userId?: string | null;
+    threadId?: string;
+    /**
+     * Whether to save incremental data (deltas) from streaming responses.
+     * Defaults to false.
+     * If false, it will not save any deltas to the database.
+     * If true, it will save deltas with {@link DEFAULT_STREAMING_OPTIONS}.
+     *
+     * Regardless of this option, when streaming you are able to use this
+     * `streamText` function as you would with the "ai" package's version:
+     * iterating over the text, streaming it over HTTP, etc.
+     */
+    saveStreamDeltas?: boolean | StreamingOptions;
+    agentForToolCtx?: Agent;
+  },
+): Promise<StreamTextResult<TOOLS, OUTPUT> & GenerationOutputMetadata> {
+  const { threadId } = options ?? {};
+  const { args, userId, order, stepOrder, promptMessageId, ...call } =
+    await startGeneration(ctx, component, streamTextArgs, options);
+
+  const steps: StepResult<TOOLS>[] = [];
+
+  const streamer =
+    threadId && options.saveStreamDeltas
+      ? new DeltaStreamer(
+          component,
+          ctx,
+          {
+            throttleMs:
+              typeof options.saveStreamDeltas === "object"
+                ? options.saveStreamDeltas.throttleMs
+                : undefined,
+            onAsyncAbort: call.fail,
+            compress: compressUIMessageChunks,
+            abortSignal: args.abortSignal,
+          },
+          {
+            threadId,
+            userId,
+            agentName: options?.agentName,
+            model: getModelName(args.model),
+            provider: getProviderName(args.model),
+            providerOptions: args.providerOptions,
+            format: "UIMessageChunk",
+            order,
+            stepOrder,
+          },
+        )
+      : undefined;
+
+  const result = streamTextAi({
+    ...args,
+    abortSignal: streamer?.abortController.signal ?? args.abortSignal,
+    experimental_transform: mergeTransforms(
+      options?.saveStreamDeltas,
+      streamTextArgs.experimental_transform,
+    ),
+    onError: async (error) => {
+      console.error("onError", error);
+      await call.fail(errorToString(error.error));
+      await streamer?.fail(errorToString(error.error));
+      return streamTextArgs.onError?.(error);
+    },
+    prepareStep: async (options) => {
+      const result = await streamTextArgs.prepareStep?.(options);
+      if (result) {
+        const model = result.model ?? options.model;
+        call.updateModel(model);
+        // streamer?.updateMetadata({
+        //   model: getModelName(model),
+        //   provider: getProviderName(model),
+        //   providerOptions: options.messages.at(-1)?.providerOptions,
+        // });
+        return result;
+      }
+      return undefined;
+    },
+    onStepFinish: async (step) => {
+      steps.push(step);
+      const createPendingMessage = await willContinue(steps, args.stopWhen);
+      await call.save({ step }, createPendingMessage);
+      return args.onStepFinish?.(step);
+    },
+  }) as StreamTextResult<TOOLS, OUTPUT>;
+  const stream = streamer?.consumeStream(
+    result.toUIMessageStream<AIUIMessage<TOOLS>>(),
+  );
+  if (
+    (typeof options?.saveStreamDeltas === "object" &&
+      !options.saveStreamDeltas.returnImmediately) ||
+    options?.saveStreamDeltas === true
+  ) {
+    await stream;
+    await result.consumeStream();
+  }
+  const metadata: GenerationOutputMetadata = {
+    promptMessageId,
+    order,
+    savedMessages: call.getSavedMessages(),
+    messageId: promptMessageId,
+  };
+  return Object.assign(result, metadata);
+}

package/src/client/streaming.test.ts

@@ -0,0 +1,186 @@
+import { beforeEach, describe, expect, test } from "vitest";
+import { createThread } from "./index.js";
+import type { GenericSchema, SchemaDefinition } from "convex/server";
+import { streamText } from "ai";
+import { components, initConvexTest } from "./setup.test.js";
+import { mockModel } from "./mockModel.js";
+import { compressUIMessageChunks, DeltaStreamer } from "./streaming.js";
+import { getParts } from "../deltas.js";
+import type { TestConvex } from "convex-test";
+
+const defaultTestOptions = {
+  throttleMs: 0,
+  abortSignal: undefined,
+  compress: null,
+  onAsyncAbort: async () => {
+    throw new Error("unexpected");
+  },
+};
+
+const testMetadata = {
+  order: 0,
+  stepOrder: 0,
+  agentName: "test agent",
+  model: "test model",
+  provider: "test provider",
+  providerOptions: {},
+  format: "UIMessageChunk" as const,
+};
+
+describe("DeltaStreamer", () => {
+  let t: TestConvex<SchemaDefinition<GenericSchema, boolean>>;
+  let threadId: string;
+  beforeEach(async () => {
+    t = initConvexTest();
+    await t.run(async (ctx) => {
+      threadId = await createThread(ctx, components.agent, {});
+    });
+  });
+  test("should save chunks via DeltaStreamer", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer(
+        components.agent,
+        ctx,
+        { ...defaultTestOptions },
+        { ...testMetadata, threadId },
+      );
+      const result = streamText({
+        model: mockModel(),
+        prompt: "Test prompt",
+      });
+      await streamer.consumeStream(result.toUIMessageStream());
+      const streamId = streamer.streamId!;
+      expect(streamId).toBeDefined();
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
+      const { parts } = getParts(deltas);
+      const stream = result.toUIMessageStream();
+      for await (const part of stream) {
+        const expected = parts.shift();
+        expect(part).toEqual(expected);
+      }
+    });
+  });
+  test("should save all parts when throttleMs is 0", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer(
+        components.agent,
+        ctx,
+        { ...defaultTestOptions, throttleMs: 0 },
+        { ...testMetadata, threadId },
+      );
+      const result = streamText({
+        model: mockModel({
+          content: [
+            // The mockModel splits these into deltas based on spaces
+            { type: "text", text: "A B C" },
+            { type: "reasoning", text: "D E F" },
+          ],
+        }),
+        prompt: "Test prompt",
+      });
+      await streamer.consumeStream(result.toUIMessageStream());
+      const streamId = streamer.streamId!;
+      expect(streamId).toBeDefined();
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
+      const { parts } = getParts(deltas);
+      const expected = [
+        { type: "start" },
+        { type: "start-step" },
+        { type: "text-start" },
+        { type: "text-delta", delta: "A" },
+        { type: "text-delta", delta: " B" },
+        { type: "text-delta", delta: " C" },
+        { type: "text-end" },
+        { type: "reasoning-start" },
+        { type: "reasoning-delta", delta: "D" },
+        { type: "reasoning-delta", delta: " E" },
+        { type: "reasoning-delta", delta: " F" },
+        { type: "reasoning-end" },
+        { type: "finish-step" },
+        { type: "finish" },
+      ];
+      for (const expectedPart of expected) {
+        const part = parts.shift();
+        expect(part).toBeDefined();
+        expect(part).toMatchObject(expectedPart);
+      }
+    });
+  });
+
+  test("should save compressed parts via DeltaStreamer", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer(
+        components.agent,
+        ctx,
+        {
+          throttleMs: 1000,
+          abortSignal: undefined,
+          compress: compressUIMessageChunks,
+          onAsyncAbort: async () => {
+            throw new Error("async abort");
+          },
+        },
+        {
+          ...testMetadata,
+          threadId,
+        },
+      );
+      const result = streamText({
+        model: mockModel({
+          content: [
+            // The mockModel splits these into deltas based on spaces
+            { type: "text", text: "A B C" },
+            { type: "text", text: "D E F" },
+            { type: "reasoning", text: "J K L" },
+            { type: "text", text: "M N O" },
+          ],
+        }),
+        prompt: "Test prompt",
+        // experimental_transform: smoothStream({ chunking: "line" }),
+        onError: (error) => {
+          console.error(error);
+        },
+      });
+      await streamer.consumeStream(result.toUIMessageStream());
+      const streamId = streamer.streamId!;
+      expect(streamId).toBeDefined();
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
+      const { parts } = getParts(deltas);
+      const expected = [
+        { type: "start" },
+        { type: "start-step" },
+        { type: "text-start" },
+        // These are collapsed into a single delta
+        { type: "text-delta", delta: "A B C" },
+        { type: "text-end" },
+        { type: "text-start" },
+        { type: "text-delta", delta: "D E F" },
+        { type: "text-end" },
+        { type: "reasoning-start" },
+        { type: "reasoning-delta", delta: "J K L" },
+        { type: "reasoning-end" },
+        { type: "text-start" },
+        { type: "text-delta", delta: "M N O" },
+        { type: "text-end" },
+        { type: "finish-step" },
+        { type: "finish" },
+      ];
+      for (const expectedPart of expected) {
+        const part = parts.shift();
+        expect(part).toBeDefined();
+        expect(part).toMatchObject(expectedPart);
+      }
+    });
+  });
+  // TODO: test errors & aborted states
+  // TODO: test fetching partial stream data - syncStreams w/ cursors
+});