Record what every model call spends, and show it
Nothing anywhere recorded what a model call cost: the harness used a provider's usage numbers only to decide whether a turn had finished, then dropped them. So a reply could not say what it took to make, one model could not be compared with another on price, and there was no way to ask what a day, a week or a month of the crew's work came to. - `token_usages` records one row per call: user, bot, workspace, session, job and message, the harness, the provider and model that actually answered, the prompt and completion counts, and the cost when the provider reports one. - `harness-usage-scope.ts` carries those ids in an AsyncLocalStorage scope that `run-bot-job.ts` enters once, so a run's turns, its session summary, its transcriptions, images and embeddings are all attributed without threading ids through every harness function. Nested runs merge rather than replace, and a null genuinely clears a field (a bot's own chat has no workspace). - Work with no run behind it is charged to its owner instead: `run-with-fallback-scope.ts` reaches the workspace's creator else the bot's, and deliberately does not override a real run, so an ingest a bot asked for is not misattributed to the bot's owner. - All six harnesses record through `withCandidates`'s `on_success`. Recording never throws: a model that answered must not lose its answer to a failed insert. - A job's calls land before its message exists, so `complete-bot-job.ts` back-fills `message_id` once the message is inserted. - The message's actions row shows what it cost, beside the timestamp and the play button, with the breakdown in its title. - `/admin/usage` totals the install's spending for 7, 30 or 90 days: the total against the previous period, a day-by-day chart, a split per kind of call, and a table per model with the price per thousand tokens worked out from cost and tokens, which is what makes two models comparable. Crew managers only, since it is the whole install's work rather than one user's conversation. It sits with Routines and Activity in the user popover rather than under Settings, because it is a record of what happened, not something to configure. Days are grouped in local time (`date(created_at / 1000, 'unixepoch', 'localtime')`), matching the local midnights the window builds and the `YYYY-MM-DD` buckets it looks up. Grouping by the bare UTC day put a late evening's calls under a name no bucket carried, so they vanished and the chart's total silently disagreed with the totals card; a test pins it. `canSeeRailItem` now decides who sees a popover or rail entry. The two each filtered with their own inline rule and only one understood `can_view`, so a gated page could be offered in the popover to everyone.
This commit is contained in:
@@ -2,7 +2,9 @@ import grabHarnessSetting from "@/src/functions/backend/db/harness/grab-harness-
|
||||
import grabHarnessAdapter from "../providers/grab-harness-adapter";
|
||||
import resolveHarnessCandidates from "../resolve/resolve-harness-candidates";
|
||||
import withCandidates from "../resolve/with-candidates";
|
||||
import type { HarnessCandidate } from "../types";
|
||||
import type { HarnessAudioStream, HarnessCandidate } from "../types";
|
||||
import buildUsageCall from "../usage/build-usage-call";
|
||||
import recordHarnessUsage from "../usage/record-harness-usage";
|
||||
|
||||
type Params = {
|
||||
text: string;
|
||||
@@ -39,9 +41,22 @@ export default async function speakText({
|
||||
grabHarnessSetting({ harness: "speech", bot_id, workspace_id }),
|
||||
]);
|
||||
|
||||
const { result } = await withCandidates({
|
||||
const { result } = await withCandidates<HarnessAudioStream>({
|
||||
harness: "speech",
|
||||
candidates,
|
||||
/**
|
||||
* Speech is billed per character and streamed, so there is no
|
||||
* completion count to read: the call is recorded as it is, with
|
||||
* the text's length standing in for what was spoken.
|
||||
*/
|
||||
on_success: async ({ candidate }) =>
|
||||
await recordHarnessUsage({
|
||||
call: buildUsageCall({
|
||||
harness: "speech",
|
||||
candidate,
|
||||
usage: { prompt_tokens: text.length, completion_tokens: 0 },
|
||||
}),
|
||||
}),
|
||||
run: async (candidate) =>
|
||||
await grabHarnessAdapter({ kind: candidate.provider.kind }).speak({
|
||||
candidate,
|
||||
|
||||
@@ -4,7 +4,9 @@ import HarnessError from "../harness-error";
|
||||
import grabHarnessAdapter from "../providers/grab-harness-adapter";
|
||||
import resolveHarnessCandidates from "../resolve/resolve-harness-candidates";
|
||||
import withCandidates from "../resolve/with-candidates";
|
||||
import type { HarnessAudioInput } from "../types";
|
||||
import type { HarnessAudioInput, HarnessUsage } from "../types";
|
||||
import buildUsageCall from "../usage/build-usage-call";
|
||||
import recordHarnessUsage from "../usage/record-harness-usage";
|
||||
|
||||
type Params = {
|
||||
audio: HarnessAudioInput;
|
||||
@@ -46,15 +48,23 @@ export default async function transcribeAudio({
|
||||
|
||||
const final_signal = signal || new AbortController().signal;
|
||||
|
||||
const { result } = await withCandidates({
|
||||
const { result } = await withCandidates<{ text: string; usage?: HarnessUsage }>({
|
||||
harness: "transcription",
|
||||
candidates,
|
||||
on_success: async ({ candidate, result: transcript }): Promise<void> =>
|
||||
await recordHarnessUsage({
|
||||
call: buildUsageCall({
|
||||
harness: "transcription",
|
||||
candidate,
|
||||
usage: transcript.usage,
|
||||
}),
|
||||
}),
|
||||
run: async (candidate) => {
|
||||
const adapter = grabHarnessAdapter({
|
||||
kind: candidate.provider.kind,
|
||||
});
|
||||
|
||||
const { text } = await adapter.transcribe({
|
||||
const { text, usage } = await adapter.transcribe({
|
||||
candidate,
|
||||
audio,
|
||||
language,
|
||||
@@ -69,9 +79,9 @@ export default async function transcribeAudio({
|
||||
});
|
||||
}
|
||||
|
||||
return text.trim();
|
||||
return { text: text.trim(), usage };
|
||||
},
|
||||
});
|
||||
|
||||
return result;
|
||||
return result.text;
|
||||
}
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import grabHarnessAdapter from "../providers/grab-harness-adapter";
|
||||
import resolveHarnessCandidates from "../resolve/resolve-harness-candidates";
|
||||
import withCandidates from "../resolve/with-candidates";
|
||||
import type { HarnessEmbedResult } from "../types";
|
||||
import buildUsageCall from "../usage/build-usage-call";
|
||||
import recordHarnessUsage from "../usage/record-harness-usage";
|
||||
import fitVectorDimensions from "./fit-vector-dimensions";
|
||||
|
||||
type Params = {
|
||||
@@ -52,9 +55,17 @@ export default async function embedText({
|
||||
workspace_id,
|
||||
});
|
||||
|
||||
const { result } = await withCandidates({
|
||||
const { result } = await withCandidates<HarnessEmbedResult>({
|
||||
harness: "embedding",
|
||||
candidates,
|
||||
on_success: async ({ candidate, result: embedded }): Promise<void> =>
|
||||
await recordHarnessUsage({
|
||||
call: buildUsageCall({
|
||||
harness: "embedding",
|
||||
candidate,
|
||||
usage: embedded.usage,
|
||||
}),
|
||||
}),
|
||||
run: async (candidate) =>
|
||||
await grabHarnessAdapter({ kind: candidate.provider.kind }).embed({
|
||||
candidate,
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import grabHarnessAdapter from "../providers/grab-harness-adapter";
|
||||
import resolveHarnessCandidates from "../resolve/resolve-harness-candidates";
|
||||
import withCandidates from "../resolve/with-candidates";
|
||||
import type { HarnessImageOutput, HarnessUsage } from "../types";
|
||||
import buildUsageCall from "../usage/build-usage-call";
|
||||
import recordHarnessUsage from "../usage/record-harness-usage";
|
||||
|
||||
type Params = {
|
||||
prompt: string;
|
||||
@@ -35,9 +38,17 @@ export default async function generateImage({
|
||||
}: Params) {
|
||||
const candidates = await resolveHarnessCandidates({ harness: "image", bot_id, workspace_id });
|
||||
|
||||
const { result } = await withCandidates({
|
||||
const { result } = await withCandidates<{ images: HarnessImageOutput[]; usage?: HarnessUsage }>({
|
||||
harness: "image",
|
||||
candidates,
|
||||
on_success: async ({ candidate, result: made }): Promise<void> =>
|
||||
await recordHarnessUsage({
|
||||
call: buildUsageCall({
|
||||
harness: "image",
|
||||
candidate,
|
||||
usage: made.usage,
|
||||
}),
|
||||
}),
|
||||
run: async (candidate) =>
|
||||
await grabHarnessAdapter({
|
||||
kind: candidate.provider.kind,
|
||||
|
||||
@@ -3,7 +3,10 @@ import HarnessError from "../harness-error";
|
||||
import grabHarnessAdapter from "../providers/grab-harness-adapter";
|
||||
import resolveHarnessCandidates from "../resolve/resolve-harness-candidates";
|
||||
import withCandidates from "../resolve/with-candidates";
|
||||
import type { HarnessUsage } from "../types";
|
||||
import grabChatParams from "../text/grab-chat-params";
|
||||
import buildUsageCall from "../usage/build-usage-call";
|
||||
import recordHarnessUsage from "../usage/record-harness-usage";
|
||||
|
||||
type Params = {
|
||||
/**
|
||||
@@ -42,9 +45,17 @@ export default async function readImage({ url, focus, signal, bot_id, workspace_
|
||||
workspace_id,
|
||||
});
|
||||
|
||||
const { result } = await withCandidates({
|
||||
const { result } = await withCandidates<{ text: string; usage?: HarnessUsage }>({
|
||||
harness: "vision",
|
||||
candidates,
|
||||
on_success: async ({ candidate, result: answer }): Promise<void> =>
|
||||
await recordHarnessUsage({
|
||||
call: buildUsageCall({
|
||||
harness: "vision",
|
||||
candidate,
|
||||
usage: answer.usage,
|
||||
}),
|
||||
}),
|
||||
run: async (candidate) => {
|
||||
const answer = await grabHarnessAdapter({
|
||||
kind: candidate.provider.kind,
|
||||
@@ -83,9 +94,9 @@ export default async function readImage({ url, focus, signal, bot_id, workspace_
|
||||
});
|
||||
}
|
||||
|
||||
return answer.text.trim();
|
||||
return { text: answer.text.trim(), usage: answer.usage };
|
||||
},
|
||||
});
|
||||
|
||||
return result;
|
||||
return result.text;
|
||||
}
|
||||
|
||||
@@ -13,6 +13,11 @@ const embed: HarnessAdapter["embed"] = async ({ candidate, input, signal }) => {
|
||||
const { provider, model } = candidate;
|
||||
const preset = grabProviderPreset({ kind: provider.kind });
|
||||
const vectors: number[][] = [];
|
||||
/**
|
||||
* What the provider said the batch cost. Most report it once per
|
||||
* request; the rows are opened before the loop so each batch counts.
|
||||
*/
|
||||
let prompt_tokens = 0;
|
||||
|
||||
/**
|
||||
* Every OpenAI-compatible API puts embeddings at `/embeddings`.
|
||||
@@ -47,6 +52,8 @@ const embed: HarnessAdapter["embed"] = async ({ candidate, input, signal }) => {
|
||||
});
|
||||
}
|
||||
|
||||
prompt_tokens += Number(json?.usage?.prompt_tokens) || 0;
|
||||
|
||||
for (const item of items) {
|
||||
const vector = Array.isArray(item) ? item : item?.embedding;
|
||||
|
||||
@@ -62,7 +69,12 @@ const embed: HarnessAdapter["embed"] = async ({ candidate, input, signal }) => {
|
||||
}
|
||||
}
|
||||
|
||||
return { vectors };
|
||||
return {
|
||||
vectors,
|
||||
...(prompt_tokens
|
||||
? { usage: { prompt_tokens, completion_tokens: 0 } }
|
||||
: {}),
|
||||
};
|
||||
};
|
||||
|
||||
export default embed;
|
||||
|
||||
@@ -81,7 +81,20 @@ const generateImage: HarnessAdapter["generateImage"] = async ({
|
||||
});
|
||||
}
|
||||
|
||||
return { images };
|
||||
/**
|
||||
* Image models bill per image, not per token: only a provider that
|
||||
* volunteers a usage figure is recorded with one.
|
||||
*/
|
||||
const usage =
|
||||
json?.usage && typeof json.usage == "object"
|
||||
? {
|
||||
prompt_tokens: Number(json.usage.prompt_tokens) || 0,
|
||||
completion_tokens: Number(json.usage.completion_tokens) || 0,
|
||||
...(typeof json.usage.cost == "number" ? { cost: json.usage.cost } : {}),
|
||||
}
|
||||
: undefined;
|
||||
|
||||
return { images, ...(usage ? { usage } : {}) };
|
||||
};
|
||||
|
||||
export default generateImage;
|
||||
|
||||
@@ -50,7 +50,22 @@ const transcribe: HarnessAdapter["transcribe"] = async ({
|
||||
});
|
||||
}
|
||||
|
||||
return { text: String(json?.text || "").trim() };
|
||||
/**
|
||||
* Providers that bill by audio seconds sometimes report a token count;
|
||||
* the OpenAI shape puts it under `usage`. Left out when absent, so the
|
||||
* row records that the call happened without a token count.
|
||||
*/
|
||||
const usage =
|
||||
json?.usage && typeof json.usage == "object"
|
||||
? {
|
||||
prompt_tokens: Number(json.usage.prompt_tokens ?? json.usage.input_tokens) || 0,
|
||||
completion_tokens:
|
||||
Number(json.usage.completion_tokens ?? json.usage.output_tokens) || 0,
|
||||
...(typeof json.usage.cost == "number" ? { cost: json.usage.cost } : {}),
|
||||
}
|
||||
: undefined;
|
||||
|
||||
return { text: String(json?.text || "").trim(), ...(usage ? { usage } : {}) };
|
||||
};
|
||||
|
||||
export default transcribe;
|
||||
|
||||
@@ -14,6 +14,15 @@ type Params<T> = {
|
||||
* uses, so it doesn't switch models mid-task without a reason.
|
||||
*/
|
||||
prefer_key?: string;
|
||||
/**
|
||||
* Called with what the candidate that worked returned, before the call
|
||||
* comes back to the caller: where a harness writes its usage down.
|
||||
* Never fails the call (the caller passes a recorder that can't throw).
|
||||
*/
|
||||
on_success?: (params: {
|
||||
candidate: HarnessCandidate;
|
||||
result: T;
|
||||
}) => Promise<void> | void;
|
||||
cooldowns?: HarnessCooldowns;
|
||||
now?: () => number;
|
||||
};
|
||||
@@ -42,6 +51,7 @@ export default async function withCandidates<T>({
|
||||
candidates,
|
||||
run,
|
||||
prefer_key,
|
||||
on_success,
|
||||
cooldowns = DefaultCooldowns,
|
||||
now = Date.now,
|
||||
}: Params<T>): Promise<{ result: T; candidate: HarnessCandidate }> {
|
||||
@@ -93,7 +103,11 @@ export default async function withCandidates<T>({
|
||||
if (!isReady(candidate)) continue;
|
||||
|
||||
try {
|
||||
return { result: await run(candidate), candidate };
|
||||
const result = await run(candidate);
|
||||
|
||||
await on_success?.({ candidate, result });
|
||||
|
||||
return { result, candidate };
|
||||
} catch (error: unknown) {
|
||||
if (!(error instanceof HarnessError)) throw error;
|
||||
if (error.code == "aborted") throw error;
|
||||
|
||||
@@ -8,6 +8,8 @@ import type {
|
||||
HarnessMessage,
|
||||
HarnessToolDefinition,
|
||||
} from "../types";
|
||||
import buildUsageCall from "../usage/build-usage-call";
|
||||
import recordHarnessUsage from "../usage/record-harness-usage";
|
||||
import grabChatParams from "./grab-chat-params";
|
||||
import prepareMessagesForModel from "./prepare-messages-for-model";
|
||||
import stripThinkTags from "./strip-think-tags";
|
||||
@@ -47,6 +49,14 @@ export type GenerateTextResult = HarnessChatResult & {
|
||||
candidate: HarnessCandidate;
|
||||
};
|
||||
|
||||
/**
|
||||
* One turn as `withCandidates` carries it: the answer the provider gave,
|
||||
* with the reasoning tags stripped and an empty answer turned into null.
|
||||
*/
|
||||
type ChatAnswer = Omit<HarnessChatResult, "text"> & {
|
||||
text: string | null;
|
||||
};
|
||||
|
||||
/**
|
||||
* One model turn on the text harness, falling back across its models.
|
||||
* Media the model can't take is read by the helper harnesses first.
|
||||
@@ -71,10 +81,18 @@ export default async function generateText({
|
||||
? undefined
|
||||
: await grabChatParams({ harness: "text", bot_id, workspace_id });
|
||||
|
||||
const { result, candidate } = await withCandidates({
|
||||
const { result, candidate } = await withCandidates<ChatAnswer>({
|
||||
harness: "text",
|
||||
candidates: resolved,
|
||||
prefer_key,
|
||||
on_success: async ({ candidate: answered, result: answer }): Promise<void> =>
|
||||
await recordHarnessUsage({
|
||||
call: buildUsageCall({
|
||||
harness: "text",
|
||||
candidate: answered,
|
||||
usage: answer.usage,
|
||||
}),
|
||||
}),
|
||||
run: async (candidate) => {
|
||||
const answer = await grabHarnessAdapter({
|
||||
kind: candidate.provider.kind,
|
||||
|
||||
+25
-2
@@ -138,6 +138,18 @@ export type HarnessUsage = {
|
||||
cost?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* What a harness does with a model, for its usage rows and the totals
|
||||
* that read them. Mirrors the `harness` column of `token_usages`.
|
||||
*/
|
||||
export type HarnessUsageKind =
|
||||
| "text"
|
||||
| "vision"
|
||||
| "transcription"
|
||||
| "speech"
|
||||
| "image"
|
||||
| "embedding";
|
||||
|
||||
export type HarnessChatParams = {
|
||||
candidate: HarnessCandidate;
|
||||
messages: HarnessMessage[];
|
||||
@@ -238,11 +250,22 @@ export type HarnessEmbedResult = {
|
||||
*/
|
||||
export type HarnessAdapter = {
|
||||
chat: (params: HarnessChatParams) => Promise<HarnessChatResult>;
|
||||
transcribe: (params: HarnessTranscribeParams) => Promise<{ text: string }>;
|
||||
/**
|
||||
* `usage` when the provider says how many tokens it heard, which
|
||||
* OpenAI-style speech-to-text endpoints do only sometimes.
|
||||
*/
|
||||
transcribe: (params: HarnessTranscribeParams) => Promise<{
|
||||
text: string;
|
||||
usage?: HarnessUsage;
|
||||
}>;
|
||||
speak: (params: HarnessSpeakParams) => Promise<HarnessAudioStream>;
|
||||
/**
|
||||
* Image models report what they made rather than how many tokens it
|
||||
* took, so `usage` is only whatever the provider volunteers.
|
||||
*/
|
||||
generateImage: (
|
||||
params: HarnessImageParams,
|
||||
) => Promise<{ images: HarnessImageOutput[] }>;
|
||||
) => Promise<{ images: HarnessImageOutput[]; usage?: HarnessUsage }>;
|
||||
embed: (params: HarnessEmbedParams) => Promise<HarnessEmbedResult>;
|
||||
listModels: (params: {
|
||||
provider: HarnessProviderConfig;
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import type { HarnessCandidate } from "../types";
|
||||
import buildUsageCall from "./build-usage-call";
|
||||
|
||||
function model(overrides: { row_id?: number } = {}) {
|
||||
return {
|
||||
slug: "deepseek/deepseek-chat",
|
||||
name: "DeepSeek Chat",
|
||||
input_modalities: ["text"],
|
||||
output_modalities: ["text"],
|
||||
supports_tools: true,
|
||||
...overrides,
|
||||
} as HarnessCandidate["model"];
|
||||
}
|
||||
|
||||
function candidate(overrides: Partial<HarnessCandidate> = {}): HarnessCandidate {
|
||||
return {
|
||||
provider: {
|
||||
id: 2,
|
||||
name: "OpenRouter",
|
||||
kind: "openrouter",
|
||||
} as HarnessCandidate["provider"],
|
||||
key: "openrouter:deepseek/deepseek-chat",
|
||||
source: "auto",
|
||||
model: model({ row_id: 41 }),
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe("buildUsageCall", () => {
|
||||
test("names the candidate that really answered and what it spent", () => {
|
||||
const call = buildUsageCall({
|
||||
harness: "text",
|
||||
candidate: candidate(),
|
||||
usage: { prompt_tokens: 1200, completion_tokens: 340, cost: 0.0042 },
|
||||
});
|
||||
|
||||
expect(call).toEqual({
|
||||
harness: "text",
|
||||
provider_id: 2,
|
||||
provider_name: "OpenRouter",
|
||||
model_id: 41,
|
||||
model_slug: "deepseek/deepseek-chat",
|
||||
prompt_tokens: 1200,
|
||||
completion_tokens: 340,
|
||||
cost: 0.0042,
|
||||
});
|
||||
});
|
||||
|
||||
test("a provider with no name reads as its kind", () => {
|
||||
const call = buildUsageCall({
|
||||
harness: "image",
|
||||
candidate: candidate({
|
||||
provider: { id: 4, name: "", kind: "openai" } as HarnessCandidate["provider"],
|
||||
}),
|
||||
usage: { prompt_tokens: 0, completion_tokens: 0 },
|
||||
});
|
||||
|
||||
expect(call.provider_name).toBe("openai");
|
||||
});
|
||||
|
||||
test("a provider that reports no usage still gets a row", () => {
|
||||
const call = buildUsageCall({ harness: "transcription", candidate: candidate() });
|
||||
|
||||
expect(call.prompt_tokens).toBe(0);
|
||||
expect(call.completion_tokens).toBe(0);
|
||||
expect(call.cost).toBeUndefined();
|
||||
});
|
||||
|
||||
test("a candidate with no catalog row has no model id to point at", () => {
|
||||
const call = buildUsageCall({
|
||||
harness: "text",
|
||||
candidate: candidate({
|
||||
model: model({ slug: "made-up-model" } as any),
|
||||
}),
|
||||
usage: { prompt_tokens: 10, completion_tokens: 5 },
|
||||
});
|
||||
|
||||
expect(call.model_id).toBeUndefined();
|
||||
expect(call.model_slug).toBe("made-up-model");
|
||||
});
|
||||
|
||||
test("a cost the provider did not report is left out, not zeroed", () => {
|
||||
const call = buildUsageCall({
|
||||
harness: "text",
|
||||
candidate: candidate(),
|
||||
usage: { prompt_tokens: 10, completion_tokens: 5, cost: Number.NaN },
|
||||
});
|
||||
|
||||
expect(call.cost).toBeUndefined();
|
||||
});
|
||||
|
||||
test("a free provider's zero cost is kept as zero", () => {
|
||||
const call = buildUsageCall({
|
||||
harness: "text",
|
||||
candidate: candidate(),
|
||||
usage: { prompt_tokens: 10, completion_tokens: 5, cost: 0 },
|
||||
});
|
||||
|
||||
expect(call.cost).toBe(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,37 @@
|
||||
import type { HarnessCandidate, HarnessUsage } from "../types";
|
||||
import type { HarnessUsageCall } from "./types";
|
||||
|
||||
type Params = {
|
||||
harness: HarnessUsageCall["harness"];
|
||||
candidate: HarnessCandidate;
|
||||
usage?: HarnessUsage;
|
||||
};
|
||||
|
||||
/**
|
||||
* The candidate that really answered a call, and what it reported, as the
|
||||
* row `recordHarnessUsage` writes. Providers that report no usage at all
|
||||
* still get a row (no tokens), so a call is never invisible.
|
||||
*/
|
||||
export default function buildUsageCall({ harness, candidate, usage }: Params): HarnessUsageCall {
|
||||
return {
|
||||
harness,
|
||||
provider_id: candidate.provider.id,
|
||||
provider_name: candidate.provider.name || candidate.provider.kind,
|
||||
model_id: modelIdOf(candidate),
|
||||
model_slug: candidate.model.slug,
|
||||
prompt_tokens: Number(usage?.prompt_tokens) || 0,
|
||||
completion_tokens: Number(usage?.completion_tokens) || 0,
|
||||
cost: typeof usage?.cost == "number" && Number.isFinite(usage.cost) ? usage.cost : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The `harness_models` row id, when the candidate carries the catalog row
|
||||
* it was resolved from (`withCandidates` doesn't drop it; a candidate made
|
||||
* by hand may not have one).
|
||||
*/
|
||||
function modelIdOf(candidate: HarnessCandidate) {
|
||||
const row_id = (candidate.model as { row_id?: number }).row_id;
|
||||
|
||||
return typeof row_id == "number" ? row_id : undefined;
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { grabHarnessUsageScope, runWithHarnessUsageScope } from "./harness-usage-scope";
|
||||
|
||||
describe("harness usage scope", () => {
|
||||
test("nothing is attributed outside a run", () => {
|
||||
expect(grabHarnessUsageScope()).toBeUndefined();
|
||||
});
|
||||
|
||||
test("a run's scope is visible to everything it calls", async () => {
|
||||
const scope = await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7, bot_id: 3, session_id: 11 },
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
});
|
||||
|
||||
expect(scope).toEqual({ user_id: 7, bot_id: 3, session_id: 11 });
|
||||
});
|
||||
|
||||
test("a nested run keeps what the outer one knew", async () => {
|
||||
const scope = await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7, bot_id: 3, session_id: 11, job_id: 5 },
|
||||
task: () =>
|
||||
runWithHarnessUsageScope({
|
||||
scope: { message_id: 42 },
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
}),
|
||||
});
|
||||
|
||||
expect(scope).toEqual({
|
||||
user_id: 7,
|
||||
bot_id: 3,
|
||||
session_id: 11,
|
||||
job_id: 5,
|
||||
message_id: 42,
|
||||
});
|
||||
});
|
||||
|
||||
test("a nested run overrides only what it sets", async () => {
|
||||
const scope = await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7, bot_id: 3, workspace_id: 9 },
|
||||
task: () =>
|
||||
runWithHarnessUsageScope({
|
||||
scope: { bot_id: 4 },
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
}),
|
||||
});
|
||||
|
||||
expect(scope).toEqual({ user_id: 7, bot_id: 4, workspace_id: 9 });
|
||||
});
|
||||
|
||||
test("an undefined field leaves the outer value alone", async () => {
|
||||
const scope = await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7, bot_id: 3, job_id: 5 },
|
||||
task: () =>
|
||||
runWithHarnessUsageScope({
|
||||
scope: { job_id: undefined, message_id: 42 },
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
}),
|
||||
});
|
||||
|
||||
expect(scope).toEqual({ user_id: 7, bot_id: 3, job_id: 5, message_id: 42 });
|
||||
});
|
||||
|
||||
test("a null field clears the outer value, since it means 'none'", async () => {
|
||||
const scope = await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7, bot_id: 3, workspace_id: 9, job_id: 5 },
|
||||
task: () =>
|
||||
runWithHarnessUsageScope({
|
||||
scope: { workspace_id: null, job_id: null },
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
}),
|
||||
});
|
||||
|
||||
expect(scope).toEqual({ user_id: 7, bot_id: 3, workspace_id: null, job_id: null });
|
||||
});
|
||||
|
||||
test("the scope is gone once the run returns", async () => {
|
||||
await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7 },
|
||||
task: async () => undefined,
|
||||
});
|
||||
|
||||
expect(grabHarnessUsageScope()).toBeUndefined();
|
||||
});
|
||||
|
||||
test("two concurrent runs never see each other's scope", async () => {
|
||||
const [first, second] = await Promise.all([
|
||||
runWithHarnessUsageScope({
|
||||
scope: { user_id: 1 },
|
||||
task: async () => {
|
||||
await Bun.sleep(5);
|
||||
return grabHarnessUsageScope();
|
||||
},
|
||||
}),
|
||||
runWithHarnessUsageScope({
|
||||
scope: { user_id: 2 },
|
||||
task: async () => {
|
||||
await Bun.sleep(1);
|
||||
return grabHarnessUsageScope();
|
||||
},
|
||||
}),
|
||||
]);
|
||||
|
||||
expect(first).toEqual({ user_id: 1 });
|
||||
expect(second).toEqual({ user_id: 2 });
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,63 @@
|
||||
import { AsyncLocalStorage } from "node:async_hooks";
|
||||
import type { HarnessUsageScope } from "./types";
|
||||
|
||||
/**
|
||||
* The attribution a run's model calls carry.
|
||||
*
|
||||
* A bot's run is entered once (`runWithHarnessUsageScope`) and everything
|
||||
* it calls underneath - the text loop's turns, a session summary, a
|
||||
* transcription, an image, an embedding - is written down for that user,
|
||||
* bot, session and job without threading the ids through every harness
|
||||
* function. One of those calls is itself entered by a `run` that has to
|
||||
* swap the scope it reports under (`withHarnessUsageScope`): the session
|
||||
* namer runs outside the job's scope, on purpose, since the job it named
|
||||
* is already complete.
|
||||
*
|
||||
* The store is merged, not replaced: a call that only knows the message id
|
||||
* (the text loop attaching a turn's usage) keeps the run's user, bot and
|
||||
* job.
|
||||
*/
|
||||
const storage = new AsyncLocalStorage<HarnessUsageScope>();
|
||||
|
||||
/**
|
||||
* Runs `task` with `scope` as its attribution. Ambient values are merged
|
||||
* under it, so a nested call keeps what the outer one set and overrides
|
||||
* only what it knows better (null clears a field).
|
||||
*/
|
||||
export async function runWithHarnessUsageScope<T>(
|
||||
{ scope, task }: { scope: HarnessUsageScope; task: () => Promise<T> },
|
||||
): Promise<T> {
|
||||
return await storage.run(mergeScopes({ base: storage.getStore(), overlay: scope }), task);
|
||||
}
|
||||
|
||||
/**
|
||||
* The attribution a call right now carries, if any.
|
||||
*/
|
||||
export function grabHarnessUsageScope(): HarnessUsageScope | undefined {
|
||||
return storage.getStore();
|
||||
}
|
||||
|
||||
function mergeScopes({
|
||||
base,
|
||||
overlay,
|
||||
}: {
|
||||
base?: HarnessUsageScope;
|
||||
overlay: HarnessUsageScope;
|
||||
}): HarnessUsageScope {
|
||||
const merged: HarnessUsageScope = { ...base };
|
||||
|
||||
for (const [key, value] of Object.entries(overlay) as [
|
||||
keyof HarnessUsageScope,
|
||||
HarnessUsageScope[keyof HarnessUsageScope],
|
||||
][]) {
|
||||
if (value === undefined) continue;
|
||||
|
||||
/**
|
||||
* A null is a real value here: it says "none", e.g. a bot's own
|
||||
* chat has no workspace, and must clear what the outer scope had.
|
||||
*/
|
||||
(merged as Record<string, unknown>)[key] = value;
|
||||
}
|
||||
|
||||
return merged;
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
import BunSQLite from "@moduletrace/bun-sqlite";
|
||||
import type {
|
||||
BUN_SQLITE_COBBLE_TOKEN_USAGES,
|
||||
BunSQLiteTables,
|
||||
} from "@/db/types/db";
|
||||
import { logError } from "@/src/functions/backend/logger";
|
||||
import { grabHarnessUsageScope } from "./harness-usage-scope";
|
||||
import type { HarnessUsageCall, HarnessUsageScope } from "./types";
|
||||
|
||||
type Params = {
|
||||
/**
|
||||
* The scope the call ran under. Left out, the ambient one
|
||||
* (`grabHarnessUsageScope`) is used, which is what every entry point
|
||||
* does: the ids came in through `runWithHarnessUsageScope`.
|
||||
*/
|
||||
scope?: HarnessUsageScope;
|
||||
/**
|
||||
* Which harness and which model answered, with what the provider
|
||||
* reported it cost.
|
||||
*/
|
||||
call: HarnessUsageCall;
|
||||
};
|
||||
|
||||
/**
|
||||
* Writes one model call into `token_usages`: the tokens it spent, on which
|
||||
* model and provider, for whom, and which message or job it was for. This
|
||||
* is what a message's token line, a bot's and a provider's totals and the
|
||||
* usage page all read.
|
||||
*
|
||||
* Never throws: a model that answered must not have its answer lost
|
||||
* because the bookkeeping failed. A call without an owner (no scope, no
|
||||
* user) is not written at all, since nothing could show it.
|
||||
*/
|
||||
export default async function recordHarnessUsage({ scope, call }: Params) {
|
||||
const resolved = scope || grabHarnessUsageScope();
|
||||
const user_id = Number(resolved?.user_id) || null;
|
||||
|
||||
if (!user_id) return;
|
||||
|
||||
const row: BUN_SQLITE_COBBLE_TOKEN_USAGES = {
|
||||
user_id,
|
||||
bot_id: numberOrNull(resolved?.bot_id),
|
||||
workspace_id: numberOrNull(resolved?.workspace_id),
|
||||
session_id: numberOrNull(resolved?.session_id),
|
||||
message_id: numberOrNull(resolved?.message_id),
|
||||
job_id: numberOrNull(resolved?.job_id),
|
||||
harness: call.harness,
|
||||
provider_id: call.provider_id,
|
||||
provider_name: call.provider_name,
|
||||
model_id: call.model_id,
|
||||
model_slug: call.model_slug,
|
||||
prompt_tokens: Math.max(0, Math.round(call.prompt_tokens || 0)),
|
||||
completion_tokens: Math.max(0, Math.round(call.completion_tokens || 0)),
|
||||
cost: typeof call.cost == "number" && Number.isFinite(call.cost) ? call.cost : undefined,
|
||||
};
|
||||
|
||||
try {
|
||||
const res = await BunSQLite.insert<
|
||||
BUN_SQLITE_COBBLE_TOKEN_USAGES,
|
||||
(typeof BunSQLiteTables)[number]
|
||||
>({
|
||||
table: "token_usages",
|
||||
data: [row],
|
||||
});
|
||||
|
||||
if (!res.success) {
|
||||
throw new Error(res.msg || `The usage wasn't saved.`);
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
await logError({
|
||||
title: "Model usage not recorded",
|
||||
error,
|
||||
metadata: {
|
||||
harness: call.harness,
|
||||
model_slug: call.model_slug || null,
|
||||
provider_name: call.provider_name || null,
|
||||
user_id,
|
||||
},
|
||||
}).catch(() => {});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* SQLite stores NULL for what isn't given; `Number(undefined)` is NaN,
|
||||
* which bun-sqlite would refuse. An absent id is `undefined`, not `null`:
|
||||
* bun-sqlite's row types spell "no value" that way.
|
||||
*/
|
||||
function numberOrNull(value: unknown): number | undefined {
|
||||
const number = Number(value);
|
||||
|
||||
return Number.isFinite(number) && number > 0 ? number : undefined;
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { grabHarnessUsageScope, runWithHarnessUsageScope } from "./harness-usage-scope";
|
||||
import runWithFallbackScope from "./run-with-fallback-scope";
|
||||
|
||||
const owner = { user_id: 9, bot_id: 3, session_id: null, message_id: null, job_id: null };
|
||||
|
||||
describe("runWithFallbackScope", () => {
|
||||
test("the enclosing run keeps the work it already owns", async () => {
|
||||
const seen = await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7, bot_id: 3, job_id: 5 },
|
||||
task: () =>
|
||||
runWithFallbackScope({
|
||||
scope: owner,
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
}),
|
||||
});
|
||||
|
||||
expect(seen).toEqual({ user_id: 7, bot_id: 3, job_id: 5 });
|
||||
});
|
||||
|
||||
test("work with no run is attributed to its owner", async () => {
|
||||
const seen = await runWithFallbackScope({
|
||||
scope: owner,
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
});
|
||||
|
||||
expect(seen).toEqual(owner);
|
||||
});
|
||||
|
||||
test("nobody to charge leaves the work unattributed rather than failing", async () => {
|
||||
const seen = await runWithFallbackScope({
|
||||
scope: null,
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
});
|
||||
|
||||
expect(seen).toBeUndefined();
|
||||
});
|
||||
|
||||
test("a scope with no user is not an owner", async () => {
|
||||
const seen = await runWithFallbackScope({
|
||||
scope: { ...owner, user_id: undefined },
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
});
|
||||
|
||||
expect(seen).toBeUndefined();
|
||||
});
|
||||
|
||||
test("work that must be charged on its own overrides the enclosing run", async () => {
|
||||
const seen = await runWithHarnessUsageScope({
|
||||
scope: { user_id: 7, bot_id: 3, job_id: 5 },
|
||||
task: () =>
|
||||
runWithFallbackScope({
|
||||
scope: owner,
|
||||
mode: "own",
|
||||
task: async () => grabHarnessUsageScope(),
|
||||
}),
|
||||
});
|
||||
|
||||
expect(seen).toEqual(owner);
|
||||
});
|
||||
|
||||
test("the work's result is passed through", async () => {
|
||||
const result = await runWithFallbackScope({
|
||||
scope: owner,
|
||||
task: async () => "the answer",
|
||||
});
|
||||
|
||||
expect(result).toBe("the answer");
|
||||
});
|
||||
|
||||
test("a failure inside the work is not swallowed", async () => {
|
||||
expect(
|
||||
runWithFallbackScope({
|
||||
scope: owner,
|
||||
task: async () => {
|
||||
throw new Error("the provider refused");
|
||||
},
|
||||
}),
|
||||
).rejects.toThrow("the provider refused");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,41 @@
|
||||
import { grabHarnessUsageScope, runWithHarnessUsageScope } from "./harness-usage-scope";
|
||||
import type { HarnessUsageScope } from "./types";
|
||||
|
||||
type Params<T> = {
|
||||
/**
|
||||
* What to attribute the work to when nothing already owns it: a file
|
||||
* ingested by the queue has no run behind it, but its owner must still
|
||||
* see what the embedding cost. Null when there is nobody to charge.
|
||||
*/
|
||||
scope: HarnessUsageScope | null;
|
||||
task: () => Promise<T>;
|
||||
/**
|
||||
* How the work is attributed: `enclosing` keeps whatever run it is in,
|
||||
* `own` always reports under `scope` (the session namer, whose job is
|
||||
* already complete and must not be charged to it).
|
||||
*/
|
||||
mode?: "enclosing" | "own";
|
||||
};
|
||||
|
||||
/**
|
||||
* Runs work under an attribution the run may not have. Nothing is entered
|
||||
* when a run is already in charge, so an ingest a bot asked for is charged
|
||||
* to that bot's user and not to the file's owner.
|
||||
*/
|
||||
export default async function runWithFallbackScope<T>({
|
||||
scope,
|
||||
task,
|
||||
mode = "enclosing",
|
||||
}: Params<T>): Promise<T> {
|
||||
if (mode == "enclosing" && grabHarnessUsageScope()?.user_id) {
|
||||
return await task();
|
||||
}
|
||||
|
||||
const user_id = Number(scope?.user_id);
|
||||
|
||||
if (!Number.isFinite(user_id) || user_id <= 0 || user_id === null) {
|
||||
return await task();
|
||||
}
|
||||
|
||||
return await runWithHarnessUsageScope({ scope: { ...scope, user_id }, task });
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
/**
|
||||
* Who and what a model call was for. Every harness call that reports its
|
||||
* usage is written down with this (`/api/auth/crud` no: `token_usages`),
|
||||
* so a message, a session, a bot or a provider can be totaled later.
|
||||
*
|
||||
* It is deliberately shaped slightly differently from the table's
|
||||
* columns: `user_id` is required here (a row without an owner can't be
|
||||
* shown or totaled for anyone), the rest are optional.
|
||||
*/
|
||||
export type HarnessUsageScope = {
|
||||
/**
|
||||
* Whose work it was: the user whose conversation the call ran for.
|
||||
*/
|
||||
user_id?: number | string;
|
||||
bot_id?: number | string | null;
|
||||
workspace_id?: number | string | null;
|
||||
session_id?: number | string | null;
|
||||
/**
|
||||
* The queued job it ran for, when it ran for one.
|
||||
*/
|
||||
job_id?: number | string | null;
|
||||
/**
|
||||
* The message it produced, for the ones that produce one.
|
||||
*/
|
||||
message_id?: number | string | null;
|
||||
};
|
||||
|
||||
/**
|
||||
* What the caller knows about the call beyond its scope: which harness
|
||||
* asked, and which candidate answered. Both are filled in by the call's
|
||||
* own entry point, never by the code that runs it.
|
||||
*/
|
||||
export type HarnessUsageCall = {
|
||||
harness: "text" | "vision" | "transcription" | "speech" | "image" | "embedding";
|
||||
provider_id?: number;
|
||||
provider_name?: string;
|
||||
model_id?: number;
|
||||
model_slug?: string;
|
||||
prompt_tokens: number;
|
||||
completion_tokens: number;
|
||||
cost?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* A message's usage, as its thread shows it and as `token_usages` stores
|
||||
* it: what one answer cost, summed over its model calls.
|
||||
*/
|
||||
export type TokenUsageSummary = {
|
||||
prompt_tokens: number;
|
||||
completion_tokens: number;
|
||||
total_tokens: number;
|
||||
/**
|
||||
* In USD, when the provider reported it.
|
||||
*/
|
||||
cost?: number;
|
||||
/**
|
||||
* How many model calls it took.
|
||||
*/
|
||||
calls: number;
|
||||
};
|
||||
@@ -0,0 +1,19 @@
|
||||
import { runWithHarnessUsageScope } from "./harness-usage-scope";
|
||||
import type { HarnessUsageScope } from "./types";
|
||||
|
||||
/**
|
||||
* Runs one call the app makes on a user's behalf inside the scope its
|
||||
* usage is written down with (transcribing a dictation, speaking a
|
||||
* message, previewing a voice). The bot's run is already inside a scope of
|
||||
* its own, so this only adds what the run didn't know: no job, no message,
|
||||
* and the session the call belongs to.
|
||||
*/
|
||||
export default async function withUsageScope<T>({
|
||||
scope,
|
||||
task,
|
||||
}: {
|
||||
scope: HarnessUsageScope;
|
||||
task: () => Promise<T>;
|
||||
}): Promise<T> {
|
||||
return await runWithHarnessUsageScope({ scope, task });
|
||||
}
|
||||
Reference in New Issue
Block a user