Q65: every call path records real tokens, cache, reasoning and cost
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
1367a8900e
commit
c2f8e656f6
15 files changed
+187
-37
No files matched your search
@@ -0,0 +1,12 @@
|
||||
export const usageWindowTokensSchema = `
|
||||
alter table provider_usage_windows add column input_tokens integer not null default 0;
|
||||
alter table provider_usage_windows add column output_tokens integer not null default 0;
|
||||
alter table provider_usage_windows add column cached_input_tokens integer not null default 0;
|
||||
alter table provider_usage_windows add column reasoning_tokens integer not null default 0;
|
||||
alter table provider_usage_windows add column cost_cents real not null default 0;
|
||||
alter table provider_usage_windows add column requests_without_usage integer not null default 0;
|
||||
alter table provider_usage_windows add column unpriced_requests integer not null default 0;
|
||||
|
||||
insert into schema_migrations(version, name, applied_at)
|
||||
values (73, 'usage_window_tokens', datetime('now'));
|
||||
`;
|
||||
@@ -71,6 +71,7 @@ import { goalStallSchema } from "./migrations/069-goal-stall";
|
||||
import { runActionsSchema } from "./migrations/070-run-actions";
|
||||
import { backgroundPollFailuresSchema } from "./migrations/071-background-poll-failures";
|
||||
import { spendSchema } from "./migrations/072-spend";
|
||||
import { usageWindowTokensSchema } from "./migrations/073-usage-window-tokens";
|
||||
|
||||
interface Migration {
|
||||
version: number;
|
||||
@@ -151,6 +152,7 @@ const MIGRATIONS: Migration[] = [
|
||||
{ version: 70, name: "run_actions", sql: runActionsSchema },
|
||||
{ version: 71, name: "background_poll_failures", sql: backgroundPollFailuresSchema },
|
||||
{ version: 72, name: "spend", sql: spendSchema },
|
||||
{ version: 73, name: "usage_window_tokens", sql: usageWindowTokensSchema },
|
||||
];
|
||||
|
||||
const SCHEMA_VERSION = MIGRATIONS[MIGRATIONS.length - 1].version;
|
||||
|
||||
@@ -186,12 +186,14 @@ export const providerAdapters: ProviderAdapterRegistry = {
|
||||
label: "xAI",
|
||||
endpoint: "https://api.x.ai/v1/chat/completions",
|
||||
healthUrl: "https://api.x.ai/v1/models",
|
||||
requestStreamUsage: true,
|
||||
}),
|
||||
groq: openAiCompatibleAdapter({
|
||||
providerKey: "groq",
|
||||
label: "Groq",
|
||||
endpoint: "https://api.groq.com/openai/v1/chat/completions",
|
||||
healthUrl: "https://api.groq.com/openai/v1/models",
|
||||
requestStreamUsage: true,
|
||||
}),
|
||||
mistral: openAiCompatibleAdapter({
|
||||
providerKey: "mistral",
|
||||
@@ -367,11 +369,14 @@ function openAiCompatibleAdapter({
|
||||
label,
|
||||
endpoint,
|
||||
healthUrl,
|
||||
requestStreamUsage = false,
|
||||
}: {
|
||||
providerKey: string;
|
||||
label: string;
|
||||
endpoint: string;
|
||||
healthUrl: string;
|
||||
/** Providers that report a stream's token usage only when asked (xAI, Groq). */
|
||||
requestStreamUsage?: boolean;
|
||||
}): ProviderAdapter {
|
||||
return {
|
||||
providerKey,
|
||||
@@ -394,6 +399,7 @@ function openAiCompatibleAdapter({
|
||||
model: input.model,
|
||||
messages,
|
||||
stream: true,
|
||||
...(requestStreamUsage ? { stream_options: { include_usage: true } } : {}),
|
||||
});
|
||||
return { events: mapOpenAiChatCompletionStream(response.messages), rateLimitState: response.rateLimitState };
|
||||
},
|
||||
|
||||
@@ -170,7 +170,10 @@ export class ProviderBackgroundService {
|
||||
startedAt,
|
||||
action: "background_response_completed",
|
||||
});
|
||||
this.requestStore.recordUsage(request.org_id, provider.id, request.input_char_count, responseText.length);
|
||||
this.requestStore.settleUsage({
|
||||
orgId: request.org_id, requestId: request.id, provider, rawResponseJson: result.rawResponseJson,
|
||||
inputCharCount: request.input_char_count, outputCharCount: responseText.length,
|
||||
});
|
||||
return;
|
||||
}
|
||||
if (CANCELLED_BACKGROUND_STATUSES.has(result.status)) {
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
import type { DatabaseSync } from "node:sqlite";
|
||||
import { PROVIDER_EXECUTION_MODE, PROVIDER_REQUEST_STATE, type ProviderExecutionMode } from "../../shared/domain";
|
||||
import { createId } from "../../shared/ids";
|
||||
import type { TokenUsage } from "./spend";
|
||||
import { providerPrices, readTokenUsage, requestCostCents, type TokenUsage } from "./spend";
|
||||
|
||||
/** Providers already logged as not reporting token usage. */
|
||||
const providersWithoutUsage = new Set<string>();
|
||||
import { currentProviderUsageWindow } from "../../shared/provider-runtime-domain";
|
||||
|
||||
export interface ProviderRequestOwner {
|
||||
@@ -269,20 +272,48 @@ export class ProviderRequestStore {
|
||||
return request;
|
||||
}
|
||||
|
||||
recordUsage(orgId: string, providerId: string, inputCharCount: number, outputCharCount: number): void {
|
||||
/** Settles a successful call: reads the provider's token usage, prices it, and adds it to the request and the month's totals. */
|
||||
settleUsage(input: { orgId: string; requestId: string; provider: { id: string; provider_key: string }; rawResponseJson: string; inputCharCount: number; outputCharCount: number }): void {
|
||||
const usage = readTokenUsage(input.rawResponseJson);
|
||||
if (!usage) {
|
||||
if (!providersWithoutUsage.has(input.provider.provider_key)) console.warn(`[spend] ${input.provider.provider_key} returned no token usage; its calls cost nothing on record.`);
|
||||
providersWithoutUsage.add(input.provider.provider_key);
|
||||
this.recordUsage(input.orgId, input.provider.id, input.inputCharCount, input.outputCharCount, null);
|
||||
return;
|
||||
}
|
||||
const costCents = requestCostCents(usage, providerPrices(this.db, input.provider.id));
|
||||
this.recordSpend(input.requestId, usage, costCents);
|
||||
this.recordUsage(input.orgId, input.provider.id, input.inputCharCount, input.outputCharCount, { usage, costCents });
|
||||
}
|
||||
|
||||
/** Adds one call to the provider's monthly totals; spend is null when the provider reported no token usage. */
|
||||
recordUsage(orgId: string, providerId: string, inputCharCount: number, outputCharCount: number, spend: { usage: TokenUsage; costCents: number | null } | null): void {
|
||||
const window = currentProviderUsageWindow();
|
||||
const now = new Date().toISOString();
|
||||
const usage = spend?.usage;
|
||||
this.db.prepare(`
|
||||
insert into provider_usage_windows(
|
||||
id, org_id, provider_id, window_start, window_end, request_count,
|
||||
input_char_count, output_char_count, created_at, updated_at
|
||||
) values (?, ?, ?, ?, ?, 1, ?, ?, ?, ?)
|
||||
input_char_count, output_char_count, created_at, updated_at,
|
||||
input_tokens, output_tokens, cached_input_tokens, reasoning_tokens, cost_cents, requests_without_usage, unpriced_requests
|
||||
) values (?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
on conflict(provider_id, window_start) do update set
|
||||
request_count = request_count + 1,
|
||||
input_char_count = input_char_count + excluded.input_char_count,
|
||||
output_char_count = output_char_count + excluded.output_char_count,
|
||||
input_tokens = input_tokens + excluded.input_tokens,
|
||||
output_tokens = output_tokens + excluded.output_tokens,
|
||||
cached_input_tokens = cached_input_tokens + excluded.cached_input_tokens,
|
||||
reasoning_tokens = reasoning_tokens + excluded.reasoning_tokens,
|
||||
cost_cents = cost_cents + excluded.cost_cents,
|
||||
requests_without_usage = requests_without_usage + excluded.requests_without_usage,
|
||||
unpriced_requests = unpriced_requests + excluded.unpriced_requests,
|
||||
updated_at = excluded.updated_at
|
||||
`).run(createId("puw"), orgId, providerId, window.start, window.end, inputCharCount, outputCharCount, now, now);
|
||||
`).run(
|
||||
createId("puw"), orgId, providerId, window.start, window.end, inputCharCount, outputCharCount, now, now,
|
||||
usage?.inputTokens ?? 0, usage?.outputTokens ?? 0, usage?.cachedInputTokens ?? 0, usage?.reasoningTokens ?? 0,
|
||||
spend?.costCents ?? 0, spend ? 0 : 1, spend && spend.costCents === null ? 1 : 0,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ export class ProviderResponseObserver {
|
||||
});
|
||||
// A late error on a request already cancelled or finished changes nothing.
|
||||
if (!failed) return;
|
||||
this.requestStore.recordUsage(input.orgId, input.provider.id, input.inputCharCount, 0);
|
||||
this.requestStore.recordUsage(input.orgId, input.provider.id, input.inputCharCount, 0, null);
|
||||
this.recordProviderHealth(input.provider.id, input.orgId, HEALTH_STATE_FAILED, message);
|
||||
this.operations.recordSignal({
|
||||
orgId: input.orgId,
|
||||
|
||||
@@ -17,7 +17,7 @@ import { type AgentTurnMessage, type AgentTurnResult, type ToolDefinition, type
|
||||
import { selectToolTurn, supportsToolTurn } from "./provider-tool-turns/router";
|
||||
import { resolveTurnCredential } from "./provider-tool-turns/credential-resolver";
|
||||
import { currentRunId } from "./run-control";
|
||||
import { budgetBlock, isPriced, providerPrices, readTokenUsage, requestCostCents } from "./spend";
|
||||
import { budgetBlock, isPriced, providerPrices } from "./spend";
|
||||
import {
|
||||
HEALTH_STATE_HEALTHY,
|
||||
WORKSPACE_SCOPES,
|
||||
@@ -49,9 +49,6 @@ export interface ProviderResponseRunnerDeps {
|
||||
}
|
||||
|
||||
export class ProviderResponseRunner {
|
||||
/** Providers already logged as not reporting token usage. */
|
||||
private readonly unreportedUsage = new Set<string>();
|
||||
|
||||
constructor(private readonly deps: ProviderResponseRunnerDeps) {}
|
||||
|
||||
async createResponse(input: ProviderResponseCreateInput): Promise<string> {
|
||||
@@ -87,7 +84,10 @@ export class ProviderResponseRunner {
|
||||
startedAt,
|
||||
action: "response_created",
|
||||
});
|
||||
this.deps.requestStore.recordUsage(input.orgId, execution.provider.id, execution.prompt.length, result.text.length);
|
||||
this.deps.requestStore.settleUsage({
|
||||
orgId: input.orgId, requestId, provider: execution.provider, rawResponseJson: result.rawResponseJson,
|
||||
inputCharCount: execution.prompt.length, outputCharCount: result.text.length,
|
||||
});
|
||||
return requestId;
|
||||
} catch (error) {
|
||||
this.deps.responseObserver.recordFailure({
|
||||
@@ -163,8 +163,10 @@ export class ProviderResponseRunner {
|
||||
startedAt,
|
||||
action: "response_created",
|
||||
});
|
||||
this.deps.requestStore.recordUsage(input.orgId, provider.id, promptChars, result.text.length);
|
||||
this.recordSpend(requestId, owner, result.rawResponseJson);
|
||||
this.deps.requestStore.settleUsage({
|
||||
orgId: input.orgId, requestId, provider: owner, rawResponseJson: result.rawResponseJson,
|
||||
inputCharCount: promptChars, outputCharCount: result.text.length,
|
||||
});
|
||||
return { ...result, requestId };
|
||||
} catch (error) {
|
||||
this.deps.responseObserver.recordFailure({
|
||||
@@ -180,16 +182,6 @@ export class ProviderResponseRunner {
|
||||
}
|
||||
}
|
||||
|
||||
private recordSpend(requestId: string, provider: { id: string; provider_key: string }, rawResponseJson: string): void {
|
||||
const usage = readTokenUsage(rawResponseJson);
|
||||
if (!usage) {
|
||||
if (!this.unreportedUsage.has(provider.provider_key)) console.warn(`[spend] ${provider.provider_key} returned no token usage; its calls cost nothing on record.`);
|
||||
this.unreportedUsage.add(provider.provider_key);
|
||||
return;
|
||||
}
|
||||
this.deps.requestStore.recordSpend(requestId, usage, requestCostCents(usage, providerPrices(this.deps.db, provider.id)));
|
||||
}
|
||||
|
||||
async streamResponse(input: StreamProviderResponseInput, onEvent: (event: ProviderStreamRendererEvent) => void): Promise<string> {
|
||||
const execution = this.prepareExecution(input);
|
||||
const requestId = this.createRequestRecord(execution);
|
||||
@@ -209,6 +201,7 @@ export class ProviderResponseRunner {
|
||||
for await (const event of result.events) {
|
||||
if (event.type === "error") throw new Error(event.textDelta ?? "Provider stream failed.");
|
||||
rawEvents.push(event.rawEventJson);
|
||||
if (event.type === "usage") continue;
|
||||
if (event.providerResponseId) providerResponseId = event.providerResponseId;
|
||||
if (event.textDelta) textParts.push(event.textDelta);
|
||||
const responseText = textParts.join("");
|
||||
@@ -258,7 +251,10 @@ export class ProviderResponseRunner {
|
||||
startedAt,
|
||||
action: "response_streamed",
|
||||
});
|
||||
this.deps.requestStore.recordUsage(input.orgId, execution.provider.id, execution.prompt.length, responseText.length);
|
||||
this.deps.requestStore.settleUsage({
|
||||
orgId: input.orgId, requestId, provider: execution.provider, rawResponseJson: JSON.stringify(rawEvents),
|
||||
inputCharCount: execution.prompt.length, outputCharCount: responseText.length,
|
||||
});
|
||||
return requestId;
|
||||
} catch (error) {
|
||||
this.deps.responseObserver.recordFailure({
|
||||
|
||||
@@ -3,7 +3,8 @@ import type { JsonRecord, ProviderEventStreamMessage } from "./provider-http";
|
||||
import { isRecord } from "./provider-http";
|
||||
|
||||
export interface ProviderStreamEvent {
|
||||
type: "message_start" | "text_delta" | "message_stop" | "error";
|
||||
/** usage: a chunk that only reports token usage; it is recorded, never shown. */
|
||||
type: "message_start" | "text_delta" | "message_stop" | "error" | "usage";
|
||||
textDelta: string | null;
|
||||
providerResponseId: string | null;
|
||||
rawEventJson: string;
|
||||
@@ -71,6 +72,8 @@ export async function* mapOpenAiChatCompletionStream(messages: AsyncIterable<Pro
|
||||
yield streamEvent("message_start", null, message);
|
||||
}
|
||||
yield streamEvent("message_stop", null, message);
|
||||
} else if (!delta && isRecord(message.data.usage)) {
|
||||
yield streamEvent("usage", null, message);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -84,6 +87,11 @@ export async function* mapAnthropicMessageStream(messages: AsyncIterable<Provide
|
||||
yield streamEvent("message_start", null, message);
|
||||
continue;
|
||||
}
|
||||
// Anthropic reports output tokens on message_delta, input tokens on message_start.
|
||||
if (type === "message_delta") {
|
||||
yield streamEvent("usage", null, message);
|
||||
continue;
|
||||
}
|
||||
if (type === "content_block_delta") {
|
||||
const delta = readAnthropicTextDelta(message.data);
|
||||
if (delta) {
|
||||
|
||||
@@ -19,12 +19,15 @@ export interface TokenPrices {
|
||||
/** Reads token usage from a provider's raw response: one JSON body, or the list of a stream's event payloads. */
|
||||
export function readTokenUsage(rawResponseJson: string): TokenUsage | null {
|
||||
const parsed = JSON.parse(rawResponseJson) as unknown;
|
||||
const bodies = Array.isArray(parsed) ? parsed.map(parseEvent) : [parsed];
|
||||
for (let index = bodies.length - 1; index >= 0; index -= 1) {
|
||||
const usage = usageIn(bodies[index]);
|
||||
if (usage) return usage;
|
||||
}
|
||||
return null;
|
||||
const reports = (Array.isArray(parsed) ? parsed.map(parseEvent) : [parsed])
|
||||
.map(usageIn)
|
||||
.filter((usage): usage is TokenUsage => usage !== null);
|
||||
if (reports.length === 0) return null;
|
||||
// A stream may report input early and output late (Anthropic): the last report wins, input taken from the last one that has it.
|
||||
const last = reports[reports.length - 1];
|
||||
if (last.inputTokens > 0) return last;
|
||||
const withInput = [...reports].reverse().find((usage) => usage.inputTokens > 0);
|
||||
return withInput ? { ...last, inputTokens: withInput.inputTokens, cachedInputTokens: withInput.cachedInputTokens } : last;
|
||||
}
|
||||
|
||||
/** True when calls to this provider cost money on record. */
|
||||
@@ -94,7 +97,9 @@ function usageIn(body: unknown): TokenUsage | null {
|
||||
if (!isRecord(body)) return null;
|
||||
const nested = isRecord(body.response) ? body.response : null;
|
||||
const firstChoice = Array.isArray(body.choices) && isRecord(body.choices[0]) ? body.choices[0] : null;
|
||||
const openAiStyle = record(body.usage) ?? record(nested?.usage) ?? record(isRecord(body.x_groq) ? body.x_groq.usage : null) ?? record(firstChoice?.usage);
|
||||
const message = isRecord(body.message) ? body.message : null;
|
||||
const openAiStyle = record(body.usage) ?? record(nested?.usage) ?? record(message?.usage)
|
||||
?? record(isRecord(body.x_groq) ? body.x_groq.usage : null) ?? record(firstChoice?.usage);
|
||||
if (openAiStyle) return fromOpenAiOrAnthropic(openAiStyle);
|
||||
const gemini = record(body.usageMetadata) ?? record(nested?.usageMetadata);
|
||||
if (gemini) return fromGemini(gemini);
|
||||
@@ -120,11 +125,11 @@ function fromOpenAiOrAnthropic(usage: Record<string, unknown>): TokenUsage | nul
|
||||
reasoningTokens: count(record(usage.completion_tokens_details)?.reasoning_tokens),
|
||||
};
|
||||
}
|
||||
if (typeof usage.input_tokens !== "number") return null;
|
||||
if (typeof usage.input_tokens !== "number" && typeof usage.output_tokens !== "number") return null;
|
||||
// Anthropic reports cache reads and writes beside input_tokens; Responses reports them inside it.
|
||||
const anthropicCache = count(usage.cache_read_input_tokens) + count(usage.cache_creation_input_tokens);
|
||||
return {
|
||||
inputTokens: usage.input_tokens + anthropicCache,
|
||||
inputTokens: count(usage.input_tokens) + anthropicCache,
|
||||
outputTokens: count(usage.output_tokens),
|
||||
cachedInputTokens: count(usage.cache_read_input_tokens) + count(record(usage.input_tokens_details)?.cached_tokens),
|
||||
reasoningTokens: count(record(usage.output_tokens_details)?.reasoning_tokens),
|
||||
|
||||
@@ -42,7 +42,7 @@ export interface ProviderStreamRendererEvent {
|
||||
providerId: string;
|
||||
providerKey: string;
|
||||
model: string;
|
||||
streamEventType: "message_start" | "text_delta" | "message_stop" | "error";
|
||||
streamEventType: "message_start" | "text_delta" | "message_stop" | "error" | "usage";
|
||||
textDelta: string | null;
|
||||
responseText: string;
|
||||
providerResponseId: string | null;
|
||||
|
||||
@@ -24,7 +24,7 @@ describe("domain services", () => {
|
||||
const connection = createConnection();
|
||||
|
||||
const rows = connection.db.prepare("select version from schema_migrations order by version asc").all() as Array<{ version: number }>;
|
||||
expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72]);
|
||||
expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73]);
|
||||
connection.close();
|
||||
});
|
||||
|
||||
|
||||
@@ -48,6 +48,7 @@ describe("provider streaming", () => {
|
||||
message(null, { choices: [{ delta: { content: "Hel" } }] }),
|
||||
message(null, { choices: [{ delta: { content: "lo" } }] }),
|
||||
message(null, { choices: [{ delta: {}, finish_reason: "stop" }] }),
|
||||
message(null, { choices: [], usage: { prompt_tokens: 12, completion_tokens: 3 } }),
|
||||
])));
|
||||
|
||||
expect(events.map((event) => [event.type, event.textDelta])).toEqual([
|
||||
@@ -55,6 +56,7 @@ describe("provider streaming", () => {
|
||||
["text_delta", "Hel"],
|
||||
["text_delta", "lo"],
|
||||
["message_stop", null],
|
||||
["usage", null],
|
||||
]);
|
||||
});
|
||||
|
||||
@@ -63,6 +65,7 @@ describe("provider streaming", () => {
|
||||
message("message_start", { type: "message_start" }),
|
||||
message("content_block_delta", { type: "content_block_delta", delta: { type: "text_delta", text: "Hel" } }),
|
||||
message("content_block_delta", { type: "content_block_delta", delta: { type: "text_delta", text: "lo" } }),
|
||||
message("message_delta", { type: "message_delta", usage: { output_tokens: 7 } }),
|
||||
message("message_stop", { type: "message_stop" }),
|
||||
])));
|
||||
|
||||
@@ -70,6 +73,7 @@ describe("provider streaming", () => {
|
||||
["message_start", null],
|
||||
["text_delta", "Hel"],
|
||||
["text_delta", "lo"],
|
||||
["usage", null],
|
||||
["message_stop", null],
|
||||
]);
|
||||
});
|
||||
|
||||
@@ -17,6 +17,13 @@ describe("token usage and cost", () => {
|
||||
expect(readTokenUsage(JSON.stringify(responsesStream))).toEqual({ inputTokens: 10, outputTokens: 5, cachedInputTokens: 0, reasoningTokens: 0 });
|
||||
const moonshotStream = ["{\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\",\"usage\":{\"prompt_tokens\":40,\"completion_tokens\":9}}]}"];
|
||||
expect(readTokenUsage(JSON.stringify(moonshotStream))).toEqual({ inputTokens: 40, outputTokens: 9, cachedInputTokens: 0, reasoningTokens: 0 });
|
||||
const anthropicStream = [
|
||||
"{\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":700,\"cache_read_input_tokens\":300,\"output_tokens\":1}}}",
|
||||
"{\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}",
|
||||
"{\"type\":\"message_delta\",\"usage\":{\"output_tokens\":42}}",
|
||||
"{\"type\":\"message_stop\"}",
|
||||
];
|
||||
expect(readTokenUsage(JSON.stringify(anthropicStream))).toEqual({ inputTokens: 1000, outputTokens: 42, cachedInputTokens: 300, reasoningTokens: 0 });
|
||||
expect(readTokenUsage(JSON.stringify({ output: [] }))).toBeNull();
|
||||
});
|
||||
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
import { afterEach, describe, expect, test, vi } from "vitest";
|
||||
import { providerAdapters } from "../../src/main/services/provider-adapters";
|
||||
import { makeChatCompletionsTurn } from "../../src/main/services/provider-tool-turns/turns";
|
||||
|
||||
const bodies: Array<Record<string, unknown>> = [];
|
||||
|
||||
function sseResponse(): Response {
|
||||
const text = "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
|
||||
return new Response(text, { status: 200, headers: { "content-type": "text/event-stream" } });
|
||||
}
|
||||
|
||||
afterEach(() => {
|
||||
vi.unstubAllGlobals();
|
||||
bodies.length = 0;
|
||||
});
|
||||
|
||||
describe("stream usage requests", () => {
|
||||
test("xAI and Groq streams ask for token usage; Mistral and Cerebras streams do not", async () => {
|
||||
vi.stubGlobal("fetch", async (_url: string, init: RequestInit) => {
|
||||
bodies.push(JSON.parse(String(init.body)) as Record<string, unknown>);
|
||||
return sseResponse();
|
||||
});
|
||||
for (const key of ["xai", "groq", "mistral", "cerebras"] as const) {
|
||||
const result = await providerAdapters[key].streamResponse({ authProfile: { type: "api_key", apiKey: "k" }, model: "m", prompt: "p", instructions: null });
|
||||
for await (const event of result.events) void event;
|
||||
}
|
||||
expect(bodies.map((body) => body.stream_options ?? null)).toEqual([{ include_usage: true }, { include_usage: true }, null, null]);
|
||||
|
||||
await makeChatCompletionsTurn("https://api.x.ai/v1/chat/completions", "xAI", true)({ credential: { apiKey: "k" }, model: "m", instructions: null, messages: [{ role: "user", text: "hi" }], tools: [] });
|
||||
expect(bodies.at(-1)?.stream_options).toEqual({ include_usage: true });
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,44 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { describe, expect, test } from "vitest";
|
||||
import { DatabaseConnection } from "../../src/main/database/connection";
|
||||
import { OrganizationService } from "../../src/main/services/organization-service";
|
||||
import { ProviderRequestStore } from "../../src/main/services/provider-request-store";
|
||||
import { PROVIDER_EXECUTION_MODE } from "../../src/shared/domain";
|
||||
|
||||
describe("usage settlement", () => {
|
||||
test("every settled call records real tokens and cost on the request and in the month's totals", () => {
|
||||
const connection = new DatabaseConnection(join(tmpdir(), `ac-settle-${randomUUID()}.db`));
|
||||
const organization = new OrganizationService(connection.db);
|
||||
let orgId = "";
|
||||
connection.transaction(() => {
|
||||
orgId = organization.createOrganization({ organizationName: "Acme", employeeName: "Dana", mission: "m", goal: "g" });
|
||||
});
|
||||
const provider = connection.db.prepare("select id, provider_key, data_policy, workspace_scope from providers where org_id = ? and provider_key = 'anthropic'").get(orgId) as
|
||||
{ id: string; provider_key: string; data_policy: string; workspace_scope: string };
|
||||
connection.db.prepare("update providers set input_cents_per_mtok = 300, cached_input_cents_per_mtok = 30, output_cents_per_mtok = 1500 where id = ?").run(provider.id);
|
||||
const store = new ProviderRequestStore(connection.db);
|
||||
const call = () => store.create({
|
||||
orgId, provider: provider as never, model: "claude", workspaceId: null, promptSha256: "sha", inputCharCount: 40,
|
||||
executionMode: PROVIDER_EXECUTION_MODE.FOREGROUND, createdAt: new Date().toISOString(),
|
||||
});
|
||||
const priced = call();
|
||||
store.settleUsage({
|
||||
orgId, requestId: priced, provider, inputCharCount: 40, outputCharCount: 12,
|
||||
rawResponseJson: JSON.stringify({ usage: { input_tokens: 600_000, output_tokens: 100_000, cache_read_input_tokens: 400_000 } }),
|
||||
});
|
||||
const silent = call();
|
||||
store.settleUsage({ orgId, requestId: silent, provider, inputCharCount: 10, outputCharCount: 3, rawResponseJson: JSON.stringify({ content: [] }) });
|
||||
|
||||
expect(connection.db.prepare("select input_tokens, output_tokens, cached_input_tokens, cost_cents from provider_requests where id = ?").get(priced))
|
||||
.toEqual({ input_tokens: 1_000_000, output_tokens: 100_000, cached_input_tokens: 400_000, cost_cents: 180 + 12 + 150 });
|
||||
expect(connection.db.prepare(`
|
||||
select request_count, input_tokens, output_tokens, cached_input_tokens, cost_cents, requests_without_usage, unpriced_requests
|
||||
from provider_usage_windows where provider_id = ?
|
||||
`).get(provider.id)).toEqual({
|
||||
request_count: 2, input_tokens: 1_000_000, output_tokens: 100_000, cached_input_tokens: 400_000, cost_cents: 342, requests_without_usage: 1, unpriced_requests: 0,
|
||||
});
|
||||
connection.close();
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user