Hooks

import { OpenAIProvider, createLocalAgent, tool } from 'astorlm'
import { z } from 'zod'

const findTrains = tool({
  name: 'find_trains',
  description: 'List the trains to a destination on a date, with the fare for each class.',
  schema: z.object({ to: z.string(), date: z.string() }),
  execute: async ({ to, date }) => searchTimetable(to, date), // your code
})

const bookTicket = tool({
  name: 'book_ticket',
  description: 'Book one seat on a train for the employee who is asking.',
  schema: z.object({ train: z.number().int(), seat_class: z.enum(['first', 'tourist']) }),
  execute: async ({ train, seat_class }) => reserveSeat(train, seat_class), // your code
})

const agent = await createLocalAgent({
  // Any OpenAI-compatible endpoint: OpenAI, Ollama, LM Studio, vLLM, a proxy…
  provider: new OpenAIProvider({
    baseURL: 'http://localhost:11434/v1', // e.g. Ollama's default address
    model: 'your-model', // e.g. 'llama3.1', 'gpt-4o-mini'
    apiKey: 'YOUR_API_KEY', // local servers usually ignore it
  }),
  tools: [findTrains, bookTicket],
  maxTurns: 10,
  hooks: {
    // The first booth: runs before every tool call, and decides whether it runs at all.
    beforeToolExecution: async ({ toolName, input }) => {
      const { seat_class } = input as { seat_class?: string }
      if (toolName === 'book_ticket' && seat_class === 'first') {
        // The tool never runs. The model reads this text as an error result instead.
        return { authorize: false, mockResult: 'Blocked by policy: first class needs a manager’s approval. Book tourist instead.' }
      }
      return { authorize: true }
    },
    // The second booth: runs after every tool call. What you return is what the model reads.
    afterToolExecution: async ({ output }) => output.replace(/DNI [\d.]+/g, 'DNI ***'),
  },
})

// Events only watch. By the time this fires, the hook has already stamped over the DNI.
agent.on('tool-end', ({ name, output, isError }) => console.log(name, isError ? 'refused:' : 'ok:', output))

const last = await agent.run('Book me the most comfortable seat to Mar del Plata on Friday.')
console.log(last.content)

// The agent loop with hooks, from scratch. Plain fetch, no SDK.

// Any OpenAI-compatible endpoint: OpenAI, Ollama, LM Studio, vLLM, a proxy…
const LLM = {
  baseURL: 'http://localhost:11434/v1', // e.g. Ollama's default address
  model: 'your-model', // e.g. 'llama3.1', 'gpt-4o-mini'
  apiKey: 'YOUR_API_KEY', // local servers usually ignore it
}

type ToolCall = { id: string; function: { name: string; arguments: string } }
type Message =
  | { role: 'user'; content: string }
  | { role: 'assistant'; content: string | null; tool_calls?: ToolCall[] }
  | { role: 'tool'; tool_call_id: string; content: string }

type Args = Record<string, string | number>
type ToolFn = (args: Args) => Promise<string>
const tools: Record<string, ToolFn> = { find_trains: findTrains, book_ticket: bookTicket }
const toolSchemas = [/* one JSON Schema per tool */]

// The five points where the loop lets your code in. Every one is optional.
type Hooks = {
  beforeTurn?: (turn: number, messages: Message[]) => Promise<void>
  // Return the messages to send: trim them, add context, or pass them through.
  beforeProviderCall?: (messages: Message[]) => Promise<Message[]>
  // Say no, and the tool never runs: `result` goes back to the model instead.
  beforeToolExecution?: (name: string, args: Args) => Promise<{ authorize: boolean; result?: string }>
  // Whatever you return is what the model reads.
  afterToolExecution?: (name: string, output: string) => Promise<string>
  afterTurn?: (turn: number, reply: Message) => Promise<void>
}

export async function runAgent(prompt: string, hooks: Hooks = {}, maxTurns = 10): Promise<string> {
  const messages: Message[] = [{ role: 'user', content: prompt }]

  for (let turn = 1; turn <= maxTurns; turn++) {
    await hooks.beforeTurn?.(turn, messages)
    const outgoing = (await hooks.beforeProviderCall?.(messages)) ?? messages

    const res = await fetch(`${LLM.baseURL}/chat/completions`, {
      method: 'POST',
      headers: { 'content-type': 'application/json', authorization: `Bearer ${LLM.apiKey}` },
      body: JSON.stringify({ model: LLM.model, messages: outgoing, tools: toolSchemas }),
    })
    const [choice] = (await res.json()).choices
    const reply: Message = choice.message
    messages.push(reply)

    if (choice.finish_reason !== 'tool_calls') {
      await hooks.afterTurn?.(turn, reply)
      return reply.content ?? ''
    }

    for (const call of reply.tool_calls ?? []) {
      const name = call.function.name
      const run = tools[name]
      let output = `Unknown tool: ${name}`
      try {
        const args: Args = JSON.parse(call.function.arguments)
        // Booth 1: before the tool runs.
        const gate = (await hooks.beforeToolExecution?.(name, args)) ?? { authorize: true }
        if (!gate.authorize) output = gate.result ?? 'Rejected by policy.'
        else if (run) output = await run(args)
      } catch (err) {
        output = `Error: ${err instanceof Error ? err.message : err}`
      }
      // Booth 2: before the result joins the history.
      output = (await hooks.afterToolExecution?.(name, output)) ?? output
      messages.push({ role: 'tool', tool_call_id: call.id, content: output })
    }
    await hooks.afterTurn?.(turn, reply)
  }
  throw new Error(`No answer after ${maxTurns} turns`)
}

// The two booths from the animation.
const answer = await runAgent('Book me the most comfortable seat to Mar del Plata on Friday.', {
  beforeToolExecution: async (name, args) =>
    name === 'book_ticket' && args.seat_class === 'first'
      ? { authorize: false, result: 'Blocked by policy: first class needs a manager’s approval. Book tourist instead.' }
      : { authorize: true },
  afterToolExecution: async (_name, output) => output.replace(/DNI [\d.]+/g, 'DNI ***'),
})

# The agent loop with hooks, from scratch. Standard library only, no SDK.
import json
import re
import urllib.request

# Any OpenAI-compatible endpoint: OpenAI, Ollama, LM Studio, vLLM, a proxy...
LLM = {
    "base_url": "http://localhost:11434/v1",  # e.g. Ollama's default address
    "model": "your-model",  # e.g. "llama3.1", "gpt-4o-mini"
    "api_key": "YOUR_API_KEY",  # local servers usually ignore it
}

TOOLS = {"find_trains": find_trains, "book_ticket": book_ticket}
TOOL_SCHEMAS = [...]  # one JSON Schema per tool

# The five points where the loop lets your code in. Every one is optional:
#   before_turn(turn, messages)
#   before_provider_call(messages) -> the messages to send
#   before_tool_execution(name, args) -> {"authorize": bool, "result": str}
#   after_tool_execution(name, output) -> the text the model will read
#   after_turn(turn, reply)


def chat(messages):
    request = urllib.request.Request(
        f"{LLM['base_url']}/chat/completions",
        data=json.dumps({"model": LLM["model"], "messages": messages, "tools": TOOL_SCHEMAS}).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {LLM['api_key']}"},
    )
    with urllib.request.urlopen(request) as response:
        return json.load(response)["choices"][0]


def run_agent(prompt, hooks=None, max_turns=10):
    hooks = hooks or {}

    def call_hook(point, *args):
        return hooks[point](*args) if point in hooks else None

    messages = [{"role": "user", "content": prompt}]

    for turn in range(1, max_turns + 1):
        call_hook("before_turn", turn, messages)
        outgoing = call_hook("before_provider_call", messages) or messages

        choice = chat(outgoing)
        reply = choice["message"]
        messages.append(reply)

        if choice["finish_reason"] != "tool_calls":
            call_hook("after_turn", turn, reply)
            return reply.get("content") or ""

        for call in reply.get("tool_calls", []):
            name = call["function"]["name"]
            run = TOOLS.get(name)
            try:
                args = json.loads(call["function"]["arguments"])
                # Booth 1: before the tool runs. Say no, and it never does.
                gate = call_hook("before_tool_execution", name, args) or {"authorize": True}
                if not gate["authorize"]:
                    output = gate.get("result", "Rejected by policy.")
                else:
                    output = run(**args) if run else f"Unknown tool: {name}"
            except Exception as err:
                output = f"Error: {err}"
            # Booth 2: before the result joins the history. What it returns is what the model reads.
            output = call_hook("after_tool_execution", name, output) or output
            messages.append({"role": "tool", "tool_call_id": call["id"], "content": output})

        call_hook("after_turn", turn, reply)

    raise RuntimeError(f"No answer after {max_turns} turns")


# The two booths from the animation.
def check_policy(name, args):
    if name == "book_ticket" and args.get("seat_class") == "first":
        return {"authorize": False, "result": "Blocked by policy: first class needs a manager's approval. Book tourist instead."}
    return {"authorize": True}


def hide_ids(name, output):
    return re.sub(r"DNI [\d.]+", "DNI ***", output)


answer = run_agent(
    "Book me the most comfortable seat to Mar del Plata on Friday.",
    hooks={"before_tool_execution": check_policy, "after_tool_execution": hide_ids},
)