feat: image vision tool

This commit is contained in:
afjcjsbx 2026-02-20 00:10:28 +01:00
parent 521359ed4f
commit 81d5da3a81
4 changed files with 422 additions and 3 deletions

View file

@ -105,6 +105,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 {

View file

@ -315,10 +315,20 @@ type ExecConfig struct {
CustomDenyPatterns []string `json:"custom_deny_patterns" env:"PICOCLAW_TOOLS_EXEC_CUSTOM_DENY_PATTERNS"`
}
type VisionToolConfig struct {
Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_VISION_ENABLED"`
ApiKey string `json:"api_key" env:"PICOCLAW_TOOLS_VISION_API_KEY"`
ApiURL string `json:"api_url" env:"PICOCLAW_TOOLS_VISION_API_URL"`
Model string `json:"model" env:"PICOCLAW_TOOLS_VISION_MODEL"`
Workspace string `json:"workspace" env:"PICOCLAW_TOOLS_VISION_WORKSPACE"`
Restrict bool `json:"restrict" env:"PICOCLAW_TOOLS_VISION_RESTRICT"`
}
type ToolsConfig struct {
Web WebToolsConfig `json:"web"`
Cron CronToolsConfig `json:"cron"`
Exec ExecConfig `json:"exec"`
Web WebToolsConfig `json:"web"`
Cron CronToolsConfig `json:"cron"`
Exec ExecConfig `json:"exec"`
Vision VisionToolConfig `json:"vision"`
}
func DefaultConfig() *Config {
@ -439,6 +449,13 @@ func DefaultConfig() *Config {
Exec: ExecConfig{
EnableDenyPatterns: true,
},
Vision: VisionToolConfig{
Enabled: false,
ApiURL: "",
Model: "gpt-4o-mini",
Restrict: true,
Workspace: "",
},
},
Heartbeat: HeartbeatConfig{
Enabled: true,

192
pkg/tools/vision.go Normal file
View file

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

205
pkg/tools/vision_test.go Normal file
View file

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