Build an AI Agent with Tool Use

// types.ts
import { z, type ZodSchema } from "zod";

export interface ToolDefinition<TInput = unknown, TOutput = unknown> {
  name: string;
  description: string;
  inputSchema: ZodSchema<TInput>;
  execute: (input: TInput) => Promise<ToolResult<TOutput>>;
}

export interface ToolResult<T = unknown> {
  success: boolean;
  data?: T;
  error?: string;
}

export interface Message {
  role: "system" | "user" | "assistant" | "tool";
  content: string;
  toolCallId?: string;
  toolName?: string;
}

export interface ToolCall {
  id: string;
  name: string;
  arguments: Record<string, unknown>;
}

export interface AgentConfig {
  model: string;
  maxTurns: number;
  maxRetries: number;
  systemPrompt: string;
  tools: ToolDefinition[];
}

export interface PlanStep {
  thought: string;
  action: string;
  toolName?: string;
  toolInput?: Record<string, unknown>;
}

// registry.ts
import type { ToolDefinition, ToolResult } from "./types";

export class ToolRegistry {
  private tools = new Map<string, ToolDefinition>();

  register<TInput, TOutput>(tool: ToolDefinition<TInput, TOutput>): void {
    if (this.tools.has(tool.name)) {
      throw new Error(`Tool "${tool.name}" is already registered`);
    }
    this.tools.set(tool.name, tool as ToolDefinition);
  }

  get(name: string): ToolDefinition | undefined {
    return this.tools.get(name);
  }

  async execute(name: string, rawInput: unknown): Promise<ToolResult> {
    const tool = this.tools.get(name);
    if (!tool) {
      return { success: false, error: `Unknown tool: ${name}` };
    }

    const parsed = tool.inputSchema.safeParse(rawInput);
    if (!parsed.success) {
      return {
        success: false,
        error: `Invalid input for ${name}: ${parsed.error.message}`,
      };
    }

    try {
      return await tool.execute(parsed.data);
    } catch (error) {
      const message = error instanceof Error ? error.message : "Unknown error";
      return { success: false, error: `Tool ${name} failed: ${message}` };
    }
  }

  listTools(): Array<{ name: string; description: string }> {
    return Array.from(this.tools.values()).map((t) => ({
      name: t.name,
      description: t.description,
    }));
  }

  toFunctionDefinitions(): Array<Record<string, unknown>> {
    return Array.from(this.tools.values()).map((tool) => ({
      type: "function",
      function: {
        name: tool.name,
        description: tool.description,
        parameters: zodToJsonSchema(tool.inputSchema),
      },
    }));
  }
}

// tools/search.ts
import { z } from "zod";
import type { ToolDefinition, ToolResult } from "../types";

const inputSchema = z.object({
  query: z.string().min(1).max(500),
  maxResults: z.number().int().min(1).max(10).default(5),
});

type SearchInput = z.infer<typeof inputSchema>;

interface SearchResult {
  title: string;
  url: string;
  snippet: string;
}

export const searchTool: ToolDefinition<SearchInput, SearchResult[]> = {
  name: "web_search",
  description:
    "Search the web for current information. Use for facts, news, documentation, or any question about the real world.",
  inputSchema,
  async execute(input: SearchInput): Promise<ToolResult<SearchResult[]>> {
    const response = await fetch(
      `https://api.search.brave.com/res/v1/web/search?q=${encodeURIComponent(input.query)}&count=${input.maxResults}`,
      { headers: { "X-Subscription-Token": process.env.BRAVE_SEARCH_KEY! } }
    );

    if (!response.ok) {
      return { success: false, error: `Search API returned ${response.status}` };
    }

    const data = await response.json();
    const results =
      data.web?.results?.map((r: Record<string, string>) => ({
        title: r.title,
        url: r.url,
        snippet: r.description,
      })) ?? [];

    return { success: true, data: results };
  },
};

// tools/execute-code.ts
import { z } from "zod";
import { spawn } from "child_process";
import type { ToolDefinition, ToolResult } from "../types";

const inputSchema = z.object({
  code: z.string().min(1),
  language: z.enum(["javascript", "python"]),
  timeout: z.number().int().min(1000).max(30000).default(10000),
});

type CodeInput = z.infer<typeof inputSchema>;

export const executeCodeTool: ToolDefinition<CodeInput, string> = {
  name: "execute_code",
  description:
    "Execute JavaScript or Python code and return stdout. Use for calculations, data processing, or testing logic.",
  inputSchema,
  async execute(input: CodeInput): Promise<ToolResult<string>> {
    const cmd = input.language === "javascript" ? "node" : "python3";

    return new Promise((resolve) => {
      const proc = spawn(cmd, ["-e", input.code], {
        timeout: input.timeout,
        stdio: ["pipe", "pipe", "pipe"],
      });

      let stdout = "";
      let stderr = "";

      proc.stdout.on("data", (data: Buffer) => {
        stdout += data.toString();
      });
      proc.stderr.on("data", (data: Buffer) => {
        stderr += data.toString();
      });

      proc.on("close", (code) => {
        if (code === 0) {
          resolve({ success: true, data: stdout.trim() });
        } else {
          resolve({ success: false, error: stderr.trim() || `Process exited with code ${code}` });
        }
      });

      proc.on("error", (error) => {
        resolve({ success: false, error: error.message });
      });
    });
  },
};

// agent.ts
import type { AgentConfig, Message, ToolCall, ToolResult } from "./types";
import type { ToolRegistry } from "./registry";

export class Agent {
  private messages: Message[] = [];
  private config: AgentConfig;
  private registry: ToolRegistry;
  private turnCount = 0;

  constructor(config: AgentConfig, registry: ToolRegistry) {
    this.config = config;
    this.registry = registry;
    this.messages.push({ role: "system", content: config.systemPrompt });
  }

  async run(userMessage: string): Promise<string> {
    this.messages.push({ role: "user", content: userMessage });
    this.turnCount = 0;

    while (this.turnCount < this.config.maxTurns) {
      this.turnCount++;

      const response = await this.callLLM();

      if (!response.toolCalls || response.toolCalls.length === 0) {
        this.messages.push({ role: "assistant", content: response.content });
        return response.content;
      }

      // Execute all tool calls in parallel
      this.messages.push({ role: "assistant", content: response.content });

      const results = await Promise.all(
        response.toolCalls.map(async (call) => {
          const result = await this.executeWithRetry(call);
          return { call, result };
        })
      );

      for (const { call, result } of results) {
        this.messages.push({
          role: "tool",
          content: JSON.stringify(result),
          toolCallId: call.id,
          toolName: call.name,
        });
      }
    }

    return "I've reached the maximum number of steps. Here's what I found so far based on the conversation above.";
  }

  private async executeWithRetry(call: ToolCall, attempt = 0): Promise<ToolResult> {
    const result = await this.registry.execute(call.name, call.arguments);

    if (!result.success && attempt < this.config.maxRetries) {
      console.warn(`Tool ${call.name} failed (attempt ${attempt + 1}), retrying...`);
      await sleep(Math.pow(2, attempt) * 1000);
      return this.executeWithRetry(call, attempt + 1);
    }

    return result;
  }

  private async callLLM(): Promise<{
    content: string;
    toolCalls?: ToolCall[];
  }> {
    const response = await fetch("https://api.anthropic.com/v1/messages", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "x-api-key": process.env.ANTHROPIC_API_KEY!,
        "anthropic-version": "2023-06-01",
      },
      body: JSON.stringify({
        model: this.config.model,
        max_tokens: 4096,
        system: this.config.systemPrompt,
        messages: this.formatMessages(),
        tools: this.registry.toFunctionDefinitions(),
      }),
    });

    const data = await response.json();
    return this.parseResponse(data);
  }
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

private buildSystemPrompt(): string {
  return `You are an autonomous agent with access to tools. Follow this process:

1. OBSERVE: Read the user's request carefully. Identify what information you need.
2. PLAN: Think step-by-step about how to solve this. List the tools you'll need.
3. ACT: Execute one tool at a time. Use the result to decide the next step.
4. REFLECT: After each tool result, assess: Did I get what I needed? Should I try a different approach?

Rules:
- Always verify information from multiple sources when possible.
- If a tool fails, try an alternative approach before giving up.
- Show your reasoning in your responses.
- Cite sources when presenting factual information.
- If you cannot find an answer, say so explicitly rather than guessing.`;
}

// recovery.ts
interface RecoveryStrategy {
  canHandle: (toolName: string, error: string) => boolean;
  recover: (
    toolName: string,
    input: unknown,
    error: string
  ) => {
    newTool: string;
    newInput: unknown;
  } | null;
}

const strategies: RecoveryStrategy[] = [
  {
    // If web search fails, try with a simplified query
    canHandle: (name, error) => name === "web_search" && error.includes("429"),
    recover: (_, input) => ({
      newTool: "web_search",
      newInput: {
        ...(input as Record<string, unknown>),
        maxResults: 3,
      },
    }),
  },
  {
    // If code execution times out, try with a longer timeout
    canHandle: (name, error) => name === "execute_code" && error.includes("timeout"),
    recover: (_, input) => ({
      newTool: "execute_code",
      newInput: {
        ...(input as Record<string, unknown>),
        timeout: 30000,
      },
    }),
  },
];

export function findRecovery(
  toolName: string,
  input: unknown,
  error: string
): { newTool: string; newInput: unknown } | null {
  for (const strategy of strategies) {
    if (strategy.canHandle(toolName, error)) {
      return strategy.recover(toolName, input, error);
    }
  }
  return null;
}

// memory.ts
import type { Message } from "./types";

const MAX_MESSAGES = 50;
const SUMMARY_THRESHOLD = 40;

export class ConversationMemory {
  private messages: Message[] = [];
  private summaries: string[] = [];

  add(message: Message): void {
    this.messages.push(message);

    if (this.messages.length > SUMMARY_THRESHOLD) {
      this.compact();
    }
  }

  private async compact(): Promise<void> {
    // Take the oldest half of messages and summarize them
    const cutoff = Math.floor(this.messages.length / 2);
    const toSummarize = this.messages.slice(0, cutoff);
    const summary = await this.summarize(toSummarize);

    this.summaries.push(summary);
    this.messages = this.messages.slice(cutoff);
  }

  getContext(): Message[] {
    const summaryMessage: Message | null =
      this.summaries.length > 0
        ? {
            role: "system",
            content: `Previous conversation summary:\n${this.summaries.join("\n\n")}`,
          }
        : null;

    return summaryMessage ? [summaryMessage, ...this.messages] : [...this.messages];
  }

  private async summarize(messages: Message[]): Promise<string> {
    const text = messages.map((m) => `${m.role}: ${m.content.slice(0, 200)}`).join("\n");

    // Call LLM for summarization
    const response = await fetch("https://api.anthropic.com/v1/messages", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "x-api-key": process.env.ANTHROPIC_API_KEY!,
        "anthropic-version": "2023-06-01",
      },
      body: JSON.stringify({
        model: "claude-sonnet-4-20250514",
        max_tokens: 500,
        messages: [
          {
            role: "user",
            content: `Summarize this conversation in 2-3 sentences, preserving key facts and decisions:\n\n${text}`,
          },
        ],
      }),
    });

    const data = await response.json();
    return data.content[0].text;
  }
}

// stream.ts
export async function* streamAgent(agent: Agent, message: string): AsyncGenerator<string> {
  yield `\n🔍 Processing: "${message}"\n\n`;

  const steps = agent.runStreaming(message);

  for await (const step of steps) {
    switch (step.type) {
      case "thinking":
        yield `💭 ${step.content}\n`;
        break;
      case "tool_call":
        yield `🔧 Using ${step.toolName}(${JSON.stringify(step.input).slice(0, 100)})\n`;
        break;
      case "tool_result":
        yield `  ✓ ${step.content.slice(0, 200)}\n`;
        break;
      case "response":
        yield `\n${step.content}\n`;
        break;
      case "error":
        yield `  ✗ Error: ${step.content}\n`;
        break;
    }
  }
}

// tools/read-file.ts
import { z } from "zod";
import { readFile, stat } from "fs/promises";
import type { ToolDefinition, ToolResult } from "../types";

const inputSchema = z.object({
  path: z.string().min(1),
  maxLines: z.number().int().min(1).max(500).default(100),
});

export const readFileTool: ToolDefinition<z.infer<typeof inputSchema>, string> = {
  name: "read_file",
  description:
    "Read a file from disk. Returns the file contents. Use for analyzing code, configs, or data files.",
  inputSchema,
  async execute(input): Promise<ToolResult<string>> {
    try {
      const fileStat = await stat(input.path);

      if (fileStat.size > 1_000_000) {
        return { success: false, error: "File too large (>1MB). Use a more specific tool." };
      }

      const content = await readFile(input.path, "utf-8");
      const lines = content.split("\n").slice(0, input.maxLines).join("\n");

      return { success: true, data: lines };
    } catch (error) {
      const message = error instanceof Error ? error.message : "Unknown error";
      return { success: false, error: message };
    }
  },
};

// tools/calculator.ts
import { z } from "zod";
import type { ToolDefinition, ToolResult } from "../types";

const inputSchema = z.object({
  expression: z.string().min(1).max(500),
});

export const calculatorTool: ToolDefinition<z.infer<typeof inputSchema>, number> = {
  name: "calculator",
  description:
    "Evaluate a mathematical expression. Supports basic arithmetic, exponents, and common math functions.",
  inputSchema,
  async execute(input): Promise<ToolResult<number>> {
    // Whitelist allowed characters to prevent code injection
    const sanitized = input.expression.replace(/[^0-9+\-*/().%\s^]/g, "");
    if (sanitized !== input.expression.replace(/\s/g, "").replace(/[a-z]+/gi, "")) {
      return { success: false, error: "Expression contains disallowed characters" };
    }

    try {
      const fn = new Function(`"use strict"; return (${sanitized});`);
      const result = fn() as number;

      if (typeof result !== "number" || !isFinite(result)) {
        return { success: false, error: "Expression did not evaluate to a finite number" };
      }

      return { success: true, data: result };
    } catch (error) {
      return { success: false, error: `Invalid expression: ${(error as Error).message}` };
    }
  },
};

// cli.ts
import * as readline from "readline";
import { Agent } from "./agent";
import { ToolRegistry } from "./registry";
import { searchTool } from "./tools/search";
import { executeCodeTool } from "./tools/execute-code";
import { readFileTool } from "./tools/read-file";
import { calculatorTool } from "./tools/calculator";

async function main(): Promise<void> {
  const registry = new ToolRegistry();
  registry.register(searchTool);
  registry.register(executeCodeTool);
  registry.register(readFileTool);
  registry.register(calculatorTool);

  const agent = new Agent(
    {
      model: "claude-sonnet-4-20250514",
      maxTurns: 10,
      maxRetries: 2,
      systemPrompt: "You are a helpful agent with access to tools.",
      tools: [],
    },
    registry
  );

  const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout,
  });

  console.log("Agent ready. Type your message or 'quit' to exit.\n");

  const prompt = (): void => {
    rl.question("> ", async (input) => {
      if (input.toLowerCase() === "quit") {
        rl.close();
        return;
      }

      try {
        const response = await agent.run(input);
        console.log(`\n${response}\n`);
      } catch (error) {
        console.error("Agent error:", error);
      }

      prompt();
    });
  };

  prompt();
}

main();

// agent.test.ts
import { describe, it, expect, vi } from "vitest";
import { Agent } from "./agent";
import { ToolRegistry } from "./registry";
import { calculatorTool } from "./tools/calculator";

describe("Agent", () => {
  it("uses calculator tool for math questions", async () => {
    const registry = new ToolRegistry();
    registry.register(calculatorTool);

    const agent = new Agent(
      {
        model: "claude-sonnet-4-20250514",
        maxTurns: 5,
        maxRetries: 1,
        systemPrompt: "You are a helpful agent.",
        tools: [],
      },
      registry
    );

    const response = await agent.run("What is 42 * 17?");
    expect(response).toContain("714");
  });

  it("handles tool failure gracefully", async () => {
    const registry = new ToolRegistry();
    // Register a tool that always fails
    registry.register({
      name: "failing_tool",
      description: "A tool that always fails",
      inputSchema: z.object({}),
      execute: async () => ({ success: false, error: "Intentional failure" }),
    });

    const agent = new Agent(
      {
        model: "claude-sonnet-4-20250514",
        maxTurns: 3,
        maxRetries: 0,
        systemPrompt: "You are a helpful agent.",
        tools: [],
      },
      registry
    );

    const response = await agent.run("Use the failing tool");
    expect(response).toBeTruthy();
    // Agent should gracefully handle the failure and respond
  });
});