sift/internal/digest/retrieval.go

package digest

import (
	"context"
	"strings"

	"git.wntrmute.dev/kyle/sift/internal/dcache"
)

// hybridOrBM25 scores a chunked document either hybrid (cosine + lexical overlap) or, when there
// are too many uncached chunks to embed inline, lexical-only BM25. mode is "hybrid" or "bm25".
func (e *Engine) hybridOrBM25(ctx context.Context, chunks []string, question, contentHash, callerClass string) (string, int, int, string, int, error) {
	key := func(i int) string {
		return dcache.EmbedKey(contentHash+"/sp1200", i, e.embedModel(), callerClass)
	}

	if e.Cache == nil {
		if len(chunks) <= EmbedSyncMax {
			vecs, err := e.embed(ctx, append([]string{question}, capEmbed(chunks)...))
			if err != nil {
				return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "hybrid", 0, err
			}
			chunkVecs := make([][]float64, len(chunks))
			for i := range chunks {
				chunkVecs[i] = vecs[i+1]
			}
			body, used := greedyBudget(chunks, hybridScores(chunks, question, vecs[0], chunkVecs), SmallDoc)
			return body, used, len(chunks), "hybrid", len(chunks), nil
		}
		body, used := greedyBudget(chunks, bm25Scores(chunks, question), SmallDoc)
		return body, used, len(chunks), "bm25", 0, nil
	}

	// Cache present: look up every chunk's vector.
	var cached []int
	var cachedVecs [][]float64
	var uncached []int
	for i := range chunks {
		if v, ok := e.Cache.GetEmbed(key(i)); ok {
			cached = append(cached, i)
			cachedVecs = append(cachedVecs, v)
			continue
		}
		uncached = append(uncached, i)
	}

	if len(uncached) <= EmbedSyncMax {
		in := append([]string{question}, capEmbed(uncachedChunks(chunks, uncached))...)
		vecs, err := e.embed(ctx, in)
		if err != nil {
			return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "hybrid", len(cached), err
		}
		for j, i := range uncached {
			if len(vecs[j+1]) > 0 { // a vector the server left out is not cached as empty
				e.Cache.PutEmbed(key(i), vecs[j+1])
			}
		}
		chunkVecs := make([][]float64, len(chunks))
		for k, i := range cached {
			chunkVecs[i] = cachedVecs[k]
		}
		for j, i := range uncached {
			chunkVecs[i] = vecs[j+1]
		}
		body, used := greedyBudget(chunks, hybridScores(chunks, question, vecs[0], chunkVecs), SmallDoc)
		return body, used, len(chunks), "hybrid", len(chunks), nil
	}

	body, used := greedyBudget(chunks, bm25Scores(chunks, question), SmallDoc)
	e.queueForEmbed(chunks, contentHash, callerClass)
	return body, used, len(chunks), "bm25", len(cached), nil
}

func uncachedChunks(chunks []string, uncached []int) []string {
	out := make([]string, len(uncached))
	for j, i := range uncached {
		out[j] = chunks[i]
	}
	return out
}

// hybridScores builds the hybrid scorer 0.7*cosine(q, v_i) + 0.3*lexicalOverlap(question, chunk),
// where chunkVecs[i] is chunk i's vector.
func hybridScores(chunks []string, question string, q []float64, chunkVecs [][]float64) func(int) float64 {
	return func(i int) float64 {
		return 0.7*cosine(q, chunkVecs[i]) + 0.3*lexicalOverlap(question, chunks[i])
	}
}

// bm25Scores builds the BM25 scorer, normalised by its maximum so scores land in [0, 1].
func bm25Scores(chunks []string, question string) func(int) float64 {
	scores := NewBM25(chunks).Scores(question)
	var max float64
	for _, s := range scores {
		if s > max {
			max = s
		}
	}
	if max == 0 {
		return func(int) float64 { return 0 }
	}
	return func(i int) float64 { return scores[i] / max }
}

// greedyBudget picks chunk 0 first, then the highest-scoring remaining chunks while the total
// stays within budget tokens, and returns them in document order with the count used.
func greedyBudget(chunks []string, score func(int) float64, budget int) (string, int) {
	type scored struct {
		i int
		s float64
	}
	var ss []scored
	for i := range chunks {
		ss = append(ss, scored{i: i, s: score(i)})
	}
	picked := map[int]struct{}{}
	used := 0
	if len(chunks) > 0 {
		picked[0] = struct{}{}
		used = EstTokens(chunks[0])
	}
	for len(picked) < len(ss) {
		best, bi := -1.0, -1
		for _, x := range ss {
			if _, ok := picked[x.i]; ok {
				continue
			}
			if x.s > best {
				best, bi = x.s, x.i
			}
		}
		if bi < 0 {
			break
		}
		need := EstTokens(chunks[bi])
		if used+need > budget && used > 0 {
			break
		}
		picked[bi] = struct{}{}
		used += need
	}
	var ordered []string
	for i := range chunks {
		if _, ok := picked[i]; ok {
			ordered = append(ordered, chunks[i])
		}
	}
	return strings.Join(ordered, "\n\n"), len(ordered)
}