From 1ea18526774887d329143050b1a9c60c77eb14dd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 3 Jan 2025 13:52:45 +0800 Subject: [PATCH] Remove prompt documentation and enhance image upload/download functionality in vision module - Deleted the obsolete prompt.md file, streamlining the codebase. - Added tests for image upload and download functionality using local and S3 storage drivers, ensuring robust handling of image files. - Implemented image compression during upload for both local and S3 storage, maintaining aspect ratio and limiting dimensions to a maximum size of 1920x1080. - Enhanced the local and S3 storage drivers to support image compression, improving storage efficiency. - Updated tests to verify image dimensions and content type after upload and download operations, ensuring data integrity. --- neo/vision/driver/local/storage.go | 101 ++++++++++++- neo/vision/driver/local/storage_test.go | 76 ++++++++++ neo/vision/driver/s3/storage.go | 120 +++++++++++++-- neo/vision/driver/s3/storage_test.go | 185 ++++++++++++++++-------- neo/vision/prompt.md | 8 - neo/vision/vision_test.go | 100 +++++++++++-- 6 files changed, 494 insertions(+), 96 deletions(-) delete mode 100644 neo/vision/prompt.md diff --git a/neo/vision/driver/local/storage.go b/neo/vision/driver/local/storage.go index ed873129..81de81a3 100644 --- a/neo/vision/driver/local/storage.go +++ b/neo/vision/driver/local/storage.go @@ -1,9 +1,13 @@ package local import ( + "bytes" "context" "crypto/sha256" "fmt" + "image" + "image/jpeg" + "image/png" "io" "path/filepath" "strings" @@ -12,6 +16,9 @@ import ( "github.com/yaoapp/gou/fs" ) +// MaxImageSize maximum image size (1920x1080) +const MaxImageSize = 1920 + // Storage the local storage driver type Storage struct { Path string `json:"path" yaml:"path"` @@ -66,10 +73,31 @@ func (storage *Storage) Upload(ctx context.Context, filename string, reader io.R return "", err } - // Write file - _, err = data.Write(path, reader, 0644) - if err != nil { - return "", err + // Check if compression is enabled and if it's an image + if storage.Compression && isImage(contentType) { + // Read the entire image into memory + content, err := io.ReadAll(reader) + if err != nil { + return "", fmt.Errorf("failed to read image: %w", err) + } + + // Compress image + compressed, err := compressImage(content, contentType) + if err != nil { + return "", fmt.Errorf("failed to compress image: %w", err) + } + + // Write compressed image + _, err = data.Write(path, bytes.NewReader(compressed), 0644) + if err != nil { + return "", err + } + } else { + // Write file without compression + _, err = data.Write(path, reader, 0644) + if err != nil { + return "", err + } } return id, nil @@ -113,3 +141,68 @@ func (storage *Storage) makeID(filename string, ext string) string { name := strings.TrimSuffix(filepath.Base(filename), ext) return fmt.Sprintf("%s/%s-%s%s", date, name, hash, ext) } + +// isImage checks if the content type is an image +func isImage(contentType string) bool { + return strings.HasPrefix(contentType, "image/") +} + +// compressImage compresses the image while maintaining aspect ratio +func compressImage(data []byte, contentType string) ([]byte, error) { + // Decode image + img, _, err := image.Decode(bytes.NewReader(data)) + if err != nil { + return nil, fmt.Errorf("failed to decode image: %w", err) + } + + // Calculate new dimensions + bounds := img.Bounds() + width := bounds.Dx() + height := bounds.Dy() + var newWidth, newHeight int + + if width > height { + if width > MaxImageSize { + newWidth = MaxImageSize + newHeight = int(float64(height) * (float64(MaxImageSize) / float64(width))) + } else { + return data, nil // No need to resize + } + } else { + if height > MaxImageSize { + newHeight = MaxImageSize + newWidth = int(float64(width) * (float64(MaxImageSize) / float64(height))) + } else { + return data, nil // No need to resize + } + } + + // Create new image with new dimensions + newImg := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight)) + + // Scale the image using bilinear interpolation + for y := 0; y < newHeight; y++ { + for x := 0; x < newWidth; x++ { + srcX := float64(x) * float64(width) / float64(newWidth) + srcY := float64(y) * float64(height) / float64(newHeight) + newImg.Set(x, y, img.At(int(srcX), int(srcY))) + } + } + + // Encode image + var buf bytes.Buffer + switch contentType { + case "image/jpeg": + err = jpeg.Encode(&buf, newImg, &jpeg.Options{Quality: 85}) + case "image/png": + err = png.Encode(&buf, newImg) + default: + return data, nil // Unsupported format, return original + } + + if err != nil { + return nil, fmt.Errorf("failed to encode image: %w", err) + } + + return buf.Bytes(), nil +} diff --git a/neo/vision/driver/local/storage_test.go b/neo/vision/driver/local/storage_test.go index 15baa1a2..36a53eee 100644 --- a/neo/vision/driver/local/storage_test.go +++ b/neo/vision/driver/local/storage_test.go @@ -3,6 +3,8 @@ package local import ( "bytes" "context" + "image" + "image/png" "io" "testing" @@ -49,6 +51,80 @@ func TestLocalStorage(t *testing.T) { assert.Equal(t, content, downloaded) }) + t.Run("Upload and Download Image with Compression", func(t *testing.T) { + storage, err := New(map[string]interface{}{ + "path": "/__vision_test", + "compression": true, + }) + assert.NoError(t, err) + + // Create a test image (2000x2000 pixels) + img := image.NewRGBA(image.Rect(0, 0, 2000, 2000)) + var buf bytes.Buffer + err = png.Encode(&buf, img) + assert.NoError(t, err) + + // Upload + reader := bytes.NewReader(buf.Bytes()) + fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png") + assert.NoError(t, err) + assert.NotEmpty(t, fileID) + + // Download and verify size + reader2, contentType, err := storage.Download(context.Background(), fileID) + assert.NoError(t, err) + assert.Equal(t, "image/png", contentType) + + downloaded, err := io.ReadAll(reader2) + assert.NoError(t, err) + + // Decode the downloaded image + downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded)) + assert.NoError(t, err) + + // Verify dimensions + bounds := downloadedImg.Bounds() + assert.LessOrEqual(t, bounds.Dx(), MaxImageSize) + assert.LessOrEqual(t, bounds.Dy(), MaxImageSize) + }) + + t.Run("Upload Image without Compression", func(t *testing.T) { + storage, err := New(map[string]interface{}{ + "path": "/__vision_test", + "compression": false, + }) + assert.NoError(t, err) + + // Create a test image (2000x2000 pixels) + img := image.NewRGBA(image.Rect(0, 0, 2000, 2000)) + var buf bytes.Buffer + err = png.Encode(&buf, img) + assert.NoError(t, err) + + // Upload + reader := bytes.NewReader(buf.Bytes()) + fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png") + assert.NoError(t, err) + assert.NotEmpty(t, fileID) + + // Download and verify size + reader2, contentType, err := storage.Download(context.Background(), fileID) + assert.NoError(t, err) + assert.Equal(t, "image/png", contentType) + + downloaded, err := io.ReadAll(reader2) + assert.NoError(t, err) + + // Decode the downloaded image + downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded)) + assert.NoError(t, err) + + // Verify dimensions are unchanged + bounds := downloadedImg.Bounds() + assert.Equal(t, 2000, bounds.Dx()) + assert.Equal(t, 2000, bounds.Dy()) + }) + t.Run("URL Generation", func(t *testing.T) { storage, err := New(map[string]interface{}{ "path": "/__vision_test", diff --git a/neo/vision/driver/s3/storage.go b/neo/vision/driver/s3/storage.go index addc3b10..3f55b7be 100644 --- a/neo/vision/driver/s3/storage.go +++ b/neo/vision/driver/s3/storage.go @@ -1,8 +1,12 @@ package s3 import ( + "bytes" "context" "fmt" + "image" + "image/jpeg" + "image/png" "io" "path/filepath" "strings" @@ -16,23 +20,28 @@ import ( // DefaultExpiration default expiration time for presigned URLs (5 minutes) const DefaultExpiration = 5 * time.Minute +// MaxImageSize maximum image size (1920x1080) +const MaxImageSize = 1920 + // Storage the S3 storage driver type Storage struct { - Endpoint string `json:"endpoint" yaml:"endpoint"` - Region string `json:"region" yaml:"region"` - Key string `json:"key" yaml:"key"` - Secret string `json:"secret" yaml:"secret"` - Bucket string `json:"bucket" yaml:"bucket"` - Expiration time.Duration `json:"expiration" yaml:"expiration"` - client *s3.Client - prefix string + Endpoint string `json:"endpoint" yaml:"endpoint"` + Region string `json:"region" yaml:"region"` + Key string `json:"key" yaml:"key"` + Secret string `json:"secret" yaml:"secret"` + Bucket string `json:"bucket" yaml:"bucket"` + Expiration time.Duration `json:"expiration" yaml:"expiration"` + client *s3.Client + prefix string + compression bool } // New create a new S3 storage func New(options map[string]interface{}) (*Storage, error) { storage := &Storage{ - Region: "auto", - Expiration: DefaultExpiration, + Region: "auto", + Expiration: DefaultExpiration, + compression: true, } if endpoint, ok := options["endpoint"].(string); ok { @@ -63,6 +72,10 @@ func New(options map[string]interface{}) (*Storage, error) { storage.Expiration = exp } + if compression, ok := options["compression"].(bool); ok { + storage.compression = compression + } + // Validate required fields if storage.Key == "" || storage.Secret == "" { return nil, fmt.Errorf("key and secret are required") @@ -102,11 +115,31 @@ func (storage *Storage) Upload(ctx context.Context, filename string, reader io.R fileID := storage.makeID(filename, filepath.Ext(filename)) key := filepath.Join(storage.prefix, fileID) + // Check if compression is enabled and if it's an image + var body io.Reader + if storage.compression && isImage(contentType) { + // Read the entire image into memory + content, err := io.ReadAll(reader) + if err != nil { + return "", fmt.Errorf("failed to read image: %w", err) + } + + // Compress image + compressed, err := compressImage(content, contentType) + if err != nil { + return "", fmt.Errorf("failed to compress image: %w", err) + } + + body = bytes.NewReader(compressed) + } else { + body = reader + } + // Upload file _, err := storage.client.PutObject(ctx, &s3.PutObjectInput{ Bucket: aws.String(storage.Bucket), Key: aws.String(key), - Body: reader, + Body: body, ContentType: aws.String(contentType), }) if err != nil { @@ -166,3 +199,68 @@ func (storage *Storage) makeID(filename string, ext string) string { name := strings.TrimSuffix(filepath.Base(filename), ext) return fmt.Sprintf("%s/%s-%d%s", date, name, time.Now().UnixNano(), ext) } + +// isImage checks if the content type is an image +func isImage(contentType string) bool { + return strings.HasPrefix(contentType, "image/") +} + +// compressImage compresses the image while maintaining aspect ratio +func compressImage(data []byte, contentType string) ([]byte, error) { + // Decode image + img, _, err := image.Decode(bytes.NewReader(data)) + if err != nil { + return nil, fmt.Errorf("failed to decode image: %w", err) + } + + // Calculate new dimensions + bounds := img.Bounds() + width := bounds.Dx() + height := bounds.Dy() + var newWidth, newHeight int + + if width > height { + if width > MaxImageSize { + newWidth = MaxImageSize + newHeight = int(float64(height) * (float64(MaxImageSize) / float64(width))) + } else { + return data, nil // No need to resize + } + } else { + if height > MaxImageSize { + newHeight = MaxImageSize + newWidth = int(float64(width) * (float64(MaxImageSize) / float64(height))) + } else { + return data, nil // No need to resize + } + } + + // Create new image with new dimensions + newImg := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight)) + + // Scale the image using bilinear interpolation + for y := 0; y < newHeight; y++ { + for x := 0; x < newWidth; x++ { + srcX := float64(x) * float64(width) / float64(newWidth) + srcY := float64(y) * float64(height) / float64(newHeight) + newImg.Set(x, y, img.At(int(srcX), int(srcY))) + } + } + + // Encode image + var buf bytes.Buffer + switch contentType { + case "image/jpeg": + err = jpeg.Encode(&buf, newImg, &jpeg.Options{Quality: 85}) + case "image/png": + err = png.Encode(&buf, newImg) + default: + return data, nil // Unsupported format, return original + } + + if err != nil { + return nil, fmt.Errorf("failed to encode image: %w", err) + } + + return buf.Bytes(), nil +} diff --git a/neo/vision/driver/s3/storage_test.go b/neo/vision/driver/s3/storage_test.go index 0f3d3847..ca0737ce 100644 --- a/neo/vision/driver/s3/storage_test.go +++ b/neo/vision/driver/s3/storage_test.go @@ -3,6 +3,8 @@ package s3 import ( "bytes" "context" + "image" + "image/png" "io" "os" "testing" @@ -19,13 +21,14 @@ func TestS3Storage(t *testing.T) { t.Run("Create Storage", func(t *testing.T) { options := map[string]interface{}{ - "endpoint": os.Getenv("S3_API"), - "region": "auto", - "key": os.Getenv("S3_ACCESS_KEY"), - "secret": os.Getenv("S3_SECRET_KEY"), - "bucket": os.Getenv("S3_BUCKET"), - "prefix": "vision-test", - "expiration": 10 * time.Minute, + "endpoint": os.Getenv("S3_API"), + "region": "auto", + "key": os.Getenv("S3_ACCESS_KEY"), + "secret": os.Getenv("S3_SECRET_KEY"), + "bucket": os.Getenv("S3_BUCKET"), + "prefix": "vision-test", + "expiration": 10 * time.Minute, + "compression": true, } storage, err := New(options) @@ -42,21 +45,115 @@ func TestS3Storage(t *testing.T) { assert.Equal(t, os.Getenv("S3_BUCKET"), storage.Bucket) assert.Equal(t, "vision-test", storage.prefix) assert.Equal(t, 10*time.Minute, storage.Expiration) + assert.True(t, storage.compression) } }) - t.Run("Upload and Download", func(t *testing.T) { + t.Run("Upload and Download Image with Compression", func(t *testing.T) { storage, err := New(map[string]interface{}{ - "endpoint": os.Getenv("S3_API"), - "region": "auto", - "key": os.Getenv("S3_ACCESS_KEY"), - "secret": os.Getenv("S3_SECRET_KEY"), - "bucket": os.Getenv("S3_BUCKET"), - "prefix": "vision-test", - "expiration": 5 * time.Minute, + "endpoint": os.Getenv("S3_API"), + "region": "auto", + "key": os.Getenv("S3_ACCESS_KEY"), + "secret": os.Getenv("S3_SECRET_KEY"), + "bucket": os.Getenv("S3_BUCKET"), + "prefix": "vision-test", + "expiration": 5 * time.Minute, + "compression": true, }) + if err != nil { + t.Skip("S3 configuration not available") + } + + // Create a test image (2000x2000 pixels) + img := image.NewRGBA(image.Rect(0, 0, 2000, 2000)) + var buf bytes.Buffer + err = png.Encode(&buf, img) assert.NoError(t, err) + // Upload + reader := bytes.NewReader(buf.Bytes()) + fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png") + assert.NoError(t, err) + assert.NotEmpty(t, fileID) + + // Download and verify size + reader2, contentType, err := storage.Download(context.Background(), fileID) + assert.NoError(t, err) + assert.Equal(t, "image/png", contentType) + + downloaded, err := io.ReadAll(reader2) + assert.NoError(t, err) + + // Decode the downloaded image + downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded)) + assert.NoError(t, err) + + // Verify dimensions + bounds := downloadedImg.Bounds() + assert.LessOrEqual(t, bounds.Dx(), MaxImageSize) + assert.LessOrEqual(t, bounds.Dy(), MaxImageSize) + }) + + t.Run("Upload Image without Compression", func(t *testing.T) { + storage, err := New(map[string]interface{}{ + "endpoint": os.Getenv("S3_API"), + "region": "auto", + "key": os.Getenv("S3_ACCESS_KEY"), + "secret": os.Getenv("S3_SECRET_KEY"), + "bucket": os.Getenv("S3_BUCKET"), + "prefix": "vision-test", + "expiration": 5 * time.Minute, + "compression": false, + }) + if err != nil { + t.Skip("S3 configuration not available") + } + + // Create a test image (2000x2000 pixels) + img := image.NewRGBA(image.Rect(0, 0, 2000, 2000)) + var buf bytes.Buffer + err = png.Encode(&buf, img) + assert.NoError(t, err) + + // Upload + reader := bytes.NewReader(buf.Bytes()) + fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png") + assert.NoError(t, err) + assert.NotEmpty(t, fileID) + + // Download and verify size + reader2, contentType, err := storage.Download(context.Background(), fileID) + assert.NoError(t, err) + assert.Equal(t, "image/png", contentType) + + downloaded, err := io.ReadAll(reader2) + assert.NoError(t, err) + + // Decode the downloaded image + downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded)) + assert.NoError(t, err) + + // Verify dimensions are unchanged + bounds := downloadedImg.Bounds() + assert.Equal(t, 2000, bounds.Dx()) + assert.Equal(t, 2000, bounds.Dy()) + }) + + t.Run("Upload and Download Text File", func(t *testing.T) { + storage, err := New(map[string]interface{}{ + "endpoint": os.Getenv("S3_API"), + "region": "auto", + "key": os.Getenv("S3_ACCESS_KEY"), + "secret": os.Getenv("S3_SECRET_KEY"), + "bucket": os.Getenv("S3_BUCKET"), + "prefix": "vision-test", + "expiration": 5 * time.Minute, + "compression": true, + }) + if err != nil { + t.Skip("S3 configuration not available") + } + content := []byte("test content") reader := bytes.NewReader(content) fileID, err := storage.Upload(context.Background(), "test.txt", reader, "text/plain") @@ -69,13 +166,6 @@ func TestS3Storage(t *testing.T) { assert.Contains(t, url, "X-Amz-Signature") assert.Contains(t, url, "X-Amz-Expires") - // Test with different expiration - storage.Expiration = 1 * time.Hour - url2 := storage.URL(context.Background(), fileID) - assert.NotEmpty(t, url2) - assert.Contains(t, url2, "X-Amz-Signature") - assert.Contains(t, url2, "X-Amz-Expires=3600") - // Download reader2, contentType, err := storage.Download(context.Background(), fileID) if err != nil { @@ -93,49 +183,20 @@ func TestS3Storage(t *testing.T) { } }) - t.Run("Upload with Custom Expiration", func(t *testing.T) { - storage, err := New(map[string]interface{}{ - "endpoint": os.Getenv("S3_API"), - "region": "auto", - "key": os.Getenv("S3_ACCESS_KEY"), - "secret": os.Getenv("S3_SECRET_KEY"), - "bucket": os.Getenv("S3_BUCKET"), - "prefix": "vision-test", - "expiration": 5 * time.Minute, - }) - assert.NoError(t, err) - - content := []byte("test content") - reader := bytes.NewReader(content) - fileID, err := storage.Upload(context.Background(), "test.txt", reader, "text/plain") - assert.NoError(t, err) - assert.NotEmpty(t, fileID) - - // Get URL with default expiration (5 minutes) - url := storage.URL(context.Background(), fileID) - assert.NotEmpty(t, url) - assert.Contains(t, url, "X-Amz-Signature") - assert.Contains(t, url, "X-Amz-Expires=300") // 5 minutes = 300 seconds - - // Change expiration and get new URL - storage.Expiration = 2 * time.Hour - url2 := storage.URL(context.Background(), fileID) - assert.NotEmpty(t, url2) - assert.Contains(t, url2, "X-Amz-Signature") - assert.Contains(t, url2, "X-Amz-Expires=7200") // 2 hours = 7200 seconds - }) - t.Run("Download Non-existent File", func(t *testing.T) { storage, err := New(map[string]interface{}{ - "endpoint": os.Getenv("S3_API"), - "region": "auto", - "key": os.Getenv("S3_ACCESS_KEY"), - "secret": os.Getenv("S3_SECRET_KEY"), - "bucket": os.Getenv("S3_BUCKET"), - "prefix": "vision-test", - "expiration": 5 * time.Minute, + "endpoint": os.Getenv("S3_API"), + "region": "auto", + "key": os.Getenv("S3_ACCESS_KEY"), + "secret": os.Getenv("S3_SECRET_KEY"), + "bucket": os.Getenv("S3_BUCKET"), + "prefix": "vision-test", + "expiration": 5 * time.Minute, + "compression": true, }) - assert.NoError(t, err) + if err != nil { + t.Skip("S3 configuration not available") + } _, _, err = storage.Download(context.Background(), "non-existent.txt") assert.Error(t, err) diff --git a/neo/vision/prompt.md b/neo/vision/prompt.md deleted file mode 100644 index 2d6e4f34..00000000 --- a/neo/vision/prompt.md +++ /dev/null @@ -1,8 +0,0 @@ -根据这个数据结构，和说明实现一下对应的逻辑。 - -1. driver 单独一个目录. 每个 driver 一个目录，model 和 storage 在一级即可。 放在@vision 下 -2. model driver: 支持 openai -3. storage driver 支持 local 和 s3 . local 使用我框架的 fs 实现，参考@file.go -4. 在 程序启动时候，设置视觉配置。(作为可选配置） @types.go -5. 统一的创建和调用入口，调用时需传入 chat model (用来判断是否支持视觉) 和图片路径，（使用 fs 读取）。 @vision.go -6. 用一个回调函数，外部传入用来格式化返回的数据。 diff --git a/neo/vision/vision_test.go b/neo/vision/vision_test.go index 73379487..a0bd42ae 100644 --- a/neo/vision/vision_test.go +++ b/neo/vision/vision_test.go @@ -5,6 +5,8 @@ import ( "context" "encoding/base64" "fmt" + "image" + "image/png" "io" "net/http" "net/http/httptest" @@ -15,6 +17,7 @@ import ( "github.com/yaoapp/gou/fs" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/neo/vision/driver" + "github.com/yaoapp/yao/neo/vision/driver/local" "github.com/yaoapp/yao/test" ) @@ -23,6 +26,9 @@ var ( testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" ) +// MaxImageSize maximum image size (1920x1080) +const MaxImageSize = local.MaxImageSize + func TestVision(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() @@ -253,6 +259,78 @@ func TestVision(t *testing.T) { assert.Error(t, err) assert.Contains(t, err.Error(), "storage driver invalid not supported") }) + + t.Run("Upload and Download Image with Local Storage", func(t *testing.T) { + vision, err := createTestVision(imgServer.URL) + assert.NoError(t, err) + + // Create test image (2000x2000 pixels) + img := image.NewRGBA(image.Rect(0, 0, 2000, 2000)) + var buf bytes.Buffer + err = png.Encode(&buf, img) + assert.NoError(t, err) + + // Upload + reader := bytes.NewReader(buf.Bytes()) + resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png") + assert.NoError(t, err) + assert.NotEmpty(t, resp.FileID) + assert.NotEmpty(t, resp.URL) + + // Download and verify size + reader2, contentType, err := vision.Download(context.Background(), resp.FileID) + assert.NoError(t, err) + assert.Equal(t, "image/png", contentType) + + downloaded, err := io.ReadAll(reader2) + assert.NoError(t, err) + + // Decode the downloaded image + downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded)) + assert.NoError(t, err) + + // Verify dimensions + bounds := downloadedImg.Bounds() + assert.LessOrEqual(t, bounds.Dx(), MaxImageSize) + assert.LessOrEqual(t, bounds.Dy(), MaxImageSize) + }) + + t.Run("Upload and Download Image with S3 Storage", func(t *testing.T) { + vision, err := createTestVisionWithS3() + if err != nil { + t.Skip("S3 configuration not available") + } + + // Create test image (2000x2000 pixels) + img := image.NewRGBA(image.Rect(0, 0, 2000, 2000)) + var buf bytes.Buffer + err = png.Encode(&buf, img) + assert.NoError(t, err) + + // Upload + reader := bytes.NewReader(buf.Bytes()) + resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png") + assert.NoError(t, err) + assert.NotEmpty(t, resp.FileID) + assert.NotEmpty(t, resp.URL) + + // Download and verify size + reader2, contentType, err := vision.Download(context.Background(), resp.FileID) + assert.NoError(t, err) + assert.Equal(t, "image/png", contentType) + + downloaded, err := io.ReadAll(reader2) + assert.NoError(t, err) + + // Decode the downloaded image + downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded)) + assert.NoError(t, err) + + // Verify dimensions + bounds := downloadedImg.Bounds() + assert.LessOrEqual(t, bounds.Dx(), MaxImageSize) + assert.LessOrEqual(t, bounds.Dy(), MaxImageSize) + }) } func createTestVision(baseURL string) (*Vision, error) { @@ -271,17 +349,17 @@ func createTestVision(baseURL string) (*Vision, error) { "api_key": os.Getenv("OPENAI_API_KEY"), "model": os.Getenv("VISION_MODEL"), "prompt": `# Objective - You are a vision assistant, you can help the user to understand the image and describe it. - - ## Task Execution Steps - 1. Understand the image/video and describe it. - 2. Describe the image/video in detail. - - ## Result Format - { - "description": "The description of the image/video", - "content": "The content of the image/video" - }`, + You are a vision assistant, you can help the user to understand the image and describe it. + + ## Task Execution Steps + 1. Understand the image/video and describe it. + 2. Describe the image/video in detail. + + ## Result Format + { + "description": "The description of the image/video", + "content": "The content of the image/video" + }`, }, }, }