Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/ephemlets/__bench__/context.ts

Raw
import { fileURLToPath } from "node:url";
import {
	fauxAssistantMessage,
	getCurrentSystemPrompt,
	getCurrentTools,
} from "@earendil-works/pi-ai";
import type {
	ExtensionContext,
	ExtensionUIContext,
} from "@earendil-works/pi-coding-agent";
import { createPiSession } from "../../../test/fixture/create-pi-session.ts";

async function measure(
	mode: ExtensionContext["mode"],
	enabled: boolean,
): Promise<{ system: number; tools: number }> {
	const fixture = await createPiSession({
		extensionPaths: enabled
			? [fileURLToPath(new URL("../index.ts", import.meta.url))]
			: [],
		env: { PI_EPHEMLETS_DEBUG: "0" },
	});
	try {
		await fixture.session.bindExtensions({
			mode,
			uiContext: {
				...fixture.session.extensionRunner.getUIContext(),
				custom: async () => undefined,
			} as ExtensionUIContext,
		});
		let metrics = { system: 0, tools: 0 };
		fixture.setResponses([
			(context) => {
				metrics = {
					system: Math.ceil(
						getCurrentSystemPrompt(context.messages).length / 4,
					),
					tools: Math.ceil(
						JSON.stringify(getCurrentTools(context.messages)).length / 4,
					),
				};
				return fauxAssistantMessage("OK");
			},
		]);
		await fixture.session.prompt("Reply only OK. Do not use tools.");
		return metrics;
	} finally {
		await fixture.dispose();
	}
}

for (const mode of ["print", "tui"] as const) {
	const baseline = await measure(mode, false);
	const enabled = await measure(mode, true);
	process.stdout.write(
		`${JSON.stringify({ mode, measurement: "heuristic JS characters / 4; scripted provider, not live API tokens", baseline, enabled, delta: { system: enabled.system - baseline.system, tools: enabled.tools - baseline.tools, total: enabled.system + enabled.tools - baseline.system - baseline.tools } })}\n`,
	);
}