src/piAgent.js

import { Agent } from "@earendil-works/pi-agent-core";
import { createAssistantMessageEventStream, Type } from "@earendil-works/pi-ai";
import { env, pipeline } from "@huggingface/transformers";

const MODEL_ID = "Mike0021/MiniCPM5-1B-ONNX-Web";

const LOCAL_MODEL = {
  id: MODEL_ID,
  name: "MiniCPM5-1B ONNX Web",
  api: "transformers-js",
  provider: "huggingface-transformers-js",
  baseUrl: "https://huggingface.co",
  reasoning: false,
  input: ["text"],
  cost: {
    input: 0,
    output: 0,
    cacheRead: 0,
    cacheWrite: 0,
  },
  contextWindow: 4096,
  maxTokens: 512,
};

const EMPTY_USAGE = {
  input: 0,
  output: 0,
  cacheRead: 0,
  cacheWrite: 0,
  totalTokens: 0,
  cost: {
    input: 0,
    output: 0,
    cacheRead: 0,
    cacheWrite: 0,
    total: 0,
  },
};

function textFromContent(content) {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  return content
    .filter((part) => part.type === "text")
    .map((part) => part.text)
    .join("\n");
}

function now() {
  return Date.now();
}

function createMessage(content, stopReason = "stop") {
  return {
    role: "assistant",
    content,
    api: LOCAL_MODEL.api,
    provider: LOCAL_MODEL.provider,
    model: LOCAL_MODEL.id,
    usage: EMPTY_USAGE,
    stopReason,
    timestamp: now(),
  };
}

function stringifyToolResult(message) {
  const text = textFromContent(message.content);
  return `${message.toolName}(${message.toolCallId}) ${message.isError ? "failed" : "succeeded"}:\n${text}`;
}

function buildPrompt(context) {
  const tools = (context.tools || []).map((tool) => ({
    name: tool.name,
    description: tool.description,
    parameters: tool.parameters,
  }));
  const transcript = context.messages
    .slice(-8)
    .map((message) => {
      if (message.role === "toolResult") return `TOOL_RESULT:\n${stringifyToolResult(message)}`;
      return `${message.role.toUpperCase()}:\n${textFromContent(message.content)}`;
    })
    .join("\n\n");

  return `${context.systemPrompt || ""}

You are running as a pi Agent inside a browser-only app. The sandbox is a WebContainer: it has a virtual filesystem and can spawn browser-contained Node.js processes.

Use tools by returning strict JSON only. Do not use markdown.

To call tools:
{"toolCalls":[{"tool":"write_file","args":{"path":"hello.js","content":"console.log(2 + 2)\\n"}},{"tool":"run_command","args":{"command":"node","args":["hello.js"]}}]}

To answer the user after tool results:
{"final":"Short answer that explains what happened."}

Available tools:
${JSON.stringify(tools, null, 2)}

Conversation:
${transcript}

Return JSON now.`;
}

function extractJsonPayload(text) {
  const trimmed = String(text || "").trim();
  const fence = trimmed.match(/```(?:json)?\s*([\s\S]*?)```/i);
  const candidate = fence ? fence[1].trim() : trimmed;
  const firstBrace = candidate.indexOf("{");
  const firstBracket = candidate.indexOf("[");
  const starts = [firstBrace, firstBracket].filter((index) => index >= 0);
  if (starts.length === 0) return null;
  const start = Math.min(...starts);
  const lastBrace = candidate.lastIndexOf("}");
  const lastBracket = candidate.lastIndexOf("]");
  const end = Math.max(lastBrace, lastBracket);
  if (end <= start) return null;
  try {
    return JSON.parse(candidate.slice(start, end + 1));
  } catch {
    return null;
  }
}

function normalizeToolCalls(payload) {
  if (!payload) return [];
  const rawCalls = Array.isArray(payload) ? payload : payload.toolCalls || payload.tools || payload.actions || [];
  if (!Array.isArray(rawCalls)) return [];
  return rawCalls
    .map((call, index) => ({
      type: "toolCall",
      id: `tool-${now()}-${index}`,
      name: String(call.tool || call.name || ""),
      arguments: call.args || call.arguments || {},
    }))
    .filter((call) => call.name);
}

function normalizeFinalText(payload, fallback) {
  if (payload && typeof payload.final === "string") return payload.final;
  if (payload && typeof payload.message === "string") return payload.message;
  if (payload && typeof payload.answer === "string") return payload.answer;
  return String(fallback || "").trim() || "Done.";
}

function mockPlan(context) {
  const last = context.messages[context.messages.length - 1];
  if (last?.role === "toolResult") {
    const results = context.messages.filter((message) => message.role === "toolResult").slice(-4).map(stringifyToolResult);
    return {
      final: `The sandbox work completed.\n\n${results.join("\n\n")}`,
    };
  }

  const userText = textFromContent(last?.content || "").toLowerCase();
  if (userText.includes("read")) {
    return {
      toolCalls: [{ tool: "read_file", args: { path: "hello.js" } }],
    };
  }
  if (userText.includes("list")) {
    return {
      toolCalls: [{ tool: "list_files", args: { path: "." } }],
    };
  }
  return {
    toolCalls: [
      {
        tool: "write_file",
        args: {
          path: "hello.js",
          content: 'const value = 21 * 2;\nconsole.log(`pi sandbox result: ${value}`);\n',
        },
      },
      {
        tool: "run_command",
        args: {
          command: "node",
          args: ["hello.js"],
        },
      },
    ],
  };
}

function emitFinal(stream, message, text = "") {
  stream.push({ type: "start", partial: { ...message, content: [{ type: "text", text: "" }] } });
  if (text) {
    const partial = { ...message, content: [{ type: "text", text }] };
    stream.push({ type: "text_start", contentIndex: 0, partial: { ...message, content: [{ type: "text", text: "" }] } });
    stream.push({ type: "text_delta", contentIndex: 0, delta: text, partial });
    stream.push({ type: "text_end", contentIndex: 0, content: text, partial });
  }
  if (message.stopReason === "error" || message.stopReason === "aborted") {
    stream.push({ type: "error", reason: message.stopReason, error: message });
  } else {
    stream.push({ type: "done", reason: message.stopReason, message });
  }
}

export function createPiAgent({ sandbox, modelMode, device, maxTokens, temperature, onModelStatus = () => {} }) {
  env.allowLocalModels = false;
  env.allowRemoteModels = true;
  env.backends.onnx.wasm.numThreads = Math.min(4, navigator.hardwareConcurrency || 4);

  let generatorPromise = null;
  let generatorKey = "";

  async function getGenerator() {
    const key = `${MODEL_ID}:${device()}`;
    if (!generatorPromise || generatorKey !== key) {
      generatorKey = key;
      onModelStatus(`Loading ${device()}`);
      generatorPromise = pipeline("text-generation", MODEL_ID, {
        dtype: "q4",
        device: device(),
        progress_callback: (event) => {
          if (event.status === "progress") {
            onModelStatus(`${event.file} ${Math.round(event.progress)}%`);
          } else if (event.status) {
            onModelStatus(event.status);
          }
        },
      });
    }
    return generatorPromise;
  }

  async function producePlan(context, signal) {
    if (modelMode() === "mock") {
      return JSON.stringify(mockPlan(context));
    }

    const generator = await getGenerator();
    if (signal?.aborted) throw new Error("Aborted");
    const result = await generator(buildPrompt(context), {
      max_new_tokens: Number(maxTokens()) || 128,
      temperature: Number(temperature()) || 0,
      do_sample: Number(temperature()) > 0,
      return_full_text: false,
    });
    onModelStatus("Model ready");
    return result?.[0]?.generated_text ?? "";
  }

  function streamFn(_model, context, options = {}) {
    const stream = createAssistantMessageEventStream();
    queueMicrotask(async () => {
      try {
        const generated = await producePlan(context, options.signal);
        const payload = extractJsonPayload(generated);
        const lastMessage = context.messages[context.messages.length - 1];
        const forceFinal = lastMessage?.role === "toolResult";
        const fallbackPayload = payload ? null : mockPlan(context);
        const toolCalls = forceFinal ? [] : normalizeToolCalls(payload || fallbackPayload);
        if (toolCalls.length > 0) {
          const message = createMessage([{ type: "text", text: "Using sandbox tools." }, ...toolCalls], "toolUse");
          emitFinal(stream, message, "Using sandbox tools.");
          return;
        }
        const text = normalizeFinalText(payload || fallbackPayload, generated);
        const message = createMessage([{ type: "text", text }], "stop");
        emitFinal(stream, message, text);
      } catch (error) {
        const text = error instanceof Error ? error.message : String(error);
        const message = {
          ...createMessage([{ type: "text", text }], options.signal?.aborted ? "aborted" : "error"),
          errorMessage: text,
        };
        emitFinal(stream, message, text);
      }
    });
    return stream;
  }

  const tools = [
    {
      name: "list_files",
      label: "List files",
      description: "List files in the sandbox workspace.",
      parameters: Type.Object({
        path: Type.Optional(Type.String()),
      }),
      execute: async (_id, args) => {
        const output = await sandbox.listFiles(args.path || ".");
        return { content: [{ type: "text", text: output }], details: { output } };
      },
    },
    {
      name: "read_file",
      label: "Read file",
      description: "Read a UTF-8 file from the sandbox workspace.",
      parameters: Type.Object({
        path: Type.String(),
      }),
      execute: async (_id, args) => {
        const output = await sandbox.readFile(args.path);
        return { content: [{ type: "text", text: output }], details: { path: args.path, output } };
      },
    },
    {
      name: "write_file",
      label: "Write file",
      description: "Create or replace a UTF-8 file inside the sandbox workspace.",
      parameters: Type.Object({
        path: Type.String(),
        content: Type.String(),
      }),
      execute: async (_id, args) => {
        const output = await sandbox.writeFile(args.path, args.content);
        return { content: [{ type: "text", text: output }], details: { path: args.path } };
      },
    },
    {
      name: "run_command",
      label: "Run command",
      description: "Spawn a process inside the browser-only WebContainer sandbox.",
      parameters: Type.Object({
        command: Type.String(),
        args: Type.Optional(Type.Array(Type.String())),
        timeoutMs: Type.Optional(Type.Number()),
      }),
      execute: async (_id, args) => {
        const result = await sandbox.runCommand(args.command, args.args || [], args.timeoutMs || 10000);
        const text = `$ ${result.command}\nexit ${result.exitCode}\n${result.output}`;
        return {
          content: [{ type: "text", text }],
          details: result,
        };
      },
      executionMode: "sequential",
    },
  ];

  return new Agent({
    initialState: {
      model: LOCAL_MODEL,
      systemPrompt:
        "You are Pi Web Agent. Use the sandbox tools for filesystem or command tasks, then give concise results.",
      tools,
    },
    streamFn,
    toolExecution: "sequential",
  });
}

export { MODEL_ID };