Files
bifrost/plugins/maxim/attachments_test.go
Beyhan Oğur 880f412e2c first commit
2026-04-26 21:52:23 +03:00

651 lines
20 KiB
Go

package maxim
import (
"context"
"encoding/base64"
"io"
"log"
"net/http"
"os"
"strings"
"testing"
"time"
bifrost "github.com/maximhq/bifrost/core"
"github.com/maximhq/bifrost/core/schemas"
"github.com/maximhq/maxim-go/logging"
)
// Test URLs and assets from core/internal/llmtests/utils.go for consistency across Bifrost tests.
var (
testFileURL = "https://www.berkshirehathaway.com/letters/2024ltr.pdf"
testImageURL = "https://pestworldcdn-dcf2a8gbggazaghf.z01.azurefd.net/media/561791/carpenter-ant4.jpg"
testImageBase64 = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/2wBDAQkJCQwLDBgNDRgyIRwhMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjL/wAARCAAIAAoDASIAAhEBAxEB/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAb/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwCdABmX/9k="
)
func strPtr(s string) *string { return &s }
func responsesUserRole() *schemas.ResponsesMessageRoleType {
r := schemas.ResponsesInputMessageRoleUser
return &r
}
func TestExtractAttachmentsFromRequest_ChatImageUrlHttp(t *testing.T) {
req := &schemas.BifrostRequest{
RequestType: schemas.ChatCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeImage,
ImageURLStruct: &schemas.ChatInputImage{
URL: testImageURL,
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.UrlAttachment); !ok {
t.Errorf("expected *logging.UrlAttachment, got %T", attachments[0])
}
ua := attachments[0].(*logging.UrlAttachment)
if ua.URL != testImageURL {
t.Errorf("expected URL %q, got %q", testImageURL, ua.URL)
}
}
func TestExtractAttachmentsFromRequest_ChatImageUrlData(t *testing.T) {
b64 := base64.StdEncoding.EncodeToString([]byte("fake-png-bytes"))
dataURL := "data:image/png;base64," + b64
req := &schemas.BifrostRequest{
RequestType: schemas.ChatCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeImage,
ImageURLStruct: &schemas.ChatInputImage{
URL: dataURL,
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.FileDataAttachment); !ok {
t.Errorf("expected *logging.FileDataAttachment, got %T", attachments[0])
}
fda := attachments[0].(*logging.FileDataAttachment)
if string(fda.Data) != "fake-png-bytes" {
t.Errorf("expected data 'fake-png-bytes', got %q", string(fda.Data))
}
if fda.MimeType != "image/png" {
t.Errorf("expected mime image/png, got %q", fda.MimeType)
}
}
func TestExtractAttachmentsFromRequest_ChatFileUrl(t *testing.T) {
fileURL := testFileURL
req := &schemas.BifrostRequest{
RequestType: schemas.ChatCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeFile,
File: &schemas.ChatInputFile{
FileURL: &fileURL,
Filename: strPtr("2024ltr.pdf"),
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.UrlAttachment); !ok {
t.Errorf("expected *logging.UrlAttachment, got %T", attachments[0])
}
ua := attachments[0].(*logging.UrlAttachment)
if ua.URL != fileURL {
t.Errorf("expected URL %q, got %q", fileURL, ua.URL)
}
}
func TestExtractAttachmentsFromRequest_ChatFileData(t *testing.T) {
b64 := base64.StdEncoding.EncodeToString([]byte("pdf content"))
req := &schemas.BifrostRequest{
RequestType: schemas.ChatCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeFile,
File: &schemas.ChatInputFile{
FileData: &b64,
Filename: strPtr("doc.pdf"),
FileType: strPtr("application/pdf"),
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.FileDataAttachment); !ok {
t.Errorf("expected *logging.FileDataAttachment, got %T", attachments[0])
}
fda := attachments[0].(*logging.FileDataAttachment)
if string(fda.Data) != "pdf content" {
t.Errorf("expected data 'pdf content', got %q", string(fda.Data))
}
}
func TestExtractAttachmentsFromRequest_ChatInputAudio(t *testing.T) {
b64 := base64.StdEncoding.EncodeToString([]byte("audio-bytes"))
format := "wav"
req := &schemas.BifrostRequest{
RequestType: schemas.ChatCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeInputAudio,
InputAudio: &schemas.ChatInputAudio{
Data: b64,
Format: &format,
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.FileDataAttachment); !ok {
t.Errorf("expected *logging.FileDataAttachment, got %T", attachments[0])
}
fda := attachments[0].(*logging.FileDataAttachment)
if string(fda.Data) != "audio-bytes" {
t.Errorf("expected data 'audio-bytes', got %q", string(fda.Data))
}
if fda.MimeType != "audio/wav" {
t.Errorf("expected mime audio/wav, got %q", fda.MimeType)
}
}
func TestExtractAttachmentsFromRequest_ResponsesInputImage(t *testing.T) {
req := &schemas.BifrostRequest{
RequestType: schemas.ResponsesRequest,
ResponsesRequest: &schemas.BifrostResponsesRequest{
Input: []schemas.ResponsesMessage{
{
Role: responsesUserRole(),
Content: &schemas.ResponsesMessageContent{
ContentBlocks: []schemas.ResponsesMessageContentBlock{
{
Type: schemas.ResponsesInputMessageContentBlockTypeImage,
ResponsesInputMessageContentBlockImage: &schemas.ResponsesInputMessageContentBlockImage{
ImageURL: &testImageURL,
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.UrlAttachment); !ok {
t.Errorf("expected *logging.UrlAttachment, got %T", attachments[0])
}
}
func TestExtractAttachmentsFromRequest_ResponsesInputFile(t *testing.T) {
fileURL := testFileURL
req := &schemas.BifrostRequest{
RequestType: schemas.ResponsesRequest,
ResponsesRequest: &schemas.BifrostResponsesRequest{
Input: []schemas.ResponsesMessage{
{
Role: responsesUserRole(),
Content: &schemas.ResponsesMessageContent{
ContentBlocks: []schemas.ResponsesMessageContentBlock{
{
Type: schemas.ResponsesInputMessageContentBlockTypeFile,
ResponsesInputMessageContentBlockFile: &schemas.ResponsesInputMessageContentBlockFile{
FileURL: &fileURL,
Filename: strPtr("2024ltr.pdf"),
FileType: strPtr("application/pdf"),
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.UrlAttachment); !ok {
t.Errorf("expected *logging.UrlAttachment, got %T", attachments[0])
}
}
func TestExtractAttachmentsFromRequest_ResponsesInputAudio(t *testing.T) {
b64 := base64.StdEncoding.EncodeToString([]byte("mp3-bytes"))
req := &schemas.BifrostRequest{
RequestType: schemas.ResponsesRequest,
ResponsesRequest: &schemas.BifrostResponsesRequest{
Input: []schemas.ResponsesMessage{
{
Role: responsesUserRole(),
Content: &schemas.ResponsesMessageContent{
ContentBlocks: []schemas.ResponsesMessageContentBlock{
{
Type: schemas.ResponsesInputMessageContentBlockTypeAudio,
Audio: &schemas.ResponsesInputMessageContentBlockAudio{
Format: "mp3",
Data: b64,
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.FileDataAttachment); !ok {
t.Errorf("expected *logging.FileDataAttachment, got %T", attachments[0])
}
fda := attachments[0].(*logging.FileDataAttachment)
if string(fda.Data) != "mp3-bytes" {
t.Errorf("expected data 'mp3-bytes', got %q", string(fda.Data))
}
if fda.MimeType != "audio/mpeg" {
t.Errorf("expected mime audio/mpeg, got %q", fda.MimeType)
}
}
func TestExtractAttachmentsFromRequest_NoAttachments(t *testing.T) {
req := &schemas.BifrostRequest{
RequestType: schemas.ChatCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentStr: strPtr("Hello, world!"),
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 0 {
t.Fatalf("expected 0 attachments, got %d", len(attachments))
}
}
func TestExtractAttachmentsFromRequest_TextCompletion(t *testing.T) {
req := &schemas.BifrostRequest{
RequestType: schemas.TextCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeImage,
ImageURLStruct: &schemas.ChatInputImage{
URL: "https://example.com/image.png",
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 0 {
t.Fatalf("expected 0 attachments for text completion, got %d", len(attachments))
}
}
func TestExtractAttachmentsFromRequest_ImageGenerationInputImagesUrl(t *testing.T) {
req := &schemas.BifrostRequest{
RequestType: schemas.ImageGenerationRequest,
ImageGenerationRequest: &schemas.BifrostImageGenerationRequest{
Input: &schemas.ImageGenerationInput{Prompt: "make it pop"},
Params: &schemas.ImageGenerationParameters{
InputImages: []string{testImageURL},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
ua, ok := attachments[0].(*logging.UrlAttachment)
if !ok {
t.Fatalf("expected *logging.UrlAttachment, got %T", attachments[0])
}
if ua.URL != testImageURL {
t.Errorf("expected URL %q, got %q", testImageURL, ua.URL)
}
}
func TestExtractAttachmentsFromRequest_ImageGenerationInputImagesDataURL(t *testing.T) {
req := &schemas.BifrostRequest{
RequestType: schemas.ImageGenerationStreamRequest,
ImageGenerationRequest: &schemas.BifrostImageGenerationRequest{
Input: &schemas.ImageGenerationInput{Prompt: "variation"},
Params: &schemas.ImageGenerationParameters{
InputImages: []string{testImageBase64},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.FileDataAttachment); !ok {
t.Errorf("expected *logging.FileDataAttachment, got %T", attachments[0])
}
}
func TestExtractAttachmentsFromRequest_ImageGenerationInputImagesRawBase64(t *testing.T) {
idx := strings.Index(testImageBase64, ";base64,")
if idx == -1 {
t.Fatal("invalid testImageBase64 format")
}
rawB64 := testImageBase64[idx+8:]
req := &schemas.BifrostRequest{
RequestType: schemas.ImageGenerationRequest,
ImageGenerationRequest: &schemas.BifrostImageGenerationRequest{
Input: &schemas.ImageGenerationInput{Prompt: "edit"},
Params: &schemas.ImageGenerationParameters{
InputImages: []string{rawB64},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
fda, ok := attachments[0].(*logging.FileDataAttachment)
if !ok {
t.Fatalf("expected *logging.FileDataAttachment, got %T", attachments[0])
}
if fda.MimeType != "image/jpeg" {
t.Errorf("expected image/jpeg from DetectContentType, got %q", fda.MimeType)
}
if len(fda.Data) == 0 {
t.Error("expected non-empty decoded image bytes")
}
}
func TestExtractAttachmentsFromRequest_ImageGenerationNoInputImages(t *testing.T) {
req := &schemas.BifrostRequest{
RequestType: schemas.ImageGenerationRequest,
ImageGenerationRequest: &schemas.BifrostImageGenerationRequest{
Input: &schemas.ImageGenerationInput{Prompt: "text only"},
Params: &schemas.ImageGenerationParameters{},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 0 {
t.Fatalf("expected 0 attachments, got %d", len(attachments))
}
}
func TestExtractAttachmentsFromRequest_NilRequest(t *testing.T) {
attachments := ExtractAttachmentsFromRequest(nil)
if attachments != nil {
t.Fatalf("expected nil for nil request, got %v", attachments)
}
}
// TestExtractAttachmentsFromRequest_ChatFileDataFromBase64 uses testImageBase64
// and verifies FileData extraction from base64-encoded content.
func TestExtractAttachmentsFromRequest_ChatFileDataFromBase64(t *testing.T) {
// Extract base64 from data URL (format: data:image/jpeg;base64,...)
idx := strings.Index(testImageBase64, ";base64,")
if idx == -1 {
t.Fatal("invalid testImageBase64 format")
}
b64 := testImageBase64[idx+8:]
data, err := base64.StdEncoding.DecodeString(b64)
if err != nil {
t.Fatalf("failed to decode test image base64: %v", err)
}
b64ForRequest := base64.StdEncoding.EncodeToString(data)
req := &schemas.BifrostRequest{
RequestType: schemas.ChatCompletionRequest,
ChatRequest: &schemas.BifrostChatRequest{
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeFile,
File: &schemas.ChatInputFile{
FileData: &b64ForRequest,
Filename: strPtr("grey_solid.jpg"),
FileType: strPtr("image/jpeg"),
},
},
},
},
},
},
},
}
attachments := ExtractAttachmentsFromRequest(req)
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if _, ok := attachments[0].(*logging.FileDataAttachment); !ok {
t.Errorf("expected *logging.FileDataAttachment, got %T", attachments[0])
}
fda := attachments[0].(*logging.FileDataAttachment)
if len(fda.Data) != len(data) {
t.Errorf("expected data length %d, got %d", len(data), len(fda.Data))
}
if fda.MimeType != "image/jpeg" {
t.Errorf("expected mime image/jpeg, got %q", fda.MimeType)
}
}
// requireIntegrationEnv skips the test if required env vars for real API calls are not set.
func requireIntegrationEnv(t *testing.T) {
t.Helper()
if os.Getenv("OPENAI_API_KEY") == "" {
t.Skip("OPENAI_API_KEY not set, skipping integration test")
}
if os.Getenv("MAXIM_API_KEY") == "" {
t.Skip("MAXIM_API_KEY not set, skipping integration test")
}
if os.Getenv("MAXIM_LOG_REPO_ID") == "" {
t.Skip("MAXIM_LOG_REPO_ID not set, skipping integration test")
}
}
// TestVisionWithImageUrl_Integration sends a real OpenAI vision request via Bifrost
// with the maxim plugin. The plugin extracts the image_url and logs it as an attachment
// to Maxim. Verify attachments in the Maxim dashboard after the test.
func TestVisionWithImageUrl_Integration(t *testing.T) {
requireIntegrationEnv(t)
plugin, err := getPlugin()
if err != nil {
t.Fatalf("failed to get plugin: %v", err)
}
client, err := bifrost.Init(context.Background(), schemas.BifrostConfig{
Account: &BaseAccount{},
LLMPlugins: []schemas.LLMPlugin{plugin},
Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug),
})
if err != nil {
t.Fatalf("failed to init Bifrost: %v", err)
}
defer client.Shutdown()
ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
_, bifrostErr := client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
Provider: schemas.OpenAI,
Model: "gpt-4o-mini",
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeText,
Text: bifrost.Ptr("Describe this image in one sentence."),
},
{
Type: schemas.ChatContentBlockTypeImage,
ImageURLStruct: &schemas.ChatInputImage{
URL: testImageURL,
},
},
},
},
},
},
})
if bifrostErr != nil {
t.Fatalf("ChatCompletionRequest failed: %v", bifrostErr)
}
log.Printf("Vision request with image URL completed. Check Maxim dashboard for trace with attachment.")
}
// TestVisionWithImageData_Integration fetches testImageURL, encodes it as a data URL,
// and sends to OpenAI vision via Bifrost. The maxim plugin extracts the data URL and
// logs it as a FileDataAttachment. Uses a real image (carpenter ant) since OpenAI
// rejects minimal test images like the grey solid.
func TestVisionWithImageData_Integration(t *testing.T) {
requireIntegrationEnv(t)
// Fetch real image and encode as data URL (OpenAI rejects minimal grey solid)
httpClient := &http.Client{Timeout: 20 * time.Second}
resp, err := httpClient.Get(testImageURL)
if err != nil {
t.Skipf("failed to fetch test image: %v", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Skipf("test image URL returned %d", resp.StatusCode)
}
imgData, err := io.ReadAll(resp.Body)
if err != nil {
t.Skipf("failed to read test image: %v", err)
}
dataURL := "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(imgData)
plugin, err := getPlugin()
if err != nil {
t.Fatalf("failed to get plugin: %v", err)
}
client, err := bifrost.Init(context.Background(), schemas.BifrostConfig{
Account: &BaseAccount{},
LLMPlugins: []schemas.LLMPlugin{plugin},
Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug),
})
if err != nil {
t.Fatalf("failed to init Bifrost: %v", err)
}
defer client.Shutdown()
ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
_, bifrostErr := client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
Provider: schemas.OpenAI,
Model: "gpt-4o-mini",
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentBlocks: []schemas.ChatContentBlock{
{
Type: schemas.ChatContentBlockTypeText,
Text: bifrost.Ptr("What do you see in this image? One sentence."),
},
{
Type: schemas.ChatContentBlockTypeImage,
ImageURLStruct: &schemas.ChatInputImage{
URL: dataURL,
},
},
},
},
},
},
})
if bifrostErr != nil {
t.Fatalf("ChatCompletionRequest failed: %v", bifrostErr)
}
log.Printf("Vision request with base64 image completed. Check Maxim dashboard for trace with FileData attachment.")
}