Skip to content

Commit 3c09b3e

Browse files
stephentoubstephentoub
committed
Update IChatClient with support from latest bedrock runtime / M.E.AI
- Updates to M.E.AI.Abstractions 10.2.0 - Adds support for multi-modal tool returns. - Adds support for citations with URIs. - Adds support for cached tokens in usage details. - Adds a ton of tests verifying IChatClient behavior around the underlying IAmazonBedrockRuntime.
1 parent d6ed191 commit 3c09b3e

File tree

7 files changed

+3236
-626
lines changed

7 files changed

+3236
-626
lines changed

‎extensions/src/AWSSDK.Extensions.Bedrock.MEAI/AWSSDK.Extensions.Bedrock.MEAI.NetFramework.csproj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
</Choose>
3838

3939
<ItemGroup>
40-
<PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.9.1" />
40+
<PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="10.2.0" />
4141
</ItemGroup>
4242

4343
<ItemGroup>

‎extensions/src/AWSSDK.Extensions.Bedrock.MEAI/AWSSDK.Extensions.Bedrock.MEAI.NetStandard.csproj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
</Choose>
4242

4343
<ItemGroup>
44-
<PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.9.1" />
44+
<PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="10.2.0" />
4545
</ItemGroup>
4646

4747
<ItemGroup>

‎extensions/src/AWSSDK.Extensions.Bedrock.MEAI/AWSSDK.Extensions.Bedrock.MEAI.nuspec‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,17 +15,17 @@
1515
<group targetFramework="net472">
1616
<dependency id="AWSSDK.Core" version="4.0.3.9" />
1717
<dependency id="AWSSDK.BedrockRuntime" version="4.0.14.6" />
18-
<dependency id="Microsoft.Extensions.AI.Abstractions" version="9.9.1" />
18+
<dependency id="Microsoft.Extensions.AI.Abstractions" version="10.2.0" />
1919
</group>
2020
<group targetFramework="netstandard2.0">
2121
<dependency id="AWSSDK.Core" version="4.0.3.9" />
2222
<dependency id="AWSSDK.BedrockRuntime" version="4.0.14.6" />
23-
<dependency id="Microsoft.Extensions.AI.Abstractions" version="9.9.1" />
23+
<dependency id="Microsoft.Extensions.AI.Abstractions" version="10.2.0" />
2424
</group>
2525
<group targetFramework="net8.0">
2626
<dependency id="AWSSDK.Core" version="4.0.3.9" />
2727
<dependency id="AWSSDK.BedrockRuntime" version="4.0.14.6" />
28-
<dependency id="Microsoft.Extensions.AI.Abstractions" version="9.9.1" />
28+
<dependency id="Microsoft.Extensions.AI.Abstractions" version="10.2.0" />
2929
</group>
3030
</dependencies>
3131
</metadata>

‎extensions/src/AWSSDK.Extensions.Bedrock.MEAI/BedrockChatClient.cs‎

Lines changed: 45 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -162,8 +162,9 @@ public async Task<ChatResponse> GetResponseAsync(
162162
TextContent tc = new(citations.Content[i]?.Text) { RawRepresentation = citations.Content[i] };
163163
tc.Annotations = [new CitationAnnotation()
164164
{
165+
Snippet = citations.Citations[i].SourceContent?.Select(c => c.Text).FirstOrDefault() ?? citations.Citations[i].Source,
165166
Title = citations.Citations[i].Title,
166-
Snippet = citations.Citations[i].SourceContent?.Select(c => c.Text).FirstOrDefault(),
167+
Url = Uri.TryCreate(citations.Citations[i].Location?.Web?.Url, UriKind.Absolute, out Uri? uri) ? uri : null,
167168
}];
168169
result.Contents.Add(tc);
169170
}
@@ -422,15 +423,11 @@ private static UsageDetails CreateUsageDetails(TokenUsage usage)
422423
UsageDetails ud = new()
423424
{
424425
InputTokenCount = usage.InputTokens,
426+
CachedInputTokenCount = usage.CacheReadInputTokens,
425427
OutputTokenCount = usage.OutputTokens,
426428
TotalTokenCount = usage.TotalTokens,
427429
};
428430

429-
if (usage.CacheReadInputTokens is int cacheReadTokens)
430-
{
431-
(ud.AdditionalCounts ??= []).Add(nameof(usage.CacheReadInputTokens), cacheReadTokens);
432-
}
433-
434431
if (usage.CacheWriteInputTokens is int cacheWriteTokens)
435432
{
436433
(ud.AdditionalCounts ??= []).Add(nameof(usage.CacheWriteInputTokens), cacheWriteTokens);
@@ -465,8 +462,7 @@ private static List<SystemContentBlock> CreateSystem(List<SystemContentBlock>? r
465462
});
466463
}
467464

468-
foreach (var message in messages
469-
.Where(m => m.Role == ChatRole.System && m.Contents.Any(c => c is TextContent)))
465+
foreach (var message in messages.Where(m => m.Role == ChatRole.System && m.Contents.Any(c => c is TextContent)))
470466
{
471467
system.Add(new SystemContentBlock()
472468
{
@@ -567,6 +563,10 @@ private static List<ContentBlock> CreateContents(ChatMessage message)
567563
{
568564
switch (content)
569565
{
566+
case AIContent when content.RawRepresentation is ContentBlock cb:
567+
contents.Add(cb);
568+
break;
569+
570570
case TextContent tc:
571571
if (message.Role == ChatRole.Assistant)
572572
{
@@ -649,32 +649,54 @@ private static List<ContentBlock> CreateContents(ChatMessage message)
649649
break;
650650

651651
case FunctionResultContent frc:
652-
Document result = frc.Result switch
653-
{
654-
int i => i,
655-
long l => l,
656-
float f => f,
657-
double d => d,
658-
string s => s,
659-
bool b => b,
660-
JsonElement json => ToDocument(json),
661-
{ } other => ToDocument(JsonSerializer.SerializeToElement(other, BedrockJsonContext.DefaultOptions.GetTypeInfo(other.GetType()))),
662-
_ => default,
663-
};
664-
665652
contents.Add(new()
666653
{
667654
ToolResult = new()
668655
{
669656
ToolUseId = frc.CallId,
670-
Content = [new() { Json = new Document(new Dictionary<string, Document>() { ["result"] = result }) }],
657+
Content = ToToolResultContentBlocks(frc.Result),
671658
},
672659
});
673660
break;
674661
}
675662

663+
static List<ToolResultContentBlock> ToToolResultContentBlocks(object? result) =>
664+
result switch
665+
{
666+
AIContent aic => [ToolResultContentBlockFromAIContent(aic)],
667+
IEnumerable<AIContent> aics => [.. aics.Select(ToolResultContentBlockFromAIContent)],
668+
string s => [new () { Text = s }],
669+
_ => [new()
670+
{
671+
Json = new Document(new Dictionary<string, Document>()
672+
{
673+
["result"] = result switch
674+
{
675+
int i => i,
676+
long l => l,
677+
float f => f,
678+
double d => d,
679+
bool b => b,
680+
JsonElement json => ToDocument(json),
681+
{ } other => ToDocument(JsonSerializer.SerializeToElement(other, BedrockJsonContext.DefaultOptions.GetTypeInfo(other.GetType()))),
682+
_ => default,
683+
}
684+
})
685+
}],
686+
};
687+
688+
static ToolResultContentBlock ToolResultContentBlockFromAIContent(AIContent aic) =>
689+
aic switch
690+
{
691+
TextContent tc => new() { Text = tc.Text },
692+
TextReasoningContent trc => new() { Text = trc.Text },
693+
DataContent dc when GetImageFormat(dc.MediaType) is { } imageFormat => new() { Image = new() { Source = new() { Bytes = new(dc.Data.ToArray()) }, Format = imageFormat } },
694+
DataContent dc when GetVideoFormat(dc.MediaType) is { } videoFormat => new() { Video = new() { Source = new() { Bytes = new(dc.Data.ToArray()) }, Format = videoFormat } },
695+
DataContent dc when GetDocumentFormat(dc.MediaType) is { } docFormat => new() { Document = new() { Source = new() { Bytes = new(dc.Data.ToArray()) }, Format = docFormat, Name = dc.Name ?? "file" } },
696+
_ => ToToolResultContentBlocks(JsonSerializer.SerializeToElement(aic, BedrockJsonContext.DefaultOptions.GetTypeInfo(typeof(object)))).First(),
697+
};
676698

677-
if (content.AdditionalProperties?.TryGetValue(nameof(ContentBlock.CachePoint), out var maybeCachePoint) == true)
699+
if (content.AdditionalProperties?.TryGetValue(nameof(ContentBlock.CachePoint), out var maybeCachePoint) is true)
678700
{
679701
if (maybeCachePoint is CachePointBlock cachePointBlock)
680702
{

0 commit comments

Comments
 (0)