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
)

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

	payload, err := json.Marshal(openAICompatibleRequest{
		Model: e.cfg.Model,
		Input: texts,
	})
	if err != nil {
		return nil, fmt.Errorf("embedding: marshal request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.cfg.APIURL, 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)
	}

	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
}

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