@lobehub/chat/packages/model-runtime/src/core/usageConverters/google-ai.ts

Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.

import {
  GenerateContentResponseUsageMetadata,
  MediaModality,
  ModalityTokenCount,
} from '@google/genai';
import { ModelUsage } from '@lobechat/types';
import type { Pricing } from 'model-bank';

import { withUsageCost } from './utils/withUsageCost';

const getTokenCount = (details: ModalityTokenCount[] | undefined, modality: MediaModality) => {
  return details?.find((detail) => detail?.modality === modality)?.tokenCount;
};

export const convertGoogleAIUsage = (
  usage: GenerateContentResponseUsageMetadata,
  pricing?: Pricing,
): ModelUsage => {
  const inputCacheMissTokens =
    usage.promptTokenCount && usage.cachedContentTokenCount
      ? usage.promptTokenCount - usage.cachedContentTokenCount
      : undefined;

  const reasoningTokens = usage.thoughtsTokenCount;
  const candidatesDetails = usage.candidatesTokensDetails;
  const totalCandidatesTokens =
    usage.candidatesTokenCount ??
    candidatesDetails?.reduce((sum, detail) => sum + (detail?.tokenCount ?? 0), 0) ??
    0;

  const outputImageTokens = getTokenCount(candidatesDetails, MediaModality.IMAGE) ?? 0;
  const textTokensFromDetails = getTokenCount(candidatesDetails, MediaModality.TEXT);
  const outputTextTokens =
    typeof textTokensFromDetails === 'number' && textTokensFromDetails > 0
      ? textTokensFromDetails
      : Math.max(0, totalCandidatesTokens - outputImageTokens);
  const totalOutputTokens = totalCandidatesTokens + (reasoningTokens ?? 0);

  const normalizedUsage = {
    inputAudioTokens: getTokenCount(usage.promptTokensDetails, MediaModality.AUDIO),
    inputCacheMissTokens,
    inputCachedTokens: usage.cachedContentTokenCount,
    inputImageTokens: getTokenCount(usage.promptTokensDetails, MediaModality.IMAGE),
    inputTextTokens: getTokenCount(usage.promptTokensDetails, MediaModality.TEXT),
    outputImageTokens,
    outputReasoningTokens: reasoningTokens,
    outputTextTokens,
    totalInputTokens: usage.promptTokenCount,
    totalOutputTokens,
    totalTokens: usage.totalTokenCount,
  } satisfies ModelUsage;

  return withUsageCost(normalizedUsage, pricing);
};