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

package/src/client/streaming.ts

@@ -0,0 +1,166 @@
+import { ChunkDetector, smoothStream, StreamTextTransform, ToolSet } from "ai";
+import { ProviderOptions, StreamDelta, TextStreamPart } from "../validators";
+import type { AgentComponent, MessageDoc } from "./index";
+import { RunActionCtx } from "./types";
+
+export type StreamingOptions = {
+  /**
+   * The minimum granularity of deltas to save.
+   * Note: this is not a guarantee that every delta will be exactly one line.
+   * E.g. if "line" is specified, it won't save any deltas until it encounters
+   * a newline character.
+   * Defaults to "word".
+   */
+  chunking?: "word" | "line" | RegExp | ChunkDetector;
+  /**
+   * The minimum number of milliseconds to wait between saving deltas.
+   * Defaults to 250.
+   */
+  throttleMs?: number;
+};
+export const DEFAULT_STREAMING_OPTIONS = {
+  chunking: "word",
+  throttleMs: 250,
+} satisfies StreamingOptions;
+
+export function mergeTransforms<TOOLS extends ToolSet>(
+  options: StreamingOptions | boolean | undefined,
+  existing:
+    | StreamTextTransform<TOOLS>
+    | Array<StreamTextTransform<TOOLS>>
+    | undefined
+) {
+  if (!options) {
+    return existing;
+  }
+  const chunking =
+    typeof options === "boolean"
+      ? DEFAULT_STREAMING_OPTIONS.chunking
+      : options.chunking;
+  const transforms = Array.isArray(existing)
+    ? existing
+    : existing
+      ? [existing]
+      : [];
+  transforms.push(smoothStream({ delayInMs: null, chunking }));
+  return transforms;
+}
+
+export class DeltaStreamer {
+  public streamId: string | undefined;
+  public readonly options: Required<StreamingOptions>;
+  #nextParts: TextStreamPart[] = [];
+  #nextOrder: number;
+  #nextStepOrder: number;
+  #latestWrite: number = 0;
+  #ongoingWrite: Promise<void> | undefined;
+  #cursor: number = 0;
+
+  constructor(
+    public readonly component: AgentComponent,
+    public readonly ctx: RunActionCtx,
+    options: true | StreamingOptions,
+    public readonly metadata: {
+      threadId: string;
+      agentName: string | undefined;
+      model: string | undefined;
+      provider: string | undefined;
+      providerOptions: ProviderOptions | undefined;
+      userId: string | undefined;
+      order: number | undefined;
+      stepOrder: number | undefined;
+    }
+  ) {
+    this.options =
+      typeof options === "boolean"
+        ? DEFAULT_STREAMING_OPTIONS
+        : {
+            ...DEFAULT_STREAMING_OPTIONS,
+            ...options,
+          };
+    this.metadata = metadata;
+    this.#nextParts = [];
+    this.#nextOrder = metadata.order ?? 0;
+    this.#nextStepOrder = (metadata.stepOrder ?? 0) + 1;
+  }
+  public async addParts(parts: TextStreamPart[]) {
+    if (!this.streamId) {
+      this.streamId = await this.ctx.runMutation(
+        this.component.streams.create,
+        {
+          ...this.metadata,
+          order: this.#nextOrder,
+          stepOrder: this.#nextStepOrder,
+        }
+      );
+    }
+    this.#nextParts.push(...parts);
+    if (
+      !this.#ongoingWrite &&
+      Date.now() - this.#latestWrite >= this.options.throttleMs
+    ) {
+      this.#ongoingWrite = this.#sendDelta();
+    }
+  }
+
+  async #sendDelta() {
+    const delta = this.#createDelta();
+    this.#latestWrite = Date.now();
+    await this.ctx.runMutation(this.component.streams.addDelta, delta);
+    // Now that we've sent the delta, check if we need to send another one.
+    if (
+      this.#nextParts.length > 0 &&
+      Date.now() - this.#latestWrite >= this.options.throttleMs
+    ) {
+      // We send again immediately with the accumulated deltas.
+      this.#ongoingWrite = this.#sendDelta();
+      console.debug("Sent delta immediately");
+    } else {
+      this.#ongoingWrite = undefined;
+    }
+  }
+
+  #createDelta(): StreamDelta {
+    const start = this.#cursor;
+    const end = start + this.#nextParts.length;
+    this.#cursor = end;
+    const parts = this.#nextParts;
+    this.#nextParts = [];
+    if (!this.streamId) {
+      throw new Error("Creating a delta before the stream is created");
+    }
+    return {
+      streamId: this.streamId,
+      start,
+      end,
+      parts,
+    };
+  }
+
+  public async finish(messages: MessageDoc[]) {
+    if (this.#ongoingWrite) {
+      await this.#ongoingWrite;
+      this.#ongoingWrite = undefined;
+    }
+    if (!this.streamId) {
+      throw new Error("Finish called before stream is created");
+    }
+    const lastMessage = messages.at(-1);
+    if (lastMessage) {
+      this.#nextOrder = lastMessage.order;
+      this.#nextStepOrder = lastMessage.stepOrder + 1;
+    } else {
+      console.warn("Step finished without generating a message");
+    }
+    const finalDelta =
+      this.#nextParts.length > 0 ? this.#createDelta() : undefined;
+    this.#nextParts = [];
+    const streamId = this.streamId;
+    this.streamId = undefined;
+    this.#cursor = 0;
+    await this.ctx.runMutation(this.component.streams.finish, {
+      streamId,
+      finalDelta,
+    });
+  }
+}

package/src/client/types.ts

@@ -1,13 +1,438 @@
-import {
-  GenericQueryCtx,
-  GenericDataModel,
-  GenericMutationCtx,
+import type {
+  CoreMessage,
+  DeepPartial,
+  GenerateObjectResult,
+  GenerateTextResult,
+  JSONValue,
+  RepairTextFunction,
+  StreamObjectResult,
+  StreamTextResult,
+  TelemetrySettings,
+  ToolChoice,
+  ToolSet,
+} from "ai";
+import type {
+  generateObject,
+  generateText,
+  LanguageModelV1,
+  streamObject,
+  streamText,
+} from "ai";
+import type {
   Expand,
   FunctionReference,
   GenericActionCtx,
+  GenericDataModel,
+  GenericMutationCtx,
+  GenericQueryCtx,
 } from "convex/server";
-import { GenericId } from "convex/values";
+import type { GenericId } from "convex/values";
+import type {
+  CallSettings,
+  ProviderMetadata,
+  ProviderOptions,
+  StreamDelta,
+  StreamMessage,
+  Usage,
+} from "../validators.js";
+import type { Mounts } from "../component/_generated/api.js";
+import type { Schema } from "zod";
+import type { StreamingOptions } from "./streaming.js";
+
+/**
+ * Options to configure what messages are fetched as context,
+ * automatically with thread.generateText, or directly via search.
+ */
+export type ContextOptions = {
+  /** @deprecated Use excludeToolMessages instead. */
+  includeToolCalls?: boolean;
+  /**
+   * Whether to include tool messages in the context.
+   * By default, tool calls and results are not included.
+   */
+  excludeToolMessages?: boolean;
+  /**
+   * How many recent messages to include. These are added after the search
+   * messages, and do not count against the search limit.
+   * Default: 100
+   */
+  recentMessages?: number;
+  /**
+   * Options for searching messages.
+   */
+  searchOptions?: {
+    /**
+     * The maximum number of messages to fetch. Default is 10.
+     */
+    limit: number;
+    /**
+     * Whether to use text search to find messages. Default is false.
+     */
+    textSearch?: boolean;
+    /**
+     * Whether to use vector search to find messages. Default is false.
+     * At least one of textSearch or vectorSearch must be true.
+     */
+    vectorSearch?: boolean;
+    /**
+     * What messages around the search results to include.
+     * Default: { before: 2, after: 1 }
+     * (two before, and one after each message found in the search)
+     * Note, this is after the limit is applied.
+     * By default this will quadruple the number of messages fetched.
+     */
+    messageRange?: { before: number; after: number };
+  };
+  /**
+   * Whether to search across other threads for relevant messages.
+   * By default, only the current thread is searched.
+   */
+  searchOtherThreads?: boolean;
+};
+
+/**
+ * Options to configure the automatic saving of messages
+ * when generating text / objects in a thread.
+ */
+export type StorageOptions = {
+  /**
+   * Defaults to false, allowing you to pass in arbitrary context that will
+   * be in addition to automatically fetched content.
+   * Pass true to have all input messages saved to the thread history.
+   */
+  saveAllInputMessages?: boolean;
+  /** Defaults to true, saving the prompt, or last message passed to generateText. */
+  saveAnyInputMessages?: boolean;
+  /** Defaults to true. Whether to save messages generated while chatting. */
+  saveOutputMessages?: boolean;
+};
+
+export type GenerationOutputMetadata = { messageId?: string };
+
+export type UsageHandler = (
+  ctx: RunActionCtx,
+  args: {
+    userId: string | undefined;
+    threadId: string | undefined;
+    agentName: string | undefined;
+    usage: Usage;
+    // Often has more information, like cached token usage in the case of openai.
+    providerMetadata: ProviderMetadata | undefined;
+    model: string;
+    provider: string;
+  }
+) => void | Promise<void>;
+
+export type AgentComponent = UseApi<Mounts>;
+
+export type TextArgs<
+  AgentTools extends ToolSet,
+  TOOLS extends ToolSet | undefined = undefined,
+  OUTPUT = never,
+  OUTPUT_PARTIAL = never,
+> = Omit<
+  Parameters<
+    typeof generateText<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      OUTPUT,
+      OUTPUT_PARTIAL
+    >
+  >[0],
+  "toolChoice" | "tools" | "model"
+> & {
+  /**
+   * 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 tool calls. This will override the model specified
+   * in the Agent constructor.
+   */
+  model?: LanguageModelV1;
+  /**
+   * The tools to use for the tool calls. This will override tools specified
+   * in the Agent constructor or createThread / continueThread.
+   */
+  tools?: TOOLS;
+  /**
+   * The tool choice to use for the tool calls. This must be one of the tools
+   * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
+   */
+  toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
+};
+
+export type StreamingTextArgs<
+  AgentTools extends ToolSet,
+  TOOLS extends ToolSet | undefined = undefined,
+  OUTPUT = never,
+  OUTPUT_PARTIAL = never,
+> = Omit<
+  Parameters<
+    typeof streamText<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      OUTPUT,
+      OUTPUT_PARTIAL
+    >
+  >[0],
+  "toolChoice" | "tools" | "model"
+> & {
+  /**
+   * The model to use for the tool calls. This will override the model specified
+   * in the Agent constructor.
+   */
+  model?: LanguageModelV1;
+  /**
+   * The tools to use for the tool calls. This will override tools specified
+   * in the Agent constructor or createThread / continueThread.
+   */
+  tools?: TOOLS;
+  /**
+   * The tool choice to use for the tool calls. This must be one of the tools
+   * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
+   */
+  toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
+};
+
+type BaseGenerateObjectOptions = CallSettings & {
+  /**
+   * The model to use for the object generation. This will override the model
+   * specified in the Agent constructor.
+   */
+  model?: LanguageModelV1;
+  /**
+   * The system prompt to use for the object generation. This will override the
+   * system prompt specified in the Agent constructor.
+   */
+  system?: string;
+  /**
+   * The prompt to the LLM to use for the object generation.
+   * Specify this or messages, but not both.
+   */
+  prompt?: string;
+  /**
+   * The messages to use for the object generation.
+   * Note: recent messages are automatically added based on the thread it's
+   * associated with and your contextOptions.
+   */
+  messages?: CoreMessage[];
+  /**
+   * The message to use as the "prompt" for the object generation.
+   * If this is provided, it will be used 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 object generation in another call.
+   */
+  promptMessageId?: string;
+  experimental_repairText?: RepairTextFunction;
+  experimental_telemetry?: TelemetrySettings;
+  providerOptions?: ProviderOptions;
+  experimental_providerMetadata?: ProviderMetadata;
+};
+
+type GenerateObjectObjectOptions<T extends Record<string, unknown>> =
+  BaseGenerateObjectOptions & {
+    output?: "object";
+    mode?: "auto" | "json" | "tool";
+    schema: Schema<T>;
+    schemaName?: string;
+    schemaDescription?: string;
+  };
+
+type GenerateObjectArrayOptions<T> = BaseGenerateObjectOptions & {
+  output: "array";
+  mode?: "auto" | "json" | "tool";
+  schema: Schema<T>;
+  schemaName?: string;
+  schemaDescription?: string;
+};
+
+type GenerateObjectWithEnumOptions<T extends string> =
+  BaseGenerateObjectOptions & {
+    output: "enum";
+    enum: Array<T>;
+    mode?: "auto" | "json" | "tool";
+  };
+
+type GenerateObjectNoSchemaOptions = BaseGenerateObjectOptions & {
+  schema?: undefined;
+  mode?: "json";
+};
+
+// TODO: simplify this to just use the generateObject args, with an optional
+// model and tool/toolChoice types
+type GenerateObjectArgs<T> =
+  T extends Record<string, unknown>
+    ? GenerateObjectObjectOptions<T>
+    : T extends Array<unknown>
+      ? GenerateObjectArrayOptions<T>
+      : T extends string
+        ? GenerateObjectWithEnumOptions<T>
+        : GenerateObjectNoSchemaOptions;
+
+type StreamObjectArgs<T> =
+  T extends Record<string, unknown>
+    ? GenerateObjectObjectOptions<T>
+    : T extends Array<unknown>
+      ? GenerateObjectArrayOptions<T>
+      : GenerateObjectNoSchemaOptions;
+
+export type OurObjectArgs<T> = GenerateObjectArgs<T> &
+  Pick<
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    Parameters<typeof generateObject<any>>[0],
+    "experimental_repairText" | "abortSignal"
+  >;
+
+export type OurStreamObjectArgs<T> = StreamObjectArgs<T> &
+  Pick<
+    Parameters<typeof streamObject<T>>[0],
+    "onError" | "onFinish" | "abortSignal"
+  >;
+
+type ThreadOutputMetadata = GenerationOutputMetadata & {
+  messageId: string;
+};
+
+/**
+ * The interface for a thread returned from {@link createThread} or {@link continueThread}.
+ * This is contextual to a thread and/or user.
+ */
+export interface Thread<DefaultTools extends ToolSet> {
+  /**
+   * The target threadId, from the startThread or continueThread initializers.
+   */
+  threadId: string;
+  /**
+   * This behaves like {@link generateText} 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).
+   * @param args The arguments to the generateText function, along with extra controls
+   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @returns The result of the generateText function.
+   */
+  generateText<
+    TOOLS extends ToolSet | undefined = undefined,
+    OUTPUT = never,
+    OUTPUT_PARTIAL = never,
+  >(
+    args: TextArgs<
+      TOOLS extends undefined ? DefaultTools : TOOLS,
+      TOOLS,
+      OUTPUT,
+      OUTPUT_PARTIAL
+    >,
+    options?: Options
+  ): Promise<
+    GenerateTextResult<TOOLS extends undefined ? DefaultTools : TOOLS, OUTPUT> &
+      ThreadOutputMetadata
+  >;
+
+  /**
+   * 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).
+   * @param args The arguments to the streamText function, along with extra controls
+   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @returns The result of the streamText function.
+   */
+  streamText<
+    TOOLS extends ToolSet | undefined = undefined,
+    OUTPUT = never,
+    PARTIAL_OUTPUT = never,
+  >(
+    args: StreamingTextArgs<
+      TOOLS extends undefined ? DefaultTools : TOOLS,
+      TOOLS,
+      OUTPUT,
+      PARTIAL_OUTPUT
+    >,
+    options?: Options & {
+      /**
+       * 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;
+    }
+  ): Promise<
+    StreamTextResult<
+      TOOLS extends undefined ? DefaultTools : TOOLS,
+      PARTIAL_OUTPUT
+    > &
+      ThreadOutputMetadata
+  >;
+  /**
+   * This behaves like {@link generateObject} 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. This overload is for objects, arrays, and enums.
+   * Use {@link continueThread} to get a version of this function already scoped
+   * to a thread (and optionally userId).
+   * @param args The arguments to the generateObject function, along with extra controls
+   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @returns The result of the generateObject function.
+   */
+  generateObject<T>(
+    args: OurObjectArgs<T>,
+    options?: Options
+  ): Promise<GenerateObjectResult<T> & ThreadOutputMetadata>;
+  /**
+   * This behaves like {@link generateObject} 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. This overload is for when there's no schema.
+   * Use {@link continueThread} to get a version of this function already scoped
+   * to a thread (and optionally userId).
+   * @param args The arguments to the generateObject function, along with extra controls
+   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @returns The result of the generateObject function.
+   */
+  generateObject(
+    args: GenerateObjectNoSchemaOptions,
+    options?: Options
+  ): Promise<GenerateObjectResult<JSONValue> & ThreadOutputMetadata>;
+  /**
+   * This behaves like {@link streamObject} 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).
+   * @param args The arguments to the streamObject function, along with extra controls
+   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @returns The result of the streamObject function.
+   */
+  streamObject<T>(
+    args: OurStreamObjectArgs<T>,
+    options?: Options
+  ): Promise<
+    StreamObjectResult<DeepPartial<T>, T, never> & ThreadOutputMetadata
+  >;
+}
+
+export type Options = {
+  /**
+   * The context options to use for passing in message history to the LLM.
+   */
+  contextOptions?: ContextOptions;
+  /**
+   * The storage options to use for saving the input and output messages to the thread.
+   */
+  storageOptions?: StorageOptions;
+};
 
+export type SyncStreamsReturnValue =
+  | { kind: "list"; messages: StreamMessage[] }
+  | { kind: "deltas"; deltas: StreamDelta[] }
+  | undefined;
 
 /* Type utils follow */
 export type RunQueryCtx = {