Vectrum
Sign In
104 lines · 4.6 KB · TypeScript
// Tool parameters, declared once: each gives the JSON Schema that tools/list advertises and checks
// the value an agent sends. Bad arguments come back as tool errors, so the agent can correct itself.

export class ToolError extends Error {}

type Schema = Record<string, unknown>;

export type Param<T> = { schema: Schema; required: boolean; parse: (value: unknown, name: string) => T };

export type Props = Record<string, Param<unknown>>;

export type Args<P extends Props> = { [K in keyof P]: P[K] extends Param<infer T> ? T : never };

/** A required parameter, with .optional() and .default() variants. Agents often send null for "not given". */
function param<T>(schema: Schema, check: (value: unknown, name: string) => T) {
  return {
    schema,
    required: true,
    parse: check,
    optional(): Param<T | undefined> {
      return { schema, required: false, parse: (v, name) => (v === undefined || v === null ? undefined : check(v, name)) };
    },
    default(fallback: T): Param<T> {
      return { schema: { ...schema, default: fallback }, required: false, parse: (v, name) => (v === undefined || v === null ? fallback : check(v, name)) };
    },
  };
}

export function string(description: string, opts: { max?: number; min?: number } = {}) {
  const { max = 1000, min = 0 } = opts;
  return param({ type: "string", description, ...(min ? { minLength: min } : {}), maxLength: max }, (v, name) => {
    if (typeof v !== "string") throw new ToolError(`${name} must be a string.`);
    if (v.length > max) throw new ToolError(`${name} can be at most ${max} characters.`);
    if (v.trim().length < min) throw new ToolError(`${name} can't be empty.`);
    return v;
  });
}

export function oneOf<const V extends string>(description: string, values: readonly V[]) {
  return param({ type: "string", description, enum: values }, (v, name) => {
    if (!values.includes(v as V)) throw new ToolError(`${name} must be one of ${values.join(", ")}.`);
    return v as V;
  });
}

export function integer(description: string, opts: { min?: number; max?: number } = {}) {
  const { min = 0, max = Number.MAX_SAFE_INTEGER } = opts;
  return param({ type: "integer", description, minimum: min, ...(opts.max !== undefined ? { maximum: max } : {}) }, (v, name) => {
    // Some agents send numbers as strings.
    const n = typeof v === "string" && /^-?\d+$/.test(v) ? Number(v) : v;
    if (typeof n !== "number" || !Number.isInteger(n)) throw new ToolError(`${name} must be a whole number.`);
    if (n < min || n > max) throw new ToolError(`${name} must be between ${min} and ${max}.`);
    return n;
  });
}

export function boolean(description: string) {
  return param({ type: "boolean", description }, (v, name) => {
    if (typeof v !== "boolean") throw new ToolError(`${name} must be true or false.`);
    return v;
  });
}

export function array<T>(description: string, item: Param<T>, opts: { max?: number } = {}) {
  const { max = 1000 } = opts;
  return param({ type: "array", description, items: item.schema, maxItems: max }, (v, name) => {
    if (!Array.isArray(v)) throw new ToolError(`${name} must be an array.`);
    if (v.length > max) throw new ToolError(`${name} can have at most ${max} items.`);
    return v.map((x, i) => item.parse(x, `${name}[${i}]`));
  });
}

export function object<P extends Props>(description: string, props: P) {
  return param(objectSchema(props, description), (v, name) => parseArgs(props, v, `${name}.`));
}

export function objectSchema(props: Props, description?: string): Schema {
  return {
    type: "object",
    ...(description ? { description } : {}),
    properties: Object.fromEntries(Object.entries(props).map(([k, p]) => [k, p.schema])),
    required: Object.entries(props)
      .filter(([, p]) => p.required)
      .map(([k]) => k),
    additionalProperties: false,
  };
}

/** Checks arguments against their parameters. Unknown names are refused, so typos don't go unnoticed. */
export function parseArgs<P extends Props>(props: P, value: unknown, prefix = ""): Args<P> {
  const input = value ?? {};
  if (typeof input !== "object" || Array.isArray(input)) throw new ToolError(`${prefix || "arguments "}must be an object.`);
  const unknown = Object.keys(input).filter((k) => !(k in props));
  if (unknown.length) {
    throw new ToolError(`Unknown argument ${unknown.map((k) => prefix + k).join(", ")}. Expected: ${Object.keys(props).join(", ")}.`);
  }
  const out: Record<string, unknown> = {};
  for (const [key, p] of Object.entries(props)) {
    const v = (input as Record<string, unknown>)[key];
    if (p.required && (v === undefined || v === null)) throw new ToolError(`${prefix}${key} is required.`);
    out[key] = p.parse(v, prefix + key);
  }
  return out as Args<P>;
}