@convex-dev/agent 0.1.18 → 0.2.0

package/src/client/types.ts

@@ -1,48 +1,50 @@
 import type {
-  CoreMessage,
-  DeepPartial,
-  generateObject,
+  InferSchema,
+  ProviderOptions,
+  Schema,
+} from "@ai-sdk/provider-utils";
+import type { JSONValue } from "@ai-sdk/provider";
+import type {
   GenerateObjectResult,
   generateText,
   GenerateTextResult,
-  JSONValue,
   LanguageModelRequestMetadata,
   LanguageModelResponseMetadata,
-  LanguageModelV1,
-  RepairTextFunction,
+  LanguageModelUsage,
+  LanguageModel,
   streamObject,
-  StreamObjectResult,
   streamText,
   StreamTextResult,
-  TelemetrySettings,
   ToolChoice,
   ToolSet,
+  RepairTextFunction,
+  TelemetrySettings,
+  CallSettings,
+  Prompt,
 } from "ai";
 import type {
   Auth,
   Expand,
-  FunctionReference,
-  StorageActionWriter,
-  StorageReader,
-  WithoutSystemFields,
   FunctionArgs,
+  FunctionReference,
   FunctionReturnType,
   GenericActionCtx,
   GenericDataModel,
+  StorageActionWriter,
+  StorageReader,
+  WithoutSystemFields,
 } from "convex/server";
 import type { GenericId } from "convex/values";
-import type { Schema } from "zod";
 import type { Mounts } from "../component/_generated/api.js";
-import type { ThreadDoc } from "../component/schema.js";
+import type { MessageDoc, ThreadDoc } from "../component/schema.js";
 import type {
-  CallSettings,
   ProviderMetadata,
-  ProviderOptions,
   StreamDelta,
   StreamMessage,
-  Usage,
 } from "../validators.js";
 import type { StreamingOptions } from "./streaming.js";
+import type * as z3 from "zod/v3";
+import type * as z4 from "zod/v4";
 
 /**
  * Options to configure what messages are fetched as context,
@@ -114,7 +116,28 @@ export type StorageOptions = {
   saveMessages?: "all" | "none" | "promptAndOutput";
 };
 
-export type GenerationOutputMetadata = { messageId?: string; order?: number };
+export type GenerationOutputMetadata = {
+  /**
+   * The ID of the prompt message for the generation.
+   */
+  promptMessageId?: string;
+  /**
+   * The order of the prompt message and responses for the generation.
+   * Each order starts with a user message, then followed by agent responses.
+   * If a promptMessageId is provided, that dictates the order.
+   */
+  order?: number;
+  /**
+   * The messages saved for the generation - both saved input and output.
+   * If you passed promptMessageId, it will not include that message.
+   */
+  savedMessages?: MessageDoc[];
+  /**
+   * @deprecated Use promptMessageId instead.
+   * The ID of the prompt message for the generation.
+   */
+  messageId?: string;
+};
 
 export type UsageHandler = (
   ctx: RunActionCtx,
@@ -122,7 +145,7 @@ export type UsageHandler = (
     userId: string | undefined;
     threadId: string | undefined;
     agentName: string | undefined;
-    usage: Usage;
+    usage: LanguageModelUsage;
     // Often has more information, like cached token usage in the case of openai.
     providerMetadata: ProviderMetadata | undefined;
     model: string;
@@ -169,7 +192,7 @@ export type TextArgs<
    * The model to use for the LLM calls. This will override the model specified
    * in the Agent constructor.
    */
-  model?: LanguageModelV1;
+  model?: LanguageModel;
   /**
    * The tools to use for the tool calls. This will override tools specified
    * in the Agent constructor or createThread / continueThread.
@@ -208,7 +231,7 @@ export type StreamingTextArgs<
    * The model to use for the tool calls. This will override the model specified
    * in the Agent constructor.
    */
-  model?: LanguageModelV1;
+  model?: LanguageModel;
   /**
    * The tools to use for the tool calls. This will override tools specified
    * in the Agent constructor or createThread / continueThread.
@@ -221,100 +244,142 @@ export type StreamingTextArgs<
   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;
-};
+export type ObjectMode = "object" | "array" | "enum" | "no-schema";
+export type ObjectSchema = z3.Schema | z4.core.$ZodType | Schema;
+export type DefaultObjectSchema = z4.core.$ZodType<JSONValue>;
 
-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;
-};
+/**
+ * Due to some issues with the type inference of the ai sdk, we need to
+ * manually type the arguments for the generateObject function.
+ * This is a workaround to allow the model to be optional.
+ */
+export type GenerateObjectArgs<
+  SCHEMA extends ObjectSchema = DefaultObjectSchema,
+  OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+    ? "enum"
+    : "object",
+  RESULT = OUTPUT extends "array"
+    ? Array<InferSchema<SCHEMA>>
+    : InferSchema<SCHEMA>,
+> = Omit<CallSettings, "stopSequences"> &
+  Prompt &
+  (OUTPUT extends "enum"
+    ? {
+        /**
+The enum values that the model should use.
+*/
+        enum: Array<RESULT>;
+        mode?: "json";
+        output: "enum";
+      }
+    : OUTPUT extends "no-schema"
+      ? // eslint-disable-next-line @typescript-eslint/no-empty-object-type
+        {}
+      : {
+          /**
+The schema of the object that the model should generate.
+*/
+          schema: SCHEMA;
+          /**
+Optional name of the output that should be generated.
+Used by some providers for additional LLM guidance, e.g.
+via tool or schema name.
+*/
+          schemaName?: string;
+          /**
+Optional description of the output that should be generated.
+Used by some providers for additional LLM guidance, e.g.
+via tool or schema description.
+*/
+          schemaDescription?: string;
+          /**
+The mode to use for object generation.
 
-type GenerateObjectWithEnumOptions<T extends string> =
-  BaseGenerateObjectOptions & {
-    output: "enum";
-    enum: Array<T>;
-    mode?: "auto" | "json" | "tool";
-  };
+The schema is converted into a JSON schema and used in one of the following ways
 
-type GenerateObjectNoSchemaOptions = BaseGenerateObjectOptions & {
-  schema?: undefined;
-  mode?: "json";
-};
+- 'auto': The provider will choose the best mode for the model.
+- 'tool': A tool with the JSON schema as parameters is provided and the provider is instructed to use it.
+- 'json': The JSON schema and an instruction are injected into the prompt. If the provider supports JSON mode, it is enabled. If the provider supports JSON grammars, the grammar is used.
 
-// 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;
+Please note that most providers do not support all modes.
 
-type StreamObjectArgs<T> =
-  T extends Record<string, unknown>
-    ? GenerateObjectObjectOptions<T>
-    : T extends Array<unknown>
-      ? GenerateObjectArrayOptions<T>
-      : GenerateObjectNoSchemaOptions;
+Default and recommended: 'auto' (best mode for the model).
+*/
+          mode?: "auto" | "json" | "tool";
+        }) & {
+    output?: OUTPUT;
+    /**
+The language model to use.
+   */
+    model?: LanguageModel;
+    /**
+A function that attempts to repair the raw output of the model
+to enable JSON parsing.
+   */
+    experimental_repairText?: RepairTextFunction;
+    /**
+Optional telemetry configuration (experimental).
+     */
+    experimental_telemetry?: TelemetrySettings;
+    /**
+Additional provider-specific options. They are passed through
+to the provider from the AI SDK and enable provider-specific
+functionality that can be fully encapsulated in the provider.
+*/
+    providerOptions?: ProviderOptions;
+    /**
+     * Internal. For test use only. May change without notice.
+     */
+    _internal?: { generateId?: () => string; currentDate?: () => Date };
+  };
 
-export type OurObjectArgs<T> = GenerateObjectArgs<T> &
+export type StreamObjectArgs<
+  SCHEMA extends ObjectSchema = DefaultObjectSchema,
+  OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+    ? "enum"
+    : "object",
+  RESULT = OUTPUT extends "array"
+    ? Array<InferSchema<SCHEMA>>
+    : InferSchema<SCHEMA>,
+> = GenerateObjectArgs<SCHEMA, OUTPUT, RESULT> &
   Pick<
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    Parameters<typeof generateObject<any>>[0],
-    "experimental_repairText" | "abortSignal"
+    Parameters<typeof streamObject<SCHEMA, OUTPUT, RESULT>>[0],
+    "onError" | "onFinish" | "_internal"
   >;
 
-export type OurStreamObjectArgs<T> = StreamObjectArgs<T> &
-  Pick<
-    Parameters<typeof streamObject<T>>[0],
-    "onError" | "onFinish" | "abortSignal"
-  >;
+export type MaybeCustomCtx<
+  CustomCtx,
+  DataModel extends GenericDataModel,
+  AgentTools extends ToolSet,
+> =
+  CustomCtx extends Record<string, unknown>
+    ? {
+        /**
+         * If you have a custom ctx that you use with the Agent
+         * (e.g. new Agent<{ orgId: string }>(...))
+         * you need to provide this function to add any extra fields.
+         * e.g.
+         * ```ts
+         * const myAgent = new Agent<{ orgId: string }>(...);
+         * const myAction = myAgent.asTextAction({
+         *   customCtx: (ctx: ActionCtx, target, llmArgs) => {
+         *     const orgId = await lookupOrgId(ctx, target.threadId);
+         *     return { orgId };
+         *   },
+         * });
+         * ```
+         * Then, in your tools, you can
+         */
+        customCtx: (
+          ctx: GenericActionCtx<DataModel>,
+          target: {
+            userId?: string | undefined;
+            threadId?: string | undefined;
+          },
+          llmArgs: TextArgs<AgentTools>,
+        ) => CustomCtx;
+      }
+    : { customCtx?: never };
 
 type ThreadOutputMetadata = Required<GenerationOutputMetadata>;
 
@@ -352,7 +417,7 @@ export interface Thread<DefaultTools extends ToolSet> {
     OUTPUT = never,
     OUTPUT_PARTIAL = never,
   >(
-    args: TextArgs<
+    generateTextArgs: TextArgs<
       TOOLS extends undefined ? DefaultTools : TOOLS,
       TOOLS,
       OUTPUT,
@@ -379,7 +444,7 @@ export interface Thread<DefaultTools extends ToolSet> {
     OUTPUT = never,
     PARTIAL_OUTPUT = never,
   >(
-    args: StreamingTextArgs<
+    streamTextArgs: StreamingTextArgs<
       TOOLS extends undefined ? DefaultTools : TOOLS,
       TOOLS,
       OUTPUT,
@@ -415,24 +480,35 @@ export interface Thread<DefaultTools extends ToolSet> {
    * 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,
+  generateObject<
+    SCHEMA extends ObjectSchema = DefaultObjectSchema,
+    OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+      ? "enum"
+      : "object",
+    RESULT = OUTPUT extends "array"
+      ? Array<InferSchema<SCHEMA>>
+      : InferSchema<SCHEMA>,
+  >(
+    generateObjectArgs: GenerateObjectArgs<SCHEMA, OUTPUT, RESULT> & {
+      /**
+       * 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.
+       */
+    },
     options?: Options,
-  ): Promise<GenerateObjectResult<JSONValue> & ThreadOutputMetadata>;
+  ): Promise<GenerateObjectResult<RESULT> & 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
@@ -443,11 +519,40 @@ export interface Thread<DefaultTools extends ToolSet> {
    * for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the streamObject function.
    */
-  streamObject<T>(
-    args: OurStreamObjectArgs<T>,
+  streamObject<
+    SCHEMA extends ObjectSchema = DefaultObjectSchema,
+    OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+      ? "enum"
+      : "object",
+    RESULT = OUTPUT extends "array"
+      ? Array<InferSchema<SCHEMA>>
+      : InferSchema<SCHEMA>,
+  >(
+    /**
+     * The same arguments you'd pass to "ai" sdk {@link streamObject}.
+     */
+    streamObjectArgs: StreamObjectArgs<SCHEMA, OUTPUT, RESULT> & {
+      /**
+       * 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.
+       */
+    },
     options?: Options,
   ): Promise<
-    StreamObjectResult<DeepPartial<T>, T, never> & ThreadOutputMetadata
+    ReturnType<typeof streamObject<SCHEMA, OUTPUT, RESULT>> &
+      ThreadOutputMetadata
   >;
 }
 
@@ -460,6 +565,11 @@ export type Options = {
    * The storage options to use for saving the input and output messages to the thread.
    */
   storageOptions?: StorageOptions;
+  /**
+   * The usage handler to use for this thread. Overrides any handler
+   * set in the agent constructor.
+   */
+  usageHandler?: UsageHandler;
 };
 
 export type SyncStreamsReturnValue =
@@ -491,9 +601,7 @@ export type ActionCtx = RunActionCtx & {
   auth: Auth;
   storage: StorageActionWriter;
 };
-export type QueryCtx = RunQueryCtx & {
-  storage: StorageReader;
-};
+export type QueryCtx = RunQueryCtx & { storage: StorageReader };
 
 export type OpaqueIds<T> =
   T extends GenericId<infer _T>