fix(llma): price override should not multiply by tokens

.changeset/clean-moons-kiss.md

@@ -0,0 +1,5 @@
+---
+'@posthog/ai': patch
+---
+
+fix price override params

packages/ai/src/gemini/index.ts

@@ -70,9 +70,7 @@ export class WrappedModels {
         usage: {
           inputTokens: metadata?.promptTokenCount ?? 0,
           outputTokens: metadata?.candidatesTokenCount ?? 0,
-          reasoningTokens:
-            (metadata as GenerateContentResponseUsageMetadata & { thoughtsTokenCount?: number })?.thoughtsTokenCount ??
-            0,
+          reasoningTokens: metadata?.thoughtsTokenCount ?? 0,
           cacheReadInputTokens: metadata?.cachedContentTokenCount ?? 0,
         },
         tools: availableTools,
@@ -170,9 +168,7 @@ export class WrappedModels {
           usage = {
             inputTokens: metadata.promptTokenCount ?? 0,
             outputTokens: metadata.candidatesTokenCount ?? 0,
-            reasoningTokens:
-              (metadata as GenerateContentResponseUsageMetadata & { thoughtsTokenCount?: number }).thoughtsTokenCount ??
-              0,
+            reasoningTokens: metadata.thoughtsTokenCount ?? 0,
             cacheReadInputTokens: metadata.cachedContentTokenCount ?? 0,
           }
         }

packages/ai/src/types.ts

@@ -55,7 +55,7 @@ export interface FormattedMessage {
 export interface TokenUsage {
   inputTokens?: number
   outputTokens?: number
-  reasoningTokens?: unknown // Use unknown since various providers return different types
-  cacheReadInputTokens?: unknown // Use unknown for provider flexibility
-  cacheCreationInputTokens?: unknown // Use unknown for provider flexibility
+  reasoningTokens?: number
+  cacheReadInputTokens?: number
+  cacheCreationInputTokens?: number
 }

packages/ai/src/utils.ts

@@ -412,16 +412,20 @@ export const sendEventToPosthog = async ({
   const safeError = sanitizeValues(error)
 
   let errorData = {}
+
   if (isError) {
     errorData = {
       $ai_is_error: true,
       $ai_error: safeError,
     }
   }
+
   let costOverrideData = {}
+
   if (params.posthogCostOverride) {
-    const inputCostUSD = (params.posthogCostOverride.inputCost ?? 0) * (usage.inputTokens ?? 0)
-    const outputCostUSD = (params.posthogCostOverride.outputCost ?? 0) * (usage.outputTokens ?? 0)
+    const inputCostUSD = params.posthogCostOverride.inputCost ?? 0
+    const outputCostUSD = params.posthogCostOverride.outputCost ?? 0
+
     costOverrideData = {
       $ai_input_cost_usd: inputCostUSD,
       $ai_output_cost_usd: outputCostUSD,

packages/ai/src/vercel/middleware.ts

@@ -1,4 +1,4 @@
-import { wrapLanguageModel } from 'ai'
+import { JSONValue, wrapLanguageModel } from 'ai'
 import type {
   LanguageModelV2,
   LanguageModelV2Content,
@@ -12,6 +12,7 @@ import { CostOverride, sendEventToPosthog, truncate, MAX_OUTPUT_SIZE, extractAva
 import { Buffer } from 'buffer'
 import { redactBase64DataUrl } from '../sanitization'
 import { isString } from '../typeGuards'
+import { TokenUsage } from '../types'
 
 interface ClientOptions {
   posthogDistinctId?: string
@@ -277,21 +278,19 @@ export const createInstrumentationMiddleware = (
         const baseURL = '' // cannot currently get baseURL from vercel
         const content = mapVercelOutput(result.content)
         const latency = (Date.now() - startTime) / 1000
-        const providerMetadata = result.providerMetadata
-        const additionalTokenValues = {
-          ...(providerMetadata?.anthropic
-            ? {
-                cacheCreationInputTokens: providerMetadata.anthropic.cacheCreationInputTokens,
-              }
-            : {}),
-        }
-        const usage = {
+
+        const anthropicCacheCreationInputTokens: JSONValue | undefined =
+          result.providerMetadata?.anthropic?.cacheCreationInputTokens
+
+        const usage: TokenUsage = {
           inputTokens: result.usage.inputTokens,
           outputTokens: result.usage.outputTokens,
           reasoningTokens: result.usage.reasoningTokens,
           cacheReadInputTokens: result.usage.cachedInputTokens,
-          ...additionalTokenValues,
+          cacheCreationInputTokens:
+            typeof anthropicCacheCreationInputTokens === 'number' ? anthropicCacheCreationInputTokens : undefined,
         }
+
         await sendEventToPosthog({
           client: phClient,
           distinctId: options.posthogDistinctId,
@@ -341,13 +340,7 @@ export const createInstrumentationMiddleware = (
       const startTime = Date.now()
       let generatedText = ''
       let reasoningText = ''
-      let usage: {
-        inputTokens?: number
-        outputTokens?: number
-        reasoningTokens?: any
-        cacheReadInputTokens?: any
-        cacheCreationInputTokens?: any
-      } = {}
+      let usage: TokenUsage = {}
       const mergedParams = {
         ...options,
         ...mapVercelParams(params),
@@ -410,20 +403,16 @@ export const createInstrumentationMiddleware = (
             }
 
             if (chunk.type === 'finish') {
-              const providerMetadata = chunk.providerMetadata
-              const additionalTokenValues = {
-                ...(providerMetadata?.anthropic
-                  ? {
-                      cacheCreationInputTokens: providerMetadata.anthropic.cacheCreationInputTokens,
-                    }
-                  : {}),
-              }
+              const anthropicCacheCreationInputTokens: JSONValue | undefined =
+                chunk.providerMetadata?.anthropic?.cacheCreationInputTokens
+
               usage = {
                 inputTokens: chunk.usage?.inputTokens,
                 outputTokens: chunk.usage?.outputTokens,
                 reasoningTokens: chunk.usage?.reasoningTokens,
                 cacheReadInputTokens: chunk.usage?.cachedInputTokens,
-                ...additionalTokenValues,
+                cacheCreationInputTokens:
+                  typeof anthropicCacheCreationInputTokens === 'number' ? anthropicCacheCreationInputTokens : undefined,
               }
             }
             controller.enqueue(chunk)