Support for ChatOptions.ResponseFormat in AWSSDK.Extensions.Bedrock.MEAI

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

@@ -18,8 +18,11 @@
 using Amazon.Runtime.Internal.Util;
 using Microsoft.Extensions.AI;
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Globalization;
+using System.IO;
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Text;
@@ -35,6 +38,11 @@ internal sealed partial class BedrockChatClient : IChatClient
     /// <summary>A default logger to use.</summary>
     private static readonly ILogger DefaultLogger = Logger.GetLogger(typeof(BedrockChatClient));
 
+    /// <summary>The name used for the synthetic tool that enforces response format.</summary>
+    private const string ResponseFormatToolName = "generate_response";
+    /// <summary>The description used for the synthetic tool that enforces response format.</summary>
+    private const string ResponseFormatToolDescription = "Generate response in specified format";
+
     /// <summary>The wrapped <see cref="IAmazonBedrockRuntime"/> instance.</summary>
     private readonly IAmazonBedrockRuntime _runtime;
     /// <summary>Default model ID to use when no model is specified in the request.</summary>
@@ -63,6 +71,12 @@ public void Dispose()
     }
 
     /// <inheritdoc />
+    /// <remarks>
+    /// When <see cref="ChatOptions.ResponseFormat"/> is specified, the model must support
+    /// the ToolChoice feature. Models without this support will throw <see cref="NotSupportedException"/>.
+    /// If the model fails to return the expected structured output, <see cref="InvalidOperationException"/>
+    /// is thrown.
+    /// </remarks>
     public async Task<ChatResponse> GetResponseAsync(
         IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
     {
@@ -79,7 +93,29 @@ public async Task<ChatResponse> GetResponseAsync(
         request.InferenceConfig = CreateInferenceConfiguration(request.InferenceConfig, options);
         request.AdditionalModelRequestFields = CreateAdditionalModelRequestFields(request.AdditionalModelRequestFields, options);
 
-        var response = await _runtime.ConverseAsync(request, cancellationToken).ConfigureAwait(false);
+        ConverseResponse response;
+        try
+        {
+            response = await _runtime.ConverseAsync(request, cancellationToken).ConfigureAwait(false);
+        }
+        // Transforms ValidationException to NotSupportedException when error message indicates model lacks tool use support (required for ResponseFormat).
+        // This detection relies on error message text which may change in future Bedrock API versions.
+        catch (AmazonBedrockRuntimeException ex) when (options?.ResponseFormat is ChatResponseFormatJson)
+        {
+            // Detect unsupported model: ValidationException with specific tool support error messages
+            if (ex.ErrorCode == "ValidationException" &&
+                (ex.Message.IndexOf("toolChoice is not supported by this model", StringComparison.OrdinalIgnoreCase) >= 0 ||
+                 ex.Message.IndexOf("This model doesn't support tool use", StringComparison.OrdinalIgnoreCase) >= 0))
+            {
+                throw new NotSupportedException(
+                    $"The model '{request.ModelId}' does not support ResponseFormat. " +
+                    $"ResponseFormat requires ToolChoice support, which is only available in Claude 3+ and Mistral Large models. " +
+                    $"See: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html",
+                    ex);
+            }
+
+            throw;
+        }
 
         ChatMessage result = new()
         {
@@ -89,6 +125,42 @@ public async Task<ChatResponse> GetResponseAsync(
             MessageId = Guid.NewGuid().ToString("N"),
         };
 
+        // Check if ResponseFormat was used and extract structured content
+        bool usingResponseFormat = options?.ResponseFormat is ChatResponseFormatJson;
+        if (usingResponseFormat)
+        {
+            var structuredContent = ExtractResponseFormatContent(response.Output?.Message);
+            if (structuredContent is not null)
+            {
+                // Replace the content with the extracted JSON as a TextContent
+                result.Contents.Add(new TextContent(structuredContent) { RawRepresentation = response.Output?.Message });
+
+                // Skip normal content processing since we've extracted the structured response
+                if (DocumentToDictionary(response.AdditionalModelResponseFields) is { } responseFieldsDict)
+                {
+                    result.AdditionalProperties = new(responseFieldsDict);
+                }
+
+                return new(result)
+                {
+                    CreatedAt = result.CreatedAt,
+                    FinishReason = response.StopReason is not null ? GetChatFinishReason(response.StopReason) : null,
+                    Usage = response.Usage is TokenUsage tokenUsage ? CreateUsageDetails(tokenUsage) : null,
+                    RawRepresentation = response,
+                };
+            }
+            else
+            {
+                // Model succeeded but did not return expected structured output
+                throw new InvalidOperationException(
+                    $"Model '{request.ModelId}' did not return structured output as requested. " +
+                    $"This may indicate the model refused to follow the tool use instruction, " +
+                    $"the schema was too complex, or the prompt conflicted with the requirement. " +
+                    $"StopReason: {response.StopReason?.Value ?? "unknown"}.");
+            }
+        }
+
+        // Normal content processing when not using ResponseFormat or extraction failed
         if (response.Output?.Message?.Content is { } contents)
         {
             foreach (var content in contents)
@@ -182,6 +254,14 @@ public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
             throw new ArgumentNullException(nameof(messages));
         }
 
+        // Check if ResponseFormat is set - not supported for streaming yet
+        if (options?.ResponseFormat is ChatResponseFormatJson)
+        {
+            throw new NotSupportedException(
+                "ResponseFormat is not yet supported for streaming responses with Amazon Bedrock. " +
+                "Please use GetResponseAsync for structured output.");
+        }
+
         ConverseStreamRequest request = options?.RawRepresentationFactory?.Invoke(this) as ConverseStreamRequest ?? new();
         request.ModelId ??= options?.ModelId ?? _modelId;
         request.Messages = CreateMessages(request.Messages, messages);
@@ -794,7 +874,11 @@ private static Document ToDocument(JsonElement json)
         }
     }
 
-    /// <summary>Creates an <see cref="ToolConfiguration"/> from the specified options.</summary>
+    /// <summary>Creates a <see cref="ToolConfiguration"/> from the specified options.</summary>
+    /// <remarks>
+    /// When ResponseFormat is specified, creates a synthetic tool to enforce structured output.
+    /// This conflicts with user-provided tools as Bedrock only supports a single ToolChoice value.
+    /// </remarks>
     private static ToolConfiguration? CreateToolConfig(ToolConfiguration? toolConfig, ChatOptions? options)
     {
         if (options?.Tools is { Count: > 0 } tools)
@@ -857,6 +941,56 @@ private static Document ToDocument(JsonElement json)
             }
         }
 
+        // Handle ResponseFormat by creating a synthetic tool
+        if (options?.ResponseFormat is ChatResponseFormatJson jsonFormat)
+        {
+            // Check for conflict with user-provided tools
+            if (toolConfig?.Tools?.Count > 0)
+            {
+                throw new ArgumentException(
+                    "ResponseFormat cannot be used with Tools in Amazon Bedrock. " +
+                    "ResponseFormat uses Bedrock's tool mechanism for structured output, " +
+                    "which conflicts with user-provided tools.");
+            }
+
+            // Create the synthetic tool with the schema from ResponseFormat
+            toolConfig ??= new();
+            toolConfig.Tools ??= [];
+
+            // Parse the schema if provided, otherwise create an empty object schema
+            Document schemaDoc;
+            if (jsonFormat.Schema.HasValue)
+            {
+                // Schema is already a JsonElement (parsed JSON), convert directly to Document
+                schemaDoc = ToDocument(jsonFormat.Schema.Value);
+            }
+            else
+            {
+                // For JSON mode without schema, create a generic object schema
+                schemaDoc = new Document(new Dictionary<string, Document>
+                {
+                    ["type"] = new Document("object"),
+                    ["additionalProperties"] = new Document(true)
+                });
+            }
+
+            toolConfig.Tools.Add(new Tool
+            {
+                ToolSpec = new ToolSpecification
+                {
+                    Name = ResponseFormatToolName,
+                    Description = jsonFormat.SchemaDescription ?? ResponseFormatToolDescription,
+                    InputSchema = new ToolInputSchema
+                    {
+                        Json = schemaDoc
+                    }
+                }
+            });
+
+            // Force the model to use the synthetic tool
+            toolConfig.ToolChoice = new ToolChoice { Tool = new() { Name = ResponseFormatToolName } };
+        }
+
         if (toolConfig?.Tools is { Count: > 0 } && toolConfig.ToolChoice is null)
         {
             switch (options!.ToolMode)
@@ -872,6 +1006,43 @@ private static Document ToDocument(JsonElement json)
         return toolConfig;
     }
 
+    /// <summary>Extracts JSON content from the synthetic ResponseFormat tool use, if present.</summary>
+    private static string? ExtractResponseFormatContent(Message? message)
+    {
+        if (message?.Content is null)
+        {
+            return null;
+        }
+
+        foreach (var content in message.Content)
+        {
+            if (content.ToolUse is ToolUseBlock toolUse &&
+                toolUse.Name == ResponseFormatToolName &&
+                toolUse.Input != default)
+            {
+                // Convert the Document back to JSON string
+                return DocumentToJsonString(toolUse.Input);
+            }
+        }
+
+        return null;
+    }
+
+    /// <summary>
+    /// Converts a <see cref="Document"/> to a JSON string using the SDK's standard DocumentMarshaller.
+    /// Note: Document is a struct (value type), so circular references are structurally impossible.
+    /// </summary>
+    private static string DocumentToJsonString(Document document)
+    {
+        using var stream = new MemoryStream();
+        using (var writer = new Utf8JsonWriter(stream, new JsonWriterOptions { Indented = false }))
+        {
+            Amazon.Runtime.Documents.Internal.Transform.DocumentMarshaller.Instance.Write(writer, document);
+        }
+        return Encoding.UTF8.GetString(stream.ToArray());
+    }
+
+
     /// <summary>Creates an <see cref="InferenceConfiguration"/> from the specified options.</summary>
     private static InferenceConfiguration CreateInferenceConfiguration(InferenceConfiguration config, ChatOptions? options)
     {