aps-agent/poc/pi-aps/e2e/mock-llm.mjs

/**
 * mock-llm.mjs - scripted OpenAI-compatible LLM endpoint for real Pi E2E.
 *
 * This server only scripts the model layer. It does NOT mock APS data. The APS
 * request is made by the real Pi extension against the real APS HTTP server.
 *
 * Behaviour:
 * 1. First chat completion request (no tool result yet): ask the model to call
 *    the chosen APS tool (aps_order_pool if present, otherwise aps_health).
 * 2. Second request (tool result present): return a Chinese final message.
 */
import http from "node:http";
import { appendFileSync } from "node:fs";
import path from "node:path";

const RESERVED = new Set(["aps_health", "aps_summary", "aps_invoke"]);

function args() {
  const out = { port: 0, portFile: "" };
  const argv = process.argv.slice(2);
  for (let i = 0; i < argv.length; i += 1) {
    if (argv[i] === "--port") out.port = Number(argv[i + 1]);
    if (argv[i] === "--port-file") out.portFile = String(argv[i + 1] || "");
  }
  return out;
}

function chunk(delta, finishReason = null) {
  const payload = {
    id: "chatcmpl-aps-e2e-mock",
    object: "chat.completion.chunk",
    created: Math.floor(Date.now() / 1000),
    model: "mock-local",
    choices: [{ index: 0, delta, finish_reason: finishReason }],
  };
  return `data: ${JSON.stringify(payload)}\n\n`;
}

function toolCallSse(toolName, callIndex) {
  const id = `call_aps_e2e_${Date.now()}_${callIndex}`;
  const parts = [
    chunk({ role: "assistant", content: null }),
    chunk({
      tool_calls: [{
        index: callIndex,
        id,
        type: "function",
        function: { name: toolName, arguments: "" },
      }],
    }),
    chunk({
      tool_calls: [{
        index: callIndex,
        function: { arguments: "{}" },
      }],
    }),
    chunk({}, "tool_calls"),
    chunk({}, null),
    "data: [DONE]\n\n",
  ];
  return parts.join("");
}

function finalTextSse(text) {
  const parts = [chunk({ role: "assistant", content: "" })];
  for (let i = 0; i < text.length; i += 40) {
    parts.push(chunk({ content: text.slice(i, i + 40) }));
  }
  parts.push(chunk({}, "stop"));
  parts.push("data: [DONE]\n\n");
  return parts.join("");
}

function pickTool(body) {
  const tools = (body.tools || [])
    .map((t) => t?.function?.name)
    .filter((name) => typeof name === "string" && name.startsWith("aps_"));
  const forced = process.env.APS_E2E_TOOL;
  if (forced && tools.includes(forced)) return forced;
  const dynamic = tools.find((name) => name !== "aps_health" && name !== "aps_summary" && name !== "aps_invoke");
  return dynamic || tools.find((name) => name === "aps_health") || tools[0];
}

function logRequest(summary) {
  const logFile = process.env.APS_MOCK_REQUEST_LOG;
  if (!logFile) return;
  try {
    appendFileSync(logFile, `${JSON.stringify({ ...summary, at: new Date().toISOString() })}\n`, "utf8");
  } catch {
    // Logging must never break the E2E run.
  }
}

function readBody(req) {
  return new Promise((resolve, reject) => {
    const chunks = [];
    req.on("data", (c) => chunks.push(c));
    req.on("end", () => {
      try {
        resolve(JSON.parse(Buffer.concat(chunks).toString("utf8") || "{}"));
      } catch (err) {
        reject(err);
      }
    });
    req.on("error", reject);
  });
}

function sendSse(res, payload, status = 200) {
  res.writeHead(status, {
    "Content-Type": "text/event-stream",
    "Cache-Control": "no-cache",
    Connection: "close",
  });
  res.end(payload);
}

export async function startMockLlmServer(options = {}) {
  const port = options.port ?? 0;
  const portFile = options.portFile ?? "";
  const { writeFileSync } = await import("node:fs");
  const server = http.createServer(async (req, res) => {
    if (req.method === "POST" && req.url.replace(/\/+$/, "").endsWith("/chat/completions")) {
      logRequest({
        route: "chat/completions/request_seen",
        url: req.url,
        contentLength: Number(req.headers["content-length"] || 0),
      });
      let body;
      try {
        body = await readBody(req);
      } catch {
        sendSse(res, finalTextSse("剧本 LLM 收到无法解析的请求，E2E 失败。"), 400);
        return;
      }
      const roles = (body.messages || []).map((m) => m.role);
      const hasToolResult = roles.includes("tool");
      logRequest({ route: "chat/completions", roles, hasToolResult });
      if (hasToolResult) {
        const toolMessage = (body.messages || []).find((m) => m.role === "tool");
        const toolName = toolMessage?.toolName || toolMessage?.tool_name || toolMessage?.name || "APS 工具";
        sendSse(
          res,
          finalTextSse(
            `${toolName} 已由真实 Pi 扩展调用并返回结果。我核对了工具事件，` +
            "未发现 HTTP 失败；以下是面向用户的中文总结：APS 服务请求已完成，结果已按工具返回内容如实呈现。",
          ),
        );
        return;
      }
      const toolName = pickTool(body);
      if (!toolName) {
        sendSse(
          res,
          finalTextSse("剧本 LLM 未发现 pi-aps 注册的 APS 工具，扩展加载或目录注册失败。"),
        );
        return;
      }
      sendSse(res, toolCallSse(toolName, 0));
      return;
    }
    res.writeHead(404, { "Content-Type": "application/json" });
    res.end(JSON.stringify({ error: "not found" }));
  });

  await new Promise((resolve, reject) => {
    server.once("error", reject);
    server.listen(port, "127.0.0.1", resolve);
  });
  const actualPort = server.address().port;
  if (portFile) {
    writeFileSync(portFile, `MOCK_LLM_PORT=${actualPort}\n`, "utf8");
  }
  process.stdout.write(`[mock-llm] listening on http://127.0.0.1:${actualPort}/v1\n`);
  return { server, port: actualPort, close: () => new Promise((resolve) => server.close(resolve)) };
}

const isMain = typeof process.argv[1] === "string" && path.basename(process.argv[1]) === "mock-llm.mjs";
if (isMain) {
  const { server } = await startMockLlmServer(args());
  const shutdown = () => server.close(() => process.exit(0));
  process.on("SIGINT", shutdown);
  process.on("SIGTERM", shutdown);
}