Agent-System/internal/controlcenter/embed/embed.go

// Package embed provides an OpenAI-compatible embeddings client for the
// Control Center document vector store, plus a deterministic stub for tests.
// Self-contained: modeled on internal/kb/embed.go but does not import it.
package embed

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

const (
	// Dim is the bge-m3 embedding dimension used by cc_doc_chunks.embedding.
	Dim = 1024

	maxBatchItems = 16
	httpTimeout   = 30 * time.Second
)

// Embedder generates embedding vectors for text batches.
type Embedder interface {
	Embed(ctx context.Context, texts []string) ([][]float32, error)
	Dim() int
}

// HTTPClient talks to an OpenAI-compatible /embeddings endpoint.
type HTTPClient struct {
	endpoint string
	apiKey   string
	model    string
	dim      int
	http     *http.Client
}

// NewHTTP builds a client; endpoint may already include /v1 (or the full
// /embeddings suffix). An empty endpoint yields a client whose Embed always
// fails, so callers can treat it as "embedding unavailable".
func NewHTTP(endpoint, apiKey, model string, dim int) *HTTPClient {
	if dim <= 0 {
		dim = Dim
	}
	return &HTTPClient{
		endpoint: normalizeEndpoint(endpoint),
		apiKey:   apiKey,
		model:    model,
		dim:      dim,
		http:     &http.Client{Timeout: httpTimeout},
	}
}

func (c *HTTPClient) Dim() int { return c.dim }

// Embed embeds texts in batches of at most maxBatchItems.
func (c *HTTPClient) Embed(ctx context.Context, texts []string) ([][]float32, error) {
	if c.endpoint == "" {
		return nil, fmt.Errorf("embed: endpoint not configured")
	}
	if len(texts) == 0 {
		return nil, nil
	}
	out := make([][]float32, len(texts))
	for start := 0; start < len(texts); start += maxBatchItems {
		end := start + maxBatchItems
		if end > len(texts) {
			end = len(texts)
		}
		vecs, err := c.embedBatch(ctx, texts[start:end])
		if err != nil {
			return nil, err
		}
		copy(out[start:end], vecs)
	}
	return out, nil
}

// embedBatch posts one batch; 5xx and network errors are retried once, 4xx
// is returned immediately.
func (c *HTTPClient) embedBatch(ctx context.Context, texts []string) ([][]float32, error) {
	var lastErr error
	for attempt := 0; attempt < 2; attempt++ {
		vecs, retryable, err := c.post(ctx, texts)
		if err == nil {
			return vecs, nil
		}
		if !retryable {
			return nil, err
		}
		lastErr = err
	}
	return nil, lastErr
}

func (c *HTTPClient) post(ctx context.Context, texts []string) ([][]float32, bool, error) {
	body, _ := json.Marshal(map[string]any{"model": c.model, "input": texts})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, false, err
	}
	req.Header.Set("Content-Type", "application/json")
	if c.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
	}
	resp, err := c.http.Do(req)
	if err != nil {
		return nil, true, fmt.Errorf("embed request: %w", err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if resp.StatusCode >= 300 {
		retryable := resp.StatusCode >= 500
		return nil, retryable, fmt.Errorf("embed http %d: %s", resp.StatusCode, truncateUTF8(strings.TrimSpace(string(raw)), 300))
	}
	var parsed struct {
		Data []struct {
			Embedding []float32 `json:"embedding"`
			Index     int       `json:"index"`
		} `json:"data"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, false, fmt.Errorf("embed decode: %w", err)
	}
	out := make([][]float32, len(texts))
	for _, d := range parsed.Data {
		if d.Index >= 0 && d.Index < len(out) {
			out[d.Index] = padOrTrim(d.Embedding, c.dim)
		}
	}
	for i := range out {
		if out[i] == nil && i < len(parsed.Data) {
			out[i] = padOrTrim(parsed.Data[i].Embedding, c.dim)
		}
		if out[i] == nil {
			return nil, false, fmt.Errorf("embed: missing embedding %d", i)
		}
	}
	return out, false, nil
}

func normalizeEndpoint(endpoint string) string {
	endpoint = strings.TrimRight(strings.TrimSpace(endpoint), "/")
	if endpoint == "" {
		return ""
	}
	if strings.HasSuffix(endpoint, "/embeddings") {
		return endpoint
	}
	if strings.HasSuffix(endpoint, "/v1") || strings.Contains(endpoint, "/v1/") {
		return endpoint + "/embeddings"
	}
	return endpoint + "/v1/embeddings"
}

func padOrTrim(v []float32, n int) []float32 {
	if len(v) == n {
		return v
	}
	out := make([]float32, n)
	copy(out, v)
	return out
}

func truncateUTF8(s string, n int) string {
	if n <= 0 || len(s) <= n {
		return s
	}
	for n > 0 && !utf8.RuneStart(s[n]) {
		n--
	}
	return s[:n]
}

// VectorLiteral serializes a vector into the pgvector text format
// "[0.1,0.2,...]"; bind it as a string and cast with ::vector.
func VectorLiteral(vec []float32) string {
	if len(vec) == 0 {
		return ""
	}
	var sb strings.Builder
	sb.WriteByte('[')
	for i, f := range vec {
		if i > 0 {
			sb.WriteByte(',')
		}
		sb.WriteString(strconv.FormatFloat(float64(f), 'g', -1, 32))
	}
	sb.WriteByte(']')
	return sb.String()
}