getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-15 09:01:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Replace Zod with TypeBox for schema validation

Browse source 
- Switch from Zod to TypeBox for tool parameter schemas
- TypeBox schemas can be serialized/deserialized as JSON
- Use AJV for runtime validation instead of Zod's parse
- Add StringEnum helper for Google API compatibility (avoids anyOf/const patterns)
- Export Type and Static from main package for convenience
- Update all tests and documentation to reflect TypeBox usage
This commit is contained in:
Mario Zechner 2025-09-16 01:10:40 +02:00
parent f5ac1ef521
commit e8370436d7
16 changed files with 196 additions and 121 deletions
Download patch file Download diff file Expand all files Collapse all files

62
packages/ai/README.md
View file

@ -24,17 +24,17 @@ npm install @mariozechner/pi-ai
## Quick Start
```typescript
import { getModel, stream, complete, Context, Tool, z } from '@mariozechner/pi-ai';
import { Type, getModel, stream, complete, Context, Tool, StringEnum } from '@mariozechner/pi-ai';
// Fully typed with auto-complete support for both providers and models
const model = getModel('openai', 'gpt-4o-mini');
// Define tools with Zod schemas for type safety and validation
// Define tools with TypeBox schemas for type safety and validation
const tools: Tool[] = [{
name: 'get_time',
description: 'Get the current time',
parameters: z.object({
timezone: z.string().optional().describe('Optional timezone (e.g., America/New_York)')
parameters: Type.Object({
timezone: Type.Optional(Type.String({ description: 'Optional timezone (e.g., America/New_York)' }))
})
}];
@ -133,36 +133,35 @@ for (const block of response.content) {
## Tools
Tools enable LLMs to interact with external systems. This library uses Zod schemas for type-safe tool definitions with automatic validation.
Tools enable LLMs to interact with external systems. This library uses TypeBox schemas for type-safe tool definitions with automatic validation using AJV. TypeBox schemas can be serialized and deserialized as plain JSON, making them ideal for distributed systems.
### Defining Tools
```typescript
import { z, Tool } from '@mariozechner/pi-ai';
import { Type, Tool, StringEnum } from '@mariozechner/pi-ai';
// Define tool parameters with Zod
// Define tool parameters with TypeBox
const weatherTool: Tool = {
name: 'get_weather',
description: 'Get current weather for a location',
parameters: z.object({
location: z.string().describe('City name or coordinates'),
units: z.enum(['celsius', 'fahrenheit']).default('celsius')
parameters: Type.Object({
location: Type.String({ description: 'City name or coordinates' }),
units: StringEnum(['celsius', 'fahrenheit'], { default: 'celsius' })
})
};
// Complex validation with Zod refinements
// Note: For Google API compatibility, use StringEnum helper instead of Type.Enum
// Type.Enum generates anyOf/const patterns that Google doesn't support
const bookMeetingTool: Tool = {
name: 'book_meeting',
description: 'Schedule a meeting',
parameters: z.object({
title: z.string().min(1),
startTime: z.string().datetime(),
endTime: z.string().datetime(),
attendees: z.array(z.string().email()).min(1)
}).refine(
data => new Date(data.endTime) > new Date(data.startTime),
{ message: 'End time must be after start time' }
)
parameters: Type.Object({
title: Type.String({ minLength: 1 }),
startTime: Type.String({ format: 'date-time' }),
endTime: Type.String({ format: 'date-time' }),
attendees: Type.Array(Type.String({ format: 'email' }), { minItems: 1 })
})
};
```
@ -179,7 +178,7 @@ const response = await complete(model, context);
// Check for tool calls in the response
for (const block of response.content) {
if (block.type === 'toolCall') {
// Arguments are automatically validated against the Zod schema
// Arguments are automatically validated against the TypeBox schema using AJV
// If validation fails, an error event is emitted
const result = await executeWeatherApi(block.arguments);
@ -687,19 +686,20 @@ const messages = await stream.result();
context.messages.push(...messages);
```
### Defining Tools with Zod
### Defining Tools with TypeBox
Tools use Zod schemas for runtime validation and type inference:
Tools use TypeBox schemas for runtime validation and type inference:
```typescript
import { z } from 'zod';
import { AgentTool, AgentToolResult } from '@mariozechner/pi-ai';
import { Type, Static, AgentTool, AgentToolResult, StringEnum } from '@mariozechner/pi-ai';
const weatherSchema = z.object({
city: z.string().min(1, 'City is required'),
units: z.enum(['celsius', 'fahrenheit']).default('celsius')
const weatherSchema = Type.Object({
city: Type.String({ minLength: 1 }),
units: StringEnum(['celsius', 'fahrenheit'], { default: 'celsius' })
});
type WeatherParams = Static<typeof weatherSchema>;
const weatherTool: AgentTool<typeof weatherSchema, { temp: number }> = {
label: 'Get Weather',
name: 'get_weather',
@ -718,7 +718,7 @@ const weatherTool: AgentTool<typeof weatherSchema, { temp: number }> = {
### Validation and Error Handling
Tool arguments are automatically validated using the Zod schema. Invalid arguments result in detailed error messages:
Tool arguments are automatically validated using AJV with the TypeBox schema. Invalid arguments result in detailed error messages:
```typescript
// If the LLM calls with invalid arguments:
@ -727,8 +727,8 @@ Tool arguments are automatically validated using the Zod schema. Invalid argumen
// The tool execution will fail with:
/*
Validation failed for tool "get_weather":
- city: City is required
- units: Invalid enum value. Expected 'celsius' | 'fahrenheit', received 'kelvin'
- city: must NOT have fewer than 1 characters
- units: must be equal to one of the allowed values
Received arguments:
{

8
packages/ai/src/agent/tools/calculate.ts
View file

@ -1,4 +1,4 @@
import { z } from "zod";
import { type Static, Type } from "@sinclair/typebox";
import type { AgentTool } from "../../agent";
export interface CalculateResult {
@ -15,10 +15,12 @@ export function calculate(expression: string): CalculateResult {
}
}
const calculateSchema = z.object({
expression: z.string().describe("The mathematical expression to evaluate"),
const calculateSchema = Type.Object({
expression: Type.String({ description: "The mathematical expression to evaluate" }),
});
type CalculateParams = Static<typeof calculateSchema>;
export const calculateTool: AgentTool<typeof calculateSchema, undefined> = {
label: "Calculator",
name: "calculate",

10
packages/ai/src/agent/tools/get-current-time.ts
View file

@ -1,4 +1,4 @@
import { z } from "zod";
import { type Static, Type } from "@sinclair/typebox";
import type { AgentTool } from "../../agent";
import type { AgentToolResult } from "../types";
@ -26,10 +26,14 @@ export async function getCurrentTime(timezone?: string): Promise<GetCurrentTimeR
};
}
const getCurrentTimeSchema = z.object({
timezone: z.string().optional().describe("Optional timezone (e.g., 'America/New_York', 'Europe/London')"),
const getCurrentTimeSchema = Type.Object({
timezone: Type.Optional(
Type.String({ description: "Optional timezone (e.g., 'America/New_York', 'Europe/London')" }),
),
});
type GetCurrentTimeParams = Static<typeof getCurrentTimeSchema>;
export const getCurrentTimeTool: AgentTool<typeof getCurrentTimeSchema, { utcTimestamp: number }> = {
label: "Current Time",
name: "get_current_time",

6
packages/ai/src/agent/types.ts
View file

@ -1,4 +1,4 @@
import type { ZodSchema, z } from "zod";
import type { Static, TSchema } from "@sinclair/typebox";
import type {
AssistantMessage,
AssistantMessageEvent,
@ -17,12 +17,12 @@ export interface AgentToolResult<T> {
}
// AgentTool extends Tool but adds the execute function
export interface AgentTool<TParameters extends ZodSchema = ZodSchema, TDetails = any> extends Tool<TParameters> {
export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {
// A human-readable label for the tool to be displayed in UI
label: string;
execute: (
toolCallId: string,
params: z.infer<TParameters>,
params: Static<TParameters>,
signal?: AbortSignal,
) => Promise<AgentToolResult<TDetails>>;
}

3
packages/ai/src/index.ts
View file

@ -1,4 +1,4 @@
export { z } from "zod";
export { type Static, Type } from "@sinclair/typebox";
export * from "./agent/index.js";
export * from "./models.js";
export * from "./providers/anthropic.js";
@ -6,4 +6,5 @@ export * from "./providers/google.js";
export * from "./providers/openai-completions.js";
export * from "./providers/openai-responses.js";
export * from "./stream.js";
export * from "./typebox-helpers.js";
export * from "./types.js";

3
packages/ai/src/providers/anthropic.ts
View file

@ -4,7 +4,6 @@ import type {
MessageCreateParamsStreaming,
MessageParam,
} from "@anthropic-ai/sdk/resources/messages.js";
import { zodToJsonSchema } from "zod-to-json-schema";
import { AssistantMessageEventStream } from "../event-stream.js";
import { calculateCost } from "../models.js";
import type {
@ -442,7 +441,7 @@ function convertTools(tools: Tool[]): Anthropic.Messages.Tool[] {
if (!tools) return [];
return tools.map((tool) => {
const jsonSchema = zodToJsonSchema(tool.parameters, { $refStrategy: "none" }) as any;
const jsonSchema = tool.parameters as any; // TypeBox already generates JSON Schema
return {
name: tool.name,

3
packages/ai/src/providers/google.ts
View file

@ -7,7 +7,6 @@ import {
GoogleGenAI,
type Part,
} from "@google/genai";
import { zodToJsonSchema } from "zod-to-json-schema";
import { AssistantMessageEventStream } from "../event-stream.js";
import { calculateCost } from "../models.js";
import type {
@ -394,7 +393,7 @@ function convertTools(tools: Tool[]): any[] | undefined {
functionDeclarations: tools.map((tool) => ({
name: tool.name,
description: tool.description,
parameters: zodToJsonSchema(tool.parameters, { $refStrategy: "none" }),
parameters: tool.parameters as any, // TypeBox already generates JSON Schema
})),
},
];

3
packages/ai/src/providers/openai-completions.ts
View file

@ -7,7 +7,6 @@ import type {
ChatCompletionContentPartText,
ChatCompletionMessageParam,
} from "openai/resources/chat/completions.js";
import { zodToJsonSchema } from "zod-to-json-schema";
import { AssistantMessageEventStream } from "../event-stream.js";
import { calculateCost } from "../models.js";
import type {
@ -396,7 +395,7 @@ function convertTools(tools: Tool[]): OpenAI.Chat.Completions.ChatCompletionTool
function: {
name: tool.name,
description: tool.description,
parameters: zodToJsonSchema(tool.parameters, { $refStrategy: "none" }),
parameters: tool.parameters as any, // TypeBox already generates JSON Schema
},
}));
}

3
packages/ai/src/providers/openai-responses.ts
View file

@ -10,7 +10,6 @@ import type {
ResponseOutputMessage,
ResponseReasoningItem,
} from "openai/resources/responses/responses.js";
import { zodToJsonSchema } from "zod-to-json-schema";
import { AssistantMessageEventStream } from "../event-stream.js";
import { calculateCost } from "../models.js";
import type {
@ -462,7 +461,7 @@ function convertTools(tools: Tool[]): OpenAITool[] {
type: "function",
name: tool.name,
description: tool.description,
parameters: zodToJsonSchema(tool.parameters, { $refStrategy: "none" }),
parameters: tool.parameters as any, // TypeBox already generates JSON Schema
strict: null,
}));
}

24
packages/ai/src/typebox-helpers.ts Normal file
View file

@ -0,0 +1,24 @@
import { type TUnsafe, Type } from "@sinclair/typebox";
/**
* Creates a string enum schema compatible with Google's API and other providers
* that don't support anyOf/const patterns.
*
* @example
* const OperationSchema = StringEnum(["add", "subtract", "multiply", "divide"], {
* description: "The operation to perform"
* });
*
* type Operation = Static<typeof OperationSchema>; // "add" | "subtract" | "multiply" | "divide"
*/
export function StringEnum<T extends readonly string[]>(
values: T,
options?: { description?: string; default?: T[number] },
): TUnsafe<T[number]> {
return Type.Unsafe<T[number]>({
type: "string",
enum: values as any,
...(options?.description && { description: options.description }),
...(options?.default && { default: options.default }),
});
}

4
packages/ai/src/types.ts
View file

@ -119,9 +119,9 @@ export interface ToolResultMessage<TDetails = any> {
export type Message = UserMessage | AssistantMessage | ToolResultMessage;
import type { ZodSchema } from "zod";
import type { TSchema } from "@sinclair/typebox";
export interface Tool<TParameters extends ZodSchema = ZodSchema> {
export interface Tool<TParameters extends TSchema = TSchema> {
name: string;
description: string;
parameters: TParameters;

49
packages/ai/src/validation.ts
View file

@ -1,32 +1,37 @@
import { z } from "zod";
import Ajv from "ajv";
import addFormats from "ajv-formats";
import type { Tool, ToolCall } from "./types.js";
// Create a singleton AJV instance with formats
const ajv = new Ajv({ allErrors: true, strict: false });
addFormats(ajv);
/**
* Validates tool call arguments against the tool's Zod schema
* @param tool The tool definition with Zod schema
* Validates tool call arguments against the tool's TypeBox schema
* @param tool The tool definition with TypeBox schema
* @param toolCall The tool call from the LLM
* @returns The validated arguments
* @throws ZodError with formatted message if validation fails
* @throws Error with formatted message if validation fails
*/
export function validateToolArguments(tool: Tool, toolCall: ToolCall): any {
try {
// Validate arguments with Zod schema
return tool.parameters.parse(toolCall.arguments);
} catch (e) {
if (e instanceof z.ZodError) {
// Format validation errors nicely
const errors = e.issues
.map((err) => {
const path = err.path.length > 0 ? err.path.join(".") : "root";
return ` - ${path}: ${err.message}`;
})
.join("\n");
// Compile the schema
const validate = ajv.compile(tool.parameters);
const errorMessage = `Validation failed for tool "${toolCall.name}":\n${errors}\n\nReceived arguments:\n${JSON.stringify(toolCall.arguments, null, 2)}`;
// Throw a new error with the formatted message
throw new Error(errorMessage);
}
throw e;
// Validate the arguments
if (validate(toolCall.arguments)) {
return toolCall.arguments;
}
// Format validation errors nicely
const errors =
validate.errors
?.map((err) => {
const path = err.instancePath ? err.instancePath.substring(1) : err.params.missingProperty || "root";
return ` - ${path}: ${err.message}`;
})
.join("\n") || "Unknown validation error";
const errorMessage = `Validation failed for tool "${toolCall.name}":\n${errors}\n\nReceived arguments:\n${JSON.stringify(toolCall.arguments, null, 2)}`;
throw new Error(errorMessage);
}

17
packages/ai/test/enum-test.ts Normal file
View file

@ -0,0 +1,17 @@
import { Type } from "@sinclair/typebox";
import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
import { StringEnum } from "../src/typebox-helpers.js";
// Zod version
const zodSchema = z.object({
operation: z.enum(["add", "subtract", "multiply", "divide"]),
});
// TypeBox with our StringEnum helper
const typeboxHelper = Type.Object({
operation: StringEnum(["add", "subtract", "multiply", "divide"]),
});
console.log("Zod:", JSON.stringify(zodToJsonSchema(zodSchema), null, 2));
console.log("\nTypeBox.StringEnum:", JSON.stringify(typeboxHelper, null, 2));

23
packages/ai/test/generate.test.ts
View file

@ -1,27 +1,32 @@
import { Type } from "@sinclair/typebox";
import { type ChildProcess, execSync, spawn } from "child_process";
import { readFileSync } from "fs";
import { dirname, join } from "path";
import { fileURLToPath } from "url";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { z } from "zod";
import { getModel } from "../src/models.js";
import { complete, stream } from "../src/stream.js";
import { StringEnum } from "../src/typebox-helpers.js";
import type { Api, Context, ImageContent, Model, OptionsForApi, Tool, ToolResultMessage } from "../src/types.js";
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
// Calculator tool definition (same as examples)
const calculatorTool: Tool = {
// Note: Using StringEnum helper because Google's API doesn't support anyOf/const patterns
// that Type.Enum generates. Google requires { type: "string", enum: [...] } format.
const calculatorSchema = Type.Object({
a: Type.Number({ description: "First number" }),
b: Type.Number({ description: "Second number" }),
operation: StringEnum(["add", "subtract", "multiply", "divide"], {
description: "The operation to perform. One of 'add', 'subtract', 'multiply', 'divide'.",
}),
});
const calculatorTool: Tool<typeof calculatorSchema> = {
name: "calculator",
description: "Perform basic arithmetic operations",
parameters: z.object({
a: z.number().describe("First number"),
b: z.number().describe("Second number"),
operation: z
.enum(["add", "subtract", "multiply", "divide"])
.describe("The operation to perform. One of 'add', 'subtract', 'multiply', 'divide'."),
}),
parameters: calculatorSchema,
};
async function basicTextGeneration<TApi extends Api>(model: Model<TApi>, options?: OptionsForApi<TApi>) {

12
packages/ai/test/handoff.test.ts
View file

@ -1,16 +1,18 @@
import { Type } from "@sinclair/typebox";
import { describe, expect, it } from "vitest";
import { z } from "zod";
import { getModel } from "../src/models.js";
import { complete } from "../src/stream.js";
import type { Api, AssistantMessage, Context, Message, Model, Tool, ToolResultMessage } from "../src/types.js";
// Tool for testing
const weatherTool: Tool = {
const weatherSchema = Type.Object({
location: Type.String({ description: "City name" }),
});
const weatherTool: Tool<typeof weatherSchema> = {
name: "get_weather",
description: "Get the weather for a location",
parameters: z.object({
location: z.string().describe("City name"),
}),
parameters: weatherSchema,
};
// Pre-built contexts representing typical outputs from each provider

87
packages/ai/test/tool-validation.test.ts
View file

@ -1,16 +1,20 @@
import { type Static, Type } from "@sinclair/typebox";
import Ajv from "ajv";
import addFormats from "ajv-formats";
import { describe, expect, it } from "vitest";
import { z } from "zod";
import type { AgentTool } from "../src/agent/types.js";
describe("Tool Validation with Zod", () => {
// Define a test tool with Zod schema
const testSchema = z.object({
name: z.string().min(1, "Name is required"),
age: z.number().int().min(0).max(150),
email: z.string().email("Invalid email format"),
tags: z.array(z.string()).optional(),
describe("Tool Validation with TypeBox and AJV", () => {
// Define a test tool with TypeBox schema
const testSchema = Type.Object({
name: Type.String({ minLength: 1 }),
age: Type.Integer({ minimum: 0, maximum: 150 }),
email: Type.String({ format: "email" }),
tags: Type.Optional(Type.Array(Type.String())),
});
type TestParams = Static<typeof testSchema>;
const testTool: AgentTool<typeof testSchema, void> = {
label: "Test Tool",
name: "test_tool",
@ -24,6 +28,10 @@ describe("Tool Validation with Zod", () => {
},
};
// Create AJV instance for validation
const ajv = new Ajv({ allErrors: true });
addFormats(ajv);
it("should validate correct input", () => {
const validInput = {
name: "John Doe",
@ -32,9 +40,10 @@ describe("Tool Validation with Zod", () => {
tags: ["developer", "typescript"],
};
// This should not throw
const result = testTool.parameters.parse(validInput);
expect(result).toEqual(validInput);
// Validate with AJV
const validate = ajv.compile(testTool.parameters);
const isValid = validate(validInput);
expect(isValid).toBe(true);
});
it("should reject invalid email", () => {
@ -44,7 +53,10 @@ describe("Tool Validation with Zod", () => {
email: "not-an-email",
};
expect(() => testTool.parameters.parse(invalidInput)).toThrowError(z.ZodError);
const validate = ajv.compile(testTool.parameters);
const isValid = validate(invalidInput);
expect(isValid).toBe(false);
expect(validate.errors).toBeDefined();
});
it("should reject missing required fields", () => {
@ -53,7 +65,10 @@ describe("Tool Validation with Zod", () => {
email: "john@example.com",
};
expect(() => testTool.parameters.parse(invalidInput)).toThrowError(z.ZodError);
const validate = ajv.compile(testTool.parameters);
const isValid = validate(invalidInput);
expect(isValid).toBe(false);
expect(validate.errors).toBeDefined();
});
it("should reject invalid age", () => {
@ -63,7 +78,10 @@ describe("Tool Validation with Zod", () => {
email: "john@example.com",
};
expect(() => testTool.parameters.parse(invalidInput)).toThrowError(z.ZodError);
const validate = ajv.compile(testTool.parameters);
const isValid = validate(invalidInput);
expect(isValid).toBe(false);
expect(validate.errors).toBeDefined();
});
it("should format validation errors nicely", () => {
@ -73,25 +91,23 @@ describe("Tool Validation with Zod", () => {
email: "invalid",
};
try {
testTool.parameters.parse(invalidInput);
// Should not reach here
expect(true).toBe(false);
} catch (e) {
if (e instanceof z.ZodError) {
const errors = e.issues
.map((err) => {
const path = err.path.length > 0 ? err.path.join(".") : "root";
return ` - ${path}: ${err.message}`;
})
.join("\n");
const validate = ajv.compile(testTool.parameters);
const isValid = validate(invalidInput);
expect(isValid).toBe(false);
expect(validate.errors).toBeDefined();
expect(errors).toContain("name: Name is required");
expect(errors).toContain("age: Number must be less than or equal to 150");
expect(errors).toContain("email: Invalid email format");
} else {
throw e;
}
if (validate.errors) {
const errors = validate.errors
.map((err) => {
const path = err.instancePath ? err.instancePath.substring(1) : err.params.missingProperty || "root";
return ` - ${path}: ${err.message}`;
})
.join("\n");
// AJV error messages are different from Zod
expect(errors).toContain("name: must NOT have fewer than 1 characters");
expect(errors).toContain("age: must be <= 150");
expect(errors).toContain('email: must match format "email"');
}
});
@ -103,8 +119,11 @@ describe("Tool Validation with Zod", () => {
};
// Validate and execute
const validated = testTool.parameters.parse(validInput);
const result = await testTool.execute("test-id", validated);
const validate = ajv.compile(testTool.parameters);
const isValid = validate(validInput);
expect(isValid).toBe(true);
const result = await testTool.execute("test-id", validInput as TestParams);
expect(result.output).toBe("Processed: John Doe, 30, john@example.com");
expect(result.details).toBeUndefined();