Cyrene_v2/internal/runtime/multimodal/vision/adapter.go

package vision

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// visionMaxTokens 限制 vision/OCR 响应的长度。
const visionMaxTokens = 512

// VisionAdapter 理解一张图像并返回对它的描述。它是
// 视觉工具背后的后端；真实提供方实现它，且不把图像
// 暴露给主模型上下文。
type VisionAdapter interface {
	// Describe 返回 img 的描述。prompt 是可选的指令，
	// 用于覆盖默认的「描述这张图」提问。
	Describe(ctx context.Context, img ImageRef, prompt string) (string, error)
}

// OCRAdapter 从图像中抽取文字。它是 OCR 工具背后的后端。
type OCRAdapter interface {
	// Extract 返回 img 中发现的文字。
	Extract(ctx context.Context, img ImageRef) (string, error)
}

// OpenAICompatibleAdapter 通过调用 OpenAI 兼容的 /chat/completions
// endpoint（图像作为内容部件内嵌）同时实现 VisionAdapter 与 OCRAdapter。
// 它是开箱可用的默认实现；真实的
// 云端 vision/OCR 提供方暴露相同协议。
//
// 两个方向可以跑在同一个 endpoint 的两个模型上（裁决 14.4.1 的「独立视觉
// 模型 + OCR 模型」）：Describe 用 model（LLM_VISION_MODEL），Extract 用
// ocrModel（LLM_OCR_MODEL），未配 OCR 模型时回落到同一个模型。
type OpenAICompatibleAdapter struct {
	cfg        config.LLMConfig
	model      string
	ocrModel   string
	client     *http.Client
	log        *slog.Logger
	uploadsDir string
}

// OpenAIOption 用于定制 OpenAICompatibleAdapter。
type OpenAIOption func(*OpenAICompatibleAdapter)

// WithModel 覆盖适配器的模型（默认取 cfg.Model）。
func WithModel(m string) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if m != "" {
			a.model = m
		}
	}
}

// WithHTTPClient 设置请求所用的 *http.Client；nil 表示保持默认。
func WithHTTPClient(c *http.Client) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if c != nil {
			a.client = c
		}
	}
}

// WithLogger 设置失败时使用的日志器；nil 表示保持丢弃日志器。
func WithLogger(l *slog.Logger) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if l != nil {
			a.log = l
		}
	}
}

// WithOCRModel 覆盖 OCR 方向使用的模型（默认取 cfg.OCRModel，再回落
// 到 model）。
func WithOCRModel(m string) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if m != "" {
			a.ocrModel = m
		}
	}
}

// WithUploadsDir 配置本机上传目录（<DataDir>/uploads，即网关静态路由
// /api/v1/files/ 的根），使指向本服务的上传 URL 直接走本地文件读取，
// 而不是绕一圈网络回自己的 HTTP 路由。留空时上传 URL 按普通 URL 取回，
// 行为与配置前一致。
func WithUploadsDir(dir string) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if d := strings.TrimSpace(dir); d != "" {
			a.uploadsDir = d
		}
	}
}

// NewOpenAICompatible 用 cfg 构造一个适配器，未覆盖时使用
// cfg.Model。
func NewOpenAICompatible(cfg config.LLMConfig, opts ...OpenAIOption) *OpenAICompatibleAdapter {
	a := &OpenAICompatibleAdapter{
		cfg:    cfg,
		model:  cfg.Model,
		client: defaultHTTPClient(),
		log:    slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	for _, o := range opts {
		o(a)
	}
	if a.ocrModel == "" {
		// LLM_OCR_MODEL 未配置时 OCR 与视觉共用同一个模型（向后兼容）。
		a.ocrModel = cfg.OCRModel
	}
	if a.ocrModel == "" {
		a.ocrModel = a.model
	}
	return a
}

// Describe 从底层 chat endpoint 返回图像描述。
func (a *OpenAICompatibleAdapter) Describe(ctx context.Context, img ImageRef, prompt string) (string, error) {
	if prompt == "" {
		prompt = "Describe this image in detail."
	}
	return a.ask(ctx, a.model, img, prompt)
}

// Extract 经同一 endpoint 返回图像中发现的文字；当配置了独立 OCR 模型
// （LLM_OCR_MODEL / WithOCRModel）时用它，否则与视觉模型相同。
func (a *OpenAICompatibleAdapter) Extract(ctx context.Context, img ImageRef) (string, error) {
	return a.ask(ctx, a.ocrModel, img, "Extract all text visible in this image. Return only the text.")
}

// ask 解析图像、构造 chat 请求并返回助手文本。model 是本次请求使用的模型，
// 让两个方向能在同一 endpoint 上跑不同模型。
func (a *OpenAICompatibleAdapter) ask(ctx context.Context, model string, img ImageRef, prompt string) (string, error) {
	if strings.TrimSpace(a.cfg.BaseURL) == "" {
		return "", ErrNotConfigured
	}
	if strings.TrimSpace(model) == "" {
		return "", ErrNotConfigured
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}
	imageURL, err := a.imageURL(ctx, img)
	if err != nil {
		return "", err
	}
	payload, err := json.Marshal(a.request(model, imageURL, prompt))
	if err != nil {
		return "", fmt.Errorf("vision: %w: marshal request: %v", ErrUpstream, err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(a.cfg.BaseURL, "/")+"/chat/completions", bytes.NewReader(payload))
	if err != nil {
		return "", fmt.Errorf("vision: %w: build request: %v", ErrUpstream, trace.URLCause(err))
	}
	req.Header.Set("Content-Type", "application/json")
	if a.cfg.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+a.cfg.APIKey)
	}
	resp, err := a.client.Do(req)
	if err != nil {
		return "", fmt.Errorf("vision: %w: %v", ErrUpstream, trace.URLCause(err))
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		// 错误路径**不再整段读**：这里原来沿用取图的 maxFetchedBody（50 MiB），
		// 等于「最多 50 MiB 的错误体进内存、再进 error」。错误体是攻击面，因此
		// 只读 maxErrorBodyBytes，并且一个字节都不进 error —— 只留 status /
		// contentType / bodyLen / sha256 与 code/type（trace.UpstreamError）。
		body, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBodyBytes))
		upstream := trace.UpstreamError(trace.HostOf(a.cfg.BaseURL), resp.StatusCode,
			resp.Header.Get("Content-Type"), body)
		return "", fmt.Errorf("vision: %w: %w", ErrUpstream, upstream)
	}
	return a.decode(resp.Body)
}

// imageURL 返回线上 image_url 的取值。无论来源是 http(s) URL、本地路径
// 还是 base64，送上游的字节一律由服务端自己取回并内联为 data: URI：
// 把同源（局域网）URL 原样透传给云 provider，对方根本取不到，
// 自部署形态下视觉链路会整条不通（B1 上传大图的实际意义也就落空）。
// URL 取回沿用 Resolve→fetchURL 的既有守卫（http/https 白名单 + 50MiB
// 体积上限），并加取回超时（imageFetchTimeout）。
func (a *OpenAICompatibleAdapter) imageURL(ctx context.Context, img ImageRef) (string, error) {
	ref := a.localizeRef(img)
	if ref.Source == SourceURL {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, imageFetchTimeout)
		defer cancel()
	}
	mime, data, err := Resolve(ctx, ref, a.client)
	if err != nil {
		return "", err
	}
	return DataURI(mime, data), nil
}

// localizeRef 把指向本服务上传路由的 URL 改写成 uploads 目录下的本地路径。
// uploadsDir 未配置、或 URL 不是本服务的上传形状时原样返回（此时仍会按
// 普通 URL 取回，因此行为不退化）；文件缺失等读取错误由 Resolve 如实报告。
func (a *OpenAICompatibleAdapter) localizeRef(img ImageRef) ImageRef {
	if img.Source != SourceURL || a.uploadsDir == "" {
		return img
	}
	path, ok := localUploadPath(img.Data, a.uploadsDir)
	if !ok {
		return img
	}
	return ImageRef{Source: SourcePath, Data: path}
}

// openAIVisionRequest 是发往 /chat/completions 的 JSON 请求体。
type openAIVisionRequest struct {
	Model     string          `json:"model"`
	Messages  []openAIMessage `json:"messages"`
	MaxTokens int             `json:"max_tokens"`
}

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

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

type openAIImageURL struct {
	URL string `json:"url"`
}

type openAIVisionResponse struct {
	Choices []openAIChoice `json:"choices"`
}

type openAIChoice struct {
	Message openAIResponseMessage `json:"message"`
}

type openAIResponseMessage struct {
	Content string `json:"content"`
}

// request 构造图像 + prompt 的线上请求体。
func (a *OpenAICompatibleAdapter) request(model, imageURL, prompt string) openAIVisionRequest {
	return openAIVisionRequest{
		Model: model,
		Messages: []openAIMessage{
			{Role: "user", Content: []openAIPart{
				{Type: "text", Text: prompt},
				{Type: "image_url", ImageURL: &openAIImageURL{URL: imageURL}},
			}},
		},
		MaxTokens: visionMaxTokens,
	}
}

// decode 从响应体中取出第一个 choice 的文本。
func (a *OpenAICompatibleAdapter) decode(r io.Reader) (string, error) {
	var wire openAIVisionResponse
	if err := json.NewDecoder(r).Decode(&wire); err != nil {
		return "", fmt.Errorf("vision: %w: decode response: %v", ErrUpstream, err)
	}
	if len(wire.Choices) == 0 {
		return "", fmt.Errorf("vision: %w: empty choices", ErrUpstream)
	}
	return wire.Choices[0].Message.Content, nil
}

// 接口断言：同一个具体适配器同时服务两个方向。
var _ VisionAdapter = (*OpenAICompatibleAdapter)(nil)
var _ OCRAdapter = (*OpenAICompatibleAdapter)(nil)