Skip to content

Commit 483e769

Browse files
stephentoubamanda-tarafa
stephentoub
authored and
amanda-tarafa
committed
feat: Update to M.E.AI.Abstractions 10.1.1
Additional usage token counts are strongly-typed. Also fix usage reporting in embedding generation.
1 parent 7f38b46 commit 483e769

File tree

6 files changed

+108
-42
lines changed

6 files changed

+108
-42
lines changed

‎Directory.Packages.props‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
<PackageVersion Include="Google.Protobuf" Version="3.31.1" />
2020

2121
<!-- System packages -->
22-
<PackageVersion Include="Microsoft.Bcl.AsyncInterfaces" Version="10.0.0" />
22+
<PackageVersion Include="Microsoft.Bcl.AsyncInterfaces" Version="10.0.1" />
2323

2424
<!-- MS Extensions packages -->
2525
<PackageVersion Include="Microsoft.Extensions.Http" Version="8.0.1" />
@@ -29,7 +29,7 @@
2929
<PackageVersion Include="Microsoft.Extensions.Hosting" Version="8.0.1" />
3030
<PackageVersion Include="Microsoft.Extensions.Hosting.Abstractions" Version="8.0.1" />
3131
<PackageVersion Include="Microsoft.Extensions.Logging.Console" Version= "8.0.1" />
32-
<PackageVersion Include="Microsoft.Extensions.AI.Abstractions" Version= "10.0.0" />
32+
<PackageVersion Include="Microsoft.Extensions.AI.Abstractions" Version= "10.1.1" />
3333

3434
<!-- Misc packages (might warrant a reorganization at some point) -->
3535
<PackageVersion Include="xunit" Version="2.9.3" />
@@ -39,7 +39,7 @@
3939
<PackageVersion Include="Xunit.SkippableFact" Version="1.5.23" />
4040
<PackageVersion Include="Microsoft.NET.Test.Sdk" Version="18.0.1" />
4141
<PackageVersion Include="NSubstitute" Version="5.3.0" />
42-
<PackageVersion Include="System.Linq.AsyncEnumerable" Version="10.0.0" />
42+
<PackageVersion Include="System.Linq.AsyncEnumerable" Version="10.0.1" />
4343
<PackageVersion Include="Newtonsoft.Json" Version="13.0.4" />
4444
<PackageVersion Include="YamlDotNet" Version="16.3.0" />
4545
<PackageVersion Include="CommandLineParser" Version="2.9.1" />

‎apis/Google.Cloud.VertexAI.Extensions/Google.Cloud.VertexAI.Extensions.Tests/BuildIChatClientTest.cs‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1433,8 +1433,8 @@ public async Task IChatClient_GetResponseAsync_UsageDetails_OnlyNonZeroAdditiona
14331433
PromptTokenCount = 10,
14341434
CandidatesTokenCount = 5,
14351435
TotalTokenCount = 15,
1436-
CachedContentTokenCount = 0, // Should not appear
1437-
ThoughtsTokenCount = 7 // Should appear
1436+
CachedContentTokenCount = 0,
1437+
ThoughtsTokenCount = 7
14381438
};
14391439
return response;
14401440
}
@@ -1447,10 +1447,13 @@ public async Task IChatClient_GetResponseAsync_UsageDetails_OnlyNonZeroAdditiona
14471447

14481448
Assert.NotNull(result);
14491449
Assert.NotNull(result.Usage);
1450-
Assert.NotNull(result.Usage.AdditionalCounts);
1451-
Assert.Single(result.Usage.AdditionalCounts);
1452-
Assert.Equal(7, result.Usage.AdditionalCounts["ThoughtsTokenCount"]);
1453-
Assert.False(result.Usage.AdditionalCounts.ContainsKey("CachedContentTokenCount"));
1450+
1451+
Assert.Equal(10, result.Usage.InputTokenCount);
1452+
Assert.Equal(5, result.Usage.OutputTokenCount);
1453+
Assert.Equal(15, result.Usage.TotalTokenCount);
1454+
Assert.Equal(7, result.Usage.ReasoningTokenCount);
1455+
Assert.Null(result.Usage.CachedInputTokenCount);
1456+
Assert.Null(result.Usage.AdditionalCounts);
14541457
}
14551458

14561459
[Fact]
@@ -1586,9 +1589,9 @@ public async Task IChatClient_GetResponseAsync_WithUsageMetadata()
15861589
Assert.Equal(10, result.Usage.InputTokenCount);
15871590
Assert.Equal(5, result.Usage.OutputTokenCount);
15881591
Assert.Equal(15, result.Usage.TotalTokenCount);
1589-
Assert.NotNull(result.Usage.AdditionalCounts);
1590-
Assert.Equal(2, result.Usage.AdditionalCounts["CachedContentTokenCount"]);
1591-
Assert.Equal(3, result.Usage.AdditionalCounts["ThoughtsTokenCount"]);
1592+
Assert.Equal(2, result.Usage.CachedInputTokenCount);
1593+
Assert.Equal(3, result.Usage.ReasoningTokenCount);
1594+
Assert.Null(result.Usage.AdditionalCounts);
15921595
}
15931596

15941597
[Fact]

‎apis/Google.Cloud.VertexAI.Extensions/Google.Cloud.VertexAI.Extensions.Tests/BuildIEmbeddingGeneratorTest.cs‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,83 @@ public async Task IEmbeddingGenerator_GenerateAsync_WithComplexModelPath()
500500
Assert.Equal([0.4f, 0.5f, 0.6f], result[0].Vector.ToArray());
501501
}
502502

503+
[Fact]
504+
public async Task IEmbeddingGenerator_GenerateAsync_UsageDetails_NoTokenCounts_ReturnsNull()
505+
{
506+
DelegateCallInvoker invoker = new()
507+
{
508+
OnPredictRequest = request =>
509+
{
510+
PredictResponse response = new();
511+
response.Predictions.Add(CreateEmbeddingPrediction([0.1f, 0.2f, 0.3f]));
512+
return response;
513+
}
514+
};
515+
516+
IEmbeddingGenerator<string, Embedding<float>> embeddingGenerator = CreateClientBuilder(invoker).BuildIEmbeddingGenerator("projects/test-project/locations/us-central1/publishers/google/models/mymodel");
517+
518+
GeneratedEmbeddings<Embedding<float>> result = await embeddingGenerator.GenerateAsync(["Test"]);
519+
520+
Assert.NotNull(result);
521+
Assert.Single(result);
522+
Assert.Null(result.Usage);
523+
}
524+
525+
[Fact]
526+
public async Task IEmbeddingGenerator_GenerateAsync_UsageDetails_WithTokenCounts()
527+
{
528+
DelegateCallInvoker invoker = new()
529+
{
530+
OnPredictRequest = request =>
531+
{
532+
PredictResponse response = new();
533+
response.Predictions.Add(CreateEmbeddingPredictionWithTokens([0.1f, 0.2f, 0.3f], 10));
534+
return response;
535+
}
536+
};
537+
538+
IEmbeddingGenerator<string, Embedding<float>> embeddingGenerator = CreateClientBuilder(invoker).BuildIEmbeddingGenerator("projects/test-project/locations/us-central1/publishers/google/models/mymodel");
539+
540+
GeneratedEmbeddings<Embedding<float>> result = await embeddingGenerator.GenerateAsync(["Test with tokens"]);
541+
542+
Assert.NotNull(result);
543+
Assert.Single(result);
544+
Assert.NotNull(result.Usage);
545+
Assert.Equal(10, result.Usage.InputTokenCount);
546+
Assert.Null(result.Usage.OutputTokenCount);
547+
Assert.Null(result.Usage.ReasoningTokenCount);
548+
Assert.Equal(10, result.Usage.TotalTokenCount);
549+
}
550+
551+
[Fact]
552+
public async Task IEmbeddingGenerator_GenerateAsync_UsageDetails_SumsMultipleTokenCounts()
553+
{
554+
DelegateCallInvoker invoker = new()
555+
{
556+
OnPredictRequest = request =>
557+
{
558+
PredictResponse response = new();
559+
response.Predictions.Add(CreateEmbeddingPredictionWithTokens([0.1f, 0.2f, 0.3f], 5));
560+
response.Predictions.Add(CreateEmbeddingPredictionWithTokens([0.4f, 0.5f, 0.6f], 8));
561+
response.Predictions.Add(CreateEmbeddingPredictionWithTokens([0.7f, 0.8f, 0.9f], 12));
562+
return response;
563+
}
564+
};
565+
566+
IEmbeddingGenerator<string, Embedding<float>> embeddingGenerator = CreateClientBuilder(invoker).BuildIEmbeddingGenerator("projects/test-project/locations/us-central1/publishers/google/models/mymodel");
567+
568+
GeneratedEmbeddings<Embedding<float>> result = await embeddingGenerator.GenerateAsync(["First", "Second", "Third"]);
569+
570+
Assert.NotNull(result);
571+
Assert.Equal(3, result.Count);
572+
Assert.NotNull(result.Usage);
573+
Assert.Equal(25, result.Usage.InputTokenCount);
574+
Assert.Null(result.Usage.OutputTokenCount);
575+
Assert.Null(result.Usage.CachedInputTokenCount);
576+
Assert.Null(result.Usage.ReasoningTokenCount);
577+
Assert.Equal(25, result.Usage.TotalTokenCount);
578+
}
579+
503580
private static Value CreateEmbeddingPrediction(float[] embedding) =>
504581
new()
505582
{

‎apis/Google.Cloud.VertexAI.Extensions/Google.Cloud.VertexAI.Extensions/PredictionServiceChatClient.cs‎

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -698,21 +698,13 @@ private static ChatFinishReason ConvertFinishReason(Candidate.Types.FinishReason
698698

699699
UsageDetails usage = new()
700700
{
701-
InputTokenCount = usageMetadata.PromptTokenCount,
702-
OutputTokenCount = usageMetadata.CandidatesTokenCount,
703-
TotalTokenCount = usageMetadata.TotalTokenCount,
701+
InputTokenCount = usageMetadata.PromptTokenCount != 0 ? usageMetadata.PromptTokenCount : null,
702+
CachedInputTokenCount = usageMetadata.CachedContentTokenCount != 0 ? usageMetadata.CachedContentTokenCount : null,
703+
OutputTokenCount = usageMetadata.CandidatesTokenCount != 0 ? usageMetadata.CandidatesTokenCount : null,
704+
ReasoningTokenCount = usageMetadata.ThoughtsTokenCount != 0 ? usageMetadata.ThoughtsTokenCount : null,
705+
TotalTokenCount = usageMetadata.TotalTokenCount != 0 ? usageMetadata.TotalTokenCount : null,
704706
};
705707

706-
if (usageMetadata.CachedContentTokenCount != 0)
707-
{
708-
(usage.AdditionalCounts ??= [])[nameof(usageMetadata.CachedContentTokenCount)] = usageMetadata.CachedContentTokenCount;
709-
}
710-
711-
if (usageMetadata.ThoughtsTokenCount != 0)
712-
{
713-
(usage.AdditionalCounts ??= [])[nameof(usageMetadata.ThoughtsTokenCount)] = usageMetadata.ThoughtsTokenCount;
714-
}
715-
716708
return usage;
717709
}
718710

‎apis/Google.Cloud.VertexAI.Extensions/Google.Cloud.VertexAI.Extensions/PredictionServiceEmbeddingGenerator.cs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -111,10 +111,12 @@ public async Task<GeneratedEmbeddings<Embedding<float>>> GenerateAsync(IEnumerab
111111
tokenCountValue.KindCase == Value.KindOneofCase.NumberValue)
112112
{
113113
int tokenCount = (int) tokenCountValue.NumberValue;
114-
115-
result.Usage ??= new();
116-
result.Usage.InputTokenCount += tokenCount;
117-
result.Usage.TotalTokenCount += tokenCount;
114+
if (tokenCount > 0)
115+
{
116+
result.Usage ??= new();
117+
result.Usage.InputTokenCount = (result.Usage.InputTokenCount ?? 0) + tokenCount;
118+
result.Usage.TotalTokenCount = (result.Usage.TotalTokenCount ?? 0) + tokenCount;
119+
}
118120
}
119121
}
120122
}

‎apis/Google.Cloud.VertexAI.Extensions/Google.Cloud.VertexAI.Extensions/VertexAIExtensions.cs‎

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -228,21 +228,13 @@ public static AITool AsAITool(this Tool tool)
228228

229229
UsageDetails usage = new()
230230
{
231-
InputTokenCount = usageMetadata.PromptTokenCount,
232-
OutputTokenCount = usageMetadata.CandidatesTokenCount,
233-
TotalTokenCount = usageMetadata.TotalTokenCount,
231+
InputTokenCount = usageMetadata.PromptTokenCount != 0 ? usageMetadata.PromptTokenCount : null,
232+
CachedInputTokenCount = usageMetadata.CachedContentTokenCount != 0 ? usageMetadata.CachedContentTokenCount : null,
233+
OutputTokenCount = usageMetadata.CandidatesTokenCount != 0 ? usageMetadata.CandidatesTokenCount : null,
234+
ReasoningTokenCount = usageMetadata.ThoughtsTokenCount != 0 ? usageMetadata.ThoughtsTokenCount : null,
235+
TotalTokenCount = usageMetadata.TotalTokenCount != 0 ? usageMetadata.TotalTokenCount : null,
234236
};
235237

236-
if (usageMetadata.CachedContentTokenCount != 0)
237-
{
238-
(usage.AdditionalCounts ??= [])[nameof(usageMetadata.CachedContentTokenCount)] = usageMetadata.CachedContentTokenCount;
239-
}
240-
241-
if (usageMetadata.ThoughtsTokenCount != 0)
242-
{
243-
(usage.AdditionalCounts ??= [])[nameof(usageMetadata.ThoughtsTokenCount)] = usageMetadata.ThoughtsTokenCount;
244-
}
245-
246238
if (usageMetadata.ToolUsePromptTokenCount != 0)
247239
{
248240
(usage.AdditionalCounts ??= [])[nameof(usageMetadata.ToolUsePromptTokenCount)] = usageMetadata.ToolUsePromptTokenCount;

0 commit comments

Comments
 (0)