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

@@ -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",

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"])
+	})
 }

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

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
 	}

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) {