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