mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
540e174c72
This adds experimental support for virtual models. A virtual model is a catalog entry that an extension registers with pi.registerVirtualModel(). It does not talk to a provider itself; for every request it picks a physical model and thinking level. The motivation is to make routing policies pluggable: plan on a strong model and implement on a cheap one, pick a model with a classifier, or move to a larger window when the context grows. Today this needs model switching by hand or hacks in prepareRequest. The selection stays virtual: model_change records it, resume restores it, and each response records the physical model and thinking level that produced it. Router state can be stored on the session branch, so it follows /tree and forks and survives compaction. Context limits, compaction and image sizing use the physical model a request was routed to. examples/extensions/jev-router.ts shows a full router that uses the Jev classifier to pick Sol or Terra for planning and hands off to Luna after the first edit. docs/virtual-models.md describes the API.
397 lines
16 KiB
TypeScript
397 lines
16 KiB
TypeScript
import type { AgentTool } from "@earendil-works/pi-agent-core";
|
|
import { type FauxResponseStep, fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
|
|
import { Type } from "typebox";
|
|
import { afterEach, describe, expect, it, vi } from "vitest";
|
|
import type { ExtensionContext } from "../../src/core/extensions/index.ts";
|
|
import { type ModelRoute, type ModelRouteRequest, VIRTUAL_MODEL_STATE_ENTRY } from "../../src/core/virtual-models.ts";
|
|
import { createHarness, type Harness, type HarnessOptions } from "./harness.ts";
|
|
|
|
const echoTool: AgentTool = {
|
|
name: "echo",
|
|
label: "Echo",
|
|
description: "Echo text back",
|
|
parameters: Type.Object({ text: Type.String() }),
|
|
execute: async () => ({ content: [{ type: "text", text: "echoed" }], details: {} }),
|
|
};
|
|
|
|
type Route = (request: ModelRouteRequest, ctx: ExtensionContext) => ModelRoute;
|
|
|
|
/** Router used by these tests: new user turns pick by thinking level, everything else stays put. */
|
|
const defaultRoute: Route = (request, ctx) => {
|
|
const find = (id: string) => ctx.modelRegistry.find("faux", id)!;
|
|
if (request.reason === "direct") return { model: find("large"), thinkingLevel: "low" };
|
|
const sticky = request.failed ?? request.previous;
|
|
if (request.reason !== "user" && sticky) {
|
|
return { model: sticky.model, thinkingLevel: sticky.thinkingLevel ?? "high" };
|
|
}
|
|
return request.thinkingLevel === "high"
|
|
? { model: find("large"), thinkingLevel: "high" }
|
|
: { model: find("small"), thinkingLevel: "off" };
|
|
};
|
|
|
|
describe("AgentSession virtual models", () => {
|
|
const harnesses: Harness[] = [];
|
|
|
|
afterEach(() => {
|
|
for (const harness of harnesses.splice(0)) harness.cleanup();
|
|
});
|
|
|
|
async function createRoutedHarness(route: Route = defaultRoute, options: HarnessOptions = {}) {
|
|
const requests: ModelRouteRequest[] = [];
|
|
const harness = await createHarness({
|
|
...options,
|
|
models: [
|
|
{ id: "small", contextWindow: 1000 },
|
|
{ id: "large", contextWindow: 50_000, maxTokens: 4000, reasoning: true },
|
|
],
|
|
tools: [echoTool],
|
|
extensionFactories: [
|
|
...(options.extensionFactories ?? []),
|
|
(pi) => {
|
|
pi.registerVirtualModel({
|
|
provider: "router",
|
|
id: "auto",
|
|
name: "Auto",
|
|
thinkingLevels: ["low", "high"],
|
|
contextWindow: 1000,
|
|
route(request, ctx) {
|
|
requests.push(request);
|
|
return route(request, ctx);
|
|
},
|
|
});
|
|
},
|
|
],
|
|
});
|
|
harnesses.push(harness);
|
|
// Stream through the runtime, which records the thinking level on responses.
|
|
const runtime = harness.session.modelRuntime;
|
|
harness.session.agent.streamFunction = (model, context, options) => runtime.streamSimple(model, context, options);
|
|
await harness.session.setModel(runtime.getModel("router", "auto")!);
|
|
harness.session.setThinkingLevel("high");
|
|
const reasons = () => requests.map((request) => request.reason);
|
|
/** Physical model and thinking level recorded on each response. */
|
|
const dispatched = () =>
|
|
harness.session.messages.flatMap((message) =>
|
|
message.role === "assistant" ? [`${message.provider}/${message.model}:${message.thinkingLevel}`] : [],
|
|
);
|
|
return { harness, requests, reasons, dispatched };
|
|
}
|
|
|
|
it("routes each request, including retries, while the selection stays virtual", async () => {
|
|
const { harness, requests, reasons, dispatched } = await createRoutedHarness(defaultRoute, {
|
|
settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } },
|
|
});
|
|
harness.setResponses([
|
|
fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
|
|
fauxAssistantMessage(fauxToolCall("echo", { text: "hi" }), { stopReason: "toolUse" }),
|
|
fauxAssistantMessage("done"),
|
|
]);
|
|
|
|
await harness.session.prompt("hello");
|
|
|
|
expect(reasons()).toEqual(["user", "retry", "continuation"]);
|
|
expect(requests[1].failed?.message.errorMessage).toBe("overloaded_error");
|
|
expect(requests[2].previous?.model.id).toBe("large");
|
|
expect(dispatched()).toEqual(["faux/large:high", "faux/large:high"]);
|
|
expect(harness.session.model).toMatchObject({ provider: "router", id: "auto" });
|
|
expect(harness.session.thinkingLevel).toBe("high");
|
|
// Limits come from the physical model that produced the latest response, not the virtual model.
|
|
expect(harness.session.getContextUsage()?.contextWindow).toBe(50_000);
|
|
});
|
|
|
|
it("retries the first request of a turn on the model routed for that turn", async () => {
|
|
const { harness, requests, reasons, dispatched } = await createRoutedHarness(defaultRoute, {
|
|
settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } },
|
|
});
|
|
harness.session.setThinkingLevel("low");
|
|
harness.setResponses([
|
|
fauxAssistantMessage("easy answer"),
|
|
fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
|
|
fauxAssistantMessage("hard answer"),
|
|
]);
|
|
await harness.session.prompt("easy");
|
|
harness.session.setThinkingLevel("high");
|
|
|
|
await harness.session.prompt("hard");
|
|
|
|
expect(reasons()).toEqual(["user", "user", "retry"]);
|
|
// The retry reports the failed request on large next to the small response of the previous turn.
|
|
expect(requests[2].failed?.model.id).toBe("large");
|
|
expect(requests[2].previous?.model.id).toBe("small");
|
|
expect(dispatched()).toEqual(["faux/small:off", "faux/large:high"]);
|
|
});
|
|
|
|
it("routes the compact-and-retry after a truncated response as a retry", async () => {
|
|
const { harness, requests, reasons } = await createRoutedHarness(defaultRoute, {
|
|
settings: { compaction: { keepRecentTokens: 1, reserveTokens: 0 } },
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.on("session_before_compact", async ({ preparation: { firstKeptEntryId, tokensBefore } }) => ({
|
|
compaction: { summary: "overflow compacted", firstKeptEntryId, tokensBefore },
|
|
}));
|
|
},
|
|
],
|
|
});
|
|
harness.setResponses([
|
|
() => fauxAssistantMessage("x".repeat(64), { stopReason: "length", timestamp: Date.now() + 10_000 }),
|
|
fauxAssistantMessage("done"),
|
|
]);
|
|
|
|
await harness.session.prompt("x".repeat(5000));
|
|
|
|
expect(harness.eventsOfType("compaction_start").map((event) => event.reason)).toEqual(["overflow"]);
|
|
// Compaction may fold the prompt into the summary, so the retry is not a new user turn.
|
|
expect(reasons()).toEqual(["user", "retry"]);
|
|
expect(requests[1].failed?.model.id).toBe("large");
|
|
expect(requests[1].failed?.message.stopReason).toBe("length");
|
|
});
|
|
|
|
it("routes requests after extension messages as continuations", async () => {
|
|
const { harness, reasons } = await createRoutedHarness(defaultRoute, {
|
|
extensionFactories: [
|
|
(pi) => {
|
|
let continued = false;
|
|
pi.on("agent_before_settle", () => {
|
|
if (continued) return undefined;
|
|
continued = true;
|
|
return {
|
|
entries: [{ type: "custom_message", customType: "nudge", content: "Keep going.", display: false }],
|
|
continue: true,
|
|
};
|
|
});
|
|
},
|
|
],
|
|
});
|
|
harness.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
|
|
|
|
await harness.session.prompt("hello");
|
|
|
|
// The hidden custom message becomes a user message for the model, but the user did not write it.
|
|
expect(reasons()).toEqual(["user", "continuation"]);
|
|
});
|
|
|
|
it("routes the first request of a prompt as a user turn when extension messages follow the prompt", async () => {
|
|
const { harness, reasons } = await createRoutedHarness(defaultRoute, {
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.on("before_agent_start", () => ({
|
|
message: { customType: "context", content: "Extra context.", display: false },
|
|
}));
|
|
},
|
|
],
|
|
});
|
|
harness.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
|
|
|
|
await harness.session.prompt("hello");
|
|
await harness.session.prompt("again");
|
|
|
|
// The context ends with the extension message, but the request answers the user's prompt.
|
|
expect(harness.session.messages.at(-2)?.role).toBe("custom");
|
|
expect(reasons()).toEqual(["user", "user"]);
|
|
});
|
|
|
|
it("ends the run with an error response when routing fails and keeps the last physical limits", async () => {
|
|
let fail = false;
|
|
const { harness } = await createRoutedHarness((request, ctx) => {
|
|
if (fail) throw new Error("router unavailable");
|
|
return defaultRoute(request, ctx);
|
|
});
|
|
harness.setResponses([fauxAssistantMessage("answer")]);
|
|
await harness.session.prompt("hello");
|
|
|
|
fail = true;
|
|
await harness.session.prompt("again");
|
|
|
|
expect(harness.session.messages.at(-1)).toMatchObject({
|
|
role: "assistant",
|
|
provider: "router",
|
|
model: "auto",
|
|
stopReason: "error",
|
|
errorMessage: expect.stringContaining("router unavailable"),
|
|
});
|
|
expect(harness.faux.state.callCount).toBe(1);
|
|
// The failed attempt names the virtual model, whose declared window is 1k; the large model's 50k applies.
|
|
expect(harness.session.getContextUsage()?.contextWindow).toBe(50_000);
|
|
});
|
|
|
|
it("checks compaction against the physical model that produced the response", async () => {
|
|
const { harness } = await createRoutedHarness();
|
|
harness.setResponses([fauxAssistantMessage("short answer"), fauxAssistantMessage("long answer")]);
|
|
await harness.session.prompt("hello");
|
|
|
|
// About 20k tokens exceed the virtual model's declared 1k window but fit the large model's 50k.
|
|
await harness.session.prompt("x".repeat(80_000));
|
|
|
|
expect(harness.eventsOfType("compaction_start")).toHaveLength(0);
|
|
});
|
|
|
|
it("compacts before a request routed to a model with a smaller window", async () => {
|
|
const { harness, dispatched } = await createRoutedHarness(defaultRoute, {
|
|
settings: { compaction: { keepRecentTokens: 1, reserveTokens: 0 } },
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.on("session_before_compact", async ({ preparation: { firstKeptEntryId, tokensBefore } }) => ({
|
|
compaction: { summary: "compacted", firstKeptEntryId, tokensBefore },
|
|
}));
|
|
},
|
|
],
|
|
});
|
|
let compactedBeforeSmall = false;
|
|
harness.setResponses([
|
|
fauxAssistantMessage("y".repeat(8000)),
|
|
() => {
|
|
compactedBeforeSmall = harness.eventsOfType("compaction_end").length === 1;
|
|
return fauxAssistantMessage("small answer");
|
|
},
|
|
]);
|
|
await harness.session.prompt("hello");
|
|
expect(harness.eventsOfType("compaction_start")).toHaveLength(0);
|
|
|
|
// About 2k tokens fit the large model that answered last, but not the small model's 1k window.
|
|
harness.session.setThinkingLevel("low");
|
|
await harness.session.prompt("next");
|
|
|
|
expect(harness.eventsOfType("compaction_start").map((event) => event.reason)).toEqual(["threshold"]);
|
|
expect(compactedBeforeSmall).toBe(true);
|
|
expect(dispatched().at(-1)).toBe("faux/small:off");
|
|
});
|
|
|
|
it("compacts between turns of a run when the next request is routed to a smaller window", async () => {
|
|
const route: Route = (request, ctx) =>
|
|
request.reason === "continuation"
|
|
? { model: ctx.modelRegistry.find("faux", "small")!, thinkingLevel: "off" }
|
|
: defaultRoute(request, ctx);
|
|
const { harness, dispatched } = await createRoutedHarness(route, {
|
|
settings: { compaction: { keepRecentTokens: 1, reserveTokens: 0 } },
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.on("session_before_compact", async ({ preparation: { firstKeptEntryId, tokensBefore } }) => ({
|
|
compaction: { summary: "compacted", firstKeptEntryId, tokensBefore },
|
|
}));
|
|
},
|
|
],
|
|
});
|
|
let compactedBeforeSmall = false;
|
|
harness.setResponses([
|
|
fauxAssistantMessage(fauxToolCall("echo", { text: "hi" }), { stopReason: "toolUse" }),
|
|
() => {
|
|
compactedBeforeSmall = harness.eventsOfType("compaction_end").length === 1;
|
|
return fauxAssistantMessage("small answer");
|
|
},
|
|
]);
|
|
|
|
// About 2k tokens fit the large model of the first turn, but not the small model of the second.
|
|
await harness.session.prompt("x".repeat(8000));
|
|
|
|
expect(harness.eventsOfType("compaction_start").map((event) => event.reason)).toEqual(["threshold"]);
|
|
expect(compactedBeforeSmall).toBe(true);
|
|
expect(dispatched()).toEqual(["faux/large:high", "faux/small:off"]);
|
|
});
|
|
|
|
it("projects the session once per request under a virtual selection", async () => {
|
|
const { harness } = await createRoutedHarness();
|
|
harness.setResponses([
|
|
fauxAssistantMessage(fauxToolCall("echo", { text: "hi" }), { stopReason: "toolUse" }),
|
|
fauxAssistantMessage("done"),
|
|
]);
|
|
const projections = vi.spyOn(harness.sessionManager, "buildSessionProjection");
|
|
|
|
await harness.session.prompt("hello");
|
|
|
|
// One per request, one between the turns, and one for the compaction check after the run.
|
|
expect(projections).toHaveBeenCalledTimes(4);
|
|
});
|
|
|
|
it("stores router state on the branch and passes it to later requests", async () => {
|
|
const states: unknown[] = [];
|
|
const { harness, reasons } = await createRoutedHarness(
|
|
(request, ctx) => {
|
|
states.push(request.state);
|
|
const turns = (request.state as { turns: number } | undefined)?.turns ?? 0;
|
|
const route = defaultRoute(request, ctx);
|
|
// Returning request.state keeps it without storing it again. Direct requests neither get nor store state.
|
|
if (request.reason === "continuation") return { ...route, state: request.state };
|
|
return { ...route, state: { turns: turns + 1 } };
|
|
},
|
|
{ settings: { compaction: { keepRecentTokens: 1 } } },
|
|
);
|
|
harness.setResponses([
|
|
fauxAssistantMessage(fauxToolCall("echo", { text: "hi" }), { stopReason: "toolUse" }),
|
|
fauxAssistantMessage("first"),
|
|
fauxAssistantMessage("second"),
|
|
fauxAssistantMessage("summary"),
|
|
fauxAssistantMessage("summary"),
|
|
]);
|
|
|
|
await harness.session.prompt("one");
|
|
await harness.session.prompt("two");
|
|
const stored = () =>
|
|
harness.sessionManager
|
|
.getBranch()
|
|
.flatMap((entry) =>
|
|
entry.type === "custom" && entry.customType === VIRTUAL_MODEL_STATE_ENTRY ? [entry.data] : [],
|
|
);
|
|
expect(stored()).toEqual([
|
|
{ provider: "router", modelId: "auto", state: { turns: 1 } },
|
|
{ provider: "router", modelId: "auto", state: { turns: 2 } },
|
|
]);
|
|
|
|
await harness.session.compact();
|
|
|
|
expect(reasons()).toEqual(["user", "continuation", "user", "direct"]);
|
|
expect(states).toEqual([undefined, { turns: 1 }, { turns: 1 }, undefined]);
|
|
expect(stored()).toHaveLength(2);
|
|
});
|
|
|
|
it("does not route compactions that an extension supplies", async () => {
|
|
const route: Route = (request, ctx) => {
|
|
if (request.reason === "direct") throw new Error("router unavailable");
|
|
return defaultRoute(request, ctx);
|
|
};
|
|
const { harness, reasons } = await createRoutedHarness(route, {
|
|
settings: { compaction: { keepRecentTokens: 1 } },
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.on("session_before_compact", async ({ preparation: { firstKeptEntryId, tokensBefore } }) => ({
|
|
compaction: { summary: "extension summary", firstKeptEntryId, tokensBefore },
|
|
}));
|
|
},
|
|
],
|
|
});
|
|
harness.setResponses([fauxAssistantMessage("first answer"), fauxAssistantMessage("second answer")]);
|
|
await harness.session.prompt("first");
|
|
await harness.session.prompt("second");
|
|
|
|
const result = await harness.session.compact();
|
|
|
|
expect(result.summary).toBe("extension summary");
|
|
expect(reasons()).toEqual(["user", "user"]);
|
|
});
|
|
|
|
it("routes compaction summaries before sizing them", async () => {
|
|
const summaries: string[] = [];
|
|
const { harness, reasons } = await createRoutedHarness(defaultRoute, {
|
|
settings: { compaction: { keepRecentTokens: 1 } },
|
|
});
|
|
const summary: FauxResponseStep = (_context, options, _state, model) => {
|
|
summaries.push(`${model.id}:${options?.reasoning ?? "off"}:${options?.maxTokens}`);
|
|
return fauxAssistantMessage("summary");
|
|
};
|
|
// Compaction summarizes the history and the split turn prefix with one routed model.
|
|
harness.setResponses([
|
|
fauxAssistantMessage("first answer"),
|
|
fauxAssistantMessage("second answer"),
|
|
summary,
|
|
summary,
|
|
]);
|
|
await harness.session.prompt("first");
|
|
await harness.session.prompt("second");
|
|
|
|
const result = await harness.session.compact();
|
|
|
|
expect(result.summary).toContain("summary");
|
|
expect(reasons()).toEqual(["user", "user", "direct"]);
|
|
// The router's thinking level applies, and the output budget respects the large model's 4000 tokens.
|
|
expect(summaries).toEqual(["large:low:4000", "large:low:4000"]);
|
|
});
|
|
});
|