tool-invocation.ts

import type { z } from "zod";
import type { Schema, ToolSet } from "ai";

type ValueOf<
  ObjectType,
  ValueType extends keyof ObjectType = keyof ObjectType,
> = ObjectType[ValueType];

type Parameters = z.ZodTypeAny | Schema<any>;
type inferParameters<PARAMETERS extends Parameters> =
  PARAMETERS extends Schema<any>
    ? PARAMETERS["_type"]
    : PARAMETERS extends z.ZodTypeAny
      ? z.infer<PARAMETERS>
      : never;

type ToToolsWithoutExecute<TOOLS extends ToolSet> = {
  [K in keyof TOOLS as TOOLS[K] extends {
    execute: any;
  }
    ? never
    : K]: TOOLS[K];
};

type ToToolCallObject<TOOLS extends ToolSet> = ValueOf<{
  [NAME in keyof TOOLS]: {
    state: "call";
    toolCallId: string;
    toolName: NAME & string;
    args: inferParameters<TOOLS[NAME]["parameters"]>;
  };
}>;

type ToToolPartialCallObject<TOOLS extends ToolSet> = ValueOf<{
  [NAME in keyof TOOLS]: {
    state: "partial-call";
    toolCallId: string;
    toolName: NAME & string;
    args: Partial<inferParameters<TOOLS[NAME]["parameters"]>>;
  };
}>;

export type ToolInvocationCall<TOOLS extends ToolSet> = ToToolCallObject<TOOLS>;
export type ToolInvocationPartialCall<TOOLS extends ToolSet> =
  ToToolPartialCallObject<TOOLS>;
type ToolInvocationLocalCall<TOOLS extends ToolSet> = ToToolCallObject<
  ToToolsWithoutExecute<TOOLS>
>;

type ToToolsWithExecute<TOOLS extends ToolSet> = {
  [K in keyof TOOLS as TOOLS[K] extends {
    execute: any;
  }
    ? K
    : never]: TOOLS[K];
};

type ToToolsWithDefinedExecute<TOOLS extends ToolSet> = {
  [K in keyof TOOLS as TOOLS[K]["execute"] extends undefined
    ? never
    : K]: TOOLS[K];
};

type ToToolResultObject<TOOLS extends ToolSet> = ValueOf<{
  [NAME in keyof TOOLS]: {
    state: "result";
    toolCallId: string;
    toolName: NAME & string;
    args: inferParameters<TOOLS[NAME]["parameters"]>;
    result: Awaited<ReturnType<Exclude<TOOLS[NAME]["execute"], undefined>>>;
  };
}>;

export type ToolInvocationResult<TOOLS extends ToolSet> = ToToolResultObject<
  ToToolsWithDefinedExecute<ToToolsWithExecute<TOOLS>>
>;

export type ToolInvocation<TOOLS extends ToolSet> =
  | ToolInvocationResult<TOOLS>
  | ToolInvocationCall<TOOLS>
  | ToolInvocationPartialCall<TOOLS>;