From 64a52580fe02085988e18dca12aa03b02eac495a Mon Sep 17 00:00:00 2001 From: afjcjsbx Date: Fri, 20 Feb 2026 00:10:28 +0100 Subject: [PATCH] feat: image vision tool --- pkg/agent/loop.go | 5 + pkg/tools/vision.go | 192 ++++++++++++++++++++++++++++++++++++ pkg/tools/vision_test.go | 205 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 402 insertions(+) create mode 100644 pkg/tools/vision.go create mode 100644 pkg/tools/vision_test.go diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index f8eef395a..e64253c8d 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -106,6 +106,11 @@ func registerSharedTools(cfg *config.Config, msgBus *bus.MessageBus, registry *A agent.Tools.Register(tools.NewI2CTool()) agent.Tools.Register(tools.NewSPITool()) + // Vision tool + if cfg.Tools.Vision.Enabled { + agent.Tools.Register(tools.NewAnalyzeImageTool(cfg.Tools.Vision)) + } + // Message tool messageTool := tools.NewMessageTool() messageTool.SetSendCallback(func(channel, chatID, content string) error { diff --git a/pkg/tools/vision.go b/pkg/tools/vision.go new file mode 100644 index 000000000..5b5ff394f --- /dev/null +++ b/pkg/tools/vision.go @@ -0,0 +1,192 @@ +package tools + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/config" +) + +const ( + visionMaxFileSize = 10 * 1024 * 1024 + visionMaxTokens = 1000 + visionTimeout = 60 * time.Second + visionDefaultPrompt = "Describe exactly what you see in this image in detail. If it's a screenshot, read the texts and describe the UI elements." +) + +type AnalyzeImageTool struct { + apiKey string + apiURL string + model string + workspace string + restrict bool +} + +func NewAnalyzeImageTool(opts config.VisionToolConfig) *AnalyzeImageTool { + return &AnalyzeImageTool{ + workspace: opts.Workspace, + restrict: opts.Restrict, + apiKey: opts.ApiKey, + apiURL: opts.ApiURL, + model: opts.Model, + } +} + +func (t *AnalyzeImageTool) Name() string { + return "analyze_image" +} + +func (t *AnalyzeImageTool) Description() string { + return "Analyze an image file (e.g., png, jpeg) and return a detailed textual description of its contents. Use this to understand screenshots or photos." +} + +func (t *AnalyzeImageTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "Path to the image file to analyze", + }, + "prompt": map[string]interface{}{ + "type": "string", + "description": "Optional specific question about the image (e.g. 'Read the error message', 'Where is the login button?'). Defaults to a general description.", + }, + }, + "required": []string{"path"}, + } +} + +func getMimeType(path string) (string, error) { + ext := strings.ToLower(filepath.Ext(path)) + switch ext { + case ".png": + return "image/png", nil + case ".jpg", ".jpeg": + return "image/jpeg", nil + case ".webp": + return "image/webp", nil + case ".gif": + return "image/gif", nil + default: + return "", fmt.Errorf("unsupported image extension: %s", ext) + } +} + +func (t *AnalyzeImageTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult { + if t.apiKey == "" { + return ErrorResult("analyze_image tool is not configured properly: missing Vision API key") + } + + path, ok := args["path"].(string) + if !ok { + return ErrorResult("path is required") + } + + prompt := visionDefaultPrompt + if customPrompt, ok := args["prompt"].(string); ok && customPrompt != "" { + prompt = customPrompt + } + + resolvedPath, err := validatePath(path, t.workspace, t.restrict) + if err != nil { + return ErrorResult(err.Error()) + } + + // Secure image file reading, limited to 10MB to avoid huge payloads + fileInfo, err := os.Stat(resolvedPath) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to stat image: %v", err)) + } + if fileInfo.Size() > visionMaxFileSize { + return ErrorResult("image is too large (max 10MB allowed for analysis)") + } + + imgData, err := os.ReadFile(resolvedPath) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to read image: %v", err)) + } + + contentType := http.DetectContentType(imgData) + if !strings.HasPrefix(contentType, "image/") { + return ErrorResult(fmt.Sprintf("file is not a valid image, detected type: %s", contentType)) + } + + base64Image := base64.StdEncoding.EncodeToString(imgData) + mimeType, err := getMimeType(resolvedPath) + if err != nil { + return ErrorResult(err.Error()) + } + + dataURI := fmt.Sprintf("data:%s;base64,%s", mimeType, base64Image) + + payloadBytes, err := json.Marshal(map[string]interface{}{ + "model": t.model, + "messages": []map[string]interface{}{ + { + "role": "user", + "content": []map[string]interface{}{ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": map[string]string{"url": dataURI}}, + }, + }, + }, + "max_tokens": visionMaxTokens, + }) + if err != nil { + return ErrorResult("failed to prepare vision API payload") + } + + req, err := http.NewRequestWithContext(ctx, "POST", t.apiURL, bytes.NewBuffer(payloadBytes)) + if err != nil { + return ErrorResult("failed to create request") + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+t.apiKey) + + client := &http.Client{Timeout: visionTimeout} + resp, err := client.Do(req) + if err != nil { + return ErrorResult(fmt.Sprintf("vision API request failed: %v", err)) + } + defer resp.Body.Close() + + bodyBytes, err := io.ReadAll(resp.Body) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to read vision API response body: %v", err)) + } + + if resp.StatusCode != http.StatusOK { + return ErrorResult(fmt.Sprintf("vision API returned error (%d): %s", resp.StatusCode, string(bodyBytes))) + } + + var apiResp struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + + if err := json.Unmarshal(bodyBytes, &apiResp); err != nil { + return ErrorResult("failed to parse vision API response") + } + + if len(apiResp.Choices) == 0 { + return ErrorResult("vision API returned no content") + } + + analysisResult := apiResp.Choices[0].Message.Content + finalOutput := fmt.Sprintf("[Image Analysis Result for %s]\n%s", filepath.Base(resolvedPath), analysisResult) + + return NewToolResult(finalOutput) +} diff --git a/pkg/tools/vision_test.go b/pkg/tools/vision_test.go new file mode 100644 index 000000000..69fbabda9 --- /dev/null +++ b/pkg/tools/vision_test.go @@ -0,0 +1,205 @@ +package tools + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/sipeed/picoclaw/pkg/config" +) + +// magic bytes for a valid PNG image to fool http.DetectContentType +var validPNGBody = []byte("\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xf8\x00\x00\x00\x00IEND\xaeB`\x82") + +func TestGetMimeType(t *testing.T) { + tests := []struct { + path string + expected string + expectError bool + }{ + {"test.png", "image/png", false}, + {"TEST.JPG", "image/jpeg", false}, + {"image.jpeg", "image/jpeg", false}, + {"anim.gif", "image/gif", false}, + {"pic.webp", "image/webp", false}, + {"document.pdf", "", true}, + {"no_extension", "", true}, + } + + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + mime, err := getMimeType(tt.path) + if tt.expectError && err == nil { + t.Errorf("Expected error for path %s, but got none", tt.path) + } + if !tt.expectError && err != nil { + t.Errorf("Unexpected error for path %s: %v", tt.path, err) + } + if mime != tt.expected { + t.Errorf("Expected mime %s, got %s", tt.expected, mime) + } + }) + } +} + +func TestAnalyzeImageTool_Basic(t *testing.T) { + tool := NewAnalyzeImageTool(config.VisionToolConfig{}) + + if tool.Name() != "analyze_image" { + t.Errorf("Expected name 'analyze_image', got %s", tool.Name()) + } + + if !strings.Contains(tool.Description(), "Analyze an image") { + t.Errorf("Expected description to contain 'Analyze an image', got %s", tool.Description()) + } + + params := tool.Parameters() + if params["type"] != "object" { + t.Errorf("Expected parameters to be object type") + } +} + +func TestAnalyzeImageTool_Execute_Validations(t *testing.T) { + ctx := context.Background() + + t.Run("Missing API Key", func(t *testing.T) { + tool := NewAnalyzeImageTool(config.VisionToolConfig{ApiKey: ""}) + res := tool.Execute(ctx, map[string]interface{}{"path": "test.png"}) + if !res.IsError || !strings.Contains(res.ForLLM, "missing Vision API key") { + t.Errorf("Expected missing API key error, got: %v", res.ForLLM) + } + }) + + t.Run("Missing Path", func(t *testing.T) { + tool := NewAnalyzeImageTool(config.VisionToolConfig{ApiKey: "test-key"}) + res := tool.Execute(ctx, map[string]interface{}{}) + if !res.IsError || !strings.Contains(res.ForLLM, "path is required") { + t.Errorf("Expected missing path error, got: %v", res.ForLLM) + } + }) +} + +func TestAnalyzeImageTool_Execute_FileChecks(t *testing.T) { + ctx := context.Background() + tmpDir := t.TempDir() + tool := NewAnalyzeImageTool(config.VisionToolConfig{ + ApiKey: "test-key", + Workspace: tmpDir, + Restrict: true, + }) + + t.Run("Not An Image", func(t *testing.T) { + fakeImg := filepath.Join(tmpDir, "fake.png") + os.WriteFile(fakeImg, []byte("this is plain text, not an image"), 0644) + + res := tool.Execute(ctx, map[string]interface{}{"path": fakeImg}) + if !res.IsError || !strings.Contains(res.ForLLM, "not a valid image") { + t.Errorf("Expected invalid image error, got: %v", res.ForLLM) + } + }) + + t.Run("Unsupported Extension", func(t *testing.T) { + unsupportedFile := filepath.Join(tmpDir, "real_image.txt") + // valid image, but wrong extension + os.WriteFile(unsupportedFile, validPNGBody, 0644) + + res := tool.Execute(ctx, map[string]interface{}{"path": unsupportedFile}) + if !res.IsError || !strings.Contains(res.ForLLM, "unsupported image extension") { + t.Errorf("Expected unsupported extension error, got: %v", res.ForLLM) + } + }) + + t.Run("File Too Large", func(t *testing.T) { + largeImg := filepath.Join(tmpDir, "large.png") + // file slightly larger than 10MB + f, _ := os.Create(largeImg) + f.Truncate(10*1024*1024 + 100) + f.Close() + + res := tool.Execute(ctx, map[string]interface{}{"path": largeImg}) + if !res.IsError || !strings.Contains(res.ForLLM, "image is too large") { + t.Errorf("Expected too large error, got: %v", res.ForLLM) + } + }) +} + +func TestAnalyzeImageTool_Execute_Success(t *testing.T) { + expectedAnalysis := "I see a test image with a button." + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + + if r.Header.Get("Authorization") != "Bearer mock-api-key" { + w.WriteHeader(http.StatusUnauthorized) + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{ + "choices": [{ + "message": {"content": "` + expectedAnalysis + `"} + }] + }`)) + })) + defer ts.Close() + + ctx := context.Background() + tmpDir := t.TempDir() + + tool := NewAnalyzeImageTool(config.VisionToolConfig{ + ApiKey: "mock-api-key", + ApiURL: ts.URL, + Model: "gpt-4o-mini", + Workspace: tmpDir, + Restrict: true, + }) + + testImgPath := filepath.Join(tmpDir, "success.png") + os.WriteFile(testImgPath, validPNGBody, 0644) + + res := tool.Execute(ctx, map[string]interface{}{ + "path": testImgPath, + "prompt": "What is this?", + }) + + if res.IsError { + t.Fatalf("Expected success, got error: %s", res.ForLLM) + } + + if !strings.Contains(res.ForLLM, expectedAnalysis) { + t.Errorf("Expected output to contain %q, got: %s", expectedAnalysis, res.ForLLM) + } +} + +func TestAnalyzeImageTool_Execute_APIError(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error": "Internal Server Error"}`)) + })) + defer ts.Close() + + ctx := context.Background() + tmpDir := t.TempDir() + + tool := NewAnalyzeImageTool(config.VisionToolConfig{ + ApiKey: "mock-api-key", + ApiURL: ts.URL, + Workspace: tmpDir, + }) + + testImgPath := filepath.Join(tmpDir, "error.png") + os.WriteFile(testImgPath, validPNGBody, 0644) + + res := tool.Execute(ctx, map[string]interface{}{"path": testImgPath}) + + if !res.IsError { + t.Fatal("Expected API error, got success") + } + + if !strings.Contains(res.ForLLM, "vision API returned error (500)") { + t.Errorf("Expected 500 error message, got: %s", res.ForLLM) + } +}