Mario Zechner
75 lines
2.0 KiB
TypeScript
75 lines
2.0 KiB
TypeScript
import { describe, expect, it } from "vitest";
|
|
import {
|
|
type AssistantMessage,
|
|
type Context,
|
|
calculateContextTokens,
|
|
estimateContextTokens,
|
|
estimateMessageTokens,
|
|
type Usage,
|
|
} from "../src/index.ts";
|
|
|
|
function usage(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 assistant(totalTokens: number, stopReason: AssistantMessage["stopReason"] = "stop"): AssistantMessage {
|
|
return {
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "done" }],
|
|
api: "anthropic-messages",
|
|
provider: "test-provider",
|
|
model: "test-model",
|
|
usage: usage(totalTokens),
|
|
stopReason,
|
|
timestamp: Date.now(),
|
|
};
|
|
}
|
|
|
|
describe("token estimation", () => {
|
|
it("calculates context tokens from usage", () => {
|
|
expect(calculateContextTokens({ ...usage(0), input: 10, output: 20, cacheRead: 3, cacheWrite: 4 })).toBe(37);
|
|
expect(calculateContextTokens(usage(123))).toBe(123);
|
|
});
|
|
|
|
it("estimates message tokens from content", () => {
|
|
expect(estimateMessageTokens({ role: "user", content: "x".repeat(400), timestamp: 0 })).toBe(100);
|
|
});
|
|
|
|
it("uses the latest valid assistant usage as baseline and estimates only trailing messages", () => {
|
|
const context: Context = {
|
|
systemPrompt: "system text that should already be counted by provider usage",
|
|
messages: [
|
|
{ role: "user", content: "x".repeat(10_000), timestamp: 0 },
|
|
assistant(50_000),
|
|
{ role: "user", content: "x".repeat(400), timestamp: 0 },
|
|
],
|
|
};
|
|
|
|
expect(estimateContextTokens(context)).toMatchObject({
|
|
tokens: 50_100,
|
|
usageTokens: 50_000,
|
|
trailingTokens: 100,
|
|
lastUsageIndex: 1,
|
|
});
|
|
});
|
|
|
|
it("ignores error/aborted assistant usage", () => {
|
|
const context: Context = {
|
|
messages: [assistant(50_000, "error"), { role: "user", content: "x".repeat(400), timestamp: 0 }],
|
|
};
|
|
|
|
expect(estimateContextTokens(context)).toMatchObject({
|
|
tokens: 101,
|
|
usageTokens: 0,
|
|
lastUsageIndex: null,
|
|
});
|
|
});
|
|
});
|