Files
qinglong/packages/ql3-ai/src/usage/usageLedger.ts
T

413 lines
12 KiB
TypeScript

import { Buffer } from 'node:buffer';
import { createHash } from 'node:crypto';
import type {
ModelInvocationCompletionRecord,
ModelInvocationStartRecord,
} from '../model-invocation/modelInvocation';
import {
normalizeModelInvocationCompletionRecord,
normalizeModelInvocationStartRecord,
} from '../model-invocation/modelInvocation';
// Usage owns immutable metering facts derived from completed model invocations.
export const MODEL_INVOCATION_USAGE_LEDGER_SCHEMA =
'qinglong/model-invocation-usage-ledger@v1';
export const MAX_MODEL_INVOCATION_USAGE_LEDGER_PAGE_SIZE = 128;
export const MAX_MODEL_INVOCATION_USAGE_SUMMARY_ROWS = 100_000;
export const MAX_MODEL_INVOCATION_USAGE_QUERY_WINDOW_MS =
366 * 24 * 60 * 60_000;
const IDENTITY_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$/;
const DIGEST_PATTERN = /^[0-9a-f]{64}$/;
const LEDGER_DIGEST_DOMAIN = Buffer.from(
'qinglong/model-invocation-usage-ledger-digest@v1\0',
'utf8',
);
export interface ModelInvocationUsageLedgerRecord {
readonly schema: typeof MODEL_INVOCATION_USAGE_LEDGER_SCHEMA;
readonly invocationId: string;
readonly projectId: string;
readonly runId: string;
readonly stepRunId: string;
readonly traceId: string;
readonly provider: string;
readonly model: string;
readonly policyRevision: string;
readonly completionDigest: string;
readonly outcome: ModelInvocationCompletionRecord['outcome'];
readonly settledAtMs: number;
readonly inputBytes: number;
readonly outputBytes: number;
readonly inputTokens: number;
readonly outputTokens: number;
readonly totalTokens: number;
readonly costMicros: number | null;
readonly ledgerDigest: string;
}
export interface ModelInvocationUsageLedgerCursor {
readonly settledAtMs: number;
readonly invocationId: string;
}
export interface ModelInvocationUsageLedgerQuery {
readonly projectId: string;
readonly fromMsInclusive: number;
readonly toMsExclusive: number;
readonly limit: number;
readonly after?: Readonly<ModelInvocationUsageLedgerCursor>;
}
export interface ModelInvocationUsageLedgerSummaryQuery {
readonly projectId: string;
readonly fromMsInclusive: number;
readonly toMsExclusive: number;
}
export interface ModelInvocationUsageLedgerPage {
readonly records: readonly Readonly<ModelInvocationUsageLedgerRecord>[];
readonly hasMore: boolean;
}
export interface ModelInvocationUsageLedgerSummary {
readonly invocationCount: number;
readonly inputTokens: number;
readonly outputTokens: number;
readonly totalTokens: number;
readonly knownCostMicros: number;
readonly unknownCostInvocations: number;
}
export interface ModelInvocationUsageLedgerRepository {
findUsage(
invocationId: string,
): Promise<Readonly<ModelInvocationUsageLedgerRecord> | null>;
listProjectUsage(
query: ModelInvocationUsageLedgerQuery,
): Promise<Readonly<ModelInvocationUsageLedgerPage>>;
summarizeProjectUsage(
query: ModelInvocationUsageLedgerSummaryQuery,
): Promise<Readonly<ModelInvocationUsageLedgerSummary>>;
}
export class InvalidModelInvocationUsageLedgerError extends TypeError {
readonly code = 'MODEL_INVOCATION_USAGE_LEDGER_INVALID';
constructor(message: string) {
super(`Model invocation usage ledger is invalid: ${message}`);
this.name = 'InvalidModelInvocationUsageLedgerError';
}
}
export class ModelInvocationUsageSummaryLimitExceededError extends Error {
readonly code = 'MODEL_INVOCATION_USAGE_SUMMARY_LIMIT_EXCEEDED';
constructor() {
super('Model invocation usage summary exceeds its reviewed row limit');
this.name = 'ModelInvocationUsageSummaryLimitExceededError';
}
}
function invalid(message: string): never {
throw new InvalidModelInvocationUsageLedgerError(message);
}
function dataRecord(value: unknown, label: string): Record<string, unknown> {
if (
!value ||
typeof value !== 'object' ||
Array.isArray(value) ||
Object.getPrototypeOf(value) !== Object.prototype
) {
return invalid(`${label} must be a plain object`);
}
return value as Record<string, unknown>;
}
function exactKeys(
value: object,
expected: readonly string[],
label: string,
): void {
const actual = Object.keys(value).sort();
const canonical = [...expected].sort();
if (
actual.length !== canonical.length ||
actual.some((key, index) => key !== canonical[index])
) {
invalid(`${label} shape is invalid`);
}
}
function identifier(value: unknown, label: string): string {
if (typeof value !== 'string' || !IDENTITY_PATTERN.test(value)) {
return invalid(`${label} is invalid`);
}
return value;
}
function integer(
value: unknown,
minimum: number,
maximum: number,
label: string,
): number {
if (
!Number.isSafeInteger(value) ||
(value as number) < minimum ||
(value as number) > maximum
) {
return invalid(`${label} is invalid`);
}
return value as number;
}
function digest(value: unknown, label: string): string {
if (typeof value !== 'string' || !DIGEST_PATTERN.test(value)) {
return invalid(`${label} is invalid`);
}
return value;
}
function hash(value: object): string {
return createHash('sha256')
.update(LEDGER_DIGEST_DOMAIN)
.update(JSON.stringify(value), 'utf8')
.digest('hex');
}
function withoutDigest(
value: Readonly<ModelInvocationUsageLedgerRecord>,
): Omit<ModelInvocationUsageLedgerRecord, 'ledgerDigest'> {
const { ledgerDigest: _ledgerDigest, ...unsigned } = value;
return unsigned;
}
export function normalizeModelInvocationUsageLedgerRecord(
value: ModelInvocationUsageLedgerRecord,
): Readonly<ModelInvocationUsageLedgerRecord> {
const candidate = dataRecord(value, 'ledger record');
exactKeys(
candidate,
[
'completionDigest',
'costMicros',
'inputBytes',
'inputTokens',
'invocationId',
'ledgerDigest',
'model',
'outcome',
'outputBytes',
'outputTokens',
'policyRevision',
'projectId',
'provider',
'runId',
'schema',
'settledAtMs',
'stepRunId',
'totalTokens',
'traceId',
],
'ledger record',
);
if (
value.schema !== MODEL_INVOCATION_USAGE_LEDGER_SCHEMA ||
!['succeeded', 'failed', 'timed_out', 'outcome_unknown'].includes(
value.outcome,
)
) {
invalid('schema or outcome is invalid');
}
const inputTokens = integer(
value.inputTokens,
0,
Number.MAX_SAFE_INTEGER,
'input token count',
);
const outputTokens = integer(
value.outputTokens,
0,
Number.MAX_SAFE_INTEGER,
'output token count',
);
const totalTokens = integer(
value.totalTokens,
0,
Number.MAX_SAFE_INTEGER,
'total token count',
);
if (totalTokens !== inputTokens + outputTokens) {
invalid('total token count is inconsistent');
}
const normalized = Object.freeze({
schema: MODEL_INVOCATION_USAGE_LEDGER_SCHEMA,
invocationId: identifier(value.invocationId, 'invocation id'),
projectId: identifier(value.projectId, 'Project id'),
runId: identifier(value.runId, 'Run id'),
stepRunId: identifier(value.stepRunId, 'StepRun id'),
traceId: identifier(value.traceId, 'trace id'),
provider: identifier(value.provider, 'provider'),
model: identifier(value.model, 'model'),
policyRevision: identifier(value.policyRevision, 'policy revision'),
completionDigest: digest(value.completionDigest, 'completion digest'),
outcome: value.outcome,
settledAtMs: integer(
value.settledAtMs,
0,
Number.MAX_SAFE_INTEGER,
'settlement time',
),
inputBytes: integer(value.inputBytes, 1, 256 * 1024, 'input bytes'),
outputBytes: integer(value.outputBytes, 0, 1024 * 1024, 'output bytes'),
inputTokens,
outputTokens,
totalTokens,
costMicros:
value.costMicros === null
? null
: integer(value.costMicros, 0, Number.MAX_SAFE_INTEGER, 'cost'),
ledgerDigest: digest(value.ledgerDigest, 'ledger digest'),
});
if (hash(withoutDigest(normalized)) !== normalized.ledgerDigest) {
invalid('ledger digest is invalid');
}
return normalized;
}
export function createModelInvocationUsageLedgerRecord(
startValue: ModelInvocationStartRecord,
completionValue: ModelInvocationCompletionRecord,
): Readonly<ModelInvocationUsageLedgerRecord> | null {
const start = normalizeModelInvocationStartRecord(startValue);
const completion = normalizeModelInvocationCompletionRecord(completionValue);
if (
completion.invocationId !== start.invocationId ||
completion.projectId !== start.projectId ||
completion.runId !== start.runId ||
completion.stepRunId !== start.stepRunId ||
completion.traceId !== start.traceId ||
completion.startDigest !== start.startDigest ||
completion.completedAtMs < start.admittedAtMs
) {
invalid('completion is detached from its start');
}
if (!completion.usage) return null;
const unsigned = Object.freeze({
schema: MODEL_INVOCATION_USAGE_LEDGER_SCHEMA,
invocationId: completion.invocationId,
projectId: completion.projectId,
runId: completion.runId,
stepRunId: completion.stepRunId,
traceId: completion.traceId,
provider: start.provider,
model: start.model,
policyRevision: start.policyRevision,
completionDigest: completion.completionDigest,
outcome: completion.outcome,
settledAtMs: completion.completedAtMs,
inputBytes: start.inputBytes,
outputBytes: completion.outputBytes,
inputTokens: completion.usage.inputTokens,
outputTokens: completion.usage.outputTokens,
totalTokens: completion.usage.totalTokens,
costMicros: completion.usage.costMicros ?? null,
});
return normalizeModelInvocationUsageLedgerRecord({
...unsigned,
ledgerDigest: hash(unsigned),
});
}
function normalizeWindow(value: Record<string, unknown>): Readonly<{
projectId: string;
fromMsInclusive: number;
toMsExclusive: number;
}> {
const projectId = identifier(value.projectId, 'Project id');
const fromMsInclusive = integer(
value.fromMsInclusive,
0,
Number.MAX_SAFE_INTEGER,
'window start',
);
const toMsExclusive = integer(
value.toMsExclusive,
1,
Number.MAX_SAFE_INTEGER,
'window end',
);
if (
toMsExclusive <= fromMsInclusive ||
toMsExclusive - fromMsInclusive > MAX_MODEL_INVOCATION_USAGE_QUERY_WINDOW_MS
) {
invalid('query window is invalid');
}
return Object.freeze({ projectId, fromMsInclusive, toMsExclusive });
}
function normalizeCursor(
value: Readonly<ModelInvocationUsageLedgerCursor>,
): Readonly<ModelInvocationUsageLedgerCursor> {
const cursor = dataRecord(value, 'ledger cursor');
exactKeys(cursor, ['invocationId', 'settledAtMs'], 'ledger cursor');
return Object.freeze({
settledAtMs: integer(
value.settledAtMs,
0,
Number.MAX_SAFE_INTEGER,
'cursor settlement time',
),
invocationId: identifier(value.invocationId, 'cursor invocation id'),
});
}
export function normalizeModelInvocationUsageLedgerQuery(
value: ModelInvocationUsageLedgerQuery,
): Readonly<ModelInvocationUsageLedgerQuery> {
const query = dataRecord(value, 'ledger query');
exactKeys(
query,
value.after === undefined
? ['fromMsInclusive', 'limit', 'projectId', 'toMsExclusive']
: ['after', 'fromMsInclusive', 'limit', 'projectId', 'toMsExclusive'],
'ledger query',
);
const window = normalizeWindow(query);
const after =
value.after === undefined ? undefined : normalizeCursor(value.after);
if (
after &&
(after.settledAtMs < window.fromMsInclusive ||
after.settledAtMs >= window.toMsExclusive)
) {
invalid('cursor is outside the query window');
}
return Object.freeze({
...window,
limit: integer(
value.limit,
1,
MAX_MODEL_INVOCATION_USAGE_LEDGER_PAGE_SIZE,
'query limit',
),
...(after === undefined ? {} : { after }),
});
}
export function normalizeModelInvocationUsageLedgerSummaryQuery(
value: ModelInvocationUsageLedgerSummaryQuery,
): Readonly<ModelInvocationUsageLedgerSummaryQuery> {
const query = dataRecord(value, 'ledger summary query');
exactKeys(
query,
['fromMsInclusive', 'projectId', 'toMsExclusive'],
'ledger summary query',
);
return normalizeWindow(query);
}