HomeAgent/internal/memory/vector/http_embedder.go

package vector

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"gitcode.com/JianFeeeee/HomeAgent/pkg/embedding"
)

func init() {
	// http 也是一个普通 provider：核心只按名字打开它，不知道它背后是云 API、
	// 自建服务还是别的语言写的模型。
	embedding.Register("http", func(cfg embedding.Config) (embedding.Provider, error) {
		return NewHTTPEmbedder(HTTPEmbedderConfig{
			Endpoint:    cfg.Options["endpoint"],
			APIKey:      cfg.Options["api_key"],
			Model:       cfg.Options["model"],
			Fingerprint: cfg.Options["fingerprint"],
			Dimension:   atoiOrZero(cfg.Options["dimension"]),
			Timeout:     durationOrZero(cfg.Options["timeout"]),
		})
	})
}

func atoiOrZero(s string) int {
	n := 0
	for _, r := range strings.TrimSpace(s) {
		if r < '0' || r > '9' {
			return 0
		}
		n = n*10 + int(r-'0')
	}
	return n
}

func durationOrZero(s string) time.Duration {
	d, err := time.ParseDuration(strings.TrimSpace(s))
	if err != nil {
		return 0
	}
	return d
}

// HTTPEmbedderConfig 配置一个外部多模态向量服务。
// 服务契约刻意很小：POST Endpoint，输入 modality/data/mime/side，返回 embedding。
// 任何云 API 或自建服务只需适配这一个协议，即可复用内核全部向量存储与检索链路。
type HTTPEmbedderConfig struct {
	Endpoint    string
	APIKey      string
	Model       string
	Dimension   int
	Timeout     time.Duration
	Fingerprint string
}

// HTTPEmbedder 是 MultimodalEmbedder 的外部 API 实现。
type HTTPEmbedder struct {
	cfg    HTTPEmbedderConfig
	client *http.Client
	mu     sync.Mutex
	closed bool
}

type httpEmbedRequest struct {
	Model    string `json:"model,omitempty"`
	Modality string `json:"modality"`
	Side     string `json:"side"`
	Text     string `json:"text,omitempty"`
	Data     string `json:"data,omitempty"`
	MIME     string `json:"mime,omitempty"`
}

type httpEmbedResponse struct {
	Embedding []float64 `json:"embedding"`
	Data      []struct {
		Embedding []float64 `json:"embedding"`
	} `json:"data,omitempty"`
}

func NewHTTPEmbedder(cfg HTTPEmbedderConfig) (*HTTPEmbedder, error) {
	if strings.TrimSpace(cfg.Endpoint) == "" {
		return nil, fmt.Errorf("vector: empty HTTP embedding endpoint")
	}
	if cfg.Dimension <= 0 {
		return nil, fmt.Errorf("vector: invalid HTTP embedding dimension %d", cfg.Dimension)
	}
	if cfg.Timeout <= 0 {
		cfg.Timeout = 30 * time.Second
	}
	if cfg.Fingerprint == "" {
		cfg.Fingerprint = "http:" + cfg.Model + fmt.Sprintf(":%d", cfg.Dimension)
	}
	return &HTTPEmbedder{cfg: cfg, client: &http.Client{Timeout: cfg.Timeout}}, nil
}

func (e *HTTPEmbedder) VectorizeDense(text string) ([]float64, error) {
	return e.embed(context.Background(), httpEmbedRequest{Model: e.cfg.Model, Modality: string(ModalityText), Side: "query", Text: text})
}

func (e *HTTPEmbedder) EmbedImageDense(img []byte, mime string) ([]float64, error) {
	return e.embed(context.Background(), httpEmbedRequest{Model: e.cfg.Model, Modality: string(ModalityImage), Side: "document", Data: base64.StdEncoding.EncodeToString(img), MIME: mime})
}

func (e *HTTPEmbedder) embed(ctx context.Context, payload httpEmbedRequest) ([]float64, error) {
	e.mu.Lock()
	closed := e.closed
	e.mu.Unlock()
	if closed {
		return nil, fmt.Errorf("vector: HTTP embedder closed")
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.cfg.Endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	if e.cfg.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+e.cfg.APIKey)
	}
	resp, err := e.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("vector: HTTP embedding request: %w", err)
	}
	defer resp.Body.Close()
	b, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("vector: HTTP embedding status %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
	}
	var out httpEmbedResponse
	if err := json.Unmarshal(b, &out); err != nil {
		return nil, fmt.Errorf("vector: decode HTTP embedding: %w", err)
	}
	v := out.Embedding
	if len(v) == 0 && len(out.Data) > 0 {
		v = out.Data[0].Embedding
	}
	if len(v) != e.cfg.Dimension {
		return nil, fmt.Errorf("vector: HTTP embedding dimension %d, want %d", len(v), e.cfg.Dimension)
	}
	return v, nil
}

func (e *HTTPEmbedder) Fingerprint() string { return e.cfg.Fingerprint }
func (e *HTTPEmbedder) Dim() int            { return e.cfg.Dimension }

// Embed 实现公共 provider 契约：核心只传模态与不透明字节，本实现负责把它
// 翻译成外部服务的协议。
func (e *HTTPEmbedder) Embed(ctx context.Context, in embedding.Input) ([]float64, error) {
	req := httpEmbedRequest{
		Model:    e.cfg.Model,
		Modality: string(in.Modality),
		Side:     string(in.Purpose),
		Text:     in.Text,
		MIME:     in.MIME,
	}
	if in.Modality != embedding.ModalityText {
		req.Data = base64.StdEncoding.EncodeToString(in.Data)
	}
	return e.embed(ctx, req)
}

// Info 声明本 provider 的向量空间身份。外部服务的支持模态无法在本地探测，
// 因此只声明 text/image 这两条内核真正会走到的路径。
func (e *HTTPEmbedder) Info() embedding.Info {
	return embedding.Info{
		Dimension:   e.cfg.Dimension,
		Fingerprint: e.cfg.Fingerprint,
		Modalities:  []embedding.Modality{embedding.ModalityText, embedding.ModalityImage},
	}
}
func (e *HTTPEmbedder) Loaded() bool {
	e.mu.Lock()
	defer e.mu.Unlock()
	return !e.closed
}
func (e *HTTPEmbedder) Close() {
	e.mu.Lock()
	e.closed = true
	e.mu.Unlock()
}