diff --git a/core/changelog.md b/core/changelog.md index e69de29bb2d..ee9604424b1 100644 --- a/core/changelog.md +++ b/core/changelog.md @@ -0,0 +1 @@ +[fix]: preserve documents in Bedrock tool results [@michaeldunn9](https://github.com/michaeldunn9) diff --git a/core/providers/anthropic/emptytoolresult_test.go b/core/providers/anthropic/emptytoolresult_test.go index 99137a7f20a..ba21efcb965 100644 --- a/core/providers/anthropic/emptytoolresult_test.go +++ b/core/providers/anthropic/emptytoolresult_test.go @@ -36,17 +36,24 @@ func TestConvertToolResultWithEmptyContent(t *testing.T) { } } -// TestConvertToolResultWithUnsupportedBlocks verifies tool_result content made -// solely of block types the converter does not map (e.g. document) still -// yields a serializable output. -func TestConvertToolResultWithUnsupportedBlocks(t *testing.T) { +// TestConvertToolResultWithDocumentBlock verifies a document-only tool result +// remains serializable and retains its canonical Bifrost file block. +func TestConvertToolResultWithDocumentBlock(t *testing.T) { role := schemas.ResponsesInputMessageRoleUser blocks := []AnthropicContentBlock{ { Type: AnthropicContentBlockTypeToolResult, ToolUseID: schemas.Ptr("toolu_doc"), Content: &AnthropicContent{ContentBlocks: []AnthropicContentBlock{ - {Type: AnthropicContentBlockTypeDocument}, + { + Type: AnthropicContentBlockTypeDocument, + Title: schemas.Ptr("report.pdf"), + Source: &AnthropicBlockSource{SourceObj: &AnthropicSource{ + Type: "base64", + MediaType: schemas.Ptr("application/pdf"), + Data: schemas.Ptr("JVBERi0xLjQ="), + }}, + }, }}, }, } @@ -55,6 +62,13 @@ func TestConvertToolResultWithUnsupportedBlocks(t *testing.T) { if len(msgs) != 1 { t.Fatalf("expected 1 converted message, got %d", len(msgs)) } + output := msgs[0].Output + if output == nil || len(output.ResponsesFunctionToolCallOutputBlocks) != 1 { + t.Fatalf("expected one preserved document block, got %#v", output) + } + if output.ResponsesFunctionToolCallOutputBlocks[0].Type != schemas.ResponsesInputMessageContentBlockTypeFile { + t.Fatalf("expected canonical file block, got %#v", output.ResponsesFunctionToolCallOutputBlocks[0]) + } if _, err := schemas.MarshalSorted(msgs); err != nil { t.Fatalf("converted messages must marshal, got: %v", err) } diff --git a/core/providers/anthropic/responses.go b/core/providers/anthropic/responses.go index dad2a777cbd..e8e93d2aed1 100644 --- a/core/providers/anthropic/responses.go +++ b/core/providers/anthropic/responses.go @@ -5223,30 +5223,7 @@ func convertAnthropicContentBlocksToResponsesMessagesGrouped(contentBlocks []Ant if block.Content.ContentStr != nil { bifrostMsg.ResponsesToolMessage.Output.ResponsesToolCallOutputStr = block.Content.ContentStr } else if block.Content.ContentBlocks != nil { - var toolMsgContentBlocks []schemas.ResponsesMessageContentBlock - for _, contentBlock := range block.Content.ContentBlocks { - switch contentBlock.Type { - case AnthropicContentBlockTypeText: - if contentBlock.Text != nil { - var blockType schemas.ResponsesMessageContentBlockType - if isOutputMessage { - blockType = schemas.ResponsesOutputMessageContentTypeText - } else { - blockType = schemas.ResponsesInputMessageContentBlockTypeText - } - toolMsgContentBlocks = append(toolMsgContentBlocks, schemas.ResponsesMessageContentBlock{ - Type: blockType, - Text: contentBlock.Text, - CacheControl: contentBlock.CacheControl, - }) - } - case AnthropicContentBlockTypeImage: - if contentBlock.Source != nil && contentBlock.Source.SourceObj != nil { - toolMsgContentBlocks = append(toolMsgContentBlocks, contentBlock.toBifrostResponsesImageBlock()) - } - } - } - bifrostMsg.ResponsesToolMessage.Output.ResponsesFunctionToolCallOutputBlocks = toolMsgContentBlocks + bifrostMsg.Output.ResponsesFunctionToolCallOutputBlocks = convertAnthropicToolResultContentBlocks(block.Content.ContentBlocks, isOutputMessage) } // Handle is_error from Anthropic if block.IsError != nil && *block.IsError { @@ -5623,30 +5600,7 @@ func convertAnthropicContentBlocksToResponsesMessages(ctx *schemas.BifrostContex if block.Content.ContentStr != nil { bifrostMsg.ResponsesToolMessage.Output.ResponsesToolCallOutputStr = block.Content.ContentStr } else if block.Content.ContentBlocks != nil { - var toolMsgContentBlocks []schemas.ResponsesMessageContentBlock - for _, contentBlock := range block.Content.ContentBlocks { - switch contentBlock.Type { - case AnthropicContentBlockTypeText: - if contentBlock.Text != nil { - var blockType schemas.ResponsesMessageContentBlockType - if isOutputMessage { - blockType = schemas.ResponsesOutputMessageContentTypeText - } else { - blockType = schemas.ResponsesInputMessageContentBlockTypeText - } - toolMsgContentBlocks = append(toolMsgContentBlocks, schemas.ResponsesMessageContentBlock{ - Type: blockType, - Text: contentBlock.Text, - CacheControl: contentBlock.CacheControl, - }) - } - case AnthropicContentBlockTypeImage: - if contentBlock.Source != nil && contentBlock.Source.SourceObj != nil { - toolMsgContentBlocks = append(toolMsgContentBlocks, contentBlock.toBifrostResponsesImageBlock()) - } - } - } - bifrostMsg.ResponsesToolMessage.Output.ResponsesFunctionToolCallOutputBlocks = toolMsgContentBlocks + bifrostMsg.Output.ResponsesFunctionToolCallOutputBlocks = convertAnthropicToolResultContentBlocks(block.Content.ContentBlocks, isOutputMessage) } // Handle is_error from Anthropic if block.IsError != nil && *block.IsError { @@ -7859,6 +7813,7 @@ func (block AnthropicContentBlock) toBifrostResponsesDocumentBlock() schemas.Res if src.MediaType != nil { mediaType = *src.MediaType } + resultBlock.FileType = &mediaType dataURL := *src.Data if !strings.HasPrefix(dataURL, "data:") { dataURL = "data:" + mediaType + ";base64," + *src.Data @@ -7881,6 +7836,40 @@ func (block AnthropicContentBlock) toBifrostResponsesDocumentBlock() schemas.Res return resultBlock } +// convertAnthropicToolResultContentBlocks maps the content types Anthropic permits +// inside tool_result blocks to their canonical Bifrost Responses representations. +// Both grouped and non-grouped request conversion use this boundary so nested +// content support cannot drift between routing modes. +func convertAnthropicToolResultContentBlocks(contentBlocks []AnthropicContentBlock, isOutputMessage bool) []schemas.ResponsesMessageContentBlock { + var toolMessageContentBlocks []schemas.ResponsesMessageContentBlock + for _, contentBlock := range contentBlocks { + switch contentBlock.Type { + case AnthropicContentBlockTypeText: + if contentBlock.Text == nil { + continue + } + blockType := schemas.ResponsesInputMessageContentBlockTypeText + if isOutputMessage { + blockType = schemas.ResponsesOutputMessageContentTypeText + } + toolMessageContentBlocks = append(toolMessageContentBlocks, schemas.ResponsesMessageContentBlock{ + Type: blockType, + Text: contentBlock.Text, + CacheControl: contentBlock.CacheControl, + }) + case AnthropicContentBlockTypeImage: + if contentBlock.Source != nil && contentBlock.Source.SourceObj != nil { + toolMessageContentBlocks = append(toolMessageContentBlocks, contentBlock.toBifrostResponsesImageBlock()) + } + case AnthropicContentBlockTypeDocument: + if contentBlock.Source != nil && contentBlock.Source.SourceObj != nil { + toolMessageContentBlocks = append(toolMessageContentBlocks, contentBlock.toBifrostResponsesDocumentBlock()) + } + } + } + return toolMessageContentBlocks +} + // Helper functions for MCP tool/server conversion // convertAnthropicMCPServerV2ToBifrostTool converts a new-format MCP server to a Bifrost ResponsesTool. func convertAnthropicMCPServerV2ToBifrostTool(mcpServer *AnthropicMCPServerV2) *schemas.ResponsesTool { diff --git a/core/providers/anthropic/toolresultdocument_test.go b/core/providers/anthropic/toolresultdocument_test.go new file mode 100644 index 00000000000..373207e8aef --- /dev/null +++ b/core/providers/anthropic/toolresultdocument_test.go @@ -0,0 +1,120 @@ +package anthropic + +import ( + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +func convertToolResultDocumentBlocksForTest(grouped bool, content []AnthropicContentBlock) []schemas.ResponsesMessage { + role := schemas.ResponsesInputMessageRoleUser + blocks := []AnthropicContentBlock{ + { + Type: AnthropicContentBlockTypeToolResult, + ToolUseID: schemas.Ptr("toolu_document"), + Content: &AnthropicContent{ContentBlocks: content}, + }, + } + if grouped { + return convertAnthropicContentBlocksToResponsesMessagesGrouped(blocks, &role, false) + } + return convertAnthropicContentBlocksToResponsesMessages(nil, blocks, &role, false, "") +} + +func requireToolResultDocumentBlocks(t *testing.T, messages []schemas.ResponsesMessage) []schemas.ResponsesMessageContentBlock { + t.Helper() + if len(messages) != 1 { + t.Fatalf("expected 1 converted message, got %d", len(messages)) + } + message := messages[0] + if message.Type == nil || *message.Type != schemas.ResponsesMessageTypeFunctionCallOutput { + t.Fatalf("expected function_call_output, got %#v", message.Type) + } + if message.ResponsesToolMessage == nil || message.Output == nil { + t.Fatal("expected a populated tool message output") + } + return message.Output.ResponsesFunctionToolCallOutputBlocks +} + +func TestToolResultDocumentURLPreserved(t *testing.T) { + for _, grouped := range []bool{true, false} { + mode := "non-grouped" + if grouped { + mode = "grouped" + } + t.Run(mode, func(t *testing.T) { + content := []AnthropicContentBlock{ + { + Type: AnthropicContentBlockTypeText, + Text: schemas.Ptr("Document generated"), + }, + { + Type: AnthropicContentBlockTypeDocument, + Title: schemas.Ptr("report.pdf"), + Source: &AnthropicBlockSource{SourceObj: &AnthropicSource{ + Type: "url", + URL: schemas.Ptr("https://example.com/report.pdf"), + }}, + }, + } + + blocks := requireToolResultDocumentBlocks(t, convertToolResultDocumentBlocksForTest(grouped, content)) + if len(blocks) != 2 { + t.Fatalf("expected text and document blocks, got %d", len(blocks)) + } + if blocks[0].Text == nil || *blocks[0].Text != "Document generated" { + t.Fatalf("expected text block first, got %#v", blocks[0]) + } + if blocks[1].Type != schemas.ResponsesInputMessageContentBlockTypeFile { + t.Fatalf("expected file block second, got %q", blocks[1].Type) + } + file := blocks[1].ResponsesInputMessageContentBlockFile + if file == nil { + t.Fatal("expected canonical Bifrost file block") + } + if file.FileURL == nil || *file.FileURL != "https://example.com/report.pdf" { + t.Fatalf("expected document URL to be preserved, got %#v", file.FileURL) + } + if file.Filename == nil || *file.Filename != "report.pdf" { + t.Fatalf("expected document title as filename, got %#v", file.Filename) + } + }) + } +} + +func TestToolResultDocumentBase64Preserved(t *testing.T) { + for _, grouped := range []bool{true, false} { + mode := "non-grouped" + if grouped { + mode = "grouped" + } + t.Run(mode, func(t *testing.T) { + content := []AnthropicContentBlock{ + { + Type: AnthropicContentBlockTypeDocument, + Title: schemas.Ptr("inline.pdf"), + Source: &AnthropicBlockSource{SourceObj: &AnthropicSource{ + Type: "base64", + MediaType: schemas.Ptr("application/pdf"), + Data: schemas.Ptr("JVBERi0xLjQ="), + }}, + }, + } + + blocks := requireToolResultDocumentBlocks(t, convertToolResultDocumentBlocksForTest(grouped, content)) + if len(blocks) != 1 { + t.Fatalf("expected 1 document block, got %d", len(blocks)) + } + file := blocks[0].ResponsesInputMessageContentBlockFile + if file == nil || file.FileData == nil { + t.Fatal("expected inline file data") + } + if got := *file.FileData; got != "data:application/pdf;base64,JVBERi0xLjQ=" { + t.Fatalf("expected base64 data and media type to be preserved, got %q", got) + } + if file.FileType == nil || *file.FileType != "application/pdf" { + t.Fatalf("expected application/pdf file type, got %#v", file.FileType) + } + }) + } +} diff --git a/core/providers/bedrock/responses.go b/core/providers/bedrock/responses.go index 4b8431da9c6..87052b6d27f 100644 --- a/core/providers/bedrock/responses.go +++ b/core/providers/bedrock/responses.go @@ -3,7 +3,6 @@ package bedrock import ( "bytes" "context" - "encoding/base64" "encoding/json" "fmt" "strings" @@ -3446,6 +3445,16 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage } else { resultContent = append(resultContent, BedrockContentBlock{Image: imageSource}) } + } else if block.Type == schemas.ResponsesInputMessageContentBlockTypeFile && + block.ResponsesInputMessageContentBlockFile != nil { + file := block.ResponsesInputMessageContentBlockFile + document, err := materializeBedrockDocument(ctx, file.FileData, file.FileURL, file.Filename, file.FileType) + if err != nil { + return nil, nil, fmt.Errorf("bedrock: converting tool result document: %w", err) + } + if document != nil { + resultContent = append(resultContent, BedrockContentBlock{Document: document}) + } } } } @@ -4524,76 +4533,12 @@ func convertBifrostResponsesMessageContentBlocksToBedrockContentBlocks(ctx conte continue case schemas.ResponsesInputMessageContentBlockTypeFile: if block.ResponsesInputMessageContentBlockFile != nil { - doc := &BedrockDocumentSource{ - Name: "document", // Default - Format: "pdf", // Default - Source: &BedrockDocumentSourceData{}, - } - - // Set filename (normalized for Bedrock) - if block.ResponsesInputMessageContentBlockFile.Filename != nil { - doc.Name = normalizeBedrockFilename(*block.ResponsesInputMessageContentBlockFile.Filename) - } - - // Determine format: text or PDF based on FileType - isTextFile := false - if block.ResponsesInputMessageContentBlockFile.FileType != nil { - fileType := *block.ResponsesInputMessageContentBlockFile.FileType - // Check if it's a text type - if fileType == "text/markdown" || fileType == "md" { - doc.Format = "md" - isTextFile = true - } else if fileType == "text/html" || fileType == "html" { - doc.Format = "html" - isTextFile = true - } else if fileType == "text/csv" || fileType == "csv" { - doc.Format = "csv" - isTextFile = true - } else if strings.HasPrefix(fileType, "text/") || fileType == "txt" { - doc.Format = "txt" - isTextFile = true - } else if strings.Contains(fileType, "pdf") || fileType == "pdf" { - doc.Format = "pdf" - } else if strings.Contains(fileType, "spreadsheetml") || fileType == "xlsx" { - doc.Format = "xlsx" - } else if fileType == "application/vnd.ms-excel" || fileType == "xls" { - doc.Format = "xls" - } else if strings.Contains(fileType, "wordprocessingml") || fileType == "docx" { - doc.Format = "docx" - } else if fileType == "application/msword" || fileType == "doc" { - doc.Format = "doc" - } - } - - // Handle file data - if block.ResponsesInputMessageContentBlockFile.FileData != nil { - fileData := *block.ResponsesInputMessageContentBlockFile.FileData - - // Check if it's a data URL (e.g., "data:application/pdf;base64,...") - if strings.HasPrefix(fileData, "data:") { - urlInfo := schemas.ExtractURLTypeInfo(fileData) - if urlInfo.DataURLWithoutPrefix != nil { - // PDF or other binary - keep as base64 - doc.Source.Bytes = urlInfo.DataURLWithoutPrefix - bedrockBlock.Document = doc - break - } - } - - // Not a data URL - use as-is - if isTextFile { - // bytes is necessary for bedrock - // base64 string of the text - doc.Source.Text = &fileData - encoded := base64.StdEncoding.EncodeToString([]byte(fileData)) - doc.Source.Bytes = &encoded - } else { - doc.Source.Bytes = &fileData - } - - bedrockBlock.Document = doc - + file := block.ResponsesInputMessageContentBlockFile + document, err := materializeBedrockDocument(ctx, file.FileData, file.FileURL, file.Filename, file.FileType) + if err != nil { + return nil, fmt.Errorf("failed to convert document in responses content block: %w", err) } + bedrockBlock.Document = document } default: // Don't add anything for unknown types diff --git a/core/providers/bedrock/toolresultdocument_test.go b/core/providers/bedrock/toolresultdocument_test.go new file mode 100644 index 00000000000..744e0f5e15d --- /dev/null +++ b/core/providers/bedrock/toolresultdocument_test.go @@ -0,0 +1,196 @@ +package bedrock + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/maximhq/bifrost/core/providers/anthropic" + "github.com/maximhq/bifrost/core/schemas" +) + +func toolResultDocumentMessage(blocks []schemas.ResponsesMessageContentBlock) schemas.ResponsesMessage { + return schemas.ResponsesMessage{ + Type: schemas.Ptr(schemas.ResponsesMessageTypeFunctionCallOutput), + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + CallID: schemas.Ptr("toolu_document"), + Output: &schemas.ResponsesToolMessageOutputStruct{ + ResponsesFunctionToolCallOutputBlocks: blocks, + }, + }, + } +} + +func requireBedrockToolResultDocument(t *testing.T, messages []BedrockMessage) *BedrockToolResult { + t.Helper() + if len(messages) != 1 { + t.Fatalf("expected 1 Bedrock message, got %d", len(messages)) + } + if len(messages[0].Content) != 1 || messages[0].Content[0].ToolResult == nil { + t.Fatalf("expected one tool result block, got %#v", messages[0].Content) + } + return messages[0].Content[0].ToolResult +} + +func TestToolResultDocumentInlinePreserved(t *testing.T) { + input := []schemas.ResponsesMessage{ + toolResultDocumentMessage([]schemas.ResponsesMessageContentBlock{ + { + Type: schemas.ResponsesInputMessageContentBlockTypeText, + Text: schemas.Ptr("Document generated"), + }, + { + Type: schemas.ResponsesInputMessageContentBlockTypeFile, + ResponsesInputMessageContentBlockFile: &schemas.ResponsesInputMessageContentBlockFile{ + FileData: schemas.Ptr("data:application/pdf;base64,JVBERi0xLjQ="), + Filename: schemas.Ptr("quarterly/report.pdf"), + FileType: schemas.Ptr("application/pdf"), + }, + }, + }), + } + + messages, _, err := ConvertBifrostMessagesToBedrockMessages(context.Background(), input, false) + if err != nil { + t.Fatalf("unexpected conversion error: %v", err) + } + toolResult := requireBedrockToolResultDocument(t, messages) + if len(toolResult.Content) != 2 { + t.Fatalf("expected text and document blocks, got %d", len(toolResult.Content)) + } + if toolResult.Content[0].Text == nil || *toolResult.Content[0].Text != "Document generated" { + t.Fatalf("expected text block first, got %#v", toolResult.Content[0]) + } + document := toolResult.Content[1].Document + if document == nil { + t.Fatalf("expected document block second, got %#v", toolResult.Content[1]) + } + if document.Format != "pdf" { + t.Fatalf("expected pdf format, got %q", document.Format) + } + if document.Name != "quarterly_report_pdf" { + t.Fatalf("expected normalized filename, got %q", document.Name) + } + if document.Source == nil || document.Source.Bytes == nil || *document.Source.Bytes != "JVBERi0xLjQ=" { + t.Fatalf("expected inline PDF bytes, got %#v", document.Source) + } +} + +func TestToolResultDocumentURLFetchErrorReturned(t *testing.T) { + input := []schemas.ResponsesMessage{ + toolResultDocumentMessage([]schemas.ResponsesMessageContentBlock{ + { + Type: schemas.ResponsesInputMessageContentBlockTypeFile, + ResponsesInputMessageContentBlockFile: &schemas.ResponsesInputMessageContentBlockFile{ + FileURL: schemas.Ptr("file:///tmp/report.pdf"), + Filename: schemas.Ptr("report.pdf"), + FileType: schemas.Ptr("application/pdf"), + }, + }, + }), + } + + messages, _, err := ConvertBifrostMessagesToBedrockMessages(context.Background(), input, false) + if err == nil { + t.Fatalf("expected URL materialization error, got messages %#v", messages) + } + if !strings.Contains(err.Error(), "unsupported URL scheme") { + t.Fatalf("expected bounded URL fetch error, got %v", err) + } +} + +func TestToolResultDocumentURLUsesSSRFSafeFetcher(t *testing.T) { + var reached atomic.Bool + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + reached.Store(true) + })) + defer server.Close() + + input := []schemas.ResponsesMessage{ + toolResultDocumentMessage([]schemas.ResponsesMessageContentBlock{ + { + Type: schemas.ResponsesInputMessageContentBlockTypeFile, + ResponsesInputMessageContentBlockFile: &schemas.ResponsesInputMessageContentBlockFile{ + FileURL: schemas.Ptr(server.URL + "/report.pdf"), + Filename: schemas.Ptr("report.pdf"), + FileType: schemas.Ptr("application/pdf"), + }, + }, + }), + } + + messages, _, err := ConvertBifrostMessagesToBedrockMessages(context.Background(), input, false) + if err == nil { + t.Fatalf("expected loopback URL to be rejected, got messages %#v", messages) + } + if !strings.Contains(err.Error(), "blocked connection to non-public address") { + t.Fatalf("expected SSRF-safe fetch rejection, got %v", err) + } + if reached.Load() { + t.Fatal("SSRF-safe fetcher reached the loopback server") + } +} + +func TestToolResultDocumentAnthropicToBedrockRoundTrip(t *testing.T) { + request := &anthropic.AnthropicMessageRequest{ + Model: "bedrock/anthropic.claude-3-5-sonnet-v2", + MaxTokens: 1024, + Messages: []anthropic.AnthropicMessage{ + { + Role: anthropic.AnthropicMessageRoleAssistant, + Content: anthropic.AnthropicContent{ContentBlocks: []anthropic.AnthropicContentBlock{ + { + Type: anthropic.AnthropicContentBlockTypeToolUse, + ID: schemas.Ptr("toolu_document"), + Name: schemas.Ptr("create_report"), + Input: []byte(`{"topic":"quarterly results"}`), + }, + }}, + }, + { + Role: anthropic.AnthropicMessageRoleUser, + Content: anthropic.AnthropicContent{ContentBlocks: []anthropic.AnthropicContentBlock{ + { + Type: anthropic.AnthropicContentBlockTypeToolResult, + ToolUseID: schemas.Ptr("toolu_document"), + Content: &anthropic.AnthropicContent{ContentBlocks: []anthropic.AnthropicContentBlock{ + { + Type: anthropic.AnthropicContentBlockTypeText, + Text: schemas.Ptr("Document generated"), + }, + { + Type: anthropic.AnthropicContentBlockTypeDocument, + Title: schemas.Ptr("report.pdf"), + Source: &anthropic.AnthropicBlockSource{SourceObj: &anthropic.AnthropicSource{ + Type: "base64", + MediaType: schemas.Ptr("application/pdf"), + Data: schemas.Ptr("JVBERi0xLjQ="), + }}, + }, + }}, + }, + }}, + }, + }, + } + + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + bifrostRequest := request.ToBifrostResponsesRequest(ctx) + bedrockRequest, err := ToBedrockResponsesRequest(ctx, bifrostRequest) + if err != nil { + t.Fatalf("unexpected cross-provider conversion error: %v", err) + } + if len(bedrockRequest.Messages) != 2 { + t.Fatalf("expected assistant tool use and user tool result, got %d messages", len(bedrockRequest.Messages)) + } + toolResult := bedrockRequest.Messages[1].Content[0].ToolResult + if toolResult == nil || len(toolResult.Content) != 2 { + t.Fatalf("expected text and document in final tool result, got %#v", toolResult) + } + if toolResult.Content[0].Text == nil || toolResult.Content[1].Document == nil { + t.Fatalf("expected text then document, got %#v", toolResult.Content) + } +} diff --git a/core/providers/bedrock/utils.go b/core/providers/bedrock/utils.go index b1d9e871f14..ab788bfba5c 100644 --- a/core/providers/bedrock/utils.go +++ b/core/providers/bedrock/utils.go @@ -188,6 +188,103 @@ func normalizeBedrockFilename(filename string) string { return normalized } +func bedrockDocumentFormat(fileType string) (format string, isText bool) { + normalized := strings.ToLower(strings.TrimSpace(fileType)) + if mediaType, _, err := mime.ParseMediaType(normalized); err == nil { + normalized = mediaType + } + + switch { + case normalized == "text/markdown" || normalized == "md": + return "md", true + case normalized == "text/html" || normalized == "html": + return "html", true + case normalized == "text/csv" || normalized == "csv": + return "csv", true + case strings.HasPrefix(normalized, "text/") || normalized == "txt": + return "txt", true + case normalized == "application/msword" || normalized == "doc": + return "doc", false + case strings.Contains(normalized, "wordprocessingml") || normalized == "docx": + return "docx", false + case normalized == "application/vnd.ms-excel" || normalized == "xls": + return "xls", false + case strings.Contains(normalized, "spreadsheetml") || normalized == "xlsx": + return "xlsx", false + case strings.Contains(normalized, "pdf") || normalized == "pdf": + return "pdf", false + default: + return "pdf", false + } +} + +// materializeBedrockDocument converts the canonical Bifrost file representation +// into the inline document shape required by Bedrock Converse. Remote documents +// use the shared bounded, SSRF-safe fetcher; inline data URLs are reduced to their +// base64 payload before emission. +func materializeBedrockDocument( + ctx context.Context, + fileData *string, + fileURL *string, + filename *string, + fileType *string, +) (*BedrockDocumentSource, error) { + document := &BedrockDocumentSource{ + Name: "document", + Format: "pdf", + Source: &BedrockDocumentSourceData{}, + } + if filename != nil { + document.Name = normalizeBedrockFilename(*filename) + } + + isText := false + if fileType != nil { + document.Format, isText = bedrockDocumentFormat(*fileType) + } + + if fileURL != nil && *fileURL != "" { + if ctx == nil { + ctx = context.Background() + } + fetchedMediaType, fetchedBase64, err := providerUtils.FetchAndEncodeURL(ctx, *fileURL) + if err != nil { + return nil, err + } + if fetchedMediaType != "" { + document.Format, _ = bedrockDocumentFormat(fetchedMediaType) + } + document.Source.Bytes = &fetchedBase64 + return document, nil + } + + if fileData == nil { + return nil, nil + } + + data := *fileData + if strings.HasPrefix(data, "data:") { + urlInfo := schemas.ExtractURLTypeInfo(data) + if urlInfo.DataURLWithoutPrefix == nil { + return nil, fmt.Errorf("invalid document data URL") + } + if urlInfo.MediaType != nil && *urlInfo.MediaType != "" { + document.Format, _ = bedrockDocumentFormat(*urlInfo.MediaType) + } + document.Source.Bytes = urlInfo.DataURLWithoutPrefix + return document, nil + } + + if isText { + document.Source.Text = &data + encoded := base64.StdEncoding.EncodeToString([]byte(data)) + document.Source.Bytes = &encoded + } else { + document.Source.Bytes = &data + } + return document, nil +} + // bedrockAliasToolName returns a Bedrock-safe tool name and records a reverse mapping. func bedrockAliasToolName(ctx context.Context, name string) string { if len(name) <= 64 && !bedrockUnsafeToolNameCharRegex.MatchString(name) { @@ -1108,124 +1205,20 @@ func convertContentBlock(ctx context.Context, block schemas.ChatContentBlock) ([ return nil, fmt.Errorf("file block missing file field") } - documentSource := &BedrockDocumentSource{ - Name: "document", - Format: "pdf", - Source: &BedrockDocumentSourceData{}, - } - - // Set filename (normalized for Bedrock) - if block.File.Filename != nil { - documentSource.Name = normalizeBedrockFilename(*block.File.Filename) - } - - // Convert MIME type to Bedrock format - isText := false - if block.File.FileType != nil { - fileType := *block.File.FileType - switch { - case fileType == "text/plain" || fileType == "txt": - documentSource.Format = "txt" - isText = true - case fileType == "text/markdown" || fileType == "md": - documentSource.Format = "md" - isText = true - case fileType == "text/html" || fileType == "html": - documentSource.Format = "html" - isText = true - case fileType == "text/csv" || fileType == "csv": - documentSource.Format = "csv" - isText = true - case fileType == "application/msword" || fileType == "doc": - documentSource.Format = "doc" - case strings.Contains(fileType, "wordprocessingml") || fileType == "docx": - documentSource.Format = "docx" - case fileType == "application/vnd.ms-excel" || fileType == "xls": - documentSource.Format = "xls" - case strings.Contains(fileType, "spreadsheetml") || fileType == "xlsx": - documentSource.Format = "xlsx" - case strings.Contains(fileType, "pdf") || fileType == "pdf": - documentSource.Format = "pdf" - } - } - - // URL-sourced document: fetch and inline the bytes (Bedrock Converse only - // accepts inline source bytes, not remote URLs). - if block.File.FileURL != nil && *block.File.FileURL != "" { - fetchedMediaType, fetchedB64, fetchErr := providerUtils.FetchAndEncodeURL(ctx, *block.File.FileURL) - if fetchErr != nil { - return nil, fetchErr - } - // Refine format from response Content-Type when present (more reliable - // than file extension or upstream-declared media type). Normalize to - // strip parameters (e.g. "; charset=utf-8") and lowercase the base type. - if mt, _, err := mime.ParseMediaType(fetchedMediaType); err == nil { - fetchedMediaType = mt - } - switch fetchedMediaType { - case "application/pdf": - documentSource.Format = "pdf" - case "text/plain": - documentSource.Format = "txt" - isText = true - case "text/markdown": - documentSource.Format = "md" - isText = true - case "text/html": - documentSource.Format = "html" - isText = true - case "text/csv": - documentSource.Format = "csv" - isText = true - case "application/msword": - documentSource.Format = "doc" - case "application/vnd.openxmlformats-officedocument.wordprocessingml.document": - documentSource.Format = "docx" - case "application/vnd.ms-excel": - documentSource.Format = "xls" - case "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": - documentSource.Format = "xlsx" - } - documentSource.Source.Bytes = &fetchedB64 - return []BedrockContentBlock{ - { - Document: documentSource, - }, - }, nil + document, err := materializeBedrockDocument( + ctx, + block.File.FileData, + block.File.FileURL, + block.File.Filename, + block.File.FileType, + ) + if err != nil { + return nil, err } - - // Handle file data - strip data URL prefix if present - if block.File.FileData != nil { - fileData := *block.File.FileData - - // Check if it's a data URL and extract raw base64 - if strings.HasPrefix(fileData, "data:") { - urlInfo := schemas.ExtractURLTypeInfo(fileData) - if urlInfo.DataURLWithoutPrefix != nil { - documentSource.Source.Bytes = urlInfo.DataURLWithoutPrefix - return []BedrockContentBlock{ - { - Document: documentSource, - }, - }, nil - } - } - - // Set text or bytes based on file type - if isText { - documentSource.Source.Text = &fileData // Plain text - encoded := base64.StdEncoding.EncodeToString([]byte(fileData)) - documentSource.Source.Bytes = &encoded // Also sets Bytes - } else { - documentSource.Source.Bytes = &fileData - } + if document == nil { + return nil, nil } - - return []BedrockContentBlock{ - { - Document: documentSource, - }, - }, nil + return []BedrockContentBlock{{Document: document}}, nil case schemas.ChatContentBlockTypeInputAudio: // Bedrock doesn't support audio input in Converse API return nil, fmt.Errorf("audio input not supported in Bedrock Converse API")