@convex-dev/agent 0.2.0-alpha.2 → 0.2.0

package/src/client/types.ts

@@ -1,49 +1,50 @@
-import type { LanguageModelV2 } from "@ai-sdk/provider";
 import type {
-  ModelMessage,
-  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,
-  RepairTextFunction,
+  LanguageModelUsage,
+  LanguageModel,
   streamObject,
-  StreamObjectResult,
   streamText,
   StreamTextResult,
-  TelemetrySettings,
   ToolChoice,
   ToolSet,
-  Schema,
-  LanguageModelUsage,
+  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 { z } from "zod/v3";
 import type { Mounts } from "../component/_generated/api.js";
 import type { MessageDoc, ThreadDoc } from "../component/schema.js";
 import type {
-  CallSettings,
   ProviderMetadata,
-  ProviderOptions,
   StreamDelta,
   StreamMessage,
 } 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,
@@ -119,16 +120,23 @@ export type GenerationOutputMetadata = {
   /**
    * The ID of the prompt message for the generation.
    */
-  messageId?: string;
+  promptMessageId?: string;
   /**
-   * The order of the prompt message for the generation.
+   * 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.
    */
-  messages?: MessageDoc[];
+  savedMessages?: MessageDoc[];
+  /**
+   * @deprecated Use promptMessageId instead.
+   * The ID of the prompt message for the generation.
+   */
+  messageId?: string;
 };
 
 export type UsageHandler = (
@@ -184,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?: LanguageModelV2;
+  model?: LanguageModel;
   /**
    * The tools to use for the tool calls. This will override tools specified
    * in the Agent constructor or createThread / continueThread.
@@ -223,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?: LanguageModelV2;
+  model?: LanguageModel;
   /**
    * The tools to use for the tool calls. This will override tools specified
    * in the Agent constructor or createThread / continueThread.
@@ -236,77 +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?: LanguageModelV2;
-  /**
-   * 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 | Array<ModelMessage>;
-  /**
-   * 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?: Array<ModelMessage>;
-  /**
-   * 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 StandardGenerateObjectOptions<T> = {
-  schema: z.Schema<T>;
-  schemaName?: string;
-  schemaDescription?: string;
-  output?: "object" | "array";
-  mode?: "auto" | "json" | "tool";
-};
+/**
+ * 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.
+
+The schema is converted into a JSON schema and used in one of the following ways
 
-// TODO: simplify this to just use the generateObject args, with an optional
-// model and tool/toolChoice types
-type GenerateObjectArgs<T> = BaseGenerateObjectOptions &
-  (
-    | StandardGenerateObjectOptions<T>
-    | { output: "enum"; enum: Array<T>; mode?: "auto" | "json" | "tool" }
-    | { output: "any"; 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.
 
-type StreamObjectArgs<T> = BaseGenerateObjectOptions &
-  (
-    | StandardGenerateObjectOptions<T>
-    | { output: "any"; schema: undefined; mode: "json" }
-  );
+Please note that most providers do not support all modes.
+
+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 extends Schema | z.Schema> =
-  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>;
 
@@ -344,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,
@@ -371,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,
@@ -407,10 +480,35 @@ export interface Thread<DefaultTools extends ToolSet> {
    * for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the generateObject function.
    */
-  generateObject<T = JSONValue>(
-    args: OurObjectArgs<T>,
+  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<T> & 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
@@ -421,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 extends z.Schema | Schema>(
-    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
   >;
 }
 
@@ -438,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 =
@@ -469,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>