agent build
This commit is contained in:
474
go-api/internal/knowledge/retrieve.go
Normal file
474
go-api/internal/knowledge/retrieve.go
Normal file
@@ -0,0 +1,474 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user