How to Use TypeBox Schemas for Tool Definitions in pi-ai

Use TypeBox schemas to define tool parameters in pi-ai, which are automatically converted to JSON-Schema for LLM providers and validated at runtime using AJV.

The pi-ai package in the badlogic/pi-mono repository treats tool calls as first-class messages. By leveraging TypeBox (@sinclair/typebox), you get TypeScript type safety, automatic JSON-Schema generation, and runtime validation without writing boilerplate.

Understanding the Tool Type Interface

The foundation of tool definitions resides in packages/ai/src/types.ts. The generic Tool interface accepts any TypeBox schema as its type parameter:

// packages/ai/src/types.ts
export interface Tool<TParameters extends TSchema = TSchema> {
  name: string;
  description: string;
  parameters: TParameters;
}

• name identifies the tool for the LLM.
• description guides the model on when to invoke the tool.
• parameters accepts a TypeBox schema (TSchema) that defines the expected arguments.

Defining Tool Parameters with TypeBox

TypeBox provides a fluent API for constructing JSON-Schema compliant definitions. Import the Type namespace from @sinclair/typebox to build your schemas:

import { Type } from "@sinclair/typebox";
import type { Tool } from "@/ai/src/types.js";

const echoSchema = Type.Object({
  /** Message that will be echoed back */
  message: Type.String({ description: "Message to echo back" })
});

export const echoTool: Tool<typeof echoSchema> = {
  name: "echo",
  description: "Echoes the supplied message back to the user",
  parameters: echoSchema,
};

The Tool<typeof echoSchema> generic captures the exact schema type, ensuring full type safety when handling tool arguments later in your implementation.

Helper Utilities for Complex Schemas

For enum-like string parameters, pi-ai provides a StringEnum helper in packages/ai/src/utils/typebox-helpers.ts. This generates a schema compatible with providers that lack full anyOf or const support:

import { StringEnum } from "@/ai/src/utils/typebox-helpers.js";
import { Type } from "@sinclair/typebox";

const calcSchema = Type.Object({
  op: StringEnum(["add", "subtract", "multiply", "divide"], {
    description: "Arithmetic operation to perform",
    default: "add"
  }) as any,
  a: Type.Number({ description: "First operand" }),
  b: Type.Number({ description: "Second operand" }),
});

export const calcTool: Tool<typeof calcSchema> = {
  name: "calc",
  description: "Perform basic arithmetic operations",
  parameters: calcSchema,
};

Registering and Using Tools in Requests

Tools are registered by passing them to the tools array in your request configuration. The completeSimple function in packages/ai/src/stream.js handles the streaming completion:

import { completeSimple } from "@/ai/src/stream.js";
import { getModel } from "@/ai/src/models.js";
import { echoTool } from "./my-tools.js";

async function run() {
  const model = getModel("openai", "gpt-4o-mini");

  const response = await completeSimple(
    model,
    {
      systemPrompt: "You are a helpful assistant. Use tools when appropriate.",
      messages: [
        { 
          role: "user", 
          content: "Please echo 'pi-ai rocks'", 
          timestamp: Date.now() 
        }
      ],
      tools: [echoTool],  // Register the TypeBox-defined tool
    },
    { apiKey: process.env.OPENAI_API_KEY }
  );

  console.log(response);
}

run();

When the LLM decides to invoke the tool, the response contains a toolCall message with arguments that conform to your TypeBox schema.

Runtime Validation with AJV

After the model returns a tool call, pi-ai validates the arguments against the original TypeBox schema using AJV. The validation logic resides in packages/ai/src/utils/validation.ts:

// packages/ai/src/utils/validation.ts
export function validateToolArguments(tool: Tool, toolCall: ToolCall): any {
  if (!ajv) return toolCall.arguments;  // Disabled in CSP-restricted environments
  
  const validate = ajv.compile(tool.parameters);
  const args = structuredClone(toolCall.arguments);
  
  if (validate(args)) return args;
  
  // Throws a detailed validation error if schema compliance fails
  throw new Error(`Tool validation failed: ${JSON.stringify(validate.errors)}`);
}

If validation succeeds, the parsed arguments are passed to your tool implementation. If it fails, you receive a detailed error describing which constraints were violated.

Provider Integration

Each LLM provider in pi-ai forwards the Tool.parameters object directly as JSON-Schema. For example, in packages/ai/src/providers/openai-completions.ts:

// packages/ai/src/providers/openai-completions.ts
parameters: tool.parameters as any,  // TypeBox already generates JSON Schema

This seamless conversion means you write TypeBox schemas once, and they work across OpenAI, Anthropic, Google, and other supported providers without modification.

Summary

• Define tool schemas using TypeBox (@sinclair/typebox) for type-safe JSON-Schema generation.
• Use the generic Tool<TParameters> interface from packages/ai/src/types.ts to capture schema types.
• Leverage StringEnum from packages/ai/src/utils/typebox-helpers.ts for enum-like string parameters.
• Register tools in the tools array when calling completeSimple or other completion functions.
• Runtime validation automatically occurs via AJV in packages/ai/src/utils/validation.ts.
• Providers receive the TypeBox schema as standard JSON-Schema without additional conversion.

Frequently Asked Questions

How do I handle optional parameters in TypeBox schemas?

Use Type.Optional() to mark fields that the LLM may omit. For example: Type.Optional(Type.String()) creates a schema where the property is not required. When the model calls the tool without that argument, the validation still passes and your implementation receives undefined for that field.

Can I use complex nested objects in tool schemas?

Yes. TypeBox supports deeply nested structures using Type.Object(), Type.Array(), and composition utilities like Type.Intersect() or Type.Union(). The schema is automatically serialized to JSON-Schema for the provider, and AJV validates the nested structure when the tool call returns.

What happens if the LLM returns invalid arguments?

The validateToolArguments function in packages/ai/src/utils/validation.ts throws a descriptive error containing the AJV validation errors. This prevents your tool implementation from receiving malformed data. You should wrap tool execution in try-catch blocks to handle these validation failures gracefully.

Is AJV validation required for all tool calls?

No. AJV validation is automatically disabled in Content Security Policy (CSP) restricted environments where dynamic code evaluation is blocked. In such cases, validateToolArguments returns the raw arguments without validation, and you must implement your own safety checks if needed.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}