From 2becefd06e99a01f527c5ae298c56bd2b6854a4a Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 6 Jan 2025 11:43:31 +0800 Subject: [PATCH] Enhance Vision module with flexible prompt handling and comprehensive tests - Refactored the Analyze method in the Vision and OpenAI model to accept a variadic prompt parameter, allowing for optional custom prompts while defaulting to a predefined prompt if none is provided. - Added multiple test cases in vision_test.go and model_test.go to validate image analysis with default, custom, and empty prompts, ensuring robust functionality and error handling. - Updated the Model interface to reflect the new prompt handling, improving clarity and usability. These changes enhance the flexibility of the Vision module, paving the way for improved user experience and functionality in image analysis. --- neo/vision/driver/openai/model.go | 10 ++- neo/vision/driver/openai/model_test.go | 45 +++++++++++++ neo/vision/driver/types.go | 4 +- neo/vision/vision.go | 4 +- neo/vision/vision_test.go | 90 ++++++++++++++++++++++++++ 5 files changed, 148 insertions(+), 5 deletions(-) diff --git a/neo/vision/driver/openai/model.go b/neo/vision/driver/openai/model.go index a766c380..8a55abd1 100644 --- a/neo/vision/driver/openai/model.go +++ b/neo/vision/driver/openai/model.go @@ -52,11 +52,17 @@ func New(options map[string]interface{}) (*Model, error) { } // Analyze analyze image using OpenAI vision model -func (model *Model) Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error) { +func (model *Model) Analyze(ctx context.Context, fileID string, prompt ...string) (map[string]interface{}, error) { if model.APIKey == "" { return nil, fmt.Errorf("api_key is required") } + // Use default prompt if none provided + userPrompt := model.Prompt + if len(prompt) > 0 && prompt[0] != "" { + userPrompt = prompt[0] + } + // Check if fileID is a URL or base64 data var imageURL string if strings.HasPrefix(fileID, "data:image/") { @@ -103,7 +109,7 @@ func (model *Model) Analyze(ctx context.Context, fileID string, prompt string) ( "content": []map[string]interface{}{ { "type": "text", - "text": prompt, + "text": userPrompt, }, { "type": "image_url", diff --git a/neo/vision/driver/openai/model_test.go b/neo/vision/driver/openai/model_test.go index e93d516e..332d013f 100644 --- a/neo/vision/driver/openai/model_test.go +++ b/neo/vision/driver/openai/model_test.go @@ -146,4 +146,49 @@ func TestOpenAIModel(t *testing.T) { assert.Error(t, err) assert.Contains(t, err.Error(), "OpenAI API error") }) + + t.Run("Analyze with Default Prompt", func(t *testing.T) { + model, err := New(map[string]interface{}{ + "api_key": os.Getenv("OPENAI_API_KEY"), + "model": os.Getenv("VISION_MODEL"), + "prompt": "Default test prompt", + }) + assert.NoError(t, err) + + // Use base64 image data without providing a prompt + result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result["description"]) + }) + + t.Run("Analyze with Custom Prompt Overriding Default", func(t *testing.T) { + model, err := New(map[string]interface{}{ + "api_key": os.Getenv("OPENAI_API_KEY"), + "model": os.Getenv("VISION_MODEL"), + "prompt": "Default test prompt", + }) + assert.NoError(t, err) + + // Use base64 image data with custom prompt + result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt") + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result["description"]) + }) + + t.Run("Analyze with Empty Custom Prompt", func(t *testing.T) { + model, err := New(map[string]interface{}{ + "api_key": os.Getenv("OPENAI_API_KEY"), + "model": os.Getenv("VISION_MODEL"), + "prompt": "Default test prompt", + }) + assert.NoError(t, err) + + // Use base64 image data with empty prompt (should use default) + result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "") + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result["description"]) + }) } diff --git a/neo/vision/driver/types.go b/neo/vision/driver/types.go index ccc64ee6..9cc2006e 100644 --- a/neo/vision/driver/types.go +++ b/neo/vision/driver/types.go @@ -32,7 +32,9 @@ type Storage interface { // Model the vision model interface type Model interface { - Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error) + // Analyze analyzes an image file + // If prompt is empty, it will use the default prompt from model.options.prompt + Analyze(ctx context.Context, fileID string, prompt ...string) (map[string]interface{}, error) } // Response the vision response diff --git a/neo/vision/vision.go b/neo/vision/vision.go index 0886a900..792be41a 100644 --- a/neo/vision/vision.go +++ b/neo/vision/vision.go @@ -104,7 +104,7 @@ func (v *Vision) Upload(ctx context.Context, filename string, reader io.Reader, } // Analyze analyze image using vision model -func (v *Vision) Analyze(ctx context.Context, fileID string, prompt string) (*driver.Response, error) { +func (v *Vision) Analyze(ctx context.Context, fileID string, prompt ...string) (*driver.Response, error) { if v.model == nil { return nil, fmt.Errorf("model is required") } @@ -121,7 +121,7 @@ func (v *Vision) Analyze(ctx context.Context, fileID string, prompt string) (*dr } } - result, err := v.model.Analyze(ctx, url, prompt) + result, err := v.model.Analyze(ctx, url, prompt...) if err != nil { return nil, err } diff --git a/neo/vision/vision_test.go b/neo/vision/vision_test.go index a0bd42ae..411a66a7 100644 --- a/neo/vision/vision_test.go +++ b/neo/vision/vision_test.go @@ -331,6 +331,96 @@ func TestVision(t *testing.T) { assert.LessOrEqual(t, bounds.Dx(), MaxImageSize) assert.LessOrEqual(t, bounds.Dy(), MaxImageSize) }) + + t.Run("Analyze Image with Default Prompt", func(t *testing.T) { + // Create vision service with default prompt + cfg := &driver.Config{ + Storage: driver.StorageConfig{ + Driver: "local", + Options: map[string]interface{}{ + "path": "/__vision_test", + "compression": true, + }, + }, + Model: driver.ModelConfig{ + Driver: "openai", + Options: map[string]interface{}{ + "api_key": os.Getenv("OPENAI_API_KEY"), + "model": os.Getenv("VISION_MODEL"), + "prompt": "Default test prompt", + }, + }, + } + + vision, err := New(cfg) + assert.NoError(t, err) + + // Use base64 data without providing a prompt + result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result.Description) + }) + + t.Run("Analyze Image with Custom Prompt", func(t *testing.T) { + // Create vision service with default prompt + cfg := &driver.Config{ + Storage: driver.StorageConfig{ + Driver: "local", + Options: map[string]interface{}{ + "path": "/__vision_test", + "compression": true, + }, + }, + Model: driver.ModelConfig{ + Driver: "openai", + Options: map[string]interface{}{ + "api_key": os.Getenv("OPENAI_API_KEY"), + "model": os.Getenv("VISION_MODEL"), + "prompt": "Default test prompt", + }, + }, + } + + vision, err := New(cfg) + assert.NoError(t, err) + + // Use base64 data with custom prompt + result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt") + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result.Description) + }) + + t.Run("Analyze Image with Empty Custom Prompt", func(t *testing.T) { + // Create vision service with default prompt + cfg := &driver.Config{ + Storage: driver.StorageConfig{ + Driver: "local", + Options: map[string]interface{}{ + "path": "/__vision_test", + "compression": true, + }, + }, + Model: driver.ModelConfig{ + Driver: "openai", + Options: map[string]interface{}{ + "api_key": os.Getenv("OPENAI_API_KEY"), + "model": os.Getenv("VISION_MODEL"), + "prompt": "Default test prompt", + }, + }, + } + + vision, err := New(cfg) + assert.NoError(t, err) + + // Use base64 data with empty prompt (should use default) + result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "") + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result.Description) + }) } func createTestVision(baseURL string) (*Vision, error) {