| 1 | import assert from "node:assert/strict"; |
| 2 | import { test } from "node:test"; |
| 3 | |
| 4 | import { StreamUsage, fromUsage, measure, total } from "./usage.ts"; |
| 5 | |
| 6 | const sse = (events: object[]) => events.map((e) => `event: x\ndata: ${JSON.stringify(e)}\n\n`).join(""); |
| 7 | |
| 8 | test("a stream's tokens come from message_start and the last message_delta", () => { |
| 9 | const stream = sse([ |
| 10 | { type: "message_start", message: { usage: { input_tokens: 12, cache_read_input_tokens: 4000, cache_creation_input_tokens: 300, output_tokens: 1 } } }, |
| 11 | { type: "content_block_delta", delta: { text: "hi" } }, |
| 12 | { type: "message_delta", usage: { output_tokens: 40 } }, |
| 13 | { type: "message_delta", usage: { output_tokens: 95 } }, |
| 14 | { type: "message_stop" }, |
| 15 | ]); |
| 16 | const usage = new StreamUsage(); |
| 17 | // Split mid-line, as chunks arrive. |
| 18 | usage.push(stream.slice(0, 37)); |
| 19 | usage.push(stream.slice(37)); |
| 20 | assert.deepEqual(usage.finish(), { input: 12, output: 95, cacheRead: 4000, cacheWrite: 300 }); |
| 21 | }); |
| 22 | |
| 23 | test("a whole answer's usage, and nonsense counted as nothing", () => { |
| 24 | assert.deepEqual(fromUsage({ input_tokens: 5, output_tokens: 7 }), { input: 5, output: 7, cacheRead: 0, cacheWrite: 0 }); |
| 25 | assert.equal(total(fromUsage({ input_tokens: -3, output_tokens: Number.NaN })), 0); |
| 26 | const usage = new StreamUsage(); |
| 27 | usage.push("data: {not json\n\ndata: [DONE]\n"); |
| 28 | assert.equal(total(usage.finish()), 0); |
| 29 | }); |
| 30 | |
| 31 | test("measuring passes the answer through unchanged", async () => { |
| 32 | const body = sse([{ type: "message_start", message: { usage: { input_tokens: 3 } } }, { type: "message_delta", usage: { output_tokens: 9 } }]); |
| 33 | const { response, tokens } = measure(new Response(body, { headers: { "content-type": "text/event-stream" } })); |
| 34 | assert.equal(await response.text(), body); |
| 35 | assert.deepEqual(await tokens, { input: 3, output: 9, cacheRead: 0, cacheWrite: 0 }); |
| 36 | const json = measure(new Response(JSON.stringify({ usage: { input_tokens: 2, output_tokens: 1 } }), { headers: { "content-type": "application/json" } })); |
| 37 | assert.equal(JSON.parse(await json.response.text()).usage.input_tokens, 2); |
| 38 | assert.equal(total(await json.tokens), 3); |
| 39 | const failed = measure(new Response("no", { status: 500 })); |
| 40 | assert.equal(failed.response.status, 500); |
| 41 | assert.equal(total(await failed.tokens), 0); |
| 42 | }); |