Skip to content

Add token usage tracking to TokenCapChatClient and display usage summary #15

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JesseLiberty merged 1 commit into from
Aug 24, 2026
Merged

Add token usage tracking to TokenCapChatClient and display usage summary #15

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion Program.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -190,3 +191,14 @@ async Task<HttpResponseMessage> 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("==================================");
}

32 changes: 31 additions & 1 deletion TokenCapChatClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,15 @@ namespace BlogWriter;
/// </summary>
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)
{
Expand All @@ -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.");
}

/// <summary>Cumulative token usage observed across every model round-trip so far.</summary>
public TokenUsageSnapshot UsageSnapshot => new(
Interlocked.Read(ref _inputTokens),
Interlocked.Read(ref _outputTokens),
Interlocked.Read(ref _reasoningTokens),
Interlocked.Read(ref _totalTokens));

public override async Task<ChatResponse> GetResponseAsync(
IEnumerable<ChatMessage> messages,
ChatOptions? options = null,
Expand Down Expand Up @@ -54,7 +68,20 @@ public override async IAsyncEnumerable<ChatResponseUpdate> 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);
}
Comment on lines +76 to +82

long used = usage.TotalTokenCount ?? 0;
if (used == 0)
{
return;
Expand All @@ -68,6 +95,9 @@ private void Track(UsageDetails? usage)
}
}

/// <summary>Point-in-time totals of tokens consumed across all model round-trips.</summary>
public readonly record struct TokenUsageSnapshot(long InputTokens, long OutputTokens, long ReasoningTokens, long TotalTokens);

/// <summary>
/// Thrown when cumulative model token usage exceeds the configured cap. Callers
/// catch this to shut down gracefully instead of continuing to spend tokens.
Expand Down