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

package/src/client/types.ts

@@ -1,49 +1,184 @@
-import type { LanguageModelV2 } from "@ai-sdk/provider";
 import type {
+  FlexibleSchema,
+  InferSchema,
   ModelMessage,
-  DeepPartial,
-  generateObject,
+  ProviderOptions,
+} from "@ai-sdk/provider-utils";
+import type { JSONValue } from "@ai-sdk/provider";
+import type {
+  EmbeddingModel,
   GenerateObjectResult,
   generateText,
   GenerateTextResult,
-  JSONValue,
   LanguageModelRequestMetadata,
   LanguageModelResponseMetadata,
-  RepairTextFunction,
+  LanguageModelUsage,
+  LanguageModel,
   streamObject,
-  StreamObjectResult,
   streamText,
   StreamTextResult,
-  TelemetrySettings,
-  ToolChoice,
   ToolSet,
-  Schema,
-  LanguageModelUsage,
+  CallSettings,
+  generateObject,
 } from "ai";
+
+export interface Output<_T = any, _P = any, _E = any> {
+  name: string;
+  responseFormat: any;
+  parseCompleteOutput: any;
+  parsePartialOutput: any;
+  createElementStreamTransform: any;
+}
 import type {
-  Auth,
-  Expand,
-  FunctionReference,
-  StorageActionWriter,
-  StorageReader,
-  WithoutSystemFields,
-  FunctionArgs,
-  FunctionReturnType,
   GenericActionCtx,
   GenericDataModel,
+  GenericMutationCtx,
+  GenericQueryCtx,
+  WithoutSystemFields,
 } from "convex/server";
-import type { GenericId } from "convex/values";
-import type { z } from "zod";
-import type { Mounts } from "../component/_generated/api.js";
-import type { ThreadDoc } from "../component/schema.js";
 import type {
-  CallSettings,
+  MessageDoc,
   ProviderMetadata,
-  ProviderOptions,
   StreamDelta,
   StreamMessage,
+  ThreadDoc,
 } from "../validators.js";
 import type { StreamingOptions } from "./streaming.js";
+import type { ComponentApi } from "../component/_generated/component.js";
+
+export type AgentPrompt = {
+  /**
+   * System message to include in the prompt. Overwrites Agent instructions.
+   */
+  system?: string;
+  /**
+   * A prompt. It can be either a text prompt or a list of messages.
+   * If used with `promptMessageId`, it will be used in place of that
+   * prompt message and no input messages will be saved.
+   * Otherwise, if used with the storageOptions "promptAndOutput" (default),
+   * it will be the only message saved.
+   * If a string is provided, it will be a user message.
+   */
+  prompt?: string | Array<ModelMessage> | undefined;
+  /**
+   * A list of messages to use as context before the prompt.
+   * If used with `prompt`, these will precede the prompt.
+   * If used with the storageOptions "promptAndOutput" (default),
+   * none of these messages will be saved.
+   */
+  messages?: Array<ModelMessage> | undefined;
+  /**
+   * If provided, it uses this existing message to anchor the prompt:
+   * - The specified message will be included, unless `prompt` is also
+   *   provided, in which case that will be inserted in place of this
+   *   specified message.
+   * - Recent and search messages will not include messages after this
+   *   message's order.
+   * - If there are already responses on the same order,
+   *   for example, tool calls and responses,
+   *   those will be included automatically.
+   *
+   * Note: if this is provided, no input messages will be saved by default.
+   */
+  promptMessageId?: string | undefined;
+  /**
+   * The model to use for the LLM calls. This will override the languageModel
+   * specified in the Agent config.
+   */
+  model?: LanguageModel;
+};
+
+export type Config = {
+  /**
+   * The LLM model to use for generating / streaming text and objects.
+   * e.g.
+   * import { openai } from "@ai-sdk/openai"
+   * const myAgent = new Agent(components.agent, {
+   *   languageModel: openai.chat("gpt-4o-mini"),
+   */
+  languageModel?: LanguageModel;
+  /**
+   * The model to use for text embeddings. Optional.
+   * If specified, it will use this for generating vector embeddings
+   * of chats, and can opt-in to doing vector search for automatic context
+   * on generateText, etc.
+   * e.g.
+   * import { openai } from "@ai-sdk/openai"
+   * const myAgent = new Agent(components.agent, {
+   *   ...
+   *   textEmbeddingModel: openai.embedding("text-embedding-3-small")
+   * @deprecated — Use embeddingModel instead.
+   */
+  textEmbeddingModel?: EmbeddingModel;
+  /**
+   * The model to use for text embeddings. Optional.
+   * If specified, it will use this for generating vector embeddings
+   * of chats, and can opt-in to doing vector search for automatic context
+   * on generateText, etc.
+   * e.g.
+   * import { openai } from "@ai-sdk/openai"
+   * const myAgent = new Agent(components.agent, {
+   *   ...
+   *   embeddingModel: openai.embedding("text-embedding-3-small")
+   */
+  embeddingModel?: EmbeddingModel;
+  /**
+   * Options to determine what messages are included as context in message
+   * generation. To disable any messages automatically being added, pass:
+   * { recentMessages: 0 }
+   */
+  contextOptions?: ContextOptions;
+  /**
+   * Determines whether messages are automatically stored when passed as
+   * arguments or generated.
+   */
+  storageOptions?: StorageOptions;
+  /**
+   * The usage handler to use for this agent.
+   */
+  usageHandler?: UsageHandler;
+  /**
+   * By default, messages are ordered with context in `fetchContextWithPrompt`,
+   * but you can override this by providing a context handler. Here you can
+   * filter, modify, or enrich the context messages. If provided, the default
+   * ordering will not apply. This excludes the system message / instructions.
+   */
+  contextHandler?: ContextHandler;
+  /**
+   * Called for each LLM request/response, so you can do things like
+   * log the raw request body or response headers to a table, or logs.
+   */
+  rawRequestResponseHandler?: RawRequestResponseHandler;
+  /**
+   * @deprecated Reach out if you use this. Otherwise will be removed soon.
+   * Default provider options to pass for the LLM calls.
+   * This can be overridden at each generate/stream callsite on a per-field
+   * basis. To clear a default setting, you'll need to pass `undefined`.
+   */
+  providerOptions?: ProviderOptions;
+  /**
+   * The default settings to use for the LLM calls.
+   * This can be overridden at each generate/stream callsite on a per-field
+   * basis. To clear a default setting, you'll need to pass `undefined`.
+   */
+  callSettings?: CallSettings;
+  /**
+   * The maximum number of steps to allow for a single generation.
+   *
+   * For example, if an agent wants to call a tool, that call and tool response
+   * will be one step. Generating a response based on the tool call & response
+   * will be a second step.
+   * If it runs out of steps, it will return the last step result, which may
+   * not be an assistant message.
+
+   * This becomes the default value when `stopWhen` is not specified in the
+   * Agent or generation callsite.
+   * AI SDK v5 removed the `maxSteps` argument, but this is kept here for
+   * convenience and backwards compatibility.
+   * Defaults to 1.
+   */
+  maxSteps?: number;
+};
 
 /**
  * Options to configure what messages are fetched as context,
@@ -115,10 +250,31 @@ 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,
+  ctx: ActionCtx,
   args: {
     userId: string | undefined;
     threadId: string | undefined;
@@ -131,6 +287,54 @@ export type UsageHandler = (
   },
 ) => void | Promise<void>;
 
+/**
+ * By default, messages are ordered with context in `fetchContextWithPrompt`,
+ * but you can override this by providing a context handler. Here you can filter
+ * out, add in, or reorder messages.
+ */
+export type ContextHandler = (
+  ctx: ActionCtx,
+  args: {
+    /**
+     * All messages in the default order.
+     */
+    allMessages: ModelMessage[];
+    /**
+     * The messages fetched from search.
+     */
+    search: ModelMessage[];
+    /**
+     * The recent messages already in the thread history,
+     * excluding any messages that came after promptMessageId.
+     */
+    recent: ModelMessage[];
+    /**
+     * The messages passed as the `messages` argument to e.g. generateText.
+     */
+    inputMessages: ModelMessage[];
+    /**
+     * The message(s) passed as the `prompt` argument to e.g. generateText.
+     * Otherwise, if `promptMessageId` was provided, the message at that id.
+     * `prompt` will override the message at `promptMessageId`.
+     */
+    inputPrompt: ModelMessage[];
+    /**
+     * Any messages on the same `order` as the promptMessageId message after the
+     * prompt message. These are presumably existing responses to the prompt
+     * message.
+     */
+    existingResponses: ModelMessage[];
+    /**
+     * The user associated with the generation, if any.
+     */
+    userId: string | undefined;
+    /**
+     * The thread associated with the generation, if any.
+     */
+    threadId: string | undefined;
+  },
+) => ModelMessage[] | Promise<ModelMessage[]>;
+
 export type RawRequestResponseHandler = (
   ctx: ActionCtx,
   args: {
@@ -142,157 +346,117 @@ export type RawRequestResponseHandler = (
   },
 ) => void | Promise<void>;
 
-export type AgentComponent = UseApi<Mounts>;
+export type AgentComponent = ComponentApi;
 
 export type TextArgs<
   AgentTools extends ToolSet,
   TOOLS extends ToolSet | undefined = undefined,
-  OUTPUT = never,
-  OUTPUT_PARTIAL = never,
+  OUTPUT extends Output<any, any, any> = never,
 > = Omit<
   Parameters<
     typeof generateText<
       TOOLS extends undefined ? AgentTools : TOOLS,
-      OUTPUT,
-      OUTPUT_PARTIAL
+      OUTPUT
     >
   >[0],
-  "toolChoice" | "tools" | "model"
+  "model" | "prompt" | "messages"
 > & {
-  /**
-   * 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?: LanguageModelV2;
   /**
    * 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>;
-};
+} & AgentPrompt;
 
 export type StreamingTextArgs<
   AgentTools extends ToolSet,
   TOOLS extends ToolSet | undefined = undefined,
-  OUTPUT = never,
-  OUTPUT_PARTIAL = never,
+  OUTPUT extends Output<any, any, any> = never,
 > = Omit<
   Parameters<
     typeof streamText<
       TOOLS extends undefined ? AgentTools : TOOLS,
-      OUTPUT,
-      OUTPUT_PARTIAL
+      OUTPUT
     >
   >[0],
-  "toolChoice" | "tools" | "model"
+  "model" | "prompt" | "messages"
 > & {
-  /**
-   * 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?: LanguageModelV2;
   /**
    * 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?: 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;
-};
-
-type StandardGenerateObjectOptions<T> = {
-  schema: z.Schema<T>;
-  schemaName?: string;
-  schemaDescription?: string;
-  output?: "object" | "array";
-  mode?: "auto" | "json" | "tool";
-};
+} & AgentPrompt;
 
-// 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" }
-  );
+export type ObjectMode = "object" | "array" | "enum" | "no-schema";
 
-type StreamObjectArgs<T> = BaseGenerateObjectOptions &
-  (
-    | StandardGenerateObjectOptions<T>
-    | { output: "any"; schema: undefined; mode: "json" }
-  );
+export type GenerateObjectArgs<
+  SCHEMA extends FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
+  OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+    ? "enum"
+    : "object",
+  RESULT = OUTPUT extends "array"
+    ? Array<InferSchema<SCHEMA>>
+    : InferSchema<SCHEMA>,
+> = AgentPrompt &
+  Omit<
+    Parameters<typeof generateObject<SCHEMA, OUTPUT, RESULT>>[0],
+    "model" | "prompt" | "messages"
+  > & {
+    schema?: SCHEMA;
+    enum?: Array<RESULT>;
+  };
 
-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 StreamObjectArgs<
+  SCHEMA extends FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
+  OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+    ? "enum"
+    : "object",
+  RESULT = OUTPUT extends "array"
+    ? Array<InferSchema<SCHEMA>>
+    : InferSchema<SCHEMA>,
+> = AgentPrompt &
+  Omit<
+    Parameters<typeof streamObject<SCHEMA, OUTPUT, RESULT>>[0],
+    "model" | "prompt" | "messages"
+  > & {
+    schema?: SCHEMA;
+    enum?: Array<RESULT>;
+  };
 
-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>;
 
@@ -327,15 +491,14 @@ export interface Thread<DefaultTools extends ToolSet> {
    */
   generateText<
     TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    OUTPUT_PARTIAL = never,
+    OUTPUT extends Output<any, any, any> = never,
   >(
-    args: TextArgs<
-      TOOLS extends undefined ? DefaultTools : TOOLS,
-      TOOLS,
-      OUTPUT,
-      OUTPUT_PARTIAL
-    >,
+    generateTextArgs: AgentPrompt &
+      TextArgs<
+        TOOLS extends undefined ? DefaultTools : TOOLS,
+        TOOLS,
+        OUTPUT
+      >,
     options?: Options,
   ): Promise<
     GenerateTextResult<TOOLS extends undefined ? DefaultTools : TOOLS, OUTPUT> &
@@ -354,15 +517,14 @@ export interface Thread<DefaultTools extends ToolSet> {
    */
   streamText<
     TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    PARTIAL_OUTPUT = never,
+    OUTPUT extends Output<any, any, any> = never,
   >(
-    args: StreamingTextArgs<
-      TOOLS extends undefined ? DefaultTools : TOOLS,
-      TOOLS,
-      OUTPUT,
-      PARTIAL_OUTPUT
-    >,
+    streamTextArgs: AgentPrompt &
+      StreamingTextArgs<
+        TOOLS extends undefined ? DefaultTools : TOOLS,
+        TOOLS,
+        OUTPUT
+      >,
     options?: Options & {
       /**
        * Whether to save incremental data (deltas) from streaming responses.
@@ -379,7 +541,7 @@ export interface Thread<DefaultTools extends ToolSet> {
   ): Promise<
     StreamTextResult<
       TOOLS extends undefined ? DefaultTools : TOOLS,
-      PARTIAL_OUTPUT
+      OUTPUT
     > &
       ThreadOutputMetadata
   >;
@@ -393,10 +555,19 @@ 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 FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
+    OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+      ? "enum"
+      : "object",
+    RESULT = OUTPUT extends "array"
+      ? Array<InferSchema<SCHEMA>>
+      : InferSchema<SCHEMA>,
+  >(
+    generateObjectArgs: AgentPrompt &
+      GenerateObjectArgs<SCHEMA, OUTPUT, RESULT>,
     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
@@ -407,11 +578,23 @@ 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 FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
+    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: AgentPrompt & StreamObjectArgs<SCHEMA, OUTPUT, RESULT>,
     options?: Options,
   ): Promise<
-    StreamObjectResult<DeepPartial<T>, T, never> & ThreadOutputMetadata
+    ReturnType<typeof streamObject<SCHEMA, OUTPUT, RESULT>> &
+      ThreadOutputMetadata
   >;
 }
 
@@ -424,6 +607,18 @@ 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;
+  /**
+   * By default, messages are ordered with context in `fetchContextWithPrompt`,
+   * but you can override this by providing a context handler. Here you can
+   * filter, modify, or enrich the context messages. If provided, the default
+   * ordering will not apply. This excludes the system message / instructions.
+   */
+  contextHandler?: ContextHandler;
 };
 
 export type SyncStreamsReturnValue =
@@ -432,58 +627,12 @@ export type SyncStreamsReturnValue =
   | undefined;
 
 /* Type utils follow */
-export type RunQueryCtx = {
-  runQuery: <Query extends FunctionReference<"query", "internal">>(
-    query: Query,
-    args: FunctionArgs<Query>,
-  ) => Promise<FunctionReturnType<Query>>;
-};
-export type RunMutationCtx = RunQueryCtx & {
-  runMutation: <Mutation extends FunctionReference<"mutation", "internal">>(
-    mutation: Mutation,
-    args: FunctionArgs<Mutation>,
-  ) => Promise<FunctionReturnType<Mutation>>;
-};
-export type RunActionCtx = RunMutationCtx & {
-  runAction<Action extends FunctionReference<"action", "internal">>(
-    action: Action,
-    args: FunctionArgs<Action>,
-  ): Promise<FunctionReturnType<Action>>;
-};
-export type UserActionCtx = GenericActionCtx<GenericDataModel>;
-export type ActionCtx = RunActionCtx & {
-  auth: Auth;
-  storage: StorageActionWriter;
-};
-export type QueryCtx = RunQueryCtx & {
-  storage: StorageReader;
-};
-
-export type OpaqueIds<T> =
-  T extends GenericId<infer _T>
-    ? string
-    : T extends (infer U)[]
-      ? OpaqueIds<U>[]
-      : T extends ArrayBuffer
-        ? ArrayBuffer
-        : T extends object
-          ? { [K in keyof T]: OpaqueIds<T[K]> }
-          : T;
-
-export type UseApi<API> = Expand<{
-  [mod in keyof API]: API[mod] extends FunctionReference<
-    infer FType,
-    "public",
-    infer FArgs,
-    infer FReturnType,
-    infer FComponentPath
-  >
-    ? FunctionReference<
-        FType,
-        "internal",
-        OpaqueIds<FArgs>,
-        OpaqueIds<FReturnType>,
-        FComponentPath
-      >
-    : UseApi<API[mod]>;
-}>;
+export type QueryCtx = Pick<GenericQueryCtx<GenericDataModel>, "runQuery">;
+export type MutationCtx = Pick<
+  GenericMutationCtx<GenericDataModel>,
+  "runQuery" | "runMutation"
+>;
+export type ActionCtx = Pick<
+  GenericActionCtx<GenericDataModel>,
+  "runQuery" | "runMutation" | "runAction" | "storage" | "auth"
+>;