vision-tool/main.go

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-file> [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)
}