sift/internal/search/rank.go

package search

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

type Embedder interface {
	Embed(ctx context.Context, inputs []string) ([][]float64, error)
	Acquire(ctx context.Context) (ok bool, reason string)
	Release()
}

func Cosine(a, b []float64) float64 {
	if len(a) == 0 || len(a) != len(b) {
		return 0
	}
	var dot, na, nb float64
	for i := range a {
		dot += a[i] * b[i]
		na += a[i] * a[i]
		nb += b[i] * b[i]
	}
	if na == 0 || nb == 0 {
		return 0
	}
	return dot / (math.Sqrt(na) * math.Sqrt(nb))
}

func rankHits(ctx context.Context, query string, hits []Hit, emb Embedder) (out []Hit, fallback string) {
	if emb == nil || len(hits) == 0 {
		if emb == nil {
			return hits, "no embed"
		}
		return hits, ""
	}
	ok, reason := emb.Acquire(ctx)
	if !ok {
		if reason == "" {
			reason = "queue_full"
		}
		return hits, reason
	}
	defer emb.Release()
	inputs := make([]string, 0, 1+len(hits))
	inputs = append(inputs, query)
	for _, h := range hits {
		s := strings.TrimSpace(h.Title + " " + h.Snippet)
		if s == "" {
			s = h.URL
		}
		inputs = append(inputs, s)
	}
	vecs, err := emb.Embed(ctx, inputs)
	if err != nil || len(vecs) != len(inputs) {
		return hits, "embed error"
	}
	q := vecs[0]
	type pair struct {
		h Hit
		s float64
	}
	ps := make([]pair, len(hits))
	for i := range hits {
		ps[i] = pair{h: hits[i], s: Cosine(q, vecs[i+1])}
		ps[i].h.Score = ps[i].s
	}
	for i := 1; i < len(ps); i++ {
		j := i
		for j > 0 && ps[j].s > ps[j-1].s {
			ps[j], ps[j-1] = ps[j-1], ps[j]
			j--
		}
	}
	out = make([]Hit, len(ps))
	for i := range ps {
		out[i] = ps[i].h
	}
	return out, ""
}

const rrfK = 60

type namedHits struct {
	name string
	hits []Hit
}

func rrfMerge(lists [][]Hit) []Hit {
	type acc struct {
		h Hit
		s float64
	}
	seen := map[string]*acc{}
	var order []string
	for _, list := range lists {
		for r, h := range list {
			orig := h.URL
			c := CanonicalURL(orig)
			if c == "" {
				continue
			}
			sc := 1.0 / float64(rrfK+r+1)
			if a, ok := seen[c]; ok {
				a.s += sc
				keepDate := firstDate(a.h.PublishedDate, h.PublishedDate)
				if len(h.Snippet) > len(a.h.Snippet) {
					src := a.h.Source
					keep := a.h.URL
					a.h = h
					a.h.Source = src
					a.h.URL = keep
				}
				a.h.PublishedDate = keepDate
				continue
			}
			cp := h
			seen[c] = &acc{h: cp, s: sc}
			order = append(order, c)
		}
	}
	out := make([]Hit, 0, len(seen))
	for _, c := range order {
		a := seen[c]
		a.h.Score = a.s
		out = append(out, a.h)
	}
	sort.SliceStable(out, func(i, j int) bool {
		return out[i].Score > out[j].Score
	})
	return out
}

func mergeSearch(e *Engine, lists []namedHits, doRerank bool) []Hit {
	if doRerank && e != nil && e.RerankOn {
		var hold, rest [][]Hit
		for _, l := range lists {
			if e.NoRerank[l.name] {
				hold = append(hold, l.hits)
			} else {
				rest = append(rest, l.hits)
			}
		}
		var out []Hit
		for _, h := range hold {
			out = append(out, h...)
		}
		out = append(out, rrfMerge(rest)...)
		return Dedupe(out)
	}
	var all []Hit
	for _, l := range lists {
		all = append(all, l.hits...)
	}
	all = Dedupe(all)
	sortBySourceOrder(all)
	return all
}