feat: image vision tool

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 {

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

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