pi-cache-optimizer 2.8.16 → 2.8.17

package/src/usage.ts

@@ -0,0 +1,172 @@
+import { type UnknownRecord, asRecord, getNonNegativeNumber, getNumber } from "./common.ts";
+import { getAssistantRecord } from "./model-identity.ts";
+import { type UsageSnapshot } from "./stats-store.ts";
+
+export function usageRecordFromAssistant(message: unknown): UnknownRecord | undefined {
+  return asRecord(getAssistantRecord(message)?.usage);
+}
+
+export function getNestedRecord(record: UnknownRecord | undefined, key: string): UnknownRecord | undefined {
+  return asRecord(record?.[key]);
+}
+
+export function getFirstNonNegativeNumber(...values: unknown[]): number | undefined {
+  for (const value of values) {
+    const number = getNumber(value);
+    if (number !== undefined && number >= 0) return number;
+  }
+  return undefined;
+}
+
+export function readCachedTokensFromDetails(details: UnknownRecord | undefined): number | undefined {
+  return getFirstNonNegativeNumber(details?.cached_tokens, details?.cachedTokens);
+}
+
+export function readCacheWriteFromDetails(details: UnknownRecord | undefined): number | undefined {
+  return getFirstNonNegativeNumber(details?.cache_write_tokens, details?.cacheWriteTokens);
+}
+
+// Pi normalizes provider-specific raw usage (prompt_cache_hit_tokens, cached_tokens,
+// cache_read_input_tokens, etc.) into a common shape:
+//   input     = uncached prompt portion (total prompt minus cacheRead minus cacheWrite)
+//   cacheRead = tokens read from a previously-cached prefix
+//   cacheWrite= tokens newly written into cache in this request
+//
+// We reconstruct the total prompt-token count as input + cacheRead + cacheWrite.
+// Pi guarantees that input, cacheRead, and cacheWrite are always present on
+// assistant messages processed through its provider pipeline (at least as zero).
+//
+// Only DeepSeek sets allowInputOnly=true so that a cache miss (cacheRead=0) still
+// contributes total input tokens to the denominator.
+export function getPiNormalizedUsage(message: unknown, allowInputOnly = false): UsageSnapshot | undefined {
+  const usage = usageRecordFromAssistant(message);
+  if (!usage) return undefined;
+
+  const input = getNonNegativeNumber(usage, "input");
+  const cacheRead = getNonNegativeNumber(usage, "cacheRead");
+  const cacheWrite = getNonNegativeNumber(usage, "cacheWrite");
+  const hasCacheSignal = cacheRead !== undefined || cacheWrite !== undefined;
+
+  if (!hasCacheSignal && (input === undefined || !allowInputOnly)) return undefined;
+
+  // Under healthy Pi normalization input is the uncached portion, so
+  // totalInput = input + cacheRead + cacheWrite gives the full prompt token count.
+  // Guard against degenerate reads where a broken proxy omits prompt_tokens and
+  // Pi's input falls to zero: totalInput must never be less than cacheRead + cacheWrite.
+  const computed = (input ?? 0) + (cacheRead ?? 0) + (cacheWrite ?? 0);
+  const floor = (cacheRead ?? 0) + (cacheWrite ?? 0);
+  return {
+    cacheRead: cacheRead ?? 0,
+    cacheWrite: cacheWrite ?? 0,
+    totalInput: computed >= floor ? computed : floor,
+  };
+}
+
+// Raw fallback for DeepSeek responses that still carry their native usage fields.
+// In practice Pi normalizes usage before message_end fires, so this path is only
+// reached when Pi-normalized fields are absent (e.g. custom/foreign providers).
+export function getDeepSeekRawUsage(message: unknown): UsageSnapshot | undefined {
+  const usage = usageRecordFromAssistant(message);
+  if (!usage) return undefined;
+
+  const cacheRead = getFirstNonNegativeNumber(usage.prompt_cache_hit_tokens);
+  if (cacheRead === undefined) return undefined;
+
+  const cacheMiss = getFirstNonNegativeNumber(usage.prompt_cache_miss_tokens);
+  const promptTokens = getFirstNonNegativeNumber(usage.prompt_tokens);
+  // DeepSeek guarantees prompt_tokens = prompt_cache_hit_tokens + prompt_cache_miss_tokens.
+  const totalInput = promptTokens ?? cacheRead + (cacheMiss ?? 0);
+
+  return { cacheRead, cacheWrite: 0, totalInput };
+}
+
+// Raw fallback for OpenAI-family responses that still carry their native usage fields.
+// In practice Pi normalizes usage before message_end fires, so this path is only
+// reached when Pi-normalized fields are absent (e.g. custom/foreign providers).
+export function getOpenAIRawUsage(message: unknown): UsageSnapshot | undefined {
+  const usage = usageRecordFromAssistant(message);
+  if (!usage) return undefined;
+
+  const promptDetails = getNestedRecord(usage, "prompt_tokens_details") ?? getNestedRecord(usage, "promptTokensDetails");
+  const inputDetails = getNestedRecord(usage, "input_tokens_details") ?? getNestedRecord(usage, "inputTokensDetails");
+  const cacheRead = readCachedTokensFromDetails(promptDetails) ?? readCachedTokensFromDetails(inputDetails);
+  if (cacheRead === undefined) return undefined;
+
+  const cacheWrite = readCacheWriteFromDetails(promptDetails) ?? readCacheWriteFromDetails(inputDetails) ?? 0;
+  const totalInput = getFirstNonNegativeNumber(
+    usage.prompt_tokens,
+    usage.promptTokens,
+    usage.input_tokens,
+    usage.inputTokens,
+  ) ?? cacheRead + cacheWrite;
+
+  return { cacheRead, cacheWrite, totalInput };
+}
+
+// Raw fallback for Anthropic/Claude responses that still carry their native usage fields.
+// In practice Pi normalizes usage before message_end fires, so this path is only
+// reached when Pi-normalized fields are absent (e.g. custom/foreign providers).
+export function getAnthropicRawUsage(message: unknown): UsageSnapshot | undefined {
+  const usage = usageRecordFromAssistant(message);
+  if (!usage) return undefined;
+
+  const cacheRead = getFirstNonNegativeNumber(usage.cache_read_input_tokens, usage.cacheReadInputTokens);
+  const cacheWrite = getFirstNonNegativeNumber(usage.cache_creation_input_tokens, usage.cacheCreationInputTokens);
+  if (cacheRead === undefined && cacheWrite === undefined) return undefined;
+
+  // Anthropic input_tokens = tokens after the last cache breakpoint (neither read nor written).
+  const input = getFirstNonNegativeNumber(usage.input_tokens, usage.inputTokens) ?? 0;
+  return {
+    cacheRead: cacheRead ?? 0,
+    cacheWrite: cacheWrite ?? 0,
+    totalInput: input + (cacheRead ?? 0) + (cacheWrite ?? 0),
+  };
+}
+
+// Raw fallback for Gemini/Vertex responses that still carry their native usage fields.
+// In practice Pi normalizes usage before message_end fires, so this path is only
+// reached when Pi-normalized fields are absent (e.g. custom/foreign providers).
+export function getGeminiRawUsage(message: unknown): UsageSnapshot | undefined {
+  const record = getAssistantRecord(message);
+  if (!record) return undefined;
+
+  const usage = asRecord(record.usage);
+  const metadata =
+    getNestedRecord(record, "usageMetadata") ??
+    getNestedRecord(record, "usage_metadata") ??
+    getNestedRecord(usage, "usageMetadata") ??
+    getNestedRecord(usage, "usage_metadata") ??
+    usage;
+  if (!metadata) return undefined;
+
+  const cacheRead = getFirstNonNegativeNumber(
+    metadata.cachedContentTokenCount,
+    metadata.cached_content_token_count,
+  );
+  if (cacheRead === undefined) return undefined;
+
+  const totalInput = getFirstNonNegativeNumber(
+    metadata.promptTokenCount,
+    metadata.prompt_token_count,
+    metadata.inputTokenCount,
+    metadata.input_token_count,
+    usage?.input_tokens,
+    usage?.inputTokens,
+    usage?.prompt_tokens,
+    usage?.promptTokens,
+  ) ?? cacheRead;
+
+  return { cacheRead, cacheWrite: 0, totalInput };
+}
+
+// Try Pi-normalized usage first (always present for messages that went through Pi's
+// provider pipeline). Fall back to provider-specific raw-field readers when Pi-normalized
+// fields are absent (e.g. messages from custom/foreign providers whose raw usage shape
+// matches the official API).
+export function normalizeWithFallback(
+  message: unknown,
+  rawNormalizer: (message: unknown) => UsageSnapshot | undefined,
+  options: { allowInputOnlyPiUsage?: boolean } = {},
+): UsageSnapshot | undefined {
+  return getPiNormalizedUsage(message, options.allowInputOnlyPiUsage) ?? rawNormalizer(message);
+}