fathom/internal/insights/embed.go

package insights

import (
	"context"
	"math"
	"sort"
)

// Embedder produces one vector per input string.
//
// An interface rather than a concrete client so the detectors stay pure
// functions of their input: tests supply a deterministic fake, production
// supplies a local model, and a nil Embedder is the ordinary case where
// clustering falls back to MinHash.
type Embedder interface {
	Embed(ctx context.Context, texts []string) ([][]float32, error)
}

// ClusterMethod records how a cluster was actually formed.
//
// It is reported rather than hidden because it changes what the finding means.
// MinHash clusters share vocabulary; embedding clusters may share only meaning.
// A reader comparing two machines should be able to see why they disagree.
type ClusterMethod string

const (
	MethodMinHash   ClusterMethod = "lexical"
	MethodEmbedding ClusterMethod = "semantic"
)

// embedSimilarity is the cosine threshold for embedding clusters.
//
// Higher than the MinHash threshold because cosine similarity between short
// task descriptions runs high generally — unrelated developer prompts share
// enough register to sit around 0.6, so a threshold tuned for MinHash's
// distribution would group everything.
const embedSimilarity = 0.82

// clusterPrompts groups related prompts, preferring semantics when a model is
// available and falling back to vocabulary overlap when it is not.
//
// Any failure — no model, a timeout, a wrong-sized response — returns the
// MinHash result. The fallback is silent by design: a local model going away
// should quietly cost some clustering quality, not produce an error in a
// background analysis.
func clusterPrompts(ctx context.Context, emb Embedder, prompts []string, minhashThreshold float64) ([][]int, ClusterMethod) {
	lexical := func() [][]int {
		sigs := make([]signature, len(prompts))
		for i, p := range prompts {
			sigs[i] = minhash(shingle(canonPrompt(p), shingleSize))
		}
		return clusterBySimilarity(sigs, minhashThreshold)
	}

	if emb == nil {
		return lexical(), MethodMinHash
	}

	// Only non-empty prompts are sent. Blank ones cannot cluster meaningfully
	// and would otherwise all embed to the same vector and group together.
	idx := make([]int, 0, len(prompts))
	texts := make([]string, 0, len(prompts))
	for i, p := range prompts {
		if len(canonPrompt(p)) == 0 {
			continue
		}
		idx = append(idx, i)
		texts = append(texts, p)
	}
	if len(texts) < 2 {
		return lexical(), MethodMinHash
	}

	vecs, err := emb.Embed(ctx, texts)
	if err != nil || len(vecs) != len(texts) {
		return lexical(), MethodMinHash
	}
	for _, v := range vecs {
		if len(v) == 0 {
			return lexical(), MethodMinHash
		}
	}

	groups := clusterByCosine(vecs, embedSimilarity)

	// Map back to the caller's indices.
	out := make([][]int, 0, len(groups))
	for _, g := range groups {
		members := make([]int, 0, len(g))
		for _, m := range g {
			members = append(members, idx[m])
		}
		sort.Ints(members)
		out = append(out, members)
	}
	sort.SliceStable(out, func(i, j int) bool { return len(out[i]) > len(out[j]) })
	return out, MethodEmbedding
}

// clusterByCosine mirrors clusterBySimilarity: greedy assignment against
// cluster seeds, so a cluster keeps a representative and the cost stays linear
// in the number of clusters rather than quadratic in the number of items.
func clusterByCosine(vecs [][]float32, threshold float64) [][]int {
	normed := make([][]float32, len(vecs))
	for i, v := range vecs {
		normed[i] = normalize(v)
	}

	type group struct {
		seed    []float32
		members []int
	}
	var groups []group

	for i, v := range normed {
		if v == nil {
			continue
		}
		placed := false
		for g := range groups {
			if dot(groups[g].seed, v) >= threshold {
				groups[g].members = append(groups[g].members, i)
				placed = true
				break
			}
		}
		if !placed {
			groups = append(groups, group{seed: v, members: []int{i}})
		}
	}

	out := make([][]int, 0, len(groups))
	for _, g := range groups {
		if len(g.members) < 2 {
			continue
		}
		out = append(out, g.members)
	}
	return out
}

// normalize returns a unit vector, or nil for a zero vector — which would
// otherwise have cosine similarity 1 with everything and collapse the clusters.
func normalize(v []float32) []float32 {
	var sum float64
	for _, x := range v {
		sum += float64(x) * float64(x)
	}
	if sum == 0 {
		return nil
	}
	n := float32(1 / math.Sqrt(sum))
	out := make([]float32, len(v))
	for i, x := range v {
		out[i] = x * n
	}
	return out
}

func dot(a, b []float32) float64 {
	if len(a) != len(b) {
		return 0
	}
	var sum float64
	for i := range a {
		sum += float64(a[i]) * float64(b[i])
	}
	return sum
}