vision-tool/internal/service/vision_service.go

package service

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"vision-tool/internal/config"
)

// ============================================================================
// API Types (OpenAI-compatible format used by Agnes)
// ============================================================================

type ChatMessage struct {
	Role    string        `json:"role"`
	Content []MessagePart `json:"content"`
}

type MessagePart struct {
	Type     string    `json:"type"`
	Text     string    `json:"text,omitempty"`
	ImageURL *ImageURL `json:"image_url,omitempty"`
}

type ImageURL struct {
	URL    string `json:"url"`
	Detail string `json:"detail,omitempty"` // "auto", "low", "high"
}

type ChatRequest struct {
	Model              string              `json:"model"`
	Messages           []ChatMessage       `json:"messages"`
	ChatTemplateKwargs *ChatTemplateKwargs `json:"chat_template_kwargs,omitempty"`
}

// ChatTemplateKwargs 开启 Agnes 的 thinking 模式（OpenAI 兼容格式）。
type ChatTemplateKwargs struct {
	EnableThinking bool `json:"enable_thinking"`
}

type ChatResponse struct {
	ID      string   `json:"id"`
	Choices []Choice `json:"choices"`
	Usage   *Usage   `json:"usage,omitempty"`
	Error   *APIError `json:"error,omitempty"`
}

type APIError struct {
	Message string `json:"message"`
	Type    string `json:"type"`
	Code    string `json:"code"`
}

type Choice struct {
	Index   int         `json:"index"`
	Message RespMessage `json:"message"`
}

type RespMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// AnalyzeResult 识别结果。
type AnalyzeResult struct {
	Content string
	Usage   *Usage
}

// ============================================================================
// Vision Service
// ============================================================================

// VisionService 调用 Agnes 免费视觉模型识别图片内容。
type VisionService struct {
	apiKey     string
	model      string
	baseURL    string
	httpClient *http.Client
	maxRetries int
}

func NewVisionService(cfg *config.AIConfig) *VisionService {
	return &VisionService{
		apiKey:     cfg.APIKey,
		model:      cfg.Model,
		baseURL:    cfg.BaseURL,
		maxRetries: 3,
		httpClient: &http.Client{Timeout: 120 * time.Second},
	}
}

// Analyze 分析一张或多张图片（本地路径或 http(s) URL 均可），返回模型文本与 token 用量。
// noThinking 为 true 时关闭 thinking 模式（更快、更低负载）。
func (s *VisionService) Analyze(images []string, prompt string, noThinking bool) (*AnalyzeResult, error) {
	if prompt == "" {
		if len(images) > 1 {
			prompt = "请详细描述这些图片的内容，比较它们之间的异同。请用中文回答。"
		} else {
			prompt = "请详细描述这张图片的内容。如果图片中有文字，请完整识别出来。请用中文回答。"
		}
	}

	var parts []MessagePart
	for _, path := range images {
		content, err := s.buildImageContent(path)
		if err != nil {
			return nil, fmt.Errorf("读取图片 %s: %w", path, err)
		}
		parts = append(parts, content)
	}
	parts = append(parts, MessagePart{Type: "text", Text: prompt})

	req := ChatRequest{
		Model: s.model,
		Messages: []ChatMessage{
			{Role: "user", Content: parts},
		},
		ChatTemplateKwargs: &ChatTemplateKwargs{EnableThinking: !noThinking},
	}

	resp, err := s.callWithRetry(req)
	if err != nil {
		return nil, err
	}
	if len(resp.Choices) == 0 {
		return nil, fmt.Errorf("模型未返回内容")
	}
	return &AnalyzeResult{Content: resp.Choices[0].Message.Content, Usage: resp.Usage}, nil
}

// buildImageContent 构建图片消息块：URL 直接使用，本地文件 base64 编码。
func (s *VisionService) buildImageContent(path string) (MessagePart, error) {
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		return MessagePart{
			Type:     "image_url",
			ImageURL: &ImageURL{URL: path},
		}, nil
	}

	imgData, err := os.ReadFile(path)
	if err != nil {
		return MessagePart{}, err
	}
	dataURL := fmt.Sprintf("data:%s;base64,%s", detectMimeType(path), base64.StdEncoding.EncodeToString(imgData))
	return MessagePart{
		Type:     "image_url",
		ImageURL: &ImageURL{URL: dataURL},
	}, nil
}

// callWithRetry 调用 API，对 429/5xx/rate/busy 指数退避重试（2s → 5s → 8s → 11s）。
func (s *VisionService) callWithRetry(req ChatRequest) (*ChatResponse, error) {
	var lastErr error
	for attempt := 0; attempt <= s.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, s.maxRetries)
			time.Sleep(delay)
		}

		resp, err := s.call(req)
		if err != nil {
			errStr := err.Error()
			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("重试 %d 次仍失败: %w", s.maxRetries, lastErr)
}

func (s *VisionService) call(req ChatRequest) (*ChatResponse, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("序列化请求: %w", err)
	}

	httpReq, err := http.NewRequest(http.MethodPost, s.baseURL, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("创建请求: %w", err)
	}
	httpReq.Header.Set("Authorization", "Bearer "+s.apiKey)
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := s.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("调用 API: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		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("解析响应: %w\nRaw: %s", err, string(respBody))
	}
	return &chatResp, nil
}

// detectMimeType 按扩展名判断图片 MIME 类型。
func detectMimeType(path string) string {
	switch strings.ToLower(filepath.Ext(path)) {
	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"
	case ".png":
		fallthrough
	default:
		return "image/png"
	}
}