150 lines
3.5 KiB
TypeScript
150 lines
3.5 KiB
TypeScript
import type { Subprocess } from "bun";
|
|
import type { PiEvent, MessageCost } from "./types";
|
|
import { logger } from "./logger";
|
|
|
|
export async function spawnPi(
|
|
cwd: string,
|
|
model: string,
|
|
env?: Record<string, string>
|
|
): Promise<Subprocess> {
|
|
const piProcess = Bun.spawn(
|
|
[
|
|
"pi",
|
|
"--mode",
|
|
"rpc",
|
|
"--no-extensions",
|
|
"--no-skills",
|
|
"--model",
|
|
model,
|
|
],
|
|
{
|
|
cwd,
|
|
env: {
|
|
...process.env,
|
|
...env,
|
|
},
|
|
stdin: "pipe",
|
|
stdout: "pipe",
|
|
stderr: "inherit",
|
|
}
|
|
);
|
|
|
|
return piProcess;
|
|
}
|
|
|
|
export function sendPrompt(
|
|
process: Subprocess,
|
|
prompt: string,
|
|
requestId: string
|
|
): void {
|
|
const request = {
|
|
jsonrpc: "2.0",
|
|
method: "chat",
|
|
params: {
|
|
prompt,
|
|
},
|
|
id: requestId,
|
|
};
|
|
|
|
const writer = (process.stdin as any).getWriter();
|
|
writer.write(
|
|
new TextEncoder().encode(JSON.stringify(request) + "\n")
|
|
);
|
|
writer.releaseLock();
|
|
}
|
|
|
|
export function abort(
|
|
process: Subprocess,
|
|
requestId: string
|
|
): void {
|
|
const request = {
|
|
jsonrpc: "2.0",
|
|
method: "abort",
|
|
params: {},
|
|
id: requestId,
|
|
};
|
|
|
|
const writer = (process.stdin as any).getWriter();
|
|
writer.write(
|
|
new TextEncoder().encode(JSON.stringify(request) + "\n")
|
|
);
|
|
writer.releaseLock();
|
|
}
|
|
|
|
export async function* readEvents(
|
|
process: Subprocess
|
|
): AsyncGenerator<PiEvent> {
|
|
if (!process.stdout) {
|
|
throw new Error("Pi process stdout not available");
|
|
}
|
|
|
|
const reader = (process.stdout as any).getReader();
|
|
const decoder = new TextDecoder();
|
|
let buffer = "";
|
|
|
|
try {
|
|
while (true) {
|
|
const { done, value } = await reader.read();
|
|
|
|
if (done) {
|
|
break;
|
|
}
|
|
|
|
buffer += decoder.decode(value, { stream: true });
|
|
const lines = buffer.split("\n");
|
|
|
|
buffer = lines.pop() || "";
|
|
|
|
for (const line of lines) {
|
|
if (!line.trim()) continue;
|
|
|
|
try {
|
|
const event = JSON.parse(line);
|
|
|
|
if (event.method === "text") {
|
|
yield { type: "text", text: event.params.text };
|
|
} else if (event.method === "delta") {
|
|
yield { type: "delta", text: event.params.text };
|
|
} else if (event.method === "tool:start") {
|
|
yield {
|
|
type: "tool:start",
|
|
toolCallId: event.params.toolCallId,
|
|
toolName: event.params.toolName,
|
|
toolInput: event.params.toolInput,
|
|
};
|
|
} else if (event.method === "tool:result") {
|
|
yield {
|
|
type: "tool:result",
|
|
toolCallId: event.params.toolCallId,
|
|
output: event.params.output,
|
|
isError: event.params.isError || false,
|
|
};
|
|
} else if (event.method === "result") {
|
|
const cost: MessageCost = {
|
|
inputTokens: event.params.cost?.inputTokens || 0,
|
|
outputTokens: event.params.cost?.outputTokens || 0,
|
|
totalUSD: event.params.cost?.totalUSD || 0,
|
|
};
|
|
yield { type: "result", cost };
|
|
} else if (event.method === "error") {
|
|
yield {
|
|
type: "error",
|
|
message: event.params.message || "Unknown error",
|
|
};
|
|
} else if (event.method === "stopped") {
|
|
yield { type: "stopped" };
|
|
}
|
|
} catch (err) {
|
|
logger.error("Failed to parse Pi event", { line, error: String(err) });
|
|
}
|
|
}
|
|
}
|
|
} finally {
|
|
reader.releaseLock();
|
|
}
|
|
}
|
|
|
|
export function killPi(process: Subprocess): void {
|
|
process.kill();
|
|
}
|