ai-gateway-go/internal/workbench/embedder.go

package workbench

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"
)

// Embedder 把文本批量转换成固定维度向量。知识库向量化/语义检索依赖该接口,
// Embed 失败时调用方应优雅降级(文档照常入库、embedding 置 NULL)。
type Embedder interface {
	Embed(ctx context.Context, texts []string) ([][]float32, error)
	Dim() int
}

// OllamaEmbedder 调用本地 Ollama 的 /api/embed(批量),默认模型 bge-m3(1024 维)。
// 模型尚未拉取时(/api/embed 返回 404)会先 POST /api/pull 拉取一次再重试。
type OllamaEmbedder struct {
	baseURL   string
	model     string
	dim       int
	batchSize int
	timeout   time.Duration
	client    *http.Client
}

// OllamaEmbedderConfig 是 NewOllamaEmbedder 的参数;BaseURL 与 Model 已去除首尾空白。
type OllamaEmbedderConfig struct {
	BaseURL   string
	Model     string
	Dim       int
	BatchSize int
	Timeout   time.Duration
}

func NewOllamaEmbedder(cfg OllamaEmbedderConfig) *OllamaEmbedder {
	client := &http.Client{Timeout: cfg.Timeout}
	if cfg.Dim <= 0 {
		cfg.Dim = 1024
	}
	if cfg.BatchSize <= 0 {
		cfg.BatchSize = 64
	}
	return &OllamaEmbedder{
		baseURL:   cfg.BaseURL,
		model:     cfg.Model,
		dim:       cfg.Dim,
		batchSize: cfg.BatchSize,
		timeout:   cfg.Timeout,
		client:    client,
	}
}

func (o *OllamaEmbedder) Dim() int { return o.dim }

type ollamaEmbedResponse struct {
	Embeddings [][]float32 `json:"embeddings"`
}

type ollamaErrorResponse struct {
	Error string `json:"error"`
}

// Embed 把 texts 按 batchSize 切批调用 Ollama。任何一次调用失败都会返回错误,
// 由调用方决定是否降级(知识库入库语义)。
func (o *OllamaEmbedder) Embed(ctx context.Context, texts []string) ([][]float32, error) {
	if len(texts) == 0 {
		return nil, nil
	}
	all := make([][]float32, 0, len(texts))
	for start := 0; start < len(texts); start += o.batchSize {
		end := start + o.batchSize
		if end > len(texts) {
			end = len(texts)
		}
		batch, err := o.embedBatch(ctx, texts[start:end])
		if err != nil {
			return nil, fmt.Errorf("embed batch [%d:%d]: %w", start, end, err)
		}
		if len(batch) != end-start {
			return nil, fmt.Errorf("embed batch [%d:%d] returned %d vectors for %d texts", start, end, len(batch), end-start)
		}
		for _, vector := range batch {
			if len(vector) != o.dim {
				return nil, fmt.Errorf("embedding dimension mismatch: got %d want %d", len(vector), o.dim)
			}
			all = append(all, vector)
		}
	}
	return all, nil
}

func (o *OllamaEmbedder) embedBatch(ctx context.Context, texts []string) ([][]float32, error) {
	body, err := json.Marshal(map[string]any{"model": o.model, "input": texts})
	if err != nil {
		return nil, err
	}
	resp, err := o.do(ctx, "/api/embed", body)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (o *OllamaEmbedder) do(ctx context.Context, path string, body []byte) ([][]float32, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	httpResp, err := o.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer httpResp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(httpResp.Body, 8<<20))
	if httpResp.StatusCode == http.StatusNotFound {
		// 模型未拉取:先 pull(流式关闭)再重试一次;仍失败则返回明确错误。
		if err := o.pullModel(ctx); err != nil {
			return nil, fmt.Errorf("model %s not present and pull failed: %w", o.model, err)
		}
		return o.retryOnce(ctx, path, body)
	}
	if httpResp.StatusCode != http.StatusOK {
		return nil, o.decodeError(httpResp.StatusCode, raw)
	}
	var parsed ollamaEmbedResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, fmt.Errorf("decode embed response: %w", err)
	}
	return parsed.Embeddings, nil
}

func (o *OllamaEmbedder) retryOnce(ctx context.Context, path string, body []byte) ([][]float32, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	httpResp, err := o.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer httpResp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(httpResp.Body, 8<<20))
	if httpResp.StatusCode != http.StatusOK {
		return nil, o.decodeError(httpResp.StatusCode, raw)
	}
	var parsed ollamaEmbedResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, fmt.Errorf("decode embed response after pull: %w", err)
	}
	return parsed.Embeddings, nil
}

func (o *OllamaEmbedder) pullModel(ctx context.Context) error {
	body, err := json.Marshal(map[string]any{"model": o.model, "stream": false})
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+"/api/pull", bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := o.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if resp.StatusCode != http.StatusOK {
		return o.decodeError(resp.StatusCode, raw)
	}
	return nil
}

func (o *OllamaEmbedder) decodeError(status int, raw []byte) error {
	var parsed ollamaErrorResponse
	if err := json.Unmarshal(raw, &parsed); err == nil && parsed.Error != "" {
		return errors.New(parsed.Error)
	}
	return fmt.Errorf("ollama request failed with status %d", status)
}