Skip to content

Commit 9085192

Browse files
stephentoubstephentoub
committed
Augment UsageDetails with cached / reasoning token counts (dotnet#7122)
Cached tokens are currently reported by Anthropic, Gemini, OpenAI, and AWS. Reasoning tokens are currently reported by OpenAI and Gemini.
1 parent d0d7beb commit 9085192

8 files changed

Lines changed: 259 additions & 38 deletions

File tree

‎src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2236,6 +2236,14 @@
22362236
{
22372237
"Member": "long? Microsoft.Extensions.AI.UsageDetails.TotalTokenCount { get; set; }",
22382238
"Stage": "Stable"
2239+
},
2240+
{
2241+
"Member": "long? Microsoft.Extensions.AI.UsageDetails.CachedInputTokenCount { get; set; }",
2242+
"Stage": "Stable"
2243+
},
2244+
{
2245+
"Member": "long? Microsoft.Extensions.AI.UsageDetails.ReasoningTokenCount { get; set; }",
2246+
"Stage": "Stable"
22392247
}
22402248
]
22412249
}

‎src/Libraries/Microsoft.Extensions.AI.Abstractions/UsageDetails.cs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,23 @@ public class UsageDetails
2121
/// <summary>Gets or sets the total number of tokens used to produce the response.</summary>
2222
public long? TotalTokenCount { get; set; }
2323

24+
/// <summary>
25+
/// Gets or sets the number of input tokens that were read from a cache.
26+
/// </summary>
27+
/// <remarks>
28+
/// Cached input tokens should be counted as part of <see cref="InputTokenCount"/>.
29+
/// </remarks>
30+
public long? CachedInputTokenCount { get; set; }
31+
32+
/// <summary>
33+
/// Gets or sets the number of "reasoning" / "thinking" tokens used internally
34+
/// by the model.
35+
/// </summary>
36+
/// <remarks>
37+
/// Reasoning tokens should be counted as part of <see cref="OutputTokenCount"/>.
38+
/// </remarks>
39+
public long? ReasoningTokenCount { get; set; }
40+
2441
/// <summary>Gets or sets a dictionary of additional usage counts.</summary>
2542
/// <remarks>
2643
/// All values set here are assumed to be summable. For example, when middleware makes multiple calls to an underlying
@@ -38,6 +55,8 @@ public void Add(UsageDetails usage)
3855
InputTokenCount = NullableSum(InputTokenCount, usage.InputTokenCount);
3956
OutputTokenCount = NullableSum(OutputTokenCount, usage.OutputTokenCount);
4057
TotalTokenCount = NullableSum(TotalTokenCount, usage.TotalTokenCount);
58+
CachedInputTokenCount = NullableSum(CachedInputTokenCount, usage.CachedInputTokenCount);
59+
ReasoningTokenCount = NullableSum(ReasoningTokenCount, usage.ReasoningTokenCount);
4160

4261
if (usage.AdditionalCounts is { } countsToAdd)
4362
{
@@ -80,6 +99,16 @@ internal string DebuggerDisplay
8099
parts.Add($"{nameof(TotalTokenCount)} = {total}");
81100
}
82101

102+
if (CachedInputTokenCount is { } cached)
103+
{
104+
parts.Add($"{nameof(CachedInputTokenCount)} = {cached}");
105+
}
106+
107+
if (ReasoningTokenCount is { } reasoning)
108+
{
109+
parts.Add($"{nameof(ReasoningTokenCount)} = {reasoning}");
110+
}
111+
83112
if (AdditionalCounts is { } additionalCounts)
84113
{
85114
foreach (var entry in additionalCounts)

‎src/Libraries/Microsoft.Extensions.AI.OpenAI/OpenAIChatClient.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -644,6 +644,8 @@ private static UsageDetails FromOpenAIUsage(ChatTokenUsage tokenUsage)
644644
InputTokenCount = tokenUsage.InputTokenCount,
645645
OutputTokenCount = tokenUsage.OutputTokenCount,
646646
TotalTokenCount = tokenUsage.TotalTokenCount,
647+
CachedInputTokenCount = tokenUsage.InputTokenDetails?.CachedTokenCount,
648+
ReasoningTokenCount = tokenUsage.OutputTokenDetails?.ReasoningTokenCount,
647649
AdditionalCounts = [],
648650
};
649651

@@ -653,13 +655,11 @@ private static UsageDetails FromOpenAIUsage(ChatTokenUsage tokenUsage)
653655
{
654656
const string InputDetails = nameof(ChatTokenUsage.InputTokenDetails);
655657
counts.Add($"{InputDetails}.{nameof(ChatInputTokenUsageDetails.AudioTokenCount)}", inputDetails.AudioTokenCount);
656-
counts.Add($"{InputDetails}.{nameof(ChatInputTokenUsageDetails.CachedTokenCount)}", inputDetails.CachedTokenCount);
657658
}
658659

659660
if (tokenUsage.OutputTokenDetails is ChatOutputTokenUsageDetails outputDetails)
660661
{
661662
const string OutputDetails = nameof(ChatTokenUsage.OutputTokenDetails);
662-
counts.Add($"{OutputDetails}.{nameof(ChatOutputTokenUsageDetails.ReasoningTokenCount)}", outputDetails.ReasoningTokenCount);
663663
counts.Add($"{OutputDetails}.{nameof(ChatOutputTokenUsageDetails.AudioTokenCount)}", outputDetails.AudioTokenCount);
664664
counts.Add($"{OutputDetails}.{nameof(ChatOutputTokenUsageDetails.AcceptedPredictionTokenCount)}", outputDetails.AcceptedPredictionTokenCount);
665665
counts.Add($"{OutputDetails}.{nameof(ChatOutputTokenUsageDetails.RejectedPredictionTokenCount)}", outputDetails.RejectedPredictionTokenCount);

‎src/Libraries/Microsoft.Extensions.AI.OpenAI/OpenAIResponsesChatClient.cs‎

Lines changed: 2 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1143,19 +1143,9 @@ static FunctionCallOutputResponseItem SerializeAIContent(string callId, IEnumera
11431143
InputTokenCount = usage.InputTokenCount,
11441144
OutputTokenCount = usage.OutputTokenCount,
11451145
TotalTokenCount = usage.TotalTokenCount,
1146+
CachedInputTokenCount = usage.InputTokenDetails?.CachedTokenCount,
1147+
ReasoningTokenCount = usage.OutputTokenDetails?.ReasoningTokenCount,
11461148
};
1147-
1148-
if (usage.InputTokenDetails is { } inputDetails)
1149-
{
1150-
ud.AdditionalCounts ??= [];
1151-
ud.AdditionalCounts.Add($"{nameof(usage.InputTokenDetails)}.{nameof(inputDetails.CachedTokenCount)}", inputDetails.CachedTokenCount);
1152-
}
1153-
1154-
if (usage.OutputTokenDetails is { } outputDetails)
1155-
{
1156-
ud.AdditionalCounts ??= [];
1157-
ud.AdditionalCounts.Add($"{nameof(usage.OutputTokenDetails)}.{nameof(outputDetails.ReasoningTokenCount)}", outputDetails.ReasoningTokenCount);
1158-
}
11591149
}
11601150

11611151
return ud;

‎test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/Contents/UsageContentTests.cs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,9 @@ public void Serialization_Roundtrips()
6666
{
6767
InputTokenCount = 10,
6868
OutputTokenCount = 20,
69-
TotalTokenCount = 30
69+
TotalTokenCount = 30,
70+
CachedInputTokenCount = 5,
71+
ReasoningTokenCount = 8
7072
});
7173

7274
var json = JsonSerializer.Serialize(content, AIJsonUtilities.DefaultOptions);
@@ -77,5 +79,7 @@ public void Serialization_Roundtrips()
7779
Assert.Equal(content.Details.InputTokenCount, deserializedContent.Details.InputTokenCount);
7880
Assert.Equal(content.Details.OutputTokenCount, deserializedContent.Details.OutputTokenCount);
7981
Assert.Equal(content.Details.TotalTokenCount, deserializedContent.Details.TotalTokenCount);
82+
Assert.Equal(content.Details.CachedInputTokenCount, deserializedContent.Details.CachedInputTokenCount);
83+
Assert.Equal(content.Details.ReasoningTokenCount, deserializedContent.Details.ReasoningTokenCount);
8084
}
8185
}

‎test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/UsageDetailsTests.cs‎

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
4+
using System;
5+
using System.Text.Json;
6+
using Xunit;
7+
8+
namespace Microsoft.Extensions.AI;
9+
10+
public class UsageDetailsTests
11+
{
12+
[Fact]
13+
public void Constructor_PropsDefault()
14+
{
15+
UsageDetails details = new();
16+
Assert.Null(details.InputTokenCount);
17+
Assert.Null(details.OutputTokenCount);
18+
Assert.Null(details.TotalTokenCount);
19+
Assert.Null(details.CachedInputTokenCount);
20+
Assert.Null(details.ReasoningTokenCount);
21+
Assert.Null(details.AdditionalCounts);
22+
}
23+
24+
[Fact]
25+
public void Properties_Roundtrip()
26+
{
27+
UsageDetails details = new()
28+
{
29+
InputTokenCount = 10,
30+
OutputTokenCount = 20,
31+
TotalTokenCount = 30,
32+
CachedInputTokenCount = 5,
33+
ReasoningTokenCount = 8,
34+
AdditionalCounts = new() { ["custom"] = 100 }
35+
};
36+
37+
Assert.Equal(10, details.InputTokenCount);
38+
Assert.Equal(20, details.OutputTokenCount);
39+
Assert.Equal(30, details.TotalTokenCount);
40+
Assert.Equal(5, details.CachedInputTokenCount);
41+
Assert.Equal(8, details.ReasoningTokenCount);
42+
Assert.NotNull(details.AdditionalCounts);
43+
Assert.Equal(100, details.AdditionalCounts["custom"]);
44+
}
45+
46+
[Fact]
47+
public void Add_NullUsage_Throws()
48+
{
49+
UsageDetails details = new();
50+
Assert.Throws<ArgumentNullException>("usage", () => details.Add(null!));
51+
}
52+
53+
[Fact]
54+
public void Add_SumsAllProperties()
55+
{
56+
UsageDetails details1 = new()
57+
{
58+
InputTokenCount = 10,
59+
OutputTokenCount = 20,
60+
TotalTokenCount = 30,
61+
CachedInputTokenCount = 5,
62+
ReasoningTokenCount = 8,
63+
};
64+
65+
UsageDetails details2 = new()
66+
{
67+
InputTokenCount = 15,
68+
OutputTokenCount = 25,
69+
TotalTokenCount = 40,
70+
CachedInputTokenCount = 7,
71+
ReasoningTokenCount = 12,
72+
};
73+
74+
details1.Add(details2);
75+
76+
Assert.Equal(25, details1.InputTokenCount);
77+
Assert.Equal(45, details1.OutputTokenCount);
78+
Assert.Equal(70, details1.TotalTokenCount);
79+
Assert.Equal(12, details1.CachedInputTokenCount);
80+
Assert.Equal(20, details1.ReasoningTokenCount);
81+
}
82+
83+
[Fact]
84+
public void Add_WithNullValues_HandlesCorrectly()
85+
{
86+
UsageDetails details1 = new()
87+
{
88+
InputTokenCount = 10,
89+
CachedInputTokenCount = 5,
90+
};
91+
92+
UsageDetails details2 = new()
93+
{
94+
OutputTokenCount = 25,
95+
ReasoningTokenCount = 12,
96+
};
97+
98+
details1.Add(details2);
99+
100+
Assert.Equal(10, details1.InputTokenCount);
101+
Assert.Equal(25, details1.OutputTokenCount);
102+
Assert.Null(details1.TotalTokenCount);
103+
Assert.Equal(5, details1.CachedInputTokenCount);
104+
Assert.Equal(12, details1.ReasoningTokenCount);
105+
}
106+
107+
[Fact]
108+
public void Add_FromNullToValue_SetsValue()
109+
{
110+
UsageDetails details1 = new();
111+
112+
UsageDetails details2 = new()
113+
{
114+
CachedInputTokenCount = 5,
115+
ReasoningTokenCount = 10,
116+
};
117+
118+
details1.Add(details2);
119+
120+
Assert.Equal(5, details1.CachedInputTokenCount);
121+
Assert.Equal(10, details1.ReasoningTokenCount);
122+
}
123+
124+
[Fact]
125+
public void Add_AdditionalCounts_MergesCorrectly()
126+
{
127+
UsageDetails details1 = new()
128+
{
129+
AdditionalCounts = new() { ["key1"] = 10, ["key2"] = 20 }
130+
};
131+
132+
UsageDetails details2 = new()
133+
{
134+
AdditionalCounts = new() { ["key2"] = 30, ["key3"] = 40 }
135+
};
136+
137+
details1.Add(details2);
138+
139+
Assert.NotNull(details1.AdditionalCounts);
140+
Assert.Equal(10, details1.AdditionalCounts["key1"]);
141+
Assert.Equal(50, details1.AdditionalCounts["key2"]);
142+
Assert.Equal(40, details1.AdditionalCounts["key3"]);
143+
}
144+
145+
[Fact]
146+
public void Serialization_Roundtrips()
147+
{
148+
UsageDetails details = new()
149+
{
150+
InputTokenCount = 10,
151+
OutputTokenCount = 20,
152+
TotalTokenCount = 30,
153+
CachedInputTokenCount = 5,
154+
ReasoningTokenCount = 8,
155+
AdditionalCounts = new() { ["custom"] = 100 }
156+
};
157+
158+
string json = JsonSerializer.Serialize(details, AIJsonUtilities.DefaultOptions);
159+
UsageDetails? deserialized = JsonSerializer.Deserialize<UsageDetails>(json, AIJsonUtilities.DefaultOptions);
160+
161+
Assert.NotNull(deserialized);
162+
Assert.Equal(details.InputTokenCount, deserialized.InputTokenCount);
163+
Assert.Equal(details.OutputTokenCount, deserialized.OutputTokenCount);
164+
Assert.Equal(details.TotalTokenCount, deserialized.TotalTokenCount);
165+
Assert.Equal(details.CachedInputTokenCount, deserialized.CachedInputTokenCount);
166+
Assert.Equal(details.ReasoningTokenCount, deserialized.ReasoningTokenCount);
167+
Assert.NotNull(deserialized.AdditionalCounts);
168+
Assert.Equal(100, deserialized.AdditionalCounts["custom"]);
169+
}
170+
171+
[Fact]
172+
public void Serialization_WithNullProperties_Roundtrips()
173+
{
174+
UsageDetails details = new()
175+
{
176+
InputTokenCount = 10,
177+
OutputTokenCount = 20,
178+
};
179+
180+
string json = JsonSerializer.Serialize(details, AIJsonUtilities.DefaultOptions);
181+
UsageDetails? deserialized = JsonSerializer.Deserialize<UsageDetails>(json, AIJsonUtilities.DefaultOptions);
182+
183+
Assert.NotNull(deserialized);
184+
Assert.Equal(10, deserialized.InputTokenCount);
185+
Assert.Equal(20, deserialized.OutputTokenCount);
186+
Assert.Null(deserialized.TotalTokenCount);
187+
Assert.Null(deserialized.CachedInputTokenCount);
188+
Assert.Null(deserialized.ReasoningTokenCount);
189+
}
190+
}

0 commit comments

Comments
 (0)