155 lines · 4.8 KB · TypeScript
RawDownload
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));};