Cyrene_v2/internal/runtime/embedding/api.go

package embedding

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

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

// maxResponseBytes caps how much of an /embeddings response body is read, so an
// oversized or malformed response cannot exhaust memory. Legitimate embeddings
// for a normal batch fall far below this; the cap is a guard, not a limit.
const maxResponseBytes = 16 << 20

// maxErrorBodyBytes caps how much of a non-OK response body is read for the
// error message.
const maxErrorBodyBytes = 4096

// HTTP client tuning values, matching the recommended pool in
// DevDocs/02-development-constraints.md §2.3.
const (
	httpMaxIdleConns        = 100
	httpMaxIdleConnsPerHost = 20
	httpIdleConnTimeout     = 90 * time.Second
	httpTLSHandshakeTimeout = 15 * time.Second
	httpRequestTimeout      = 180 * time.Second
)

// DefaultRequestTimeout 是**单次** /embeddings 调用的墙钟上限，独立于 HTTP 客户端
// 那一层的 180s 总超时（httpRequestTimeout）。
//
// 依据（是实测数据，不是拍脑袋的调优）：
//   - 上层「入库前审查」的单条总预算是 20s（memory.defaultReviewTimeout），embedding
//     只是它内部的第一步。2026-10-07 实测（本机 llama.cpp + Qwen3-Embedding-0.6B，
//     服务空闲）：单条文本约 60ms；**64 条一批约 2.1~2.9s**；审查最坏情况
//     （ScanLimit=200 + 新条目 1）要 4 批 ⇒ 约 8~10s，已经占掉 20s 预算的一半，
//     留给后面的 T1 灰区判定（defaultT1Timeout=8s）的余量很薄。一次卡死的 HTTP
//     就能把这点余量连同整条审查一起吃掉（实测日志里
//     `memory/review: context deadline exceeded` 正是这个形态）；
//   - 所以每次调用必须有**自己的**上限：卡住的那一批到点就失败并降级，不能把整条
//     审查的 20s 都赌在一次 HTTP 上；
//   - 5s 的取值：相对实测的 64 条一批 2.9s 留约 1.7 倍余量，能吸收并发争用带来的
//     抖动，同时远小于 20s 的审查预算。若日后批次变大或服务变慢，先调这里而不是
//     调审查总预算；
//   - 真慢就**快速失败**：调用方按**既有语义**降级（审查降级到词项余弦并如实留痕
//     degraded），**绝不静默丢记忆** —— 本常量只砍等待时间，不改任何判定规则；
//   - 已知代价：llama.cpp 首次加载模型可能超过 5s，那一次调用会失败并降级，模型
//     常驻后的下一次调用即恢复正常。降级本就是既有行为，因此可接受。
const DefaultRequestTimeout = 5 * time.Second

// openAICompatibleRequest is the JSON body sent to an OpenAI-compatible
// /embeddings endpoint.
type openAICompatibleRequest struct {
	Model string   `json:"model"`
	Input []string `json:"input"`
}

// openAICompatibleResponse is the subset of the /embeddings JSON response the
// APIEmbedder consumes.
type openAICompatibleResponse struct {
	Data []openAICompatibleDatum `json:"data"`
}

// openAICompatibleDatum is a single embedding in the response data array.
type openAICompatibleDatum struct {
	Index     int       `json:"index"`
	Embedding []float32 `json:"embedding"`
}

// ErrNoAPIURL is returned by APIEmbedder.Embed when the configured API URL is
// blank, so a mis-configured embedder fails loudly instead of issuing a request
// to an empty URL.
var ErrNoAPIURL = errors.New("embedding: api url is empty")

// APIEmbedder implements Embedder by calling an OpenAI-compatible /embeddings
// endpoint. It is the semantic-vector enhancement selected when
// config.EmbeddingConfig.APIURL is set, and it always uses the real HTTP path
// (never a stub).
//
// The vector dimension is not known until the first successful response; Dim()
// reports the last observed dimension, or 0 before any successful Embed. The
// runtime should therefore call Embed before relying on Dim().
type APIEmbedder struct {
	cfg    config.EmbeddingConfig
	client *http.Client
	log    *slog.Logger

	mu  sync.RWMutex
	dim int
}

// APIOption customizes an APIEmbedder built by NewAPI.
type APIOption func(*APIEmbedder)

// WithHTTPClient sets the *http.Client used for the /embeddings request. A nil
// client is ignored and the default is kept.
func WithHTTPClient(c *http.Client) APIOption {
	return func(e *APIEmbedder) {
		if c != nil {
			e.client = c
		}
	}
}

// WithRoundTripper sets the transport for the /embeddings request by wrapping
// it in a fresh *http.Client. Tests use this to inject a mock RoundTripper
// without starting a network server.
func WithRoundTripper(rt http.RoundTripper) APIOption {
	return func(e *APIEmbedder) {
		e.client = &http.Client{Transport: rt}
	}
}

// WithLogger sets the logger used for embed failures. A nil logger is ignored
// and the default discard logger is kept.
func WithLogger(l *slog.Logger) APIOption {
	return func(e *APIEmbedder) {
		if l != nil {
			e.log = l
		}
	}
}

// NewAPI builds an APIEmbedder from cfg using the default HTTP client.
func NewAPI(cfg config.EmbeddingConfig, opts ...APIOption) *APIEmbedder {
	e := &APIEmbedder{
		cfg:    cfg,
		client: defaultHTTPClient(),
		log:    slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	for _, opt := range opts {
		opt(e)
	}
	return e
}

// Dim returns the dimension observed from the last successful response, or 0
// if no response has been received yet.
func (e *APIEmbedder) Dim() int {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.dim
}

// Embed sends texts to the configured /embeddings endpoint and returns the
// returned vectors aligned by index. It returns ErrNoAPIURL for a blank URL
// and a wrapped error for any HTTP or decode failure.
//
// 单次调用带独立的、比上层审查总预算（20s）短得多的超时（DefaultRequestTimeout），
// 卡住即快速失败，把预算留给后续步骤；失败语义与降级口径不变。
func (e *APIEmbedder) Embed(ctx context.Context, texts []string) ([]Vector, error) {
	if len(texts) == 0 {
		return []Vector{}, nil
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if strings.TrimSpace(e.cfg.APIURL) == "" {
		return nil, ErrNoAPIURL
	}

	callCtx, cancel := context.WithTimeout(ctx, DefaultRequestTimeout)
	defer cancel()

	req, err := e.newEmbeddingRequest(callCtx, texts)
	if err != nil {
		return nil, err
	}

	resp, err := e.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("embedding: request: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		msg := readErrorBody(resp.Body)
		return nil, fmt.Errorf("embedding: status %d: %s", resp.StatusCode, msg)
	}

	data, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
	if err != nil {
		return nil, fmt.Errorf("embedding: read response: %w", err)
	}

	var decoded openAICompatibleResponse
	if err := json.Unmarshal(data, &decoded); err != nil {
		return nil, fmt.Errorf("embedding: decode response: %w", err)
	}
	if len(decoded.Data) != len(texts) {
		return nil, fmt.Errorf("embedding: got %d vectors for %d texts", len(decoded.Data), len(texts))
	}

	vectors := make([]Vector, len(texts))
	dim := 0
	for _, d := range decoded.Data {
		if d.Index < 0 || d.Index >= len(vectors) {
			return nil, fmt.Errorf("embedding: index %d out of range", d.Index)
		}
		vec := make(Vector, len(d.Embedding))
		copy(vec, d.Embedding)
		vectors[d.Index] = vec
		if dim == 0 {
			dim = len(vec)
		}
	}

	e.mu.Lock()
	e.dim = dim
	e.mu.Unlock()

	return vectors, nil
}

// newEmbeddingRequest 组装一次 /embeddings 调用：序列化请求体、**规范化端点主机**
// （URL 解析之后、拨号之前；只改主机，path/query 一字不动）、构造 POST 请求并打上
// 鉴权头。它只负责组装，不负责发送。
func (e *APIEmbedder) newEmbeddingRequest(ctx context.Context, texts []string) (*http.Request, error) {
	payload, err := json.Marshal(openAICompatibleRequest{
		Model: e.cfg.Model,
		Input: texts,
	})
	if err != nil {
		return nil, fmt.Errorf("embedding: marshal request: %w", err)
	}

	endpoint, err := normalizeEndpointURL(e.cfg.APIURL)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
	if err != nil {
		return nil, fmt.Errorf("embedding: build request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	if e.cfg.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+e.cfg.APIKey)
	}
	return req, nil
}

// readErrorBody drains a small prefix of a non-OK response body into a message,
// capped at maxErrorBodyBytes so an error response cannot exhaust memory.
func readErrorBody(r io.Reader) string {
	raw, err := io.ReadAll(io.LimitReader(r, maxErrorBodyBytes))
	if err != nil {
		return err.Error()
	}
	msg := strings.TrimSpace(string(raw))
	if msg == "" {
		return "bad status"
	}
	return msg
}

// defaultHTTPClient builds the HTTP client configuration described in
// DevDocs/02-development-constraints.md §2.3: tuned idle/TLS timeouts plus a
// total request timeout. A fresh client is returned per call (never a package
// global, which the constraints forbid).
func defaultHTTPClient() *http.Client {
	return &http.Client{
		Transport: &http.Transport{
			MaxIdleConns:        httpMaxIdleConns,
			MaxIdleConnsPerHost: httpMaxIdleConnsPerHost,
			IdleConnTimeout:     httpIdleConnTimeout,
			TLSHandshakeTimeout: httpTLSHandshakeTimeout,
		},
		Timeout: httpRequestTimeout,
	}
}