diff --git a/message/datauri.go b/message/datauri.go index b336c0a6..2a6730d4 100644 --- a/message/datauri.go +++ b/message/datauri.go @@ -68,6 +68,31 @@ func parseDataURI(uri string) (*dataURI, error) { }, nil } +// DecodeDataURI parses a data URI (RFC 2397) and returns the decoded bytes +// together with its media type. It is intended for providers that must send +// inline binary data (rather than a URI reference) when handling a URIContent +// whose URI is a data: URI. +func DecodeDataURI(uri string) (data []byte, mediaType string, err error) { + parsed, err := parseDataURI(uri) + if err != nil { + return nil, "", err + } + if !parsed.IsBase64 { + // Non-base64 data URIs carry percent-encoded payloads; return the + // unescaped bytes directly instead of round-tripping through base64. + unescaped, err := url.PathUnescape(parsed.Data) + if err != nil { + unescaped = parsed.Data + } + return []byte(unescaped), parsed.MediaType, nil + } + decoded, err := base64.StdEncoding.DecodeString(parsed.Data) + if err != nil { + return nil, "", fmt.Errorf("invalid data URI format: failed to decode data: %w", err) + } + return decoded, parsed.MediaType, nil +} + // data returns the raw data portion of the data URI as a base64-encoded string. func (d *dataURI) data() string { if d.IsBase64 { diff --git a/message/datauri_test.go b/message/datauri_test.go new file mode 100644 index 00000000..ed43babf --- /dev/null +++ b/message/datauri_test.go @@ -0,0 +1,63 @@ +// Copyright (c) Microsoft. All rights reserved. + +package message_test + +import ( + "testing" + + "github.com/microsoft/agent-framework-go/message" +) + +func TestDecodeDataURI(t *testing.T) { + tests := []struct { + name string + uri string + wantData string + wantMediaType string + wantErr bool + }{ + { + name: "percent-encoded payload", + uri: "data:text/plain,hello%20world", + wantData: "hello world", + wantMediaType: "text/plain", + }, + { + name: "percent-encoded default media type", + uri: "data:,hello%2Cworld", + wantData: "hello,world", + wantMediaType: "text/plain;charset=US-ASCII", + }, + { + name: "base64 payload", + uri: "data:text/plain;base64,aGVsbG8gd29ybGQ=", + wantData: "hello world", + wantMediaType: "text/plain", + }, + { + name: "missing scheme", + uri: "text/plain,hello", + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, mediaType, err := message.DecodeDataURI(tt.uri) + if tt.wantErr { + if err == nil { + t.Fatalf("DecodeDataURI(%q) = nil error, want error", tt.uri) + } + return + } + if err != nil { + t.Fatalf("DecodeDataURI(%q) unexpected error: %v", tt.uri, err) + } + if string(data) != tt.wantData { + t.Errorf("data = %q, want %q", string(data), tt.wantData) + } + if mediaType != tt.wantMediaType { + t.Errorf("mediaType = %q, want %q", mediaType, tt.wantMediaType) + } + }) + } +} diff --git a/provider/geminiprovider/agent.go b/provider/geminiprovider/agent.go index 3dfd88d3..9c049f42 100644 --- a/provider/geminiprovider/agent.go +++ b/provider/geminiprovider/agent.go @@ -372,12 +372,33 @@ func buildRequestParts(msg *message.Message, callIDToName map[string]string) ([] }, }) case *message.URIContent: - parts = append(parts, &genai.Part{ - FileData: &genai.FileData{ - FileURI: c.URI, - MIMEType: c.MediaType, - }, - }) + if len(c.URI) >= len("data:") && strings.EqualFold(c.URI[:len("data:")], "data:") { + // A data: URI carries the bytes inline. Gemini's FileData.FileURI + // requires an external reference (gs:// or https://), so a data: URI + // would be silently dropped. Decode it into InlineData instead, + // mirroring the DataContent handling above and the Python SDK + // (from_bytes for data: URIs, from_uri otherwise). + data, mt, err := message.DecodeDataURI(c.URI) + if err != nil { + return nil, fmt.Errorf("geminiprovider: failed to decode data URI content: %w", err) + } + if c.MediaType != "" { + mt = c.MediaType + } + parts = append(parts, &genai.Part{ + InlineData: &genai.Blob{ + Data: data, + MIMEType: mt, + }, + }) + } else { + parts = append(parts, &genai.Part{ + FileData: &genai.FileData{ + FileURI: c.URI, + MIMEType: c.MediaType, + }, + }) + } case *message.HostedFileContent: parts = append(parts, &genai.Part{ FileData: &genai.FileData{ diff --git a/provider/geminiprovider/agent_test.go b/provider/geminiprovider/agent_test.go index dfa5c05b..642ac89e 100644 --- a/provider/geminiprovider/agent_test.go +++ b/provider/geminiprovider/agent_test.go @@ -1002,6 +1002,97 @@ func TestURIAndHostedFileInRequest(t *testing.T) { } } +// TestDataURIContentInRequest verifies that a URIContent whose URI is a data: +// URI is decoded into inlineData (Gemini's fileData.fileUri only accepts +// external references), while a true external URI still maps to fileData. +func TestDataURIContentInRequest(t *testing.T) { + const b64 = "iVBORw0KGgo=" + + t.Run("data uri decodes to inlineData", func(t *testing.T) { + bodyCh := make(chan []byte, 1) + server := httptest.NewServer(captureAndRespond(t, bodyCh, "application/json", minimalTextResponse("ok"))) + defer server.Close() + + a := newTestClient(t, server) + messages := []*message.Message{{ + Role: message.RoleUser, + Contents: []message.Content{ + &message.URIContent{URI: "data:image/png;base64," + b64, MediaType: "image/png"}, + }, + }} + if _, err := a.Run(t.Context(), messages).Collect(); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + part := firstPart(t, <-bodyCh) + if _, ok := part["fileData"]; ok { + t.Errorf("expected no fileData for data URI, got %v", part["fileData"]) + } + inlineData, _ := part["inlineData"].(map[string]any) + if inlineData == nil { + t.Fatal("expected inlineData for data URI") + } + if mime, _ := inlineData["mimeType"].(string); mime != "image/png" { + t.Errorf("inlineData.mimeType = %q, want %q", mime, "image/png") + } + // genai marshals the decoded bytes back to base64, so the payload + // round-trips to the original data URI base64. + if data, _ := inlineData["data"].(string); data != b64 { + t.Errorf("inlineData.data = %q, want %q", data, b64) + } + }) + + t.Run("external uri maps to fileData", func(t *testing.T) { + bodyCh := make(chan []byte, 1) + server := httptest.NewServer(captureAndRespond(t, bodyCh, "application/json", minimalTextResponse("ok"))) + defer server.Close() + + a := newTestClient(t, server) + messages := []*message.Message{{ + Role: message.RoleUser, + Contents: []message.Content{ + &message.URIContent{URI: "https://example.com/x.png", MediaType: "image/png"}, + }, + }} + if _, err := a.Run(t.Context(), messages).Collect(); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + part := firstPart(t, <-bodyCh) + if _, ok := part["inlineData"]; ok { + t.Errorf("expected no inlineData for external URI, got %v", part["inlineData"]) + } + fileData, _ := part["fileData"].(map[string]any) + if fileData == nil { + t.Fatal("expected fileData for external URI") + } + if fileURI, _ := fileData["fileUri"].(string); fileURI != "https://example.com/x.png" { + t.Errorf("fileData.fileUri = %q, want %q", fileURI, "https://example.com/x.png") + } + }) +} + +// firstPart unmarshals a captured Gemini request body and returns the first +// part of its single content. +func firstPart(t *testing.T, body []byte) map[string]any { + t.Helper() + var req map[string]any + if err := json.Unmarshal(body, &req); err != nil { + t.Fatalf("unmarshal request body: %v", err) + } + contents, _ := req["contents"].([]any) + if len(contents) != 1 { + t.Fatalf("contents length = %d, want 1", len(contents)) + } + content0, _ := contents[0].(map[string]any) + parts, _ := content0["parts"].([]any) + if len(parts) != 1 { + t.Fatalf("parts length = %d, want 1", len(parts)) + } + part, _ := parts[0].(map[string]any) + return part +} + func TestResponseWithFileAndInlineData(t *testing.T) { resp := map[string]any{ "candidates": []any{