mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
9e05370b29
This change makes system prompt text and tool changes part of the transcript rather than silently rewriting its starting conditions. This lets Pi record when instructions changed or tools became available, restore that state after resuming or navigating branches, and preserve cached prompt prefixes where the upstream supports it.
82 lines
2.2 KiB
TypeScript
82 lines
2.2 KiB
TypeScript
import { describe, expect, it } from "vitest";
|
|
import { buildBaseOptions } from "../src/api/simple-options.ts";
|
|
import type { AssistantMessage, Model, Usage } from "../src/types.ts";
|
|
import { estimateContextTokens } from "../src/utils/estimate.ts";
|
|
import { normalizeContext } from "../src/utils/transcript.ts";
|
|
|
|
function createUsage(totalTokens: number): Usage {
|
|
return {
|
|
input: totalTokens,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
};
|
|
}
|
|
|
|
function createAssistant(timestamp: number, totalTokens: number): AssistantMessage {
|
|
return {
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "kept" }],
|
|
api: "openai-responses",
|
|
provider: "openai",
|
|
model: "test-model",
|
|
usage: createUsage(totalTokens),
|
|
stopReason: "stop",
|
|
timestamp,
|
|
};
|
|
}
|
|
|
|
const model: Model<"openai-responses"> = {
|
|
id: "test-model",
|
|
name: "Test Model",
|
|
api: "openai-responses",
|
|
provider: "openai",
|
|
baseUrl: "https://api.openai.com/v1",
|
|
reasoning: false,
|
|
input: ["text"],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 10_000,
|
|
maxTokens: 8_000,
|
|
};
|
|
|
|
describe("context token estimation", () => {
|
|
it("ignores stale assistant usage after a newer message is inserted before it", () => {
|
|
const context = normalizeContext({
|
|
systemPrompt: "system",
|
|
messages: [
|
|
{ role: "user", content: "summary", timestamp: 200 },
|
|
createAssistant(100, 9_500),
|
|
{ role: "user", content: "x".repeat(4_000), timestamp: 300 },
|
|
],
|
|
});
|
|
|
|
expect(estimateContextTokens(context)).toEqual({
|
|
tokens: 1_005,
|
|
usageTokens: 0,
|
|
trailingTokens: 1_005,
|
|
lastUsageIndex: null,
|
|
});
|
|
expect(buildBaseOptions(model, context).maxTokens).toBe(4_899);
|
|
});
|
|
|
|
it("uses assistant usage again after a response to the inserted context", () => {
|
|
const context = normalizeContext({
|
|
messages: [
|
|
{ role: "user", content: "summary", timestamp: 200 },
|
|
createAssistant(100, 9_500),
|
|
{ role: "user", content: "new prompt", timestamp: 300 },
|
|
createAssistant(400, 2_000),
|
|
{ role: "user", content: "tail", timestamp: 500 },
|
|
],
|
|
});
|
|
|
|
expect(estimateContextTokens(context)).toEqual({
|
|
tokens: 2_001,
|
|
usageTokens: 2_000,
|
|
trailingTokens: 1,
|
|
lastUsageIndex: 3,
|
|
});
|
|
});
|
|
});
|