krow_backend/go-api/internal/knowledge/embed.go

package knowledge

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"math"
	"net/http"
	"strings"
	"time"
	"unicode"
)

// The dense half of hybrid retrieval.
//
// §5 requires dense + BM25 fused with RRF, and says not to drop to dense-only
// for convenience. The reverse is the same sin, so the architecture below is
// hybrid from the first line even though the deployment this was built on has
// no embedding credential.
//
// That constraint is met with an interface and two implementations, and the
// difference between them is stated bluntly rather than smoothed over:
//
//   - Voyage is the real one. Semantic: "time off" retrieves a paragraph about
//     annual leave that never uses either word.
//   - Lexical is a deterministic stand-in that hashes terms into a vector. It
//     is NOT semantic. It captures term overlap and nothing else, so hybrid
//     retrieval running on it is two flavours of keyword search wearing a
//     trenchcoat. It exists so the ACL pre-filter, the fusion and the whole
//     retrieval path are testable without a network or a key, and it refuses to
//     run in production.
//
// Anthropic does not serve embeddings; Voyage is the documented partner. The
// interface is what matters — swapping in another provider is one file.

// Kind distinguishes a document from a query.
//
// Modern embedding models are asymmetric: they encode "what is our lateness
// policy?" and "Staff arriving more than ten minutes after..." differently on
// purpose, and a retriever that embeds both the same way loses accuracy for no
// reason. The interface carries it so a provider that cares can use it and one
// that does not can ignore it.
type Kind string

const (
	KindDocument Kind = "document"
	KindQuery    Kind = "query"
)

// Embedder turns text into vectors.
//
// Implementations MUST return unit-normalised vectors. The schema's similarity
// function is a plain dot product, which equals cosine similarity only for unit
// vectors — an implementation that skipped normalisation would produce a
// ranking dominated by whichever chunks happened to have the largest magnitude,
// and it would not error, it would just quietly rank badly.
type Embedder interface {
	Embed(ctx context.Context, texts []string, kind Kind) ([][]float32, error)

	// Model names the vectors this embedder produces. Stored on every chunk,
	// because vectors from two models are not comparable and a half-migrated
	// corpus returns nonsense rather than failing.
	Model() string

	// Dimensions is the vector length. Fixed per model.
	Dimensions() int
}

/* ── Voyage ─────────────────────────────────────────────────────────────── */

// VoyageEmbedder calls Voyage AI.
type VoyageEmbedder struct {
	APIKey  string
	ModelI  string
	Dims    int
	HTTP    *http.Client
	BaseURL string
}

// DefaultVoyageModel is the general-purpose retrieval model.
const (
	DefaultVoyageModel = "voyage-3.5"
	DefaultVoyageDims  = 1024
	defaultVoyageURL   = "https://api.voyageai.com/v1/embeddings"
)

// NewVoyage builds an embedder over the Voyage API.
func NewVoyage(apiKey, model string, dims int) *VoyageEmbedder {
	if model == "" {
		model = DefaultVoyageModel
	}
	if dims <= 0 {
		dims = DefaultVoyageDims
	}
	return &VoyageEmbedder{
		APIKey: apiKey, ModelI: model, Dims: dims,
		HTTP:    &http.Client{Timeout: 30 * time.Second},
		BaseURL: defaultVoyageURL,
	}
}

func (v *VoyageEmbedder) Model() string   { return v.ModelI }
func (v *VoyageEmbedder) Dimensions() int { return v.Dims }

func (v *VoyageEmbedder) Embed(ctx context.Context, texts []string, kind Kind) ([][]float32, error) {
	if v.APIKey == "" {
		// Structured rather than a bare string, and raised here rather than at
		// startup: this service boots and serves everything that is not
		// retrieval without an embedding key, and a refusal to start would
		// make the knowledge layer's absence take the whole API with it.
		return nil, &Error{Code: ErrNotConfigured, Message: "no embedding credential is configured"}
	}
	if len(texts) == 0 {
		return nil, nil
	}

	inputType := "document"
	if kind == KindQuery {
		inputType = "query"
	}
	body, err := json.Marshal(map[string]any{
		"input":      texts,
		"model":      v.ModelI,
		"input_type": inputType,
		// Unit-normalised at the source where the provider offers it, so the
		// dot product in SQL is cosine similarity without a second pass.
		"output_dimension": v.Dims,
	})
	if err != nil {
		return nil, &Error{Code: ErrEmbedFailed, Message: "the embedding request could not be encoded", Cause: err}
	}

	url := v.BaseURL
	if url == "" {
		url = defaultVoyageURL
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return nil, &Error{Code: ErrEmbedFailed, Message: "the embedding request could not be built", Cause: err}
	}
	req.Header.Set("Authorization", "Bearer "+v.APIKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := v.HTTP.Do(req)
	if err != nil {
		return nil, &Error{Code: ErrEmbedUnavailable, Message: "the embedding service could not be reached", Cause: err}
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		code := ErrEmbedFailed
		if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
			code = ErrEmbedUnavailable
		}
		// The response body is deliberately not included. It is provider text,
		// it can echo the input, and the input is tenant content.
		return nil, &Error{Code: code, Message: fmt.Sprintf("the embedding service answered %d", resp.StatusCode)}
	}

	var decoded struct {
		Data []struct {
			Index     int       `json:"index"`
			Embedding []float32 `json:"embedding"`
		} `json:"data"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&decoded); err != nil {
		return nil, &Error{Code: ErrEmbedFailed, Message: "the embedding response could not be read", Cause: err}
	}
	if len(decoded.Data) != len(texts) {
		return nil, &Error{Code: ErrEmbedFailed, Message: fmt.Sprintf(
			"asked for %d embeddings and got %d", len(texts), len(decoded.Data))}
	}

	// Placed by the index the provider reports rather than by arrival order. A
	// mis-ordered batch would attach every chunk's vector to its neighbour,
	// which produces a corpus that retrieves confidently and wrongly.
	out := make([][]float32, len(texts))
	for _, d := range decoded.Data {
		if d.Index < 0 || d.Index >= len(out) {
			return nil, &Error{Code: ErrEmbedFailed, Message: "the embedding response was mis-indexed"}
		}
		out[d.Index] = normalise(d.Embedding)
	}
	for i, vec := range out {
		if len(vec) == 0 {
			return nil, &Error{Code: ErrEmbedFailed, Message: fmt.Sprintf("no embedding came back for input %d", i)}
		}
	}
	return out, nil
}

/* ── Lexical stand-in ───────────────────────────────────────────────────── */

// LexicalEmbedder hashes terms into a fixed-width vector.
//
// READ THIS BEFORE USING IT. It is not a semantic embedder and does not
// approximate one. It hashes each term to a dimension and counts it, so two
// texts are "similar" here exactly when they share vocabulary — "annual leave"
// and "time off" are orthogonal. Running hybrid retrieval on it gives you BM25
// twice, and any evaluation of retrieval QUALITY against it is measuring
// nothing.
//
// It exists for one reason: the ACL pre-filter, the fusion, the citation path
// and the prompt assembly all need to be exercised and asserted, and none of
// them should require a network call or a credential to test. Those properties
// are independent of whether the vectors mean anything.
//
// It refuses outside development, so it cannot become the thing that shipped.
type LexicalEmbedder struct {
	Dims int
	// AllowInProduction is the deliberate override, and there is no reason to
	// set it. It exists so the refusal below is a decision someone had to make
	// in code rather than a flag they could set in an environment.
	AllowInProduction bool
	Production        bool
}

// NewLexical builds the stand-in embedder.
func NewLexical(dims int) *LexicalEmbedder {
	if dims <= 0 {
		dims = 256
	}
	return &LexicalEmbedder{Dims: dims}
}

func (l *LexicalEmbedder) Model() string   { return fmt.Sprintf("lexical-hash-%d", l.Dims) }
func (l *LexicalEmbedder) Dimensions() int { return l.Dims }

func (l *LexicalEmbedder) Embed(_ context.Context, texts []string, _ Kind) ([][]float32, error) {
	if l.Production && !l.AllowInProduction {
		return nil, &Error{
			Code: ErrNotConfigured,
			Message: "the lexical stand-in embedder cannot run in production; it is not semantic, " +
				"and a corpus indexed with it would retrieve on word overlap alone",
		}
	}

	out := make([][]float32, len(texts))
	for i, text := range texts {
		vec := make([]float32, l.Dims)
		for _, term := range terms(text) {
			h := fnv.New32a()
			h.Write([]byte(term))
			d := int(h.Sum32()) % l.Dims
			if d < 0 {
				d += l.Dims
			}
			// A second hash decides the sign, so unrelated terms colliding on a
			// dimension tend to cancel rather than reinforce. Cheap, and it
			// keeps a small vector from saturating.
			s := fnv.New32()
			s.Write([]byte(term))
			if s.Sum32()%2 == 0 {
				vec[d] += 1
			} else {
				vec[d] -= 1
			}
		}
		out[i] = normalise(vec)
	}
	return out, nil
}

// terms splits text into lower-cased word tokens.
func terms(text string) []string {
	fields := strings.FieldsFunc(strings.ToLower(text), func(r rune) bool {
		return !unicode.IsLetter(r) && !unicode.IsDigit(r)
	})
	out := make([]string, 0, len(fields))
	for _, f := range fields {
		if len(f) > 1 {
			out = append(out, f)
		}
	}
	return out
}

/* ── Shared ─────────────────────────────────────────────────────────────── */

// normalise scales a vector to unit length.
//
// The schema's similarity function is a dot product, which is cosine similarity
// only for unit vectors. A zero vector — a chunk of pure punctuation, or a
// provider returning zeros — is returned unchanged rather than divided by zero;
// it scores 0 against everything, which is the right answer for text with no
// content.
func normalise(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
}

/* ── Ollama ─────────────────────────────────────────────────────────────── */

// OllamaEmbedder calls a model running on this machine.
//
// The third option, and for a workforce corpus often the right one. It is a
// real semantic embedder — "time off" finds "annual leave" — with three
// properties the hosted one does not have:
//
//   - **No credential.** Nothing to provision, rotate, or leak.
//   - **No per-token cost.** Re-embedding the whole corpus after a chunking
//     change is free, which is the difference between tuning retrieval and
//     being afraid to.
//   - **No tenant text leaving the machine.** For handbooks and worker notes
//     that is a substantive argument, not a preference.
//
// The cost is quality: `nomic-embed-text` is genuinely good and still behind
// the best hosted models on subtle retrieval over a large messy corpus. For a
// policy library it is not the limiting factor.
type OllamaEmbedder struct {
	BaseURL string
	ModelI  string
	Dims    int
	HTTP    *http.Client
}

const (
	// DefaultOllamaModel is a retrieval-tuned embedding model that runs
	// comfortably on a laptop.
	DefaultOllamaModel = "nomic-embed-text"
	// DefaultOllamaDims is that model's output width.
	DefaultOllamaDims = 768
	defaultOllamaURL  = "http://localhost:11434"
)

// NewOllama builds an embedder over a local Ollama.
func NewOllama(baseURL, model string, dims int) *OllamaEmbedder {
	if baseURL == "" {
		baseURL = defaultOllamaURL
	}
	if model == "" {
		model = DefaultOllamaModel
	}
	if dims <= 0 {
		dims = DefaultOllamaDims
	}
	return &OllamaEmbedder{
		BaseURL: strings.TrimRight(baseURL, "/"),
		ModelI:  model,
		Dims:    dims,
		// Longer than the hosted client's: a local model that has just been
		// pulled loads into memory on the first request, and that first call
		// can take tens of seconds on a cold start. Timing it out would make
		// the very first ingest look broken.
		HTTP: &http.Client{Timeout: 120 * time.Second},
	}
}

// Model names the vectors this embedder produces.
//
// Prefixed, so a corpus embedded by a local `nomic-embed-text` is never
// mistaken for one embedded by a hosted model of the same name. Vectors from
// two models are not comparable, and the model name on the chunk row is the
// only thing standing between that and confident nonsense.
func (o *OllamaEmbedder) Model() string   { return "ollama/" + o.ModelI }
func (o *OllamaEmbedder) Dimensions() int { return o.Dims }

func (o *OllamaEmbedder) Embed(ctx context.Context, texts []string, _ Kind) ([][]float32, error) {
	if len(texts) == 0 {
		return nil, nil
	}

	// Ollama's embedding endpoint takes no input_type, so the document/query
	// asymmetry the hosted models use is simply not available here. Ignored
	// rather than faked: prefixing the text with "query:" is a convention some
	// models are trained on and this one is not, and applying it anyway would
	// degrade retrieval while looking like a refinement.
	body, err := json.Marshal(map[string]any{
		"model": o.ModelI,
		"input": texts,
	})
	if err != nil {
		return nil, &Error{Code: ErrEmbedFailed, Message: "the embedding request could not be encoded", Cause: err}
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.BaseURL+"/api/embed", bytes.NewReader(body))
	if err != nil {
		return nil, &Error{Code: ErrEmbedFailed, Message: "the embedding request could not be built", Cause: err}
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := o.HTTP.Do(req)
	if err != nil {
		// The common case by a distance: Ollama is not running. Said plainly,
		// with the command to fix it, because the alternative is an operator
		// reading "connection refused" and going looking for a network problem.
		return nil, &Error{
			Code: ErrEmbedUnavailable,
			Message: fmt.Sprintf(
				"no embedding model is answering at %s — start Ollama and run "+
					"`ollama pull %s`", o.BaseURL, o.ModelI),
			Cause: err,
		}
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		if resp.StatusCode == http.StatusNotFound {
			// Ollama is up but has never seen this model. A different problem
			// from being down, and a different fix.
			return nil, &Error{
				Code: ErrNotConfigured,
				Message: fmt.Sprintf("Ollama does not have %q — run `ollama pull %s`",
					o.ModelI, o.ModelI),
			}
		}
		code := ErrEmbedFailed
		if resp.StatusCode >= 500 {
			code = ErrEmbedUnavailable
		}
		return nil, &Error{Code: code, Message: fmt.Sprintf(
			"the embedding model answered %d", resp.StatusCode)}
	}

	var decoded struct {
		Embeddings [][]float32 `json:"embeddings"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&decoded); err != nil {
		return nil, &Error{Code: ErrEmbedFailed, Message: "the embedding response could not be read", Cause: err}
	}
	if len(decoded.Embeddings) != len(texts) {
		return nil, &Error{Code: ErrEmbedFailed, Message: fmt.Sprintf(
			"asked for %d embeddings and got %d", len(texts), len(decoded.Embeddings))}
	}

	// Normalised here rather than trusted. Ollama returns whatever the model
	// produced, and the schema's similarity function is a plain dot product —
	// which equals cosine similarity only for unit vectors. Skipping this would
	// not error; it would just rank badly, dominated by whichever chunks
	// happened to have the largest magnitude.
	out := make([][]float32, len(decoded.Embeddings))
	for i, v := range decoded.Embeddings {
		if len(v) == 0 {
			return nil, &Error{Code: ErrEmbedFailed, Message: fmt.Sprintf(
				"no embedding came back for input %d", i)}
		}
		out[i] = normalise(v)
	}
	return out, nil
}