diff --git a/Program.cs b/Program.cs index e6923fb..f5fb758 100644 --- a/Program.cs +++ b/Program.cs @@ -53,13 +53,14 @@ string GetRequired(string key) => // the innermost wrapper around the raw client — so it observes every individual // model round-trip (including the extra calls tool invocation triggers) and // enforces a hard cumulative-token budget for the whole process. +TokenCapChatClient? tokenCapChatClient = null; IChatClient llm = openAIClient .GetChatClient(modelName) .AsIChatClient() .AsBuilder() .UseFunctionInvocation() .UseOpenTelemetry(sourceName: "BlogWriter.ChatClient") - .Use(inner => new TokenCapChatClient(inner, maxTotalTokens)) + .Use(inner => tokenCapChatClient = new TokenCapChatClient(inner, maxTotalTokens)) .Build(); var chatOptions = new ChatOptions @@ -190,3 +191,14 @@ async Task PostWithRetryAsync(string requestUri, object bod Console.WriteLine($"Revision Number: {result.RevisionNumber}"); Console.WriteLine("============================="); +if (tokenCapChatClient is not null) +{ + TokenUsageSnapshot usage = tokenCapChatClient.UsageSnapshot; + Console.WriteLine("\n========== TOKEN USAGE =========="); + Console.WriteLine($"Input tokens: {usage.InputTokens}"); + Console.WriteLine($"Output tokens: {usage.OutputTokens}"); + Console.WriteLine($"Reasoning tokens: {usage.ReasoningTokens}"); + Console.WriteLine($"Total tokens: {usage.TotalTokens}"); + Console.WriteLine("=================================="); +} + diff --git a/TokenCapChatClient.cs b/TokenCapChatClient.cs index b70a609..245cbde 100644 --- a/TokenCapChatClient.cs +++ b/TokenCapChatClient.cs @@ -12,8 +12,15 @@ namespace BlogWriter; /// public sealed class TokenCapChatClient : DelegatingChatClient { + // Key used by the OpenAI connector to report reasoning tokens inside + // UsageDetails.AdditionalCounts (there is no dedicated top-level property). + private const string ReasoningTokenCountKey = "OutputTokenDetails.ReasoningTokenCount"; + private readonly long _maxTotalTokens; private long _totalTokens; + private long _inputTokens; + private long _outputTokens; + private long _reasoningTokens; public TokenCapChatClient(IChatClient innerClient, long maxTotalTokens) : base(innerClient) { @@ -22,6 +29,13 @@ public TokenCapChatClient(IChatClient innerClient, long maxTotalTokens) : base(i : throw new ArgumentOutOfRangeException(nameof(maxTotalTokens), maxTotalTokens, "Token cap must be a positive number."); } + /// Cumulative token usage observed across every model round-trip so far. + public TokenUsageSnapshot UsageSnapshot => new( + Interlocked.Read(ref _inputTokens), + Interlocked.Read(ref _outputTokens), + Interlocked.Read(ref _reasoningTokens), + Interlocked.Read(ref _totalTokens)); + public override async Task GetResponseAsync( IEnumerable messages, ChatOptions? options = null, @@ -54,7 +68,20 @@ public override async IAsyncEnumerable GetStreamingResponseA private void Track(UsageDetails? usage) { - long used = usage?.TotalTokenCount ?? 0; + if (usage is null) + { + return; + } + + Interlocked.Add(ref _inputTokens, usage.InputTokenCount ?? 0); + Interlocked.Add(ref _outputTokens, usage.OutputTokenCount ?? 0); + if (usage.AdditionalCounts is { } additionalCounts && + additionalCounts.TryGetValue(ReasoningTokenCountKey, out long reasoningTokens)) + { + Interlocked.Add(ref _reasoningTokens, reasoningTokens); + } + + long used = usage.TotalTokenCount ?? 0; if (used == 0) { return; @@ -68,6 +95,9 @@ private void Track(UsageDetails? usage) } } +/// Point-in-time totals of tokens consumed across all model round-trips. +public readonly record struct TokenUsageSnapshot(long InputTokens, long OutputTokens, long ReasoningTokens, long TotalTokens); + /// /// Thrown when cumulative model token usage exceeds the configured cap. Callers /// catch this to shut down gracefully instead of continuing to spend tokens.