Cyrene_v2/internal/runtime/embedding/api.go

package embedding

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"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
	// warmupLeft 是**一次性**冷启动放宽的开关：为真时下一个请求用 warmupTimeout，
	// 取走即置假（CompareAndSwap ⇒ 并发下也只有一个请求能拿到）。
	// 它**不是**「未成功就一直放宽」—— 失败也照样花掉这一次，因此服务挂死时
	// 只会慢一次，之后每个请求都按 DefaultRequestTimeout 快速失败（见 budget.go）。
	warmupLeft 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}
	}
}

// NewAPI builds an APIEmbedder from cfg using the default HTTP client.
//
// 本包**不记日志**（没有任何 logger 缝）：失败一律沿错误返回，由四个调用方
// （审查召回 / 整理召回 / 向量索引写入 / 检索）各自带 source 记账 —— 它们本来
// 就会把错误原文写进各自的 WARN，因此这里再记一条只会是重复的第二份。
// 单次超时也不可由选项更换：它是 budget.go 那条推演的输入，批次上限与它同源，
// 因此**只允许改常量**（改了之后批次上限自动跟着改，测试的常量断言立刻复核）。
func NewAPI(cfg config.EmbeddingConfig, opts ...APIOption) *APIEmbedder {
	e := &APIEmbedder{
		cfg:    cfg,
		client: defaultHTTPClient(),
	}
	// 冷启动放宽是一个**一次性**许可：见 APIEmbedder.warmupLeft。
	e.warmupLeft.Store(true)
	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
	}

	// 切片用的上限由**稳态**超时推出（budget.go 的 DefaultRequestTimeout）：
	// 冷启动那一次只放宽**截止时刻**，批形状不变 —— 否则首个请求会按更宽的预算
	// 切出一个巨大的批，正好把冷启动的模型顶得更久。
	items, chars := batchLimitsFor(DefaultRequestTimeout)
	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
}

// callTimeout 返回本次请求该用的单次超时，并**当场花掉**一次性冷启动放宽许可
// （成功、失败都算花掉）。因此：
//   - 本进程的第一个请求最多等 warmupTimeout（冷启动），
//   - 从第二个请求起一律 DefaultRequestTimeout —— 服务挂死时不会每次都慢。
func (e *APIEmbedder) callTimeout() time.Duration {
	if e.warmupLeft.CompareAndSwap(true, false) {
		return warmupTimeout
	}
	return DefaultRequestTimeout
}

// embedChunk 发**一次** /embeddings 请求。chunk 必须已经落在 batchLimitsFor 推出来
// 的边界内（Embed 负责切），因此它的耗时上界不超过单次超时的一半。
func (e *APIEmbedder) embedChunk(ctx context.Context, texts []string) ([]Vector, error) {
	// 先取超时（取走即花掉那次放宽），再建 ctx：放宽许可与这次请求一一对应。
	callCtx, cancel := context.WithTimeout(ctx, e.callTimeout())
	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,
	}
}