rony-chat-bot/internal/embed/embed.go

// Package embed talks to an OpenAI-compatible /v1/embeddings endpoint
// (llama-server --embedding) and provides the vector maths the RAG store
// needs.
//
// Why this exists: the SQLite FTS5 index matches words exactly. Visitors ask
// in Spanish about a portfolio written in English, so the terms that carry the
// meaning — "paga", "trabajado" — appear zero times in the corpus, which says
// "Payments: Stripe" and "worked". Measured on the real corpus, keyword search
// returned nothing at all for "¿Con qué se paga en la tienda de ropa?" while a
// multilingual embedding put all three tienda-ropa chunks on top. Keyword
// search still wins on exact proper nouns, so the store keeps both and fuses
// them.
package embed

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
	"time"
)

// Config describes the embedding endpoint.
type Config struct {
	BaseURL   string // e.g. http://localhost:9200/v1
	Model     string // sent as "model"; llama-server ignores it
	BatchSize int    // texts per request (0 → 8)
	TimeoutMS int    // per-request timeout (0 → 120s)
}

// Client is a minimal embeddings client.
type Client struct {
	baseURL   string
	model     string
	batchSize int
	http      *http.Client
}

func New(cfg Config) *Client {
	batch := cfg.BatchSize
	if batch <= 0 {
		batch = 8
	}
	timeout := time.Duration(cfg.TimeoutMS) * time.Millisecond
	if timeout <= 0 {
		timeout = 120 * time.Second
	}
	return &Client{
		baseURL:   cfg.BaseURL,
		model:     cfg.Model,
		batchSize: batch,
		http:      &http.Client{Timeout: timeout},
	}
}

type embedRequest struct {
	Model string   `json:"model,omitempty"`
	Input []string `json:"input"`
}

type embedResponse struct {
	Data []struct {
		Index     int       `json:"index"`
		Embedding []float32 `json:"embedding"`
	} `json:"data"`
}

// Embed returns one unit-length vector per input, in input order.
//
// Vectors are normalised here, once, so similarity at query time is a plain
// dot product instead of a cosine with two square roots per candidate.
func (c *Client) Embed(ctx context.Context, texts []string) ([][]float32, error) {
	out := make([][]float32, 0, len(texts))
	for start := 0; start < len(texts); start += c.batchSize {
		end := min(start+c.batchSize, len(texts))
		vecs, err := c.embedBatch(ctx, texts[start:end])
		if err != nil {
			return nil, err
		}
		out = append(out, vecs...)
	}
	return out, nil
}

func (c *Client) embedBatch(ctx context.Context, texts []string) ([][]float32, error) {
	body, err := json.Marshal(embedRequest{Model: c.model, Input: texts})
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("embeddings request: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		msg, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("embeddings API %d: %s", resp.StatusCode, string(msg))
	}

	var parsed embedResponse
	if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
		return nil, fmt.Errorf("decode embeddings: %w", err)
	}
	if len(parsed.Data) != len(texts) {
		return nil, fmt.Errorf("embeddings returned %d vectors for %d inputs", len(parsed.Data), len(texts))
	}

	// The API documents an `index` field rather than guaranteeing order.
	out := make([][]float32, len(texts))
	for _, d := range parsed.Data {
		if d.Index < 0 || d.Index >= len(out) {
			return nil, fmt.Errorf("embeddings returned out-of-range index %d", d.Index)
		}
		out[d.Index] = Normalize(d.Embedding)
	}
	for i, v := range out {
		if v == nil {
			return nil, fmt.Errorf("embeddings response missing index %d", i)
		}
	}
	return out, nil
}

// Normalize scales v to unit length. A zero vector is returned unchanged —
// dividing by zero would poison every later comparison with NaN.
func Normalize(v []float32) []float32 {
	var sum float64
	for _, x := range v {
		sum += float64(x) * float64(x)
	}
	if sum == 0 {
		return v
	}
	inv := float32(1 / math.Sqrt(sum))
	out := make([]float32, len(v))
	for i, x := range v {
		out[i] = x * inv
	}
	return out
}

// Similarity is the dot product, which equals cosine similarity for the
// unit-length vectors this package produces. Mismatched lengths score 0 so a
// stale row from a different embedding model can never outrank a real hit.
func Similarity(a, b []float32) float64 {
	if len(a) != len(b) || len(a) == 0 {
		return 0
	}
	var sum float64
	for i := range a {
		sum += float64(a[i]) * float64(b[i])
	}
	return sum
}

// Encode serialises a vector as little-endian float32 for a SQLite BLOB.
func Encode(v []float32) []byte {
	buf := make([]byte, 4*len(v))
	for i, x := range v {
		binary.LittleEndian.PutUint32(buf[4*i:], math.Float32bits(x))
	}
	return buf
}

// Decode reverses Encode. A blob whose length isn't a multiple of 4 is
// corrupt and yields nil rather than a truncated vector.
func Decode(b []byte) []float32 {
	if len(b)%4 != 0 {
		return nil
	}
	out := make([]float32, len(b)/4)
	for i := range out {
		out[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[4*i:]))
	}
	return out
}