diff --git a/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Chat/ChatCompletionProvider.cs b/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Chat/ChatCompletionProvider.cs index 4843f7105..626e517ca 100644 --- a/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Chat/ChatCompletionProvider.cs +++ b/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Chat/ChatCompletionProvider.cs @@ -126,7 +126,8 @@ public async Task GetChatCompletions(Agent agent, List GetChatCompletionsAsync(Agent agent, List GetChatCompletionsStreamingAsync(Agent agent, Prompt = prompt, Provider = Provider, Model = _model, - TextInputTokens = tokenUsage?.PromptTokenCount ?? 0, + TextInputTokens = ProviderHelper.GetInputTokens(tokenUsage), + CachedTextInputTokens = ProviderHelper.GetCachedInputTokens(tokenUsage), TextOutputTokens = ProviderHelper.GetOutputTokens(tokenUsage) }); } diff --git a/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/ProviderHelper.cs b/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/ProviderHelper.cs index 91d38ae09..f1807abca 100644 --- a/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/ProviderHelper.cs +++ b/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/ProviderHelper.cs @@ -26,4 +26,18 @@ public static long GetOutputTokens(GenerativeAI.Types.UsageMetadata? usage) { return (usage?.CandidatesTokenCount ?? 0) + (usage?.ThoughtsTokenCount ?? 0); } + + /// + /// Gemini's prompt count includes the tokens served from the context cache, + /// which are billed at a separate rate, so they are subtracted here. + /// + public static long GetInputTokens(GenerativeAI.Types.UsageMetadata? usage) + { + return Math.Max((usage?.PromptTokenCount ?? 0) - GetCachedInputTokens(usage), 0); + } + + public static long GetCachedInputTokens(GenerativeAI.Types.UsageMetadata? usage) + { + return usage?.CachedContentTokenCount ?? 0; + } } diff --git a/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Text/TextCompletionProvider.cs b/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Text/TextCompletionProvider.cs index 6e5c267dc..6576c80e0 100644 --- a/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Text/TextCompletionProvider.cs +++ b/src/Plugins/BotSharp.Plugin.GoogleAI/Providers/Text/TextCompletionProvider.cs @@ -62,7 +62,8 @@ public async Task GetCompletion(string text, string agentId, string mess Prompt = text, Provider = Provider, Model = _model, - TextInputTokens = response?.UsageMetadata?.PromptTokenCount ?? 0, + TextInputTokens = ProviderHelper.GetInputTokens(response?.UsageMetadata), + CachedTextInputTokens = ProviderHelper.GetCachedInputTokens(response?.UsageMetadata), TextOutputTokens = ProviderHelper.GetOutputTokens(response?.UsageMetadata) }); } diff --git a/tests/BotSharp.Core.UnitTests/Plugins/GoogleAiUsageTests.cs b/tests/BotSharp.Core.UnitTests/Plugins/GoogleAiUsageTests.cs index c6142996d..b9ea42aa2 100644 --- a/tests/BotSharp.Core.UnitTests/Plugins/GoogleAiUsageTests.cs +++ b/tests/BotSharp.Core.UnitTests/Plugins/GoogleAiUsageTests.cs @@ -20,4 +20,22 @@ public void Output_tokens_without_thinking_are_the_candidates_count() Assert.Equal(20, ProviderHelper.GetOutputTokens(new UsageMetadata { CandidatesTokenCount = 20 })); Assert.Equal(0, ProviderHelper.GetOutputTokens(null)); } + + [Fact] + public void Cached_tokens_are_not_counted_as_fresh_input() + { + var usage = new UsageMetadata { PromptTokenCount = 1000, CachedContentTokenCount = 800 }; + + Assert.Equal(200, ProviderHelper.GetInputTokens(usage)); + Assert.Equal(800, ProviderHelper.GetCachedInputTokens(usage)); + } + + [Fact] + public void Input_tokens_without_cache_are_the_prompt_count() + { + Assert.Equal(50, ProviderHelper.GetInputTokens(new UsageMetadata { PromptTokenCount = 50 })); + Assert.Equal(0, ProviderHelper.GetCachedInputTokens(new UsageMetadata { PromptTokenCount = 50 })); + Assert.Equal(0, ProviderHelper.GetInputTokens(null)); + Assert.Equal(0, ProviderHelper.GetCachedInputTokens(null)); + } }