#第 02 章:agent loop —— 让模型调用工具

const TOOL_DEFS = [
  {
    type: "function",
    function: {
      name: "bash",
      description:
        "Execute a bash command in the current working directory. Returns stdout and stderr. Optionally provide a timeout in seconds.",
      parameters: {
        type: "object",
        properties: {
          command: { type: "string", description: "Bash command to execute" },
          timeout: { type: "number", description: "Timeout in seconds (optional, no default timeout)" },
        },
        required: ["command"],
      },
    },
  },
];

import { exec } from "node:child_process";
import * as readline from "node:readline";
import { promisify } from "node:util";

const CONFIG = {
  baseUrl: (process.env.MINI_PI_BASE_URL ?? process.env.OPENAI_BASE_URL ?? "https://api.openai.com/v1").replace(
    /\/$/,
    "",
  ),
  apiKey: process.env.MINI_PI_API_KEY ?? process.env.OPENAI_API_KEY ?? "",
  model: process.env.MINI_PI_MODEL ?? "gpt-4o-mini",
};

// ── 消息类型:OpenAI wire 格式 ──
interface ToolCall {
  id: string;
  type: "function";
  function: { name: string; arguments: string };
}

interface Message {
  role: "system" | "user" | "assistant" | "tool";
  content: string | null;
  tool_calls?: ToolCall[];
  tool_call_id?: string;
}

// ── LLM 调用(非流式,带工具) ──
interface AssistantResult {
  content: string;
  toolCalls: ToolCall[];
  finishReason: "stop" | "tool_calls" | "length";
}

async function chatWithTools(messages: Message[]): Promise<AssistantResult> {
  const res = await fetch(`${CONFIG.baseUrl}/chat/completions`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${CONFIG.apiKey}`,
    },
    body: JSON.stringify({ model: CONFIG.model, messages, tools: TOOL_DEFS }),
  });
  if (!res.ok) throw new Error(`HTTP ${res.status}: ${await res.text()}`);
  const data = (await res.json()) as any;
  const msg = data.choices[0].message;
  return {
    content: msg.content ?? "",
    toolCalls: msg.tool_calls ?? [],
    finishReason: data.choices[0].finish_reason ?? "stop",
  };
}

// ── 工具定义与执行 ──
const TOOL_DEFS = [
  {
    type: "function",
    function: {
      name: "bash",
      description:
        "Execute a bash command in the current working directory. Returns stdout and stderr. Optionally provide a timeout in seconds.",
      parameters: {
        type: "object",
        properties: {
          command: { type: "string", description: "Bash command to execute" },
          timeout: { type: "number", description: "Timeout in seconds (optional, no default timeout)" },
        },
        required: ["command"],
      },
    },
  },
];

const execAsync = promisify(exec);

async function toolBash(args: { command: string; timeout?: number }): Promise<string> {
  try {
    const { stdout, stderr } = await execAsync(args.command, {
      cwd: process.cwd(),
      shell: "/bin/bash",
      timeout: args.timeout !== undefined ? args.timeout * 1000 : undefined,
      maxBuffer: 10 * 1024 * 1024,
    });
    return [stdout, stderr].filter(Boolean).join("\n") || "(no output)";
  } catch (err: any) {
    // 命令非零退出不算工具失败:把输出和退出码一起给模型,让它自己判断
    const out = [err.stdout, err.stderr].filter(Boolean).join("\n");
    if (out) return `${out}\n[command exited with code ${err.code}]`;
    throw err;
  }
}

/** 统一入口:任何失败都转成文本返回给模型,而不是让 loop 崩溃 */
async function executeTool(name: string, argsJson: string): Promise<{ content: string; isError: boolean }> {
  try {
    let args: any;
    try {
      args = JSON.parse(argsJson || "{}");
    } catch {
      throw new Error(`Invalid JSON arguments: ${argsJson.slice(0, 200)}`);
    }
    switch (name) {
      case "bash":
        return { content: await toolBash(args), isError: false };
      default:
        throw new Error(`Unknown tool: ${name}`);
    }
  } catch (err) {
    return { content: `Error: ${err instanceof Error ? err.message : String(err)}`, isError: true };
  }
}

// ── agent loop:一个 turn = 一次 LLM 调用 + 由它触发的所有工具执行 ──
async function runAgent(messages: Message[]): Promise<void> {
  while (true) {
    const result = await chatWithTools(messages);
    if (result.content) console.log(result.content);

    const assistantMsg: Message = { role: "assistant", content: result.content || null };
    if (result.toolCalls.length > 0) assistantMsg.tool_calls = result.toolCalls;
    messages.push(assistantMsg);

    if (result.toolCalls.length === 0) return; // 模型没有开新工单,loop 结束

    for (const call of result.toolCalls) {
      console.log(`\x1b[36m[${call.function.name}]\x1b[0m \x1b[2m${call.function.arguments}\x1b[0m`);
      const toolResult = await executeTool(call.function.name, call.function.arguments);
      console.log(`  \x1b[2m${toolResult.content.split("\n").slice(0, 5).join("\n  ")}\x1b[0m`);
      messages.push({ role: "tool", tool_call_id: call.id, content: toolResult.content });
    }
  }
}

// ── REPL ──
function main(): void {
  if (!CONFIG.apiKey) {
    console.error("Error: set MINI_PI_API_KEY (or OPENAI_API_KEY) first.");
    process.exit(1);
  }
  const messages: Message[] = [
    { role: "system", content: "You are a coding assistant with access to a bash tool. Be concise." },
  ];
  const rl = readline.createInterface({ input: process.stdin, output: process.stdout, prompt: "\n> " });

  let running = false;
  rl.on("line", (line) => {
    const text = line.trim();
    if (!text || running) return rl.prompt();
    running = true;
    messages.push({ role: "user", content: text });
    runAgent(messages)
      .catch((err) => console.error(`[error] ${err.message}`))
      .finally(() => {
        running = false;
        rl.prompt();
      });
  });
  rl.prompt();
}

main();