Update BedrockChatClient with support for Instructions/RawRepresentationFactory (#3906)

* Update BedrockChatClient with support for Instructions/RawRepresentationFactory

- ChatOptions.RawRepresentationFactory, to support "breaking glass" if someone wants to set Converse{Stream}Request properties that aren't available via messages/options directly.
- ChatOptions.Instructions, as a non-messages way to set a system prompt.

Author: stephentoub

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

@@ -37,7 +37,7 @@
   </Choose>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.5.0" />
+    <PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.7.0" />
   </ItemGroup>
 
   <ItemGroup>

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

@@ -41,7 +41,7 @@
   </Choose>
 
   <ItemGroup>
-  <PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.5.0" />
+  <PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.7.0" />
   </ItemGroup>
 
   <ItemGroup>

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

@@ -3,7 +3,7 @@
   <metadata> 
     <id>AWSSDK.Extensions.Bedrock.MEAI</id>
     <title>AWSSDK - Bedrock integration with Microsoft.Extensions.AI.</title>
-    <version>4.0.1.1</version>
+    <version>4.0.2.0</version>
     <authors>Amazon Web Services</authors>
     <description>Implementations of Microsoft.Extensions.AI's abstractions for Bedrock.</description>
     <language>en-US</language>
@@ -15,17 +15,17 @@
       <group targetFramework="net472">
         <dependency id="AWSSDK.Core" version="4.0.0.4" />
         <dependency id="AWSSDK.BedrockRuntime" version="4.0.0.3" />
-        <dependency id="Microsoft.Extensions.AI.Abstractions" version="9.5.0" />
+        <dependency id="Microsoft.Extensions.AI.Abstractions" version="9.7.0" />
       </group>
       <group targetFramework="netstandard2.0">
         <dependency id="AWSSDK.Core" version="4.0.0.4" />
         <dependency id="AWSSDK.BedrockRuntime" version="4.0.0.3" />
-        <dependency id="Microsoft.Extensions.AI.Abstractions" version="9.5.0" />
+        <dependency id="Microsoft.Extensions.AI.Abstractions" version="9.7.0" />
       </group>
       <group targetFramework="net8.0">
         <dependency id="AWSSDK.Core" version="4.0.0.4" />
         <dependency id="AWSSDK.BedrockRuntime" version="4.0.0.3" />
-        <dependency id="Microsoft.Extensions.AI.Abstractions" version="9.5.0" />
+        <dependency id="Microsoft.Extensions.AI.Abstractions" version="9.7.0" />
       </group>
     </dependencies>
   </metadata> 

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

@@ -71,15 +71,13 @@ public async Task<ChatResponse> GetResponseAsync(
             throw new ArgumentNullException(nameof(messages));
         }
 
-        ConverseRequest request = new()
-        {
-            ModelId = options?.ModelId ?? _modelId,
-            Messages = CreateMessages(messages),
-            System = CreateSystem(messages),
-            ToolConfig = CreateToolConfig(options),
-            InferenceConfig = CreateInferenceConfiguration(options),
-            AdditionalModelRequestFields = CreateAdditionalModelRequestFields(options),
-        };
+        ConverseRequest request = options?.RawRepresentationFactory?.Invoke(this) as ConverseRequest ?? new();
+        request.ModelId ??= options?.ModelId ?? _modelId;
+        request.Messages = CreateMessages(request.Messages, messages);
+        request.System = CreateSystem(request.System, messages, options);
+        request.ToolConfig = CreateToolConfig(request.ToolConfig, options);
+        request.InferenceConfig = CreateInferenceConfiguration(request.InferenceConfig, options);
+        request.AdditionalModelRequestFields = CreateAdditionalModelRequestFields(request.AdditionalModelRequestFields, options);
 
         var response = await _runtime.ConverseAsync(request, cancellationToken).ConfigureAwait(false);
 
@@ -162,15 +160,13 @@ public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
             throw new ArgumentNullException(nameof(messages));
         }
 
-        ConverseStreamRequest request = new()
-        {
-            ModelId = options?.ModelId ?? _modelId,
-            Messages = CreateMessages(messages),
-            System = CreateSystem(messages),
-            ToolConfig = CreateToolConfig(options),
-            InferenceConfig = CreateInferenceConfiguration(options),
-            AdditionalModelRequestFields = CreateAdditionalModelRequestFields(options),
-        };
+        ConverseStreamRequest request = options?.RawRepresentationFactory?.Invoke(this) as ConverseStreamRequest ?? new();
+        request.ModelId ??= options?.ModelId ?? _modelId;
+        request.Messages = CreateMessages(request.Messages, messages);
+        request.System = CreateSystem(request.System, messages, options);
+        request.ToolConfig = CreateToolConfig(request.ToolConfig, options);
+        request.InferenceConfig = CreateInferenceConfiguration(request.InferenceConfig, options);
+        request.AdditionalModelRequestFields = CreateAdditionalModelRequestFields(request.AdditionalModelRequestFields, options);
 
         var result = await _runtime.ConverseStreamAsync(request, cancellationToken).ConfigureAwait(false);
 
@@ -356,11 +352,21 @@ private static ChatFinishReason GetChatFinishReason(StopReason stopReason) =>
         };
 
     /// <summary>Creates a list of <see cref="SystemContentBlock"/> from the system messages in the provided <paramref name="messages"/>.</summary>
-    private static List<SystemContentBlock> CreateSystem(IEnumerable<ChatMessage> messages) =>
-        messages
+    private static List<SystemContentBlock> CreateSystem(List<SystemContentBlock>? rawMessages, IEnumerable<ChatMessage> messages, ChatOptions? options)
+    {
+        List<SystemContentBlock> system = rawMessages ?? [];
+
+        if (options?.Instructions is { } instructions)
+        {
+            system.Add(new SystemContentBlock() { Text = instructions });
+        }
+
+        system.AddRange(messages
             .Where(m => m.Role == ChatRole.System && m.Contents.Any(c => c is TextContent))
-            .Select(m => new SystemContentBlock() { Text = string.Concat(m.Contents.OfType<TextContent>()) })
-            .ToList();
+            .Select(m => new SystemContentBlock() { Text = string.Concat(m.Contents.OfType<TextContent>()) }));
+
+        return system;
+    }
 
     /// <summary>Parses JSON tool input into a <see cref="Dictionary{String, Object}"/>.</summary>
     private static Dictionary<string, object?>? ParseToolInputs(string? jsonInput, out Exception? parseError)
@@ -382,9 +388,9 @@ private static List<SystemContentBlock> CreateSystem(IEnumerable<ChatMessage> me
     }
 
     /// <summary>Creates a list of <see cref="Message"/> from the provided <paramref name="chatMessages"/>.</summary>
-    private static List<Message> CreateMessages(IEnumerable<ChatMessage> chatMessages)
+    private static List<Message> CreateMessages(List<Message>? rawMessages, IEnumerable<ChatMessage> chatMessages)
     {
-        List<Message> messages = [];
+        List<Message> messages = rawMessages ?? [];
 
         foreach (ChatMessage chatMessage in chatMessages)
         {
@@ -681,100 +687,110 @@ private static Document ToDocument(JsonElement json)
     }
 
     /// <summary>Creates an <see cref="ToolConfiguration"/> from the specified options.</summary>
-    private static ToolConfiguration? CreateToolConfig(ChatOptions? options)
+    private static ToolConfiguration? CreateToolConfig(ToolConfiguration? toolConfig, ChatOptions? options)
     {
-        List<Tool>? tools = options?.Tools?.OfType<AIFunction>().Select(f =>
+        if (options?.Tools is { Count: > 0 } tools)
         {
-            Document inputs = default;
-            List<Document> required = [];
-
-            if (f.JsonSchema.TryGetProperty("properties", out JsonElement properties))
+            foreach (AITool tool in tools)
             {
-                foreach (JsonProperty parameter in properties.EnumerateObject())
+                if (tool is not AIFunction f)
                 {
-                    inputs.Add(parameter.Name, ToDocument(parameter.Value));
+                    continue;
                 }
-            }
 
-            if (f.JsonSchema.TryGetProperty("required", out JsonElement requiredProperties))
-            {
-                foreach (JsonElement requiredProperty in requiredProperties.EnumerateArray())
+                Document inputs = default;
+                List<Document> required = [];
+
+                if (f.JsonSchema.TryGetProperty("properties", out JsonElement properties))
                 {
-                    required.Add(requiredProperty.GetString());
+                    foreach (JsonProperty parameter in properties.EnumerateObject())
+                    {
+                        inputs.Add(parameter.Name, ToDocument(parameter.Value));
+                    }
                 }
-            }
 
-            var schemaDictionary = new Dictionary<string, Document>()
-            {
-                ["type"] = new Document("object"),
-            };
+                if (f.JsonSchema.TryGetProperty("required", out JsonElement requiredProperties))
+                {
+                    foreach (JsonElement requiredProperty in requiredProperties.EnumerateArray())
+                    {
+                        required.Add(requiredProperty.GetString());
+                    }
+                }
 
-            if (inputs != default)
-            {
-                schemaDictionary["properties"] = inputs;
-            }
+                Dictionary<string, Document> schemaDictionary = new()
+                {
+                    ["type"] = new Document("object"),
+                };
 
-            if (required.Count > 0)
-            {
-                schemaDictionary["required"] = new Document(required);
-            }
+                if (inputs != default)
+                {
+                    schemaDictionary["properties"] = inputs;
+                }
 
-            return new Tool()
-            {
-                ToolSpec = new ToolSpecification()
+                if (required.Count > 0)
+                {
+                    schemaDictionary["required"] = new Document(required);
+                }
+
+                toolConfig ??= new();
+                toolConfig.Tools ??= [];
+                toolConfig.Tools.Add(new()
                 {
-                    Name = f.Name,
-                    Description = !string.IsNullOrEmpty(f.Description) ? f.Description : f.Name,
-                    InputSchema = new()
+                    ToolSpec = new ToolSpecification()
                     {
-                        Json = new(schemaDictionary)
+                        Name = f.Name,
+                        Description = !string.IsNullOrEmpty(f.Description) ? f.Description : f.Name,
+                        InputSchema = new()
+                        {
+                            Json = new(schemaDictionary)
+                        },
                     },
-                },
-            };
-        }).ToList();
+                });
+            }
+        }
 
-        ToolChoice? choice = null;
-        if (tools is { Count: > 0 })
+        if (toolConfig?.Tools is { Count: > 0 } && toolConfig.ToolChoice is null)
         {
             switch (options!.ToolMode)
             {
-                case AutoChatToolMode:
-                case null:
-                    choice = new ToolChoice() { Auto = new() };
-                    break;
-
                 case RequiredChatToolMode r:
-                    choice = !string.IsNullOrWhiteSpace(r.RequiredFunctionName) ?
+                    toolConfig.ToolChoice = !string.IsNullOrWhiteSpace(r.RequiredFunctionName) ?
                         new ToolChoice() { Tool = new() { Name = r.RequiredFunctionName } } :
                         new ToolChoice() { Any = new() };
                     break;
             }
-
-            return new()
-            {
-                ToolChoice = choice,
-                Tools = tools,
-            };
         }
 
-        return null;
+        return toolConfig;
     }
 
     /// <summary>Creates an <see cref="InferenceConfiguration"/> from the specified options.</summary>
-    private static InferenceConfiguration CreateInferenceConfiguration(ChatOptions? options) =>
-        new()
+    private static InferenceConfiguration CreateInferenceConfiguration(InferenceConfiguration config, ChatOptions? options)
+    {
+        config ??= new();
+        
+        config.MaxTokens ??= options?.MaxOutputTokens;
+        config.Temperature ??= options?.Temperature;
+        config.TopP ??= options?.TopP;
+
+        if (options?.StopSequences is { Count: > 0 } stopOptions)
         {
-            MaxTokens = options?.MaxOutputTokens,
-            StopSequences = options?.StopSequences?.ToList(),
-            Temperature = options?.Temperature,
-            TopP = options?.TopP,
-        };
+            if (config.StopSequences is null)
+            {
+                config.StopSequences = stopOptions.ToList();
+            }
+            else
+            {
+                config.StopSequences.AddRange(stopOptions);
+            }
+        }
+
+        return config;
+    }
 
     /// <summary>Creates a <see cref="Document"/> from the specified options to use as the additional model request options.</summary>
-    private static Document CreateAdditionalModelRequestFields(ChatOptions? options)
+    private static Document CreateAdditionalModelRequestFields(Document d, ChatOptions? options)
     {
-        Document d = default;
-
         if (options is not null)
         {
             if (options.TopK is int topK)

extensions/src/AWSSDK.Extensions.Bedrock.MEAI/BedrockEmbeddingGenerator.cs

@@ -13,6 +13,7 @@
  * permissions and limitations under the License.
  */
 
+using Amazon.BedrockRuntime.Model;
 using Microsoft.Extensions.AI;
 using System;
 using System.Collections.Generic;
@@ -89,17 +90,17 @@ public async Task<GeneratedEmbeddings<Embedding<float>>> GenerateAsync(
 
         foreach (string value in values)
         {
-            var response = await _runtime.InvokeModelAsync(new()
+            InvokeModelRequest request = options?.RawRepresentationFactory?.Invoke(this) as InvokeModelRequest ?? new();
+            request.ModelId ??= options?.ModelId ?? _modelId;
+            request.Accept ??= "application/json";
+            request.ContentType ??= "application/json";
+            request.Body ??= new MemoryStream(JsonSerializer.SerializeToUtf8Bytes(new()
             {
-                ModelId = options?.ModelId ?? _modelId,
-                Accept = "application/json",
-                ContentType = "application/json",
-                Body = new MemoryStream(JsonSerializer.SerializeToUtf8Bytes(new EmbeddingRequest()
-                {
-                    InputText = value,
-                    Dimensions = options?.Dimensions ?? _dimensions,
-                }, BedrockJsonContext.Default.EmbeddingRequest)),
-            }, cancellationToken).ConfigureAwait(false);
+                InputText = value,
+                Dimensions = options?.Dimensions ?? _dimensions,
+            }, BedrockJsonContext.Default.EmbeddingRequest));
+
+            var response = await _runtime.InvokeModelAsync(request, cancellationToken).ConfigureAwait(false);
 
             var er = JsonSerializer.Deserialize(response.Body, BedrockJsonContext.Default.EmbeddingResponse);
             if (er?.Embedding is not null)

extensions/test/BedrockMEAITests/BedrockMEAITests.NetFramework.csproj

@@ -18,7 +18,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.5.0" />
+    <PackageReference Include="Microsoft.Extensions.AI.Abstractions" Version="9.7.0" />
     <PackageReference Include="xunit" Version="2.9.2" />
     <PackageReference Include="xunit.runner.visualstudio" Version="2.8.2" />
   </ItemGroup>