vision-tool/vision.go

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"
	}
}