import type { Logger } from "../../core/logger";import type { LifetimeMetricsStore } from "../system/metrics-store";import type { InferenceRequestRecord, InferenceRequestStore,} from "../../stores/inference-request-store";import { Effect } from "effect"; interface InferenceAccountingStores { lifetimeMetricsStore: Pick< LifetimeMetricsStore, "addCompletionTokens" | "addPromptTokens" | "addRequests" | "addTokens" >; inferenceRequestStore: Pick<InferenceRequestStore, "record">;} interface InferenceAccountingOptions { logger: Pick<Logger, "warn">; stores: InferenceAccountingStores;} interface InferenceUsageInput { prompt_tokens?: number; completion_tokens?: number; reasoning_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; prompt_tokens_details?: Record<string, number>; completion_tokens_details?: Record<string, number>;} interface InferenceUsageTotals { promptTokens: number; completionTokens: number; reasoningTokens: number; cacheReadTokens: number; cacheWriteTokens: number;} interface NonStreamingInferenceRecordInput { usage: InferenceUsageInput | undefined; record: Omit< InferenceRequestRecord, | "cache_read_tokens" | "cache_write_tokens" | "completion_tokens" | "prompt_tokens" | "reasoning_tokens" | "streamed" >;} interface StreamingInferenceRecordInput { usage: InferenceUsageInput; record: Omit< InferenceRequestRecord, | "cache_read_tokens" | "cache_write_tokens" | "completion_tokens" | "prompt_tokens" | "reasoning_tokens" | "streamed" >;} const hasBillableTokens = (totals: InferenceUsageTotals): boolean => totals.promptTokens > 0 || totals.completionTokens > 0; const readUsageTotals = (usage: InferenceUsageInput): InferenceUsageTotals => { const promptDetails = usage.prompt_tokens_details; const completionDetails = usage.completion_tokens_details; return { promptTokens: usage.prompt_tokens ?? 0, completionTokens: usage.completion_tokens ?? 0, reasoningTokens: usage.reasoning_tokens ?? completionDetails?.["reasoning_tokens"] ?? 0, cacheReadTokens: promptDetails?.["cached_tokens"] ?? usage.cache_read_tokens ?? 0, cacheWriteTokens: usage.cache_write_tokens ?? 0, };}; const addLifetimeUsage = ( stores: InferenceAccountingStores, totals: InferenceUsageTotals,): Effect.Effect<void, unknown> => Effect.all( [ ...(totals.promptTokens > 0 ? [ stores.lifetimeMetricsStore.addPromptTokens(totals.promptTokens), stores.lifetimeMetricsStore.addTokens(totals.promptTokens), ] : []), ...(totals.completionTokens > 0 ? [ stores.lifetimeMetricsStore.addCompletionTokens(totals.completionTokens), stores.lifetimeMetricsStore.addTokens(totals.completionTokens), ] : []), ...(hasBillableTokens(totals) ? [stores.lifetimeMetricsStore.addRequests(1)] : []), ], { concurrency: 1, discard: true }, ); const tryRecordInference = ( options: InferenceAccountingOptions, record: InferenceRequestRecord,): Effect.Effect<void> => options.stores.inferenceRequestStore .record(record) .pipe( Effect.catch((recordError) => Effect.sync(() => options.logger.warn(`Failed to record inference request: ${String(recordError)}`), ), ), ); export const recordNonStreamingInferenceUsage = ( options: InferenceAccountingOptions, input: NonStreamingInferenceRecordInput,): Effect.Effect<InferenceUsageTotals | null, unknown> => { if (!input.usage) return Effect.succeed(null); const totals = readUsageTotals(input.usage); const record = hasBillableTokens(totals) ? tryRecordInference(options, { ...input.record, prompt_tokens: totals.promptTokens, completion_tokens: totals.completionTokens, reasoning_tokens: totals.reasoningTokens, cache_read_tokens: totals.cacheReadTokens, cache_write_tokens: totals.cacheWriteTokens, streamed: false, }) : Effect.void; return addLifetimeUsage(options.stores, totals).pipe(Effect.andThen(record), Effect.as(totals));}; export const recordStreamingInferenceUsage = ( options: InferenceAccountingOptions, input: StreamingInferenceRecordInput,): Effect.Effect<InferenceUsageTotals, unknown> => { const totals = readUsageTotals(input.usage); const record = hasBillableTokens(totals) ? tryRecordInference(options, { ...input.record, prompt_tokens: totals.promptTokens, completion_tokens: totals.completionTokens, reasoning_tokens: totals.reasoningTokens, cache_read_tokens: totals.cacheReadTokens, cache_write_tokens: totals.cacheWriteTokens, streamed: true, }) : Effect.void; return addLifetimeUsage(options.stores, totals).pipe(Effect.andThen(record), Effect.as(totals));};
import type { Logger } from "../../core/logger";
import type { LifetimeMetricsStore } from "../system/metrics-store";
import type {
InferenceRequestRecord,
InferenceRequestStore,
} from "../../stores/inference-request-store";
import { Effect } from "effect";
interface InferenceAccountingStores {
lifetimeMetricsStore: Pick<
LifetimeMetricsStore,
"addCompletionTokens" | "addPromptTokens" | "addRequests" | "addTokens"
>;
inferenceRequestStore: Pick<InferenceRequestStore, "record">;
}
interface InferenceAccountingOptions {
logger: Pick<Logger, "warn">;
stores: InferenceAccountingStores;
}
interface InferenceUsageInput {
prompt_tokens?: number;
completion_tokens?: number;
reasoning_tokens?: number;
cache_read_tokens?: number;
cache_write_tokens?: number;
prompt_tokens_details?: Record<string, number>;
completion_tokens_details?: Record<string, number>;
}
interface InferenceUsageTotals {
promptTokens: number;
completionTokens: number;
reasoningTokens: number;
cacheReadTokens: number;
cacheWriteTokens: number;
}
interface NonStreamingInferenceRecordInput {
usage: InferenceUsageInput | undefined;
record: Omit<
InferenceRequestRecord,
| "cache_read_tokens"
| "cache_write_tokens"
| "completion_tokens"
| "prompt_tokens"
| "reasoning_tokens"
| "streamed"
>;
}
interface StreamingInferenceRecordInput {
usage: InferenceUsageInput;
record: Omit<
InferenceRequestRecord,
| "cache_read_tokens"
| "cache_write_tokens"
| "completion_tokens"
| "prompt_tokens"
| "reasoning_tokens"
| "streamed"
>;
}
const hasBillableTokens = (totals: InferenceUsageTotals): boolean =>
totals.promptTokens > 0 || totals.completionTokens > 0;
const readUsageTotals = (usage: InferenceUsageInput): InferenceUsageTotals => {
const promptDetails = usage.prompt_tokens_details;
const completionDetails = usage.completion_tokens_details;
return {
promptTokens: usage.prompt_tokens ?? 0,
completionTokens: usage.completion_tokens ?? 0,
reasoningTokens: usage.reasoning_tokens ?? completionDetails?.["reasoning_tokens"] ?? 0,
cacheReadTokens: promptDetails?.["cached_tokens"] ?? usage.cache_read_tokens ?? 0,
cacheWriteTokens: usage.cache_write_tokens ?? 0,
};
};
const addLifetimeUsage = (
stores: InferenceAccountingStores,
totals: InferenceUsageTotals,
): Effect.Effect<void, unknown> =>
Effect.all(
[
...(totals.promptTokens > 0
? [
stores.lifetimeMetricsStore.addPromptTokens(totals.promptTokens),
stores.lifetimeMetricsStore.addTokens(totals.promptTokens),
]
: []),
...(totals.completionTokens > 0
? [
stores.lifetimeMetricsStore.addCompletionTokens(totals.completionTokens),
stores.lifetimeMetricsStore.addTokens(totals.completionTokens),
]
: []),
...(hasBillableTokens(totals) ? [stores.lifetimeMetricsStore.addRequests(1)] : []),
],
{ concurrency: 1, discard: true },
);
const tryRecordInference = (
options: InferenceAccountingOptions,
record: InferenceRequestRecord,
): Effect.Effect<void> =>
options.stores.inferenceRequestStore
.record(record)
.pipe(
Effect.catch((recordError) =>
Effect.sync(() =>
options.logger.warn(`Failed to record inference request: ${String(recordError)}`),
),
),
);
export const recordNonStreamingInferenceUsage = (
options: InferenceAccountingOptions,
input: NonStreamingInferenceRecordInput,
): Effect.Effect<InferenceUsageTotals | null, unknown> => {
if (!input.usage) return Effect.succeed(null);
const totals = readUsageTotals(input.usage);
const record = hasBillableTokens(totals)
? tryRecordInference(options, {
...input.record,
prompt_tokens: totals.promptTokens,
completion_tokens: totals.completionTokens,
reasoning_tokens: totals.reasoningTokens,
cache_read_tokens: totals.cacheReadTokens,
cache_write_tokens: totals.cacheWriteTokens,
streamed: false,
})
: Effect.void;
return addLifetimeUsage(options.stores, totals).pipe(Effect.andThen(record), Effect.as(totals));
};
export const recordStreamingInferenceUsage = (
options: InferenceAccountingOptions,
input: StreamingInferenceRecordInput,
): Effect.Effect<InferenceUsageTotals, unknown> => {
const totals = readUsageTotals(input.usage);
const record = hasBillableTokens(totals)
? tryRecordInference(options, {
...input.record,
prompt_tokens: totals.promptTokens,
completion_tokens: totals.completionTokens,
reasoning_tokens: totals.reasoningTokens,
cache_read_tokens: totals.cacheReadTokens,
cache_write_tokens: totals.cacheWriteTokens,
streamed: true,
})
: Effect.void;
return addLifetimeUsage(options.stores, totals).pipe(Effect.andThen(record), Effect.as(totals));
};