Cyrene_v2/internal/runtime/embedding/api.go

package embedding

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"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）。
//
// 它**不是**单独取值：批次上限（DefaultBatchLimit / DefaultBatchChars）与它由
// budget.go 里**同一个不等式**推出 —— 单次请求的耗时上界不得超过超时的一半。
// 改这一行就等于改了批次上限，反之亦然；两侧永远是配套的。
//
// 为什么每次调用必须有自己的上限：上层「入库前审查」的单条总预算是 20s
// （memory.defaultReviewTimeout），embedding 只是它内部的第一步，后面还有 T1 灰区
// 判定（8s）。一次卡死的 HTTP 不能把整条审查的 20s 都赌掉。真慢就**快速失败**，
// 调用方按**既有语义**降级（审查降级到词项余弦并如实留痕 degraded），
// **绝不静默丢记忆** —— 本常量只砍等待时间，不改任何判定规则。
//
// 5s 保持不变的依据：本次故障（2026-10-13）的成因不是这个数太小，而是批次太大 ——
// 64 条的实测耗时 7.5~8.5s 已经越过它，客户端放弃、服务端还在算，把整台本机
// embedding 服务顶住，连带单条请求（写向量索引）也一起超时。所以修法是**把批次
// 收进预算内**，而不是把预算放大到迁就批次（后者会让真挂死的请求也拖住整轮）。
// 冷启动那一次另有 warmupTimeout（见 budget.go），它只生效到第一次成功响应为止。
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
	// timeout 是**单次**请求的墙钟上限，与批次上限同源（见 budget.go）。
	// 默认 DefaultRequestTimeout；WithRequestTimeout 可换，换了批次上限也跟着换。
	timeout time.Duration
	// warmed 报告是否已经拿到过成功响应。未热之前用 warmupTimeout（模型加载 /
	// GPU 缓冲分配），热之后一律用 timeout。
	warmed atomic.Bool

	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
		}
	}
}

// WithRequestTimeout 换掉单次请求的墙钟上限（非正数忽略，保留默认）。
// 它同时决定批次上限 —— 两者是同一条推演（budget.go），因此这里是**唯一**的
// 调优点，不允许调用方各自去改批次。
func WithRequestTimeout(d time.Duration) APIOption {
	return func(e *APIEmbedder) {
		if d > 0 {
			e.timeout = d
		}
	}
}

// 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)),
		timeout: DefaultRequestTimeout,
	}
	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.
//
// texts 长度不受调用方控制：本方法按 budget.go 推出来的上限（条数 / 字符总量）
// **自己切片**，分多次请求顺序发出，因此「一次请求的耗时 < 单次超时」这条不变式
// 对任何调用方都成立 —— 调用方给多大的切片都不会把服务端顶住。中途任一段失败即
// 整体失败（返回已发请求的错误），绝不半途改用混合表示。
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
	}

	// 切片用的上限由**稳态**超时推出（e.timeout），与 requestTimeout() 无关：
	// 冷启动那一次只是**截止时刻**更宽（warmupTimeout），批次形状不变 —— 否则
	// 首个请求会按 30s 的预算切出一个巨大的批，正好把冷启动的模型顶得更久。
	items, chars := batchLimitsFor(e.timeout)
	out := make([]Vector, 0, len(texts))
	for _, chunk := range chunkTexts(texts, items, chars) {
		vectors, err := e.embedChunk(ctx, chunk)
		if err != nil {
			return nil, err
		}
		out = append(out, vectors...)
	}
	return out, nil
}

// requestTimeout 返回本次调用该用的单次超时：拿到第一个成功响应之前用
// warmupTimeout（冷启动：模型加载 / GPU 缓冲分配），之后一律用配置的 timeout。
func (e *APIEmbedder) requestTimeout() time.Duration {
	if e.warmed.Load() {
		return e.timeout
	}
	return warmupTimeout
}

// embedChunk 发**一次** /embeddings 请求。chunk 必须已经落在 batchLimitsFor 推出来
// 的边界内（Embed 负责切），因此它的耗时上界不超过超时的一半。
func (e *APIEmbedder) embedChunk(ctx context.Context, texts []string) ([]Vector, error) {
	callCtx, cancel := context.WithTimeout(ctx, e.requestTimeout())
	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)
	}
	// 服务端确实答了 ⇒ 冷启动这一关过了，后续调用回到常态超时。
	e.warmed.Store(true)

	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,
	}
}