commit ee03467993420f7dfe3300f9444faf75a094688f Author: Geliebte <1297754537@qq.com> Date: Mon Jun 15 23:45:52 2026 +0800 feat: Go 图片识别 CLI 工具 — 对接 Agnes-2.0-Flash 免费视觉模型 为 Claude Code 等 AI 编码助手提供图片理解能力: - 支持本地图片 base64 编码上传 - 多图对比分析 - 自动重试 + 指数退避 - 多源配置管理(命令行/环境变量/配置文件) - 关闭 thinking 模式以降低延迟 Co-Authored-By: Claude diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..36e3c0e --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +# Binary +vision-tool.exe +vision-tool + +# API keys +config.json + +# Temp files +doc_update.json diff --git a/config.go b/config.go new file mode 100644 index 0000000..3e44f43 --- /dev/null +++ b/config.go @@ -0,0 +1,106 @@ +package main + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" +) + +// Config holds the application configuration. +type Config struct { + APIKey string `json:"api_key"` + Model string `json:"model"` + BaseURL string `json:"base_url"` +} + +// DefaultConfig returns a Config with sensible defaults. +func DefaultConfig() *Config { + return &Config{ + Model: "glm-4.6v-flash", + BaseURL: "https://open.bigmodel.cn/api/paas/v4/chat/completions", + } +} + +// LoadConfig loads configuration from multiple sources in priority order: +// 1. Command-line --apikey flag (highest) +// 2. ZHIPU_API_KEY environment variable +// 3. config.json in the same directory as the executable +// 4. Default built-in key (lowest) +func LoadConfig(flagKey, configPath string) (*Config, error) { + cfg := DefaultConfig() + + // Load from config file if it exists (lowest priority) + if configPath == "" { + execDir := getExecDir() + configPath = filepath.Join(execDir, "config.json") + } + + if data, err := os.ReadFile(configPath); err == nil { + var fileCfg Config + if err := json.Unmarshal(data, &fileCfg); err == nil { + if fileCfg.APIKey != "" { + cfg.APIKey = fileCfg.APIKey + } + if fileCfg.Model != "" { + cfg.Model = fileCfg.Model + } + if fileCfg.BaseURL != "" { + cfg.BaseURL = fileCfg.BaseURL + } + } + } + + // Override with environment variable + if envKey := os.Getenv("ZHIPU_API_KEY"); envKey != "" { + cfg.APIKey = envKey + } + + // Highest priority: command-line flag + if flagKey != "" { + cfg.APIKey = flagKey + } + + if cfg.APIKey == "" { + return cfg, fmt.Errorf("API key not configured. Set it via:\n" + + " 1. --apikey flag\n" + + " 2. ZHIPU_API_KEY environment variable\n" + + " 3. config.json file in the tools directory") + } + + return cfg, nil +} + +// SaveConfigFile saves config to a JSON file. +func SaveConfigFile(path, apiKey string) error { + cfg := DefaultConfig() + cfg.APIKey = apiKey + + data, err := json.MarshalIndent(cfg, "", " ") + if err != nil { + return err + } + + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + return err + } + + return os.WriteFile(path, data, 0600) +} + +func getExecDir() string { + execPath, err := os.Executable() + if err != nil { + return "." + } + return filepath.Dir(execPath) +} + +// MaskKey masks the API key for safe display, showing only first 8 and last 4 chars. +func MaskKey(key string) string { + if len(key) <= 12 { + return strings.Repeat("*", len(key)) + } + return key[:8] + "..." + key[len(key)-4:] +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..64ecf91 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module vision-tool + +go 1.25 diff --git a/main.go b/main.go new file mode 100644 index 0000000..f4dcba3 --- /dev/null +++ b/main.go @@ -0,0 +1,118 @@ +package main + +import ( + "flag" + "fmt" + "os" + "path/filepath" +) + +func main() { + // Command-line flags + var ( + apiKey string + configPath string + prompt string + quiet bool + noThink bool + saveConfig bool + ) + + flag.StringVar(&apiKey, "apikey", "", "Zhipu AI API key (or set ZHIPU_API_KEY env var)") + flag.StringVar(&configPath, "config", "", "Path to config.json file") + flag.StringVar(&prompt, "prompt", "", "Custom prompt for image analysis") + flag.BoolVar(&quiet, "q", false, "Quiet mode — only output the model's response text") + flag.BoolVar(&noThink, "no-think", false, "Disable thinking mode (faster, lower server load)") + flag.BoolVar(&saveConfig, "save-config", false, "Save the API key to config.json and exit") + flag.Parse() + + // Handle --save-config: persist the key and exit + if saveConfig { + if apiKey == "" { + fmt.Fprintln(os.Stderr, "Error: --apikey is required with --save-config") + os.Exit(1) + } + if configPath == "" { + configPath = filepath.Join(getToolsDir(), "config.json") + } + if err := SaveConfigFile(configPath, apiKey); err != nil { + fmt.Fprintf(os.Stderr, "Failed to save config: %v\n", err) + os.Exit(1) + } + fmt.Printf("API key saved to %s\n", configPath) + return + } + + // Load configuration + cfg, err := LoadConfig(apiKey, configPath) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + if !quiet { + fmt.Fprintf(os.Stderr, "Using model: %s | API key: %s\n", cfg.Model, MaskKey(cfg.APIKey)) + } + + // Collect image paths from remaining args + imagePaths := flag.Args() + if len(imagePaths) == 0 { + fmt.Fprintln(os.Stderr, "Usage: vision-tool [flags] [image-file2 ...]") + fmt.Fprintln(os.Stderr, "") + fmt.Fprintln(os.Stderr, "Flags:") + flag.PrintDefaults() + fmt.Fprintln(os.Stderr, "") + fmt.Fprintln(os.Stderr, "Examples:") + fmt.Fprintln(os.Stderr, " vision-tool photo.jpg") + fmt.Fprintln(os.Stderr, " vision-tool -prompt \"这张图里有什么文字?\" screenshot.png") + fmt.Fprintln(os.Stderr, " vision-tool -q img1.png img2.png") + fmt.Fprintln(os.Stderr, " vision-tool --apikey YOUR_KEY photo.jpg") + fmt.Fprintln(os.Stderr, " vision-tool --save-config --apikey YOUR_KEY") + os.Exit(1) + } + + // Validate image files exist + for _, p := range imagePaths { + if _, err := os.Stat(p); os.IsNotExist(err) { + fmt.Fprintf(os.Stderr, "Error: file not found: %s\n", p) + os.Exit(1) + } + } + + client := NewVisionClient(cfg, noThink) + + if !quiet { + fmt.Fprintf(os.Stderr, "Analyzing %d image(s)...\n\n", len(imagePaths)) + } + + // Call the API + var resp *ChatResponse + if len(imagePaths) == 1 { + resp, err = client.AnalyzeImage(imagePaths[0], prompt) + } else { + resp, err = client.AnalyzeImages(imagePaths, prompt) + } + + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + // Output the result + if len(resp.Choices) > 0 { + fmt.Println(resp.Choices[0].Message.Content) + } + + if resp.Usage != nil && !quiet { + fmt.Fprintf(os.Stderr, "\n---\nTokens: prompt=%d completion=%d total=%d\n", + resp.Usage.PromptTokens, resp.Usage.CompletionTokens, resp.Usage.TotalTokens) + } +} + +func getToolsDir() string { + exe, err := os.Executable() + if err != nil { + return "." + } + return filepath.Dir(exe) +} diff --git a/vision.go b/vision.go new file mode 100644 index 0000000..eab5efc --- /dev/null +++ b/vision.go @@ -0,0 +1,304 @@ +package main + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +// ============================================================================ +// API Types (OpenAI-compatible format used by Agnes) +// ============================================================================ + +// ChatMessage represents a message in the chat API. +type ChatMessage struct { + Role string `json:"role"` + Content []MessagePart `json:"content"` +} + +// MessagePart is a part of a message — text or image_url. +type MessagePart struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + ImageURL *ImageURL `json:"image_url,omitempty"` +} + +// ImageURL holds the image URL. +type ImageURL struct { + URL string `json:"url"` + Detail string `json:"detail,omitempty"` // "auto", "low", "high" +} + +// ChatRequest is the request body for the chat completions API. +type ChatRequest struct { + Model string `json:"model"` + Messages []ChatMessage `json:"messages"` + Temperature float64 `json:"temperature,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + Stream bool `json:"stream,omitempty"` + ChatTemplateKwargs *ChatTemplateKwargs `json:"chat_template_kwargs,omitempty"` +} + +// ChatTemplateKwargs enables thinking mode in Agnes (OpenAI-compatible format). +type ChatTemplateKwargs struct { + EnableThinking bool `json:"enable_thinking"` +} + +// ChatResponse is the response from the chat completions API. +type ChatResponse struct { + ID string `json:"id"` + Choices []Choice `json:"choices"` + Usage *Usage `json:"usage,omitempty"` + Error *APIError `json:"error,omitempty"` +} + +// APIError represents an error returned by the API. +type APIError struct { + Message string `json:"message"` + Type string `json:"type"` + Code string `json:"code"` +} + +// Choice represents a single response choice. +type Choice struct { + Index int `json:"index"` + Message RespMessage `json:"message"` +} + +// RespMessage is the message content in the response. +type RespMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +// Usage holds token usage info. +type Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +// ============================================================================ +// Vision Client +// ============================================================================ + +// VisionClient handles calling the vision API. +type VisionClient struct { + apiKey string + model string + baseURL string + httpClient *http.Client + noThinking bool + maxRetries int +} + +// NewVisionClient creates a new VisionClient. +func NewVisionClient(cfg *Config, noThinking bool) *VisionClient { + return &VisionClient{ + apiKey: cfg.APIKey, + model: cfg.Model, + baseURL: cfg.BaseURL, + noThinking: noThinking, + maxRetries: 3, + httpClient: &http.Client{ + Timeout: 120 * time.Second, + }, + } +} + +// AnalyzeImage sends an image to the vision API and returns the description. +// The imagePath can be a local file path or an http/https URL. +func (c *VisionClient) AnalyzeImage(imagePath, prompt string) (*ChatResponse, error) { + imageContent, err := c.buildImageContent(imagePath) + if err != nil { + return nil, err + } + + if prompt == "" { + prompt = "请详细描述这张图片的内容。如果图片中有文字,请完整识别出来。请用中文回答。" + } + + parts := []MessagePart{imageContent, {Type: "text", Text: prompt}} + req := c.buildRequest(parts) + return c.callWithRetry(req) +} + +// AnalyzeImages sends multiple images to the vision API and returns the description. +func (c *VisionClient) AnalyzeImages(imagePaths []string, prompt string) (*ChatResponse, error) { + if prompt == "" { + prompt = "请详细描述这些图片的内容,比较它们之间的异同。请用中文回答。" + } + + var parts []MessagePart + for _, path := range imagePaths { + content, err := c.buildImageContent(path) + if err != nil { + return nil, fmt.Errorf("reading image %s: %w", path, err) + } + parts = append(parts, content) + } + parts = append(parts, MessagePart{Type: "text", Text: prompt}) + + req := c.buildRequest(parts) + return c.callWithRetry(req) +} + +// buildImageContent creates a MessagePart for an image. +// If the path is a URL, use it directly. Otherwise, base64-encode the local file. +func (c *VisionClient) buildImageContent(path string) (MessagePart, error) { + // If it's already a URL, use it directly + if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") { + return MessagePart{ + Type: "image_url", + ImageURL: &ImageURL{URL: path}, + }, nil + } + + // Local file — base64 encode + imgData, mimeType, err := readImageFile(path) + if err != nil { + return MessagePart{}, err + } + + base64Img := base64.StdEncoding.EncodeToString(imgData) + dataURL := fmt.Sprintf("data:%s;base64,%s", mimeType, base64Img) + + return MessagePart{ + Type: "image_url", + ImageURL: &ImageURL{URL: dataURL}, + }, nil +} + +// buildRequest constructs a ChatRequest with the given content parts. +func (c *VisionClient) buildRequest(parts []MessagePart) ChatRequest { + req := ChatRequest{ + Model: c.model, + Messages: []ChatMessage{ + { + Role: "user", + Content: parts, + }, + }, + } + + // Enable thinking mode if not disabled + if !c.noThinking { + req.ChatTemplateKwargs = &ChatTemplateKwargs{EnableThinking: true} + } + + return req +} + +// call sends the request to the API and parses the response. +func (c *VisionClient) call(req ChatRequest) (*ChatResponse, error) { + body, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("marshaling request: %w", err) + } + + httpReq, err := http.NewRequest("POST", c.baseURL, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("calling API: %w", err) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("reading response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + // Try to parse error response + var chatResp ChatResponse + if json.Unmarshal(respBody, &chatResp) == nil && chatResp.Error != nil { + return nil, fmt.Errorf("API error (HTTP %d): [%s] %s", + resp.StatusCode, chatResp.Error.Code, chatResp.Error.Message) + } + return nil, fmt.Errorf("API error (HTTP %d): %s", resp.StatusCode, string(respBody)) + } + + var chatResp ChatResponse + if err := json.Unmarshal(respBody, &chatResp); err != nil { + return nil, fmt.Errorf("parsing response: %w\nRaw: %s", err, string(respBody)) + } + + return &chatResp, nil +} + +// callWithRetry calls the API with automatic retry on transient errors. +func (c *VisionClient) callWithRetry(req ChatRequest) (*ChatResponse, error) { + var lastErr error + for attempt := 0; attempt <= c.maxRetries; attempt++ { + if attempt > 0 { + delay := time.Duration(2+attempt*3) * time.Second + fmt.Fprintf(os.Stderr, "Retrying in %v (attempt %d/%d)...\n", delay, attempt, c.maxRetries) + time.Sleep(delay) + } + + resp, err := c.call(req) + if err != nil { + errStr := err.Error() + // Retry on rate limiting (429) or server errors (5xx) + if strings.Contains(errStr, "429") || + strings.Contains(errStr, "500") || + strings.Contains(errStr, "502") || + strings.Contains(errStr, "503") || + strings.Contains(errStr, "rate") || + strings.Contains(errStr, "busy") { + lastErr = err + continue + } + return nil, err + } + return resp, nil + } + return nil, fmt.Errorf("all %d attempts failed, last error: %w", c.maxRetries+1, lastErr) +} + +// ============================================================================ +// Image file helpers +// ============================================================================ + +// readImageFile reads an image file from disk and returns its bytes and MIME type. +func readImageFile(path string) ([]byte, string, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, "", err + } + return data, detectMimeType(path), nil +} + +// detectMimeType detects the MIME type from file extension. +func detectMimeType(path string) string { + ext := strings.ToLower(filepath.Ext(path)) + switch ext { + case ".png": + return "image/png" + case ".jpg", ".jpeg": + return "image/jpeg" + case ".gif": + return "image/gif" + case ".webp": + return "image/webp" + case ".bmp": + return "image/bmp" + case ".tiff", ".tif": + return "image/tiff" + default: + return "image/png" + } +}