diff --git a/packages/ai-contracts/src/client.ts b/packages/ai-contracts/src/client.ts index fe567b2..3ad969e 100644 --- a/packages/ai-contracts/src/client.ts +++ b/packages/ai-contracts/src/client.ts @@ -36,6 +36,7 @@ export interface AamosResult { costUsd?: number | undefined; inputTokens?: number | undefined; outputTokens?: number | undefined; + totalTokens?: number | undefined; } /** diff --git a/packages/ai-contracts/src/gemini.ts b/packages/ai-contracts/src/gemini.ts index c1843fe..30981e8 100644 --- a/packages/ai-contracts/src/gemini.ts +++ b/packages/ai-contracts/src/gemini.ts @@ -376,7 +376,7 @@ export class GeminiAamosClient implements AamosClient { const geminiBody = (await res.json()) as { candidates?: Array<{ content?: { parts?: Array<{ text?: string }> } }>; - usageMetadata?: { promptTokenCount?: number; candidatesTokenCount?: number }; + usageMetadata?: { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number }; }; const text = geminiBody.candidates?.[0]?.content?.parts?.[0]?.text ?? ""; @@ -394,6 +394,7 @@ export class GeminiAamosClient implements AamosClient { const inputTokens = geminiBody.usageMetadata?.promptTokenCount ?? 0; const outputTokens = geminiBody.usageMetadata?.candidatesTokenCount ?? 0; + const totalTokens = geminiBody.usageMetadata?.totalTokenCount ?? 0; const costUsd = this.estimateCostUsd(inputTokens, outputTokens); await this.recordSpend(costUsd); @@ -409,6 +410,7 @@ export class GeminiAamosClient implements AamosClient { costUsd, inputTokens, outputTokens, + totalTokens, }; } @@ -471,7 +473,7 @@ export class GeminiAamosClient implements AamosClient { const geminiBody = (await res.json()) as { candidates?: Array<{ content?: { parts?: Array<{ text?: string }> } }>; - usageMetadata?: { promptTokenCount?: number; candidatesTokenCount?: number }; + usageMetadata?: { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number }; }; const text = geminiBody.candidates?.[0]?.content?.parts?.[0]?.text ?? ""; @@ -497,6 +499,7 @@ export class GeminiAamosClient implements AamosClient { const inputTokens = geminiBody.usageMetadata?.promptTokenCount ?? 0; const outputTokens = geminiBody.usageMetadata?.candidatesTokenCount ?? 0; + const totalTokens = geminiBody.usageMetadata?.totalTokenCount ?? 0; const costUsd = this.estimateCostUsd(inputTokens, outputTokens); await this.recordSpend(costUsd); @@ -512,6 +515,7 @@ export class GeminiAamosClient implements AamosClient { costUsd, inputTokens, outputTokens, + totalTokens, }; } @@ -642,7 +646,7 @@ export class GeminiAamosClient implements AamosClient { const geminiBody = (await res.json()) as { candidates?: Array<{ content?: { parts?: Array<{ text?: string }> } }>; - usageMetadata?: { promptTokenCount?: number; candidatesTokenCount?: number }; + usageMetadata?: { promptTokenCount?: number; candidatesTokenCount?: number; totalTokenCount?: number }; }; const text = geminiBody.candidates?.[0]?.content?.parts?.[0]?.text ?? ""; @@ -666,6 +670,7 @@ export class GeminiAamosClient implements AamosClient { const inputTokens = geminiBody.usageMetadata?.promptTokenCount ?? 0; const outputTokens = geminiBody.usageMetadata?.candidatesTokenCount ?? 0; + const totalTokens = geminiBody.usageMetadata?.totalTokenCount ?? 0; const costUsd = this.estimateCostUsd(inputTokens, outputTokens); await this.recordSpend(costUsd); @@ -688,6 +693,7 @@ export class GeminiAamosClient implements AamosClient { costUsd, inputTokens, outputTokens, + totalTokens, }; }