alexion
6081e3e325
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0
103 lines
4.0 KiB
TypeScript
103 lines
4.0 KiB
TypeScript
// Replace a tool for some conversations, and decorate whichever tool wins: a bash that runs inside a Python
|
|
// virtualenv, and a wrapper that times every bash call.
|
|
// Run from packages/durable:
|
|
// node --conditions=source --experimental-strip-types test/examples/30-tool-override.ts
|
|
import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context";
|
|
import type { ToolResultMessage } from "@earendil-works/pi-ai";
|
|
import { createModels } from "@earendil-works/pi-ai/models";
|
|
import { fauxAssistantMessage, fauxProvider, fauxToolCall } from "@earendil-works/pi-ai/providers/faux";
|
|
import { NodeExecutionEnv } from "../../src/env/node.ts";
|
|
import {
|
|
type Conversation,
|
|
createRegistry,
|
|
defineExtension,
|
|
Harness,
|
|
MemoryStorage,
|
|
ToolResultEntry,
|
|
wrapTool,
|
|
} from "../../src/index.ts";
|
|
import { CodingTools, createBashTool } from "../../src/tools/index.ts";
|
|
|
|
const context = BACKGROUND_CONTEXT;
|
|
|
|
// ─── Product code ───────────────────────────────────────────────────────────
|
|
|
|
// A bash with the same name: where selected after CodingTools, it replaces CodingTools' bash in place.
|
|
const Venv = defineExtension({
|
|
name: "venv",
|
|
tools: [createBashTool({ commandPrefix: "source .venv/bin/activate" })],
|
|
});
|
|
|
|
// Wraps the bash that won, whichever it is. Wrappers never capture a base tool, so reloading either bash keeps it.
|
|
const timings: number[] = [];
|
|
const Timing = defineExtension({
|
|
name: "timing",
|
|
wraps: [
|
|
wrapTool(createBashTool(), (bash) => ({
|
|
...bash,
|
|
execute: async (args, api, callContext) => {
|
|
const start = Date.now();
|
|
try {
|
|
return await bash.execute(args, api, callContext);
|
|
} finally {
|
|
timings.push(Date.now() - start);
|
|
}
|
|
},
|
|
})),
|
|
],
|
|
});
|
|
|
|
// ─── Host setup ─────────────────────────────────────────────────────────────
|
|
|
|
const project = await mkdtemp(join(tmpdir(), "pi-durable-venv-"));
|
|
await mkdir(join(project, ".venv/bin"), { recursive: true });
|
|
await writeFile(join(project, ".venv/bin/activate"), `export VIRTUAL_ENV="${project}/.venv"\n`);
|
|
|
|
const faux = fauxProvider();
|
|
const models = createModels();
|
|
models.setProvider(faux.provider);
|
|
const probe = () =>
|
|
fauxAssistantMessage(fauxToolCall("bash", { command: "echo venv: $VIRTUAL_ENV" }), {
|
|
stopReason: "toolUse",
|
|
});
|
|
faux.setResponses([probe(), fauxAssistantMessage("Done."), probe(), fauxAssistantMessage("Done.")]);
|
|
|
|
const registry = createRegistry();
|
|
registry.install(CodingTools);
|
|
registry.install(Timing);
|
|
registry.install(Venv);
|
|
const harness = await Harness.open(
|
|
new MemoryStorage(),
|
|
{
|
|
models,
|
|
registry,
|
|
// Venv is installed but not selected by default.
|
|
settings: { extensions: [CodingTools, Timing] },
|
|
env: () => new NodeExecutionEnv({ cwd: project }),
|
|
},
|
|
context,
|
|
);
|
|
const model = { provider: "faux", modelId: "faux-1" };
|
|
const plain = await harness.root(context, { agent: { model } });
|
|
const python = await harness.createConversation(
|
|
{ ownership: { kind: "ownerless" }, agent: { model, extensions: { add: [Venv] } } },
|
|
context,
|
|
);
|
|
|
|
async function probeBash(conversation: Conversation): Promise<string> {
|
|
await (await conversation.submit({ type: "input", content: "Which venv?" }, context)).wait(context);
|
|
const page = await conversation.entries({}, 10, undefined, context);
|
|
const result = page.items.find((entry) => ToolResultEntry.is(entry))?.model?.[0] as ToolResultMessage;
|
|
return result.content.map((part) => (part.type === "text" ? part.text.trim() : "")).join("");
|
|
}
|
|
|
|
console.log("plain conversation:", await probeBash(plain));
|
|
console.log("python conversation:", (await probeBash(python)).replace(project, "<project>"));
|
|
console.log("timed bash calls:", timings.length);
|
|
|
|
await harness.close(context);
|
|
await rm(project, { recursive: true, force: true });
|