Agent-System-Backend/internal/vision/images.go

package vision

import (
	"encoding/base64"
	"fmt"
	"path/filepath"
	"strings"

	agentllm "agent-native/runtime/llm"
	core "agent-system-core"
)

// PrepareImage 把上传的图片收成一条 image-ref。Ref 是标准 base64（无换行）。
// 不支持的类型和超限图片返回警告且 ok 为 false，调用方必须把警告展示出来。
func PrepareImage(filename, mime string, content []byte) (core.DocBlock, string, bool) {
	ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(filename), "."))
	kind, known := sendableExt[ext]
	if !known {
		return core.DocBlock{}, fmt.Sprintf("图片 .%s 不作为图像发送（请转为 png、jpeg、gif 或 webp）", ext), false
	}
	if len(content) == 0 {
		return core.DocBlock{}, "图片内容为空，本轮不发送图像", false
	}
	if len(content) > MaxImageBytes {
		return core.DocBlock{}, fmt.Sprintf("图片超过 %dMB，本轮不发送图像（请压缩后重传）", MaxImageBytes>>20), false
	}
	detected := sniffImage(content)
	if detected == "" || detected != kind {
		return core.DocBlock{}, "图片内容与扩展名不符，本轮不发送图像", false
	}
	if mime != "" && !strings.EqualFold(strings.TrimSpace(strings.Split(mime, ";")[0]), kind) && !strings.EqualFold(mime, "application/octet-stream") {
		return core.DocBlock{}, "图片类型与内容不符，本轮不发送图像", false
	}
	return core.DocBlock{
		Type: core.BlockImageRef,
		Ref:  base64.StdEncoding.EncodeToString(content),
		Text: filename,
	}, "图片按图像发送给模型，不转成文字", true
}

// ImagesFromDoc 取出文档里已编码的图片。MIME 来自文件名，不信任外部声明。
func ImagesFromDoc(doc core.ParsedDocument) []agentllm.Image {
	ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(doc.Meta.Filename), "."))
	kind := sendableExt[ext]
	if kind == "" {
		return nil
	}
	var out []agentllm.Image
	for _, b := range doc.Blocks {
		if b.Type != core.BlockImageRef || strings.TrimSpace(b.Ref) == "" {
			continue
		}
		out = append(out, agentllm.Image{MIME: kind, Data: b.Ref})
	}
	return out
}

// HasVisualInput 判断这一轮是否带图：已解析的图片附件，或正文里的图片链接。
func HasVisualInput(docs []core.ParsedDocument, text string) bool {
	for _, doc := range docs {
		if len(ImagesFromDoc(doc)) > 0 {
			return true
		}
	}
	return len(ImageLinks(text)) > 0
}

// Admit 按张数和总字节数收下一批图片。被拒绝的每一项都有一句说明。
func Admit(existing []agentllm.Image, incoming []agentllm.Image) (kept []agentllm.Image, notes []string) {
	count, total := usage(existing)
	for _, img := range incoming {
		n := decodedLen(img.Data)
		switch {
		case n <= 0:
			notes = append(notes, "一张图片没有可用内容，已跳过")
		case n > MaxImageBytes:
			notes = append(notes, fmt.Sprintf("一张图片超过 %dMB，已跳过", MaxImageBytes>>20))
		case count >= MaxImages:
			notes = append(notes, fmt.Sprintf("本轮图片超过 %d 张，其余未发送", MaxImages))
			return kept, notes
		case total+n > MaxTotalBytes:
			notes = append(notes, fmt.Sprintf("本轮图片合计超过 %dMB，其余未发送", MaxTotalBytes>>20))
			return kept, notes
		default:
			kept = append(kept, img)
			count++
			total += n
		}
	}
	return kept, notes
}

func usage(images []agentllm.Image) (count, total int) {
	for _, img := range images {
		n := decodedLen(img.Data)
		if n > 0 {
			count++
			total += n
		}
	}
	return count, total
}

func decodedLen(data string) int {
	raw, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return 0
	}
	return len(raw)
}

var sendableExt = map[string]string{
	"png":  "image/png",
	"jpg":  "image/jpeg",
	"jpeg": "image/jpeg",
	"gif":  "image/gif",
	"webp": "image/webp",
}

func sniffImage(b []byte) string {
	switch {
	case len(b) >= 8 && string(b[:8]) == "\x89PNG\r\n\x1a\n":
		return "image/png"
	case len(b) >= 3 && b[0] == 0xFF && b[1] == 0xD8 && b[2] == 0xFF:
		return "image/jpeg"
	case len(b) >= 6 && (string(b[:6]) == "GIF87a" || string(b[:6]) == "GIF89a"):
		return "image/gif"
	case len(b) >= 12 && string(b[:4]) == "RIFF" && string(b[8:12]) == "WEBP":
		return "image/webp"
	default:
		return ""
	}
}

// MIMEOf 返回嗅探结果，供下载后的响应校验。
func MIMEOf(content []byte) string { return sniffImage(content) }