rony-chat-bot/internal/portfolio/hybrid.go

package portfolio

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"log/slog"
	"sort"

	"github.com/VictorVargas/rony-chat-bot/internal/embed"
)

// Embedder is the subset of embed.Client the store needs. Keeping it an
// interface lets Reindex and HybridSearch be tested without an endpoint.
type Embedder interface {
	Embed(ctx context.Context, texts []string) ([][]float32, error)
}

// rrfK is the usual Reciprocal Rank Fusion constant. RRF combines rankings
// rather than scores, which matters here because BM25 (unbounded, negative)
// and cosine similarity (0–1) are not comparable on any common scale. k=60
// is the value from the original paper and is not sensitive enough to be
// worth tuning for a corpus this size.
const rrfK = 60.0

// embedText is the exact string that gets embedded for a chunk, and the only
// definition of it. EmbedChunks and VectorSearch both go through this so the
// hash they compare can never be computed over different text.
//
// The project and heading are prepended because a bare body is ambiguous out
// of context: "Tech stack: Next.js, Stripe" says nothing about which project,
// and the question naming the shop needs to land on it.
func embedText(projectID, section, content string) string {
	return projectID + " — " + section + "\n" + content
}

func hashOf(text string) string {
	sum := sha256.Sum256([]byte(text))
	return hex.EncodeToString(sum[:])
}

// VectorSearch returns the topK chunks whose embedding is closest to the
// query, scanning every stored vector. Chunks indexed before embeddings were
// enabled simply have no row and are skipped.
func (s *Store) VectorSearch(ctx context.Context, queryVec []float32, topK int) ([]SearchResult, error) {
	if len(queryVec) == 0 {
		return nil, nil
	}
	if topK <= 0 {
		topK = 5
	}

	rows, err := s.db.QueryContext(ctx, `
		SELECT c.id, c.project_id, c.kind, c.source_file, c.section, c.chunk_index, c.content,
		       v.vec, v.content_hash
		FROM portfolio_vectors v
		JOIN portfolio_chunks c ON c.id = v.chunk_id
		WHERE v.dim = ? AND c.section <> 'frontmatter'`, len(queryVec))
	if err != nil {
		return nil, fmt.Errorf("vector search: %w", err)
	}
	defer rows.Close()

	var hits []SearchResult
	var stale int
	for rows.Next() {
		var r SearchResult
		var blob []byte
		var storedHash string
		if err := rows.Scan(&r.ID, &r.ProjectID, &r.Kind, &r.SourceFile, &r.Section, &r.Index,
			&r.Content, &blob, &storedHash); err != nil {
			return nil, err
		}
		// Skip vectors describing a previous version of this chunk. Chunk
		// ids survive body edits, so without this the row would be scored
		// against text that is no longer there — quietly, with no error.
		// The chunk stays reachable through keyword search meanwhile.
		if storedHash != hashOf(embedText(r.ProjectID, r.Section, r.Content)) {
			stale++
			continue
		}
		r.Score = embed.Similarity(queryVec, embed.Decode(blob))
		hits = append(hits, r)
	}
	if stale > 0 {
		slog.Warn("ignoring stale embeddings; re-run reindex to refresh them", "chunks", stale)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	// Higher similarity is better, unlike the bm25() score used by Search.
	sort.Slice(hits, func(i, j int) bool { return hits[i].Score > hits[j].Score })
	return hits[:min(topK, len(hits))], nil
}

// HybridSearch fuses keyword and vector retrieval with Reciprocal Rank
// Fusion and returns the topK chunks.
//
// Neither retriever is sufficient alone on this corpus. Keyword search finds
// exact proper nouns ("Kubernetes", "Stripe") that embeddings can blur, but
// returns nothing when a Spanish question meets an English document. Vectors
// bridge the languages but rank a request for a specific rare token less
// sharply. RRF asks only for each side's ordering, so a chunk that both
// retrievers like rises above one that only a single retriever loved.
//
// When embedding fails — the endpoint is down, or none is configured — this
// degrades to plain keyword search rather than failing the request: a
// keyword-only answer beats no answer.
func (s *Store) HybridSearch(ctx context.Context, embedder Embedder, query string, topK int) ([]SearchResult, error) {
	if topK <= 0 {
		topK = 5
	}
	// Over-fetch from each side: a chunk ranked 8th by one retriever and 2nd
	// by the other should still be able to win the fusion.
	pool := topK * 3

	keyword, kwErr := s.Search(ctx, query, pool)
	if embedder == nil {
		return keyword, kwErr
	}

	vecs, err := embedder.Embed(ctx, []string{query})
	if err != nil || len(vecs) == 0 {
		if kwErr != nil {
			return nil, fmt.Errorf("both retrievers failed: keyword: %v; embedding: %w", kwErr, err)
		}
		return keyword, nil
	}
	semantic, vecErr := s.VectorSearch(ctx, vecs[0], pool)
	if vecErr != nil && kwErr != nil {
		return nil, fmt.Errorf("both retrievers failed: keyword: %v; vector: %w", kwErr, vecErr)
	}

	type fused struct {
		res   SearchResult
		score float64
	}
	byID := map[string]*fused{}
	add := func(list []SearchResult) {
		for rank, r := range list {
			f, ok := byID[r.ID]
			if !ok {
				f = &fused{res: r}
				byID[r.ID] = f
			}
			f.score += 1 / (rrfK + float64(rank+1))
		}
	}
	add(keyword)
	add(semantic)

	out := make([]fused, 0, len(byID))
	for _, f := range byID {
		out = append(out, *f)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].score != out[j].score {
			return out[i].score > out[j].score
		}
		return out[i].res.ID < out[j].res.ID // stable across runs
	})

	results := make([]SearchResult, 0, min(topK, len(out)))
	for _, f := range out[:min(topK, len(out))] {
		f.res.Score = f.score
		results = append(results, f.res)
	}
	return results, nil
}

// EmbedChunks computes and stores a vector for every indexed chunk. Called
// after Reindex, since it needs the chunks to exist.
//
// The heading is prepended to each chunk's text: a bare "Tech stack" body is
// a list of technologies with no hint of which project it belongs to, and the
// embedding of "tienda-ropa — Tech stack: Next.js, Stripe…" is much closer to
// a question naming the shop.
func (s *Store) EmbedChunks(ctx context.Context, embedder Embedder) (int, error) {
	if embedder == nil {
		return 0, nil
	}
	rows, err := s.db.QueryContext(ctx,
		`SELECT id, project_id, section, content FROM portfolio_chunks ORDER BY id`)
	if err != nil {
		return 0, fmt.Errorf("read chunks for embedding: %w", err)
	}
	var ids, texts []string
	for rows.Next() {
		var id, project, section, content string
		if err := rows.Scan(&id, &project, &section, &content); err != nil {
			rows.Close()
			return 0, err
		}
		ids = append(ids, id)
		texts = append(texts, embedText(project, section, content))
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return 0, err
	}
	if len(ids) == 0 {
		return 0, nil
	}

	vecs, err := embedder.Embed(ctx, texts)
	if err != nil {
		return 0, fmt.Errorf("embed chunks: %w", err)
	}
	if len(vecs) != len(ids) {
		return 0, fmt.Errorf("embedder returned %d vectors for %d chunks", len(vecs), len(ids))
	}

	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()
	if _, err := tx.ExecContext(ctx, `DELETE FROM portfolio_vectors`); err != nil {
		return 0, err
	}
	stmt, err := tx.PrepareContext(ctx,
		`INSERT INTO portfolio_vectors (chunk_id, content_hash, dim, vec) VALUES (?,?,?,?)`)
	if err != nil {
		return 0, err
	}
	defer stmt.Close()
	for i, id := range ids {
		if _, err := stmt.ExecContext(ctx, id, hashOf(texts[i]), len(vecs[i]), embed.Encode(vecs[i])); err != nil {
			return 0, fmt.Errorf("store vector %s: %w", id, err)
		}
	}
	if err := tx.Commit(); err != nil {
		return 0, err
	}
	return len(ids), nil
}