llmgw/internal/store/knowledge.go

package store

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
	"unicode"
)

// Knowledge bases (BACKLOG F): collections of documents, split into chunks
// and embedded, that a request is answered from.
//
// Which tenant may read which chunk is decided twice. The queries here filter
// by organization and by the collections the caller may read, and the table
// enforces the same with row-level security against two settings each
// transaction makes — so a query that forgot its filter returns nothing
// rather than another tenant's passages. The role llmgw connects as owns the
// table, so the policy is FORCEd; a superuser would bypass it, which is one
// more reason the store never connects as one.

// ErrKnowledgeOff is a store that cannot hold knowledge bases.
var ErrKnowledgeOff = errors.New("knowledge bases are off")

// MaxEmbeddingDims is the most dimensions pgvector can index (HNSW on vector).
const MaxEmbeddingDims = 2000

// KBPerm is what a grant allows.
type KBPerm string

// The permissions. Write includes read.
const (
	KBRead  KBPerm = "read"
	KBWrite KBPerm = "write"
)

// KBCollection is a set of documents with one owner and one embedding model.
type KBCollection struct {
	ID         string
	OrgID      string
	OwnerScope Scope
	OwnerID    string
	Name       string
	EmbedModel string
	Dims       int
	CreatedAt  time.Time
	CreatedBy  string
}

// KBGrant lets a scope, and everything below it, read or write a collection.
type KBGrant struct {
	Scope      Scope
	ScopeID    string
	Collection string
	Perm       KBPerm
	AddedAt    time.Time
	AddedBy    string
}

// KBDocument is a document as added to a collection.
type KBDocument struct {
	ID         string
	Collection string
	Name       string
	Digest     string
	Bytes      int64
	Chunks     int
	AddedAt    time.Time
	AddedBy    string
}

// KBChunk is a piece of a document and its embedding.
type KBChunk struct {
	Body      string
	Embedding []float32
}

// KBHit is a chunk a search returned, best first.
type KBHit struct {
	Collection string
	DocID      string
	Document   string
	Ord        int
	Body       string
	// Score is the fused rank (reciprocal rank fusion over the vector and
	// text searches): only the order means anything.
	Score float64
}

// KnowledgeAvailable says whether this store can hold knowledge bases, and
// why not when it cannot.
func (s *DB) KnowledgeAvailable() error {
	if s.dialect != PostgresDialect {
		return fmt.Errorf("%w: they need the PostgreSQL store (with pgvector), and this one is SQLite", ErrKnowledgeOff)
	}
	return s.kbErr
}

// ensureKnowledge creates kb_chunks when the vector extension allows it.
//
// At open rather than in a migration: a migration that needs an extension the
// database servers do not have would stop the control plane from starting,
// and one that skipped it would be recorded as applied and never run again
// once the extension arrived.
func (s *DB) ensureKnowledge(ctx context.Context) error {
	var have bool
	if err := s.db.QueryRowContext(ctx,
		`SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector')`).Scan(&have); err != nil {
		return fmt.Errorf("%w: reading the database's extensions: %v", ErrKnowledgeOff, err)
	}
	if !have {
		var available bool
		if err := s.db.QueryRowContext(ctx,
			`SELECT EXISTS (SELECT 1 FROM pg_available_extensions WHERE name = 'vector')`).Scan(&available); err != nil {
			return fmt.Errorf("%w: reading the database's extensions: %v", ErrKnowledgeOff, err)
		}
		if !available {
			return fmt.Errorf("%w: pgvector is not installed on the database servers "+
				"(package postgresql-<version>-pgvector on every one; llmgw's database setup installs it)", ErrKnowledgeOff)
		}
		// Works only where the extension is trusted or llmgw's role may
		// create it; elsewhere a superuser has to, once.
		if _, err := s.db.ExecContext(ctx, `CREATE EXTENSION IF NOT EXISTS vector`); err != nil {
			return fmt.Errorf("%w: pgvector is installed but not enabled in llmgw's database; "+
				"as postgres, run CREATE EXTENSION vector in it (%v)", ErrKnowledgeOff, err)
		}
	}
	var exists bool
	if err := s.db.QueryRowContext(ctx, `SELECT to_regclass('kb_chunks') IS NOT NULL`).Scan(&exists); err != nil {
		return fmt.Errorf("%w: %v", ErrKnowledgeOff, err)
	}
	if exists {
		return nil
	}
	err := s.inTx(ctx, func(c conn) error {
		for _, q := range []string{
			// Partitioned by organization: each one's indexes stay its own
			// size, and an organization's chunks go with one DROP.
			`CREATE TABLE kb_chunks (
			  org_id     TEXT NOT NULL,
			  collection TEXT NOT NULL,
			  doc_id     TEXT NOT NULL REFERENCES kb_documents(id) ON DELETE CASCADE,
			  ord        INTEGER NOT NULL,
			  body       TEXT NOT NULL,
			  tsv        tsvector GENERATED ALWAYS AS (to_tsvector('simple', body)) STORED,
			  -- No dimension here: collections differ. Each collection's
			  -- index casts to its own (kbVectorIndex).
			  embedding  vector NOT NULL,
			  PRIMARY KEY (org_id, doc_id, ord)
			) PARTITION BY LIST (org_id)`,
			`CREATE INDEX kb_chunks_tsv ON kb_chunks USING gin (tsv)`,
			`ALTER TABLE kb_chunks ENABLE ROW LEVEL SECURITY`,
			`ALTER TABLE kb_chunks FORCE ROW LEVEL SECURITY`,
			`CREATE POLICY kb_scope ON kb_chunks
			  USING (org_id = current_setting('llmgw.kb_org', true)
			     AND collection = ANY (string_to_array(current_setting('llmgw.kb_collections', true), ',')))
			  WITH CHECK (org_id = current_setting('llmgw.kb_org', true)
			     AND collection = ANY (string_to_array(current_setting('llmgw.kb_collections', true), ',')))`,
		} {
			if _, err := c.ExecContext(ctx, q); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		// Another control plane opening at the same moment made it first.
		if e2 := s.db.QueryRowContext(ctx, `SELECT to_regclass('kb_chunks') IS NOT NULL`).Scan(&exists); e2 == nil && exists {
			return nil
		}
		return fmt.Errorf("%w: creating kb_chunks: %v", ErrKnowledgeOff, err)
	}
	return nil
}

// kbIDPattern is what a collection id looks like (NewID("kb_")); ids are put
// into DDL and partial-index predicates, so nothing else is let through.
var kbIDPattern = regexp.MustCompile(`^kb_[a-z2-7]{16}$`)

// kbPartition is the chunk partition of an organization. Hashed, because old
// organization ids were slugs of names and could hold anything.
func kbPartition(orgID string) string {
	h := sha256.Sum256([]byte(orgID))
	return "kb_chunks_" + hex.EncodeToString(h[:12])
}

// kbVectorIndex is the name of a collection's vector index.
func kbVectorIndex(collection string) string {
	return "kb_vec_" + strings.TrimPrefix(collection, "kb_")
}

func sqlLiteral(v string) string { return "'" + strings.ReplaceAll(v, "'", "''") + "'" }

// kbScope sets, for this transaction, which rows of kb_chunks it may touch.
func kbScope(ctx context.Context, c conn, orgID string, collections []string) error {
	_, err := c.ExecContext(ctx, `SELECT set_config('llmgw.kb_org', ?, true), set_config('llmgw.kb_collections', ?, true)`,
		orgID, strings.Join(collections, ","))
	return err
}

// orgOfScope is the organization a scope sits in. A person has none of their
// own, so a person is checked against org instead: they must be a member of
// it, or of something in it.
func (s *DB) orgOfScope(ctx context.Context, c conn, scope Scope, id, org string) (string, error) {
	var q string
	switch scope {
	case ScopeOrg:
		q = `SELECT id FROM organizations WHERE id = ?`
	case ScopeTeam:
		q = `SELECT org_id FROM teams WHERE id = ?`
	case ScopeProject:
		q = `SELECT t.org_id FROM projects p JOIN teams t ON t.id = p.team_id WHERE p.id = ?`
	case ScopeKey:
		q = `SELECT t.org_id FROM api_keys k JOIN projects p ON p.id = k.project_id
			JOIN teams t ON t.id = p.team_id WHERE k.id = ?`
	case ScopeUser:
		var n int
		err := c.QueryRowContext(ctx, `SELECT COUNT(*) FROM memberships m WHERE m.user_id = ? AND (
			  (m.scope_type = 'org' AND m.scope_id = ?)
			  OR (m.scope_type = 'team' AND m.scope_id IN (SELECT id FROM teams WHERE org_id = ?))
			  OR (m.scope_type = 'project' AND m.scope_id IN (
			       SELECT p.id FROM projects p JOIN teams t ON t.id = p.team_id WHERE t.org_id = ?)))`,
			id, org, org, org).Scan(&n)
		if err != nil {
			return "", err
		}
		if n == 0 {
			return "", fmt.Errorf("%s is not a member of anything in that organization", id)
		}
		return org, nil
	default:
		return "", fmt.Errorf("a collection cannot be owned by or granted to a %s", scope)
	}
	var got string
	err := c.QueryRowContext(ctx, q, id).Scan(&got)
	if errors.Is(err, sql.ErrNoRows) {
		if scope == ScopeKey {
			return "", fmt.Errorf("no key %s in a project (a personal key reads no collection)", id)
		}
		return "", fmt.Errorf("no %s %s: %w", scope, id, ErrNotFound)
	}
	return got, err
}

// CreateCollection makes an empty collection, its organization's chunk
// partition when this is the first, and its vector index. c.OrgID is needed
// only when a person owns it; otherwise it comes from the owner.
func (s *DB) CreateCollection(ctx context.Context, c KBCollection) (KBCollection, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return KBCollection{}, err
	}
	c.Name = strings.TrimSpace(c.Name)
	switch {
	case c.Name == "":
		return KBCollection{}, errors.New("a collection needs a name")
	case c.EmbedModel == "":
		return KBCollection{}, errors.New("a collection needs the embedding model its documents are embedded with")
	case c.Dims < 1 || c.Dims > MaxEmbeddingDims:
		return KBCollection{}, fmt.Errorf("%d dimensions: an indexed embedding has 1 to %d", c.Dims, MaxEmbeddingDims)
	case c.OwnerScope == ScopeUser && c.OrgID == "":
		return KBCollection{}, errors.New("a person's collection names the organization it belongs to")
	case c.OwnerScope == ScopeKey:
		return KBCollection{}, errors.New("a key cannot own a collection; its project or its person can")
	}
	c.ID = NewID("kb_")
	if c.CreatedAt.IsZero() {
		c.CreatedAt = time.Now()
	}
	err := s.inTx(ctx, func(tx conn) error {
		org, err := s.orgOfScope(ctx, tx, c.OwnerScope, c.OwnerID, c.OrgID)
		if err != nil {
			return err
		}
		if c.OrgID != "" && c.OrgID != org {
			return fmt.Errorf("the owner is in organization %s, not %s", org, c.OrgID)
		}
		c.OrgID = org
		var n int
		if err := tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM kb_collections WHERE org_id = ? AND name = ?`,
			c.OrgID, c.Name).Scan(&n); err != nil {
			return err
		}
		if n > 0 {
			return fmt.Errorf("%w: collection %q", ErrDuplicateName, c.Name)
		}
		if _, err := tx.ExecContext(ctx, `INSERT INTO kb_collections
			(id, org_id, owner_scope, owner_id, name, embed_model, dims, created_at, created_by)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
			c.ID, c.OrgID, string(c.OwnerScope), c.OwnerID, c.Name, c.EmbedModel, c.Dims,
			millis(c.CreatedAt), c.CreatedBy); err != nil {
			return err
		}
		part := kbPartition(c.OrgID)
		if _, err := tx.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS `+part+
			` PARTITION OF kb_chunks FOR VALUES IN (`+sqlLiteral(c.OrgID)+`)`); err != nil {
			return err
		}
		// A partial index per collection, on its organization's partition
		// only: the dimension is the collection's, and a search, which
		// names one collection at a time, uses exactly this one.
		_, err = tx.ExecContext(ctx, `CREATE INDEX `+kbVectorIndex(c.ID)+` ON `+part+
			` USING hnsw ((embedding::vector(`+strconv.Itoa(c.Dims)+`)) vector_cosine_ops)`+
			` WHERE collection = `+sqlLiteral(c.ID))
		return err
	})
	if err != nil {
		return KBCollection{}, err
	}
	return c, nil
}

const kbCollectionCols = `id, org_id, owner_scope, owner_id, name, embed_model, dims, created_at, created_by`

func scanCollection(sc scanner) (KBCollection, error) {
	var c KBCollection
	var scope string
	var at int64
	err := sc.Scan(&c.ID, &c.OrgID, &scope, &c.OwnerID, &c.Name, &c.EmbedModel, &c.Dims, &at, &c.CreatedBy)
	c.OwnerScope, c.CreatedAt = Scope(scope), fromMillis(at)
	return c, err
}

// Collections lists every collection, by organization and name.
func (s *DB) Collections(ctx context.Context) ([]KBCollection, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return nil, err
	}
	return queryAll(ctx, s.db, `SELECT `+kbCollectionCols+` FROM kb_collections ORDER BY org_id, name`, nil,
		func(r *sql.Rows) (KBCollection, error) { return scanCollection(r) })
}

// Collection reads one collection.
func (s *DB) Collection(ctx context.Context, id string) (KBCollection, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return KBCollection{}, err
	}
	c, err := scanCollection(s.db.QueryRowContext(ctx, `SELECT `+kbCollectionCols+` FROM kb_collections WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return KBCollection{}, ErrNotFound
	}
	return c, err
}

// DeleteCollection removes a collection, its documents, chunks, grants and
// index.
func (s *DB) DeleteCollection(ctx context.Context, id string) error {
	if err := s.KnowledgeAvailable(); err != nil {
		return err
	}
	if !kbIDPattern.MatchString(id) {
		return ErrNotFound
	}
	return s.inTx(ctx, func(c conn) error {
		// Chunks go with their documents (a foreign-key cascade, which row
		// security does not stand in the way of), documents and grants with
		// the collection.
		res, err := c.ExecContext(ctx, `DELETE FROM kb_collections WHERE id = ?`, id)
		if err != nil {
			return err
		}
		if n, _ := res.RowsAffected(); n == 0 {
			return ErrNotFound
		}
		_, err = c.ExecContext(ctx, `DROP INDEX IF EXISTS `+kbVectorIndex(id))
		return err
	})
}

// Grant lets a scope read or write a collection; granting again replaces the
// permission. The scope must be in the collection's organization.
func (s *DB) Grant(ctx context.Context, g KBGrant) error {
	if err := s.KnowledgeAvailable(); err != nil {
		return err
	}
	if g.Perm != KBRead && g.Perm != KBWrite {
		return fmt.Errorf("permission %q: read or write", g.Perm)
	}
	if g.AddedAt.IsZero() {
		g.AddedAt = time.Now()
	}
	return s.inTx(ctx, func(c conn) error {
		var org string
		err := c.QueryRowContext(ctx, `SELECT org_id FROM kb_collections WHERE id = ?`, g.Collection).Scan(&org)
		if errors.Is(err, sql.ErrNoRows) {
			return fmt.Errorf("collection %s: %w", g.Collection, ErrNotFound)
		}
		if err != nil {
			return err
		}
		got, err := s.orgOfScope(ctx, c, g.Scope, g.ScopeID, org)
		if err != nil {
			return err
		}
		if got != org {
			return fmt.Errorf("%s %s is in another organization; a collection is never granted outside its own", g.Scope, g.ScopeID)
		}
		_, err = c.ExecContext(ctx, `INSERT INTO kb_grants (scope, scope_id, collection, perm, added_at, added_by)
			VALUES (?, ?, ?, ?, ?, ?)
			ON CONFLICT (scope, scope_id, collection) DO UPDATE SET perm = excluded.perm,
			  added_at = excluded.added_at, added_by = excluded.added_by`,
			string(g.Scope), g.ScopeID, g.Collection, string(g.Perm), millis(g.AddedAt), g.AddedBy)
		return err
	})
}

// Revoke removes a grant, reporting whether it was there.
func (s *DB) Revoke(ctx context.Context, scope Scope, scopeID, collection string) (bool, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return false, err
	}
	res, err := s.db.ExecContext(ctx, `DELETE FROM kb_grants WHERE scope = ? AND scope_id = ? AND collection = ?`,
		string(scope), scopeID, collection)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

// Grants lists a collection's grants, or every grant when collection is "".
func (s *DB) Grants(ctx context.Context, collection string) ([]KBGrant, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return nil, err
	}
	q, args := `SELECT scope, scope_id, collection, perm, added_at, added_by FROM kb_grants`, []any(nil)
	if collection != "" {
		q, args = q+` WHERE collection = ?`, []any{collection}
	}
	return queryAll(ctx, s.db, q+` ORDER BY collection, scope, scope_id`, args, func(r *sql.Rows) (KBGrant, error) {
		var g KBGrant
		var scope, perm string
		var at int64
		err := r.Scan(&scope, &g.ScopeID, &g.Collection, &perm, &at, &g.AddedBy)
		g.Scope, g.Perm, g.AddedAt = Scope(scope), KBPerm(perm), fromMillis(at)
		return g, err
	})
}

// chainScopes is a caller's place in the tree, widest first.
func chainScopes(e Effective) [][2]string {
	var out [][2]string
	add := func(s Scope, id string) {
		if id != "" {
			out = append(out, [2]string{string(s), id})
		}
	}
	add(ScopeOrg, e.OrgID)
	add(ScopeTeam, e.TeamID)
	add(ScopeProject, e.ProjectID)
	add(ScopeUser, e.UserID)
	// The chat page's caller is a session, not a key row; its id matches no
	// key grant, which is as it should be.
	add(ScopeKey, e.KeyID)
	return out
}

// CollectionsFor lists the collections a caller may read, or with write
// those it may change. A caller outside any organization — a personal key
// that named no project — may read none: joining a second team must not
// widen what every key of that person reads.
func (s *DB) CollectionsFor(ctx context.Context, e Effective, write bool) ([]KBCollection, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return nil, err
	}
	if e.OrgID == "" {
		return nil, nil
	}
	chain := chainScopes(e)
	var match []string
	args := []any{e.OrgID}
	for _, sc := range chain {
		match = append(match, `(c.owner_scope = ? AND c.owner_id = ?)`)
		args = append(args, sc[0], sc[1])
	}
	var grant []string
	for _, sc := range chain {
		grant = append(grant, `(g.scope = ? AND g.scope_id = ?)`)
		args = append(args, sc[0], sc[1])
	}
	perm := ""
	if write {
		perm = ` AND g.perm = 'write'`
	}
	q := `SELECT ` + prefixCols("c.", kbCollectionCols) + ` FROM kb_collections c WHERE c.org_id = ? AND (` +
		strings.Join(match, " OR ") + ` OR EXISTS (SELECT 1 FROM kb_grants g WHERE g.collection = c.id` + perm +
		` AND (` + strings.Join(grant, " OR ") + `))) ORDER BY c.name`
	return queryAll(ctx, s.db, q, args, func(r *sql.Rows) (KBCollection, error) { return scanCollection(r) })
}

func prefixCols(p, cols string) string {
	parts := strings.Split(cols, ", ")
	for i := range parts {
		parts[i] = p + parts[i]
	}
	return strings.Join(parts, ", ")
}

// PutDocument adds a document to a collection, or replaces the one of the
// same name, in one transaction: it is searchable whole or not at all. The
// same digest again changes nothing, and reports so.
func (s *DB) PutDocument(ctx context.Context, c KBCollection, d KBDocument, chunks []KBChunk) (KBDocument, bool, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return KBDocument{}, false, err
	}
	if !kbIDPattern.MatchString(c.ID) {
		return KBDocument{}, false, fmt.Errorf("collection %s: %w", c.ID, ErrNotFound)
	}
	for i, ch := range chunks {
		if len(ch.Embedding) != c.Dims {
			return KBDocument{}, false, fmt.Errorf("chunk %d has %d dimensions; collection %s takes %d",
				i, len(ch.Embedding), c.Name, c.Dims)
		}
	}
	d.Collection, d.Chunks = c.ID, len(chunks)
	if d.AddedAt.IsZero() {
		d.AddedAt = time.Now()
	}
	changed := true
	err := s.inTx(ctx, func(tx conn) error {
		var oldID, oldDigest string
		err := tx.QueryRowContext(ctx, `SELECT id, digest FROM kb_documents WHERE collection = ? AND name = ?`,
			c.ID, d.Name).Scan(&oldID, &oldDigest)
		switch {
		case err == nil && oldDigest == d.Digest:
			changed = false
			return nil
		case err == nil:
			if _, err := tx.ExecContext(ctx, `DELETE FROM kb_documents WHERE id = ?`, oldID); err != nil {
				return err
			}
		case !errors.Is(err, sql.ErrNoRows):
			return err
		}
		d.ID = NewID("doc_")
		if _, err := tx.ExecContext(ctx, `INSERT INTO kb_documents
			(id, collection, name, digest, bytes, chunks, added_at, added_by) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
			d.ID, c.ID, d.Name, d.Digest, d.Bytes, d.Chunks, millis(d.AddedAt), d.AddedBy); err != nil {
			return err
		}
		if err := kbScope(ctx, tx, c.OrgID, []string{c.ID}); err != nil {
			return err
		}
		const batch = 100
		for start := 0; start < len(chunks); start += batch {
			end := min(start+batch, len(chunks))
			var vals []string
			var args []any
			for i := start; i < end; i++ {
				vals = append(vals, `(?, ?, ?, ?, ?, ?::vector)`)
				args = append(args, c.OrgID, c.ID, d.ID, i, chunks[i].Body, vectorText(chunks[i].Embedding))
			}
			if _, err := tx.ExecContext(ctx, `INSERT INTO kb_chunks (org_id, collection, doc_id, ord, body, embedding) VALUES `+
				strings.Join(vals, ", "), args...); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return KBDocument{}, false, err
	}
	return d, changed, nil
}

// vectorText is pgvector's text form of an embedding.
func vectorText(v []float32) string {
	var b strings.Builder
	b.Grow(len(v) * 10)
	b.WriteByte('[')
	for i, f := range v {
		if i > 0 {
			b.WriteByte(',')
		}
		b.WriteString(strconv.FormatFloat(float64(f), 'g', -1, 32))
	}
	b.WriteByte(']')
	return b.String()
}

// Documents lists a collection's documents, by name.
func (s *DB) Documents(ctx context.Context, collection string) ([]KBDocument, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return nil, err
	}
	return queryAll(ctx, s.db, `SELECT id, collection, name, digest, bytes, chunks, added_at, added_by
		FROM kb_documents WHERE collection = ? ORDER BY name`, []any{collection},
		func(r *sql.Rows) (KBDocument, error) {
			var d KBDocument
			var at int64
			err := r.Scan(&d.ID, &d.Collection, &d.Name, &d.Digest, &d.Bytes, &d.Chunks, &at, &d.AddedBy)
			d.AddedAt = fromMillis(at)
			return d, err
		})
}

// DeleteDocument removes a document and its chunks, reporting whether it was
// there.
func (s *DB) DeleteDocument(ctx context.Context, collection, name string) (bool, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return false, err
	}
	res, err := s.db.ExecContext(ctx, `DELETE FROM kb_documents WHERE collection = ? AND name = ?`, collection, name)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

// KBSearch is one search: the caller's organization, the collections it may
// read and asked for, the query embedded by each collection's model, and the
// query's text for the full-text arm.
type KBSearch struct {
	OrgID       string
	Collections []KBCollection
	// Embeddings is the query embedded by each model, keyed by model name.
	Embeddings map[string][]float32
	Text       string
	K          int
}

// rrfK is reciprocal rank fusion's constant: 60, as in the paper and nearly
// everywhere it is used.
const rrfK = 60

// SearchChunks finds the best chunks for a query across collections: a
// nearest-neighbour search in each collection and a full-text search across
// them, fused by rank. Every collection must be in OrgID; the caller decides
// which ones the requester may read, and the table refuses rows outside them
// regardless.
func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) {
	if err := s.KnowledgeAvailable(); err != nil {
		return nil, err
	}
	if q.K <= 0 || len(q.Collections) == 0 || q.OrgID == "" {
		return nil, nil
	}
	ids := make([]string, 0, len(q.Collections))
	quoted := make([]string, 0, len(q.Collections))
	for _, c := range q.Collections {
		if c.OrgID != q.OrgID || !kbIDPattern.MatchString(c.ID) {
			return nil, fmt.Errorf("collection %s is not in organization %s", c.ID, q.OrgID)
		}
		ids = append(ids, c.ID)
		quoted = append(quoted, sqlLiteral(c.ID))
	}
	type key struct {
		doc string
		ord int
	}
	hits := map[key]*KBHit{}
	add := func(h KBHit, rank int) {
		k := key{h.DocID, h.Ord}
		if have, ok := hits[k]; ok {
			have.Score += 1.0 / float64(rrfK+rank)
			return
		}
		h.Score = 1.0 / float64(rrfK+rank)
		hits[k] = &h
	}
	scan := func(rows *sql.Rows) ([]KBHit, error) {
		defer rows.Close()
		var out []KBHit
		for rows.Next() {
			var h KBHit
			if err := rows.Scan(&h.Collection, &h.DocID, &h.Document, &h.Ord, &h.Body); err != nil {
				return nil, err
			}
			out = append(out, h)
		}
		return out, rows.Err()
	}
	err := s.inTx(ctx, func(tx conn) error {
		if err := kbScope(ctx, tx, q.OrgID, ids); err != nil {
			return err
		}
		for _, c := range q.Collections {
			v, ok := q.Embeddings[c.EmbedModel]
			if !ok {
				continue
			}
			if len(v) != c.Dims {
				return fmt.Errorf("the query has %d dimensions; collection %s takes %d", len(v), c.Name, c.Dims)
			}
			// The literal collection and cast are what let the planner use
			// that collection's partial index.
			cast := `::vector(` + strconv.Itoa(c.Dims) + `)`
			rows, err := tx.QueryContext(ctx, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body
				FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id
				WHERE k.org_id = ? AND k.collection = `+sqlLiteral(c.ID)+`
				ORDER BY k.embedding`+cast+` <=> ?`+cast+` LIMIT ?`,
				q.OrgID, vectorText(v), q.K)
			if err != nil {
				return err
			}
			got, err := scan(rows)
			if err != nil {
				return err
			}
			for i, h := range got {
				add(h, i+1)
			}
		}
		if tsq := orQuery(q.Text); tsq != "" {
			rows, err := tx.QueryContext(ctx, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body
				FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id, to_tsquery('simple', ?) q
				WHERE k.org_id = ? AND k.collection IN (`+strings.Join(quoted, ", ")+`) AND k.tsv @@ q
				ORDER BY ts_rank(k.tsv, q) DESC LIMIT ?`,
				tsq, q.OrgID, q.K)
			if err != nil {
				return err
			}
			got, err := scan(rows)
			if err != nil {
				return err
			}
			for i, h := range got {
				add(h, i+1)
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	out := make([]KBHit, 0, len(hits))
	for _, h := range hits {
		out = append(out, *h)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Score != out[j].Score {
			return out[i].Score > out[j].Score
		}
		if out[i].DocID != out[j].DocID {
			return out[i].DocID < out[j].DocID
		}
		return out[i].Ord < out[j].Ord
	})
	if len(out) > q.K {
		out = out[:q.K]
	}
	return out, nil
}

// orQuery turns a question into a tsquery matching any of its words. All of
// them (plainto_tsquery) would match almost nothing a person types as a
// question; the rank then prefers chunks with more of them.
func orQuery(text string) string {
	seen := map[string]bool{}
	var words []string
	// Not lowercased here: to_tsquery does that with the database's own
	// rules, the same ones to_tsvector used on the chunks — Go's lower case
	// of a Turkish İ is not the database's.
	for _, w := range strings.FieldsFunc(text, func(r rune) bool {
		return !unicode.IsLetter(r) && !unicode.IsDigit(r) && !unicode.IsMark(r)
	}) {
		if len([]rune(w)) < 2 || seen[w] {
			continue
		}
		seen[w] = true
		words = append(words, w)
		if len(words) == 32 {
			break
		}
	}
	return strings.Join(words, " | ")
}

// kbOwnedBy counts the collections a scope owns, for refusing to delete it
// from under them. Zero on a store without knowledge bases.
func (s *DB) kbOwnedBy(ctx context.Context, c conn, scope Scope, id string) (int, error) {
	if s.dialect != PostgresDialect {
		return 0, nil
	}
	var n int
	q := `SELECT COUNT(*) FROM kb_collections WHERE owner_scope = ? AND owner_id = ?`
	args := []any{string(scope), id}
	if scope == ScopeOrg {
		q, args = `SELECT COUNT(*) FROM kb_collections WHERE org_id = ?`, []any{id}
	}
	err := c.QueryRowContext(ctx, q, args...).Scan(&n)
	return n, err
}

// kbForgetScope removes what a deleted scope had: its grants, and for a
// person the collections they owned.
func (s *DB) kbForgetScope(ctx context.Context, c conn, scope Scope, id string) error {
	if s.dialect != PostgresDialect {
		return nil
	}
	if _, err := c.ExecContext(ctx, `DELETE FROM kb_grants WHERE scope = ? AND scope_id = ?`, string(scope), id); err != nil {
		return err
	}
	if scope != ScopeUser {
		return nil
	}
	owned, err := queryAll(ctx, c, `SELECT id FROM kb_collections WHERE owner_scope = 'user' AND owner_id = ?`,
		[]any{id}, func(r *sql.Rows) (string, error) {
			var v string
			return v, r.Scan(&v)
		})
	if err != nil {
		return err
	}
	for _, cid := range owned {
		if _, err := c.ExecContext(ctx, `DELETE FROM kb_collections WHERE id = ?`, cid); err != nil {
			return err
		}
		if kbIDPattern.MatchString(cid) {
			if _, err := c.ExecContext(ctx, `DROP INDEX IF EXISTS `+kbVectorIndex(cid)); err != nil {
				return err
			}
		}
	}
	return nil
}

// kbDropPartition drops a deleted organization's chunk partition, empty by
// then: an organization with collections is not deleted.
func (s *DB) kbDropPartition(ctx context.Context, c conn, scope Scope, id string) error {
	if s.dialect != PostgresDialect || scope != ScopeOrg || s.kbErr != nil {
		return nil
	}
	_, err := c.ExecContext(ctx, `DROP TABLE IF EXISTS `+kbPartition(id))
	return err
}