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

package knowledge

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"github.com/krow/krow-backend/go-api/internal/authctx"
	"github.com/krow/krow-backend/go-api/internal/repo"
)

// Retrieval. The one place I1 and I2 are either kept or broken.
//
// §5 asks for three things and this file does exactly those three:
//
//  1. **Hybrid.** Dense and BM25, fused with RRF. Not dense-only "for
//     convenience" — semantic search is bad at exact terms, and a workforce
//     corpus is full of them: a shift code, a certification name, a venue. Not
//     keyword-only either, which is the failure this deployment could most
//     easily have shipped, having no embedding credential.
//  2. **Permission as a pre-filter.** The same predicate goes into BOTH
//     queries' WHERE clauses. This is I2, and it is not a style choice: rank
//     first and drop afterwards and the forbidden rows leak through the shape
//     of what is left — a short result set, a top-3 with a hole in it, a
//     confidence that tracks documents the caller cannot see.
//  3. **Citable results.** Every chunk comes back with the ids needed to point
//     at it, so a response can cite and the surface can link.
//
// There is exactly one exported entry point and it will not run without a
// principal. §5: "any retrieval function that accepts a query but not a caller
// principal is wrong by construction."

// RRFConstant is the k in RRF's 1/(k + rank).
//
// 60 is the value from the original paper and the one nearly everything uses.
// It is a flattener: with k=60 the gap between rank 1 and rank 2 is small, so a
// document both retrievers rank moderately well beats one that a single
// retriever loves. That is the entire point of fusing — agreement across two
// different notions of relevance is a stronger signal than a high score in one.
const RRFConstant = 60.0

// CandidateMultiple is how many rows each retriever fetches relative to k.
//
// Fusion needs depth: a chunk ranked 8th by keyword and 9th by vector should
// win over one ranked 1st by keyword and nowhere by vector, and it cannot if
// both lists were cut at 5. Three times k is the usual compromise between that
// and reading rows nobody will see.
const CandidateMultiple = 3

// DefaultK is how many chunks a retrieval returns when the caller does not say.
const DefaultK = 8

// MaxK is the ceiling. Not a performance guard — a context guard. Retrieved
// text is prompt, prompt is money, and a caller asking for 500 chunks has made
// a mistake this should not silently honour.
const MaxK = 50

// Query is a retrieval request.
//
// Principal is a field rather than an argument so it cannot be defaulted, and
// Retrieve refuses a zero one. That is the structural half of §5's rule; the
// other half is that this package exports no other way to search.
type Query struct {
	// Text is what to search for. The caller's words, or the model's — either
	// way untrusted, and it reaches SQL only as a bind parameter.
	Text string

	// Principal is who is asking. Required.
	Principal authctx.Identity

	// Sources are the corpora this agent's spec declares. Required: an empty
	// list is not "everything", it is a spec that named no knowledge, and the
	// correct response to it is no results rather than the whole index.
	Sources []string

	K int
}

// Result is one retrieved chunk, with everything needed to cite it.
type Result struct {
	// ChunkID is the citation's address. §5: retrieved chunks flow to the model
	// WITH source ids, so a response can cite and an unsupported claim can be
	// told apart from a grounded one.
	ChunkID    string `json:"chunkId"`
	DocumentID string `json:"documentId"`

	Source  string `json:"source"`
	Title   string `json:"title"`
	URI     string `json:"uri,omitempty"`
	Heading string `json:"heading,omitempty"`
	Ordinal int    `json:"ordinal"`

	Text string `json:"text"`

	// Score is the fused RRF score. Comparable within one result set and
	// meaningless outside it — RRF scores are ranks, not similarities, so a
	// 0.03 here is not "3% relevant" and must never be shown as a percentage.
	Score float64 `json:"score"`

	// DenseRank and SparseRank are where each retriever placed this chunk, or 0
	// for "not in that list at all". Kept because they are the only way to
	// debug a bad retrieval: a result with a good dense rank and no sparse rank
	// is a semantic match with no shared vocabulary, which is either the system
	// working or the system hallucinating a connection, and you cannot tell
	// which without seeing both.
	DenseRank  int `json:"denseRank,omitempty"`
	SparseRank int `json:"sparseRank,omitempty"`

	TokenEstimate int `json:"tokenEstimate"`
}

// Results is a retrieval's outcome.
type Results struct {
	Chunks []Result `json:"chunks"`

	// DenseSkipped says the vector half did not run, and why. Surfaced rather
	// than hidden: a retrieval that quietly degraded to keyword-only answers
	// worse in a way that looks like the model getting dumber.
	DenseSkipped string `json:"denseSkipped,omitempty"`

	TotalTokens int `json:"totalTokens"`
}

// Retriever searches the index on a caller's behalf.
type Retriever struct {
	db       repo.Querier
	embedder Embedder
}

// NewRetriever builds a retriever. A nil embedder means keyword-only, reported
// on every result rather than silently.
func NewRetriever(db repo.Querier, e Embedder) *Retriever {
	return &Retriever{db: db, embedder: e}
}

// Retrieve searches, permissioned.
//
// The only exported search in this package, and it takes a principal. There is
// no convenience overload, there is no package-level helper, and there is no
// unexported one a future call site could reach for — everything below takes
// the grants as an argument it cannot construct itself.
func (r *Retriever) Retrieve(ctx context.Context, q Query) (*Results, error) {
	// The grants ARE the permission. A caller this platform does not recognise
	// — no tenant, an unlisted role — produces nil, and nil matches no row,
	// so an unknown caller retrieves nothing rather than being special-cased.
	grants := GrantsFor(q.Principal)
	if len(grants) == 0 {
		return nil, &Error{
			Code:    ErrNoPrincipal,
			Message: "retrieval needs a caller with a tenant and a recognised role",
		}
	}
	if len(q.Sources) == 0 {
		return nil, &Error{
			Code: ErrNoSources,
			Message: "retrieval needs the sources the agent's spec declares; " +
				"an empty list is a spec that named no knowledge, not permission to read all of it",
		}
	}

	text := strings.TrimSpace(q.Text)
	if text == "" {
		return &Results{Chunks: []Result{}}, nil
	}

	k := q.K
	if k <= 0 {
		k = DefaultK
	}
	if k > MaxK {
		k = MaxK
	}
	depth := k * CandidateMultiple

	// Both halves run against the same pre-filtered set. Built once so the two
	// queries cannot drift — a permission predicate that is right in one query
	// and subtly wrong in the other is the exact bug this whole file is
	// arranged to prevent.
	scope := scopeArgs{
		orgID:   q.Principal.OrgID,
		grants:  grants,
		sources: q.Sources,
	}

	sparse, err := r.sparse(ctx, scope, text, depth)
	if err != nil {
		return nil, err
	}

	dense, skipped, err := r.dense(ctx, scope, text, depth)
	if err != nil {
		return nil, err
	}

	fused := fuse(dense, sparse, k)

	out := &Results{Chunks: fused, DenseSkipped: skipped}
	for _, c := range fused {
		out.TotalTokens += c.TokenEstimate
	}
	if out.Chunks == nil {
		out.Chunks = []Result{}
	}
	return out, nil
}

/* ── The pre-filter ─────────────────────────────────────────────────────── */

// scopeArgs is the permission predicate, as parameters.
//
// Rendered identically into both queries. The three conditions are not
// interchangeable and all three are load-bearing:
//
//	org_id = $1     I5. Tenancy, never optional, never a wildcard.
//	acl && $2       I1. The caller's own grants. A chunk with no overlapping
//	                tag is not fetched, so it cannot influence a count, a rank
//	                or a summary.
//	source = ANY($3) The agent's declared corpora. An agent granted the policy
//	                library does not thereby gain the incident log.
type scopeArgs struct {
	orgID   string
	grants  []string
	sources []string
}

// where renders the predicate and its parameters.
//
// Returns SQL with $1..$3 fixed at the front, so each query appends its own
// parameters after and there is no arithmetic to get wrong.
func (s scopeArgs) where(alias string) (string, []any) {
	c := func(col string) string {
		if alias == "" {
			return col
		}
		return alias + "." + col
	}
	predicate := fmt.Sprintf(
		"%s = $1::uuid AND %s && $2::text[] AND %s = ANY($3::text[])",
		c("org_id"), c("acl"), c("source"))
	return predicate, []any{s.orgID, s.grants, s.sources}
}

/* ── The keyword half ───────────────────────────────────────────────────── */

// sparse is the BM25-ish half: Postgres full-text ranking.
//
// `ts_rank_cd` is cover-density ranking, not textbook BM25 — Postgres does not
// ship BM25 — and the difference is worth naming rather than glossing. Both
// reward term frequency and rarity; cover density additionally rewards the
// query's terms appearing CLOSE TOGETHER, which for a policy corpus is usually
// what you want. It is not the same function, and a benchmark that assumes BM25
// will not reproduce exactly.
//
// THE `&` → `|` SUBSTITUTION IS NOT A HACK, IT IS THE POINT.
//
// `websearch_to_tsquery` joins terms with AND: "lateness policy supervisor"
// becomes 'late' & 'polici' & 'supervisor' and matches only a chunk containing
// all three. That is correct for a site search box and wrong for retrieval. A
// question is a bag of words, one of which is usually the rare, discriminating
// one — and under AND, adding that rare word to a query makes the result set
// EMPTY rather than better. Every retrieval system that works ORs its terms and
// lets the ranking function sort out which matches are good.
//
// So the query is parsed by websearch_to_tsquery — which keeps quoted phrases
// as `<->` operators, and never raises on malformed input, which matters when
// the string comes from a model — and its AND operators are then rewritten to
// OR. The phrase operators survive the rewrite untouched.
//
// The one thing lost is negation: `-term` would become `| !term`, which matches
// every chunk that lacks the term, i.e. almost all of them. Hyphens are
// stripped before parsing so a negation cannot be expressed at all. That is a
// deliberate trade — a search operator nobody asked for, against a failure mode
// that turns a query inside out.
func (r *Retriever) sparse(ctx context.Context, scope scopeArgs, text string, depth int) ([]Result, error) {
	predicate, args := scope.where("c")
	args = append(args, stripNegation(text), depth)

	rows, err := r.db.Query(ctx, `
		WITH q AS (
		  SELECT replace(websearch_to_tsquery('english', $4)::text, '&', '|')::tsquery AS query
		)
		SELECT c.id::text, c.document_id::text, c.source, d.title, d.uri,
		       c.heading, c.ordinal, c.text, c.token_estimate
		  FROM knowledge_chunks c
		  JOIN knowledge_documents d ON d.id = c.document_id
		  CROSS JOIN q
		 WHERE `+predicate+`
		   AND q.query IS NOT NULL
		   AND c.tsv @@ q.query
		 ORDER BY ts_rank_cd(c.tsv, q.query) DESC, c.id
		 LIMIT $5`, args...)
	if err != nil {
		return nil, &Error{Code: ErrRetrieveFailed, Message: "the keyword search failed", Cause: err}
	}
	defer rows.Close()

	return scanResults(rows)
}

// stripNegation removes the `-term` operator from a query string.
//
// See the note on sparse: rewriting AND to OR turns a negation into a match on
// almost everything. Removing the operator before parsing is the smaller loss.
// Hyphens INSIDE a word ("part-time") are left alone — only a leading one is an
// operator.
func stripNegation(text string) string {
	fields := strings.Fields(text)
	for i, f := range fields {
		fields[i] = strings.TrimLeft(f, "-")
	}
	return strings.Join(fields, " ")
}

/* ── The dense half ─────────────────────────────────────────────────────── */

// dense is the vector half.
//
// Returns a reason rather than an error when it cannot run. A missing embedding
// credential, a provider outage or an unembedded corpus are all cases where
// keyword-only results are far better than no results — but the caller is TOLD,
// because a retrieval that silently halved its own recall presents as the agent
// getting worse for no reason anybody can find.
func (r *Retriever) dense(ctx context.Context, scope scopeArgs, text string, depth int) ([]Result, string, error) {
	if r.embedder == nil {
		return nil, "no embedder is configured; these results are keyword-only", nil
	}

	vectors, err := r.embedder.Embed(ctx, []string{text}, KindQuery)
	if err != nil || len(vectors) == 0 || len(vectors[0]) == 0 {
		// Degraded, not failed. A knowledge error here would take the whole run
		// with it over a provider hiccup.
		reason := "the embedding service was unavailable; these results are keyword-only"
		var kErr *Error
		if ok := asKnowledgeError(err, &kErr); ok && kErr.Code == ErrNotConfigured {
			reason = "no embedding credential is configured; these results are keyword-only"
		}
		return nil, reason, nil
	}

	predicate, args := scope.where("c")
	args = append(args, vectors[0], r.embedder.Model(), depth)

	// The pre-filter and the model check are both in the WHERE clause, so the
	// dot product is only ever computed over rows this caller may read and
	// vectors that are comparable to the query's. Scoring first and filtering
	// after would be I2's violation AND a wasted scan.
	//
	// `embedding IS NOT NULL` matters: knowledge_dot is STRICT, so an unembedded
	// chunk scores NULL, and NULL sorts first under DESC. Without this the top
	// of every dense ranking would be the chunks that have no vector at all.
	rows, err := r.db.Query(ctx, `
		SELECT c.id::text, c.document_id::text, c.source, d.title, d.uri,
		       c.heading, c.ordinal, c.text, c.token_estimate
		  FROM knowledge_chunks c
		  JOIN knowledge_documents d ON d.id = c.document_id
		 WHERE `+predicate+`
		   AND c.embedding IS NOT NULL
		   AND c.embedding_model = $5
		 ORDER BY knowledge_dot(c.embedding, $4::real[]) DESC, c.id
		 LIMIT $6`, args...)
	if err != nil {
		return nil, "", &Error{Code: ErrRetrieveFailed, Message: "the vector search failed", Cause: err}
	}
	defer rows.Close()

	out, scanErr := scanResults(rows)
	if scanErr != nil {
		return nil, "", scanErr
	}
	if len(out) == 0 {
		// Distinguishable from "the provider is down": the corpus itself has no
		// vectors for this model, which is a backfill nobody has run.
		return nil, "", nil
	}
	return out, "", nil
}

/* ── Fusion ─────────────────────────────────────────────────────────────── */

// fuse combines two rankings with Reciprocal Rank Fusion.
//
// RRF scores a document 1/(k + rank) in each list and sums. It uses only the
// RANKS, never the underlying scores, and that is exactly why it is the right
// choice here: `ts_rank_cd` returns a small unbounded float and cosine
// similarity returns [-1, 1]. Any scheme that combined those numbers directly
// would need normalisation, and every normalisation is a tuning parameter that
// drifts as the corpus changes. Ranks need no scale.
//
// A chunk in only one list still scores — it simply gets one term instead of
// two, which is the correct treatment of "one retriever found this and the
// other did not".
func fuse(dense, sparse []Result, k int) []Result {
	type entry struct {
		result Result
		score  float64
	}
	merged := map[string]*entry{}

	add := func(list []Result, isDense bool) {
		for i, res := range list {
			rank := i + 1
			e, ok := merged[res.ChunkID]
			if !ok {
				e = &entry{result: res}
				merged[res.ChunkID] = e
			}
			e.score += 1.0 / (RRFConstant + float64(rank))
			if isDense {
				e.result.DenseRank = rank
			} else {
				e.result.SparseRank = rank
			}
		}
	}
	add(dense, true)
	add(sparse, false)

	out := make([]Result, 0, len(merged))
	for _, e := range merged {
		e.result.Score = e.score
		out = append(out, e.result)
	}

	// Ties broken by chunk id, so the same query over the same corpus returns
	// the same order. A retrieval whose ordering wobbles between identical
	// calls makes every downstream difference impossible to attribute.
	sort.Slice(out, func(a, b int) bool {
		if out[a].Score != out[b].Score {
			return out[a].Score > out[b].Score
		}
		return out[a].ChunkID < out[b].ChunkID
	})

	if len(out) > k {
		out = out[:k]
	}
	return out
}

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

func scanResults(rows interface {
	Next() bool
	Scan(...any) error
	Err() error
}) ([]Result, error) {
	var out []Result
	for rows.Next() {
		var res Result
		if err := rows.Scan(&res.ChunkID, &res.DocumentID, &res.Source, &res.Title,
			&res.URI, &res.Heading, &res.Ordinal, &res.Text, &res.TokenEstimate); err != nil {
			return nil, &Error{Code: ErrRetrieveFailed, Message: "a result could not be read", Cause: err}
		}
		out = append(out, res)
	}
	if err := rows.Err(); err != nil {
		return nil, &Error{Code: ErrRetrieveFailed, Message: "the results could not be read", Cause: err}
	}
	return out, nil
}

// asKnowledgeError is errors.As for this package's error, without importing
// errors into every call site's line of sight.
func asKnowledgeError(err error, target **Error) bool {
	if err == nil {
		return false
	}
	if e, ok := err.(*Error); ok {
		*target = e
		return true
	}
	return false
}