diff --git a/BACKLOG.md b/BACKLOG.md index 628fd47..64054e7 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -863,6 +863,12 @@ Built (`internal/store/knowledge.go`, migration 0035 PostgreSQL only): partial HNSW index per collection on its organization's partition (`embedding::vector(dims)`, `WHERE collection = ''`). EXPLAIN on 20k rows: partition pruned, that index used, policy applied. + Later (scale, 2026-10-08): a partial GIN per collection too + (`kb_tsv_`, backfilled at open); `vec_type` per collection, `halfvec` + on pgvector ≥ 0.7 (up to 4000 dims); `SET LOCAL hnsw.ef_search` to + candidates(k) clamped 40..1000 (EXPLAIN: 40 rows without it, all with + it); the text arm takes candidates by `ts_rank` in a subquery and scores + only those. - RLS FORCEd on `kb_chunks` against `llmgw.kb_org`/`llmgw.kb_collections` (set with set_config per transaction), one role — the table owner — not a separate one. Chunks are deleted by the foreign-key cascade from their diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cc3f56..5b38ca6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,9 +64,21 @@ APIs change between minors without a deprecation period. a control plane or agent shares the machine, and the share is read again every time `server init --patroni` runs. The new cache is applied by restarting one member at a time, the leader last. +- **Knowledge-base searches hold up with many documents.** New collections + store their index as `halfvec` where pgvector is 0.7 or later (half the + memory, and embeddings up to 4000 dimensions instead of 2000); older + collections, and clusters on pgvector 0.6, keep `vector`. Each collection + has its own full-text index, as it has its own vector index, made for + existing collections when the control plane starts. The text arm ranks by + the cheap measure first and scores only those candidates. ### Fixed +- A knowledge-base search asking for more than 10 passages + (`X-Llmgw-Knowledge-K`) got at most 40 candidates per collection from the + vector index, whatever it asked for: pgvector's `hnsw.ef_search` is now + raised to the number wanted, up to 1000. + - `llmgw server backup` and `restore` against a replicated store over TLS failed looking for `~/.postgresql/root.crt`: pg_dump was given `sslmode=verify-full` without the cluster's authority. diff --git a/docs/knowledge.md b/docs/knowledge.md index 62253a1..76a1b01 100644 --- a/docs/knowledge.md +++ b/docs/knowledge.md @@ -45,6 +45,11 @@ is available to everyone who may use them. A collection stays with the model it was made with: vectors from two models cannot be compared. + With pgvector 0.7 or later a collection's index is kept at half + precision (`halfvec`), which halves its memory and takes models of up to + 4000 dimensions; with an older pgvector the limit is 2000. A collection + keeps the kind it was made with. + ## For administrators ### Make a collection diff --git a/internal/store/knowledge.go b/internal/store/knowledge.go index 515353b..84f1dda 100644 --- a/internal/store/knowledge.go +++ b/internal/store/knowledge.go @@ -29,9 +29,13 @@ import ( // 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). +// MaxEmbeddingDims is the most dimensions pgvector can index with HNSW: on +// vector; halfvec takes up to MaxHalfvecDims. const MaxEmbeddingDims = 2000 +// MaxHalfvecDims is the most dimensions an HNSW index on halfvec takes. +const MaxHalfvecDims = 4000 + // KBPerm is what a grant allows. type KBPerm string @@ -50,8 +54,11 @@ type KBCollection struct { Name string EmbedModel string Dims int - CreatedAt time.Time - CreatedBy string + // VecType is how its vectors are indexed: vector, or halfvec where + // pgvector has it (0.7 and later), at half the index's memory. + VecType string + CreatedAt time.Time + CreatedBy string } // KBGrant lets a scope, and everything below it, read or write a collection. @@ -133,11 +140,16 @@ func (s *DB) ensureKnowledge(ctx context.Context) error { "as postgres, run CREATE EXTENSION vector in it (%v)", ErrKnowledgeOff, err) } } + var version string + if err := s.db.QueryRowContext(ctx, `SELECT extversion FROM pg_extension WHERE extname = 'vector'`).Scan(&version); err == nil { + s.kbHalfvec = versionAtLeast(version, 0, 7) + } 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 { + s.kbBackfillIndexes(ctx) return nil } err := s.inTx(ctx, func(c conn) error { @@ -197,6 +209,52 @@ func kbVectorIndex(collection string) string { return "kb_vec_" + strings.TrimPrefix(collection, "kb_") } +// kbTextIndex is the name of a collection's full-text index. +func kbTextIndex(collection string) string { return "kb_tsv_" + strings.TrimPrefix(collection, "kb_") } + +// cast is the type a collection's vectors are compared as, which its index +// is built on: the query must say the same for the planner to use it. +func (c KBCollection) cast() string { + t := c.VecType + if t != "halfvec" { + t = "vector" + } + return "::" + t + "(" + strconv.Itoa(c.Dims) + ")" +} + +// versionAtLeast compares an extension's "major.minor.patch". +func versionAtLeast(v string, major, minor int) bool { + parts := strings.SplitN(v, ".", 3) + if len(parts) < 2 { + return false + } + ma, err1 := strconv.Atoi(parts[0]) + mi, err2 := strconv.Atoi(parts[1]) + if err1 != nil || err2 != nil { + return false + } + return ma > major || ma == major && mi >= minor +} + +// kbBackfillIndexes gives collections made before their full-text index +// existed one. Best effort, at open: a search still works without it. +func (s *DB) kbBackfillIndexes(ctx context.Context) { + cs, err := queryAll(ctx, s.db, `SELECT id, org_id FROM kb_collections`, nil, func(r *sql.Rows) ([2]string, error) { + var x [2]string + return x, r.Scan(&x[0], &x[1]) + }) + if err != nil { + return + } + for _, c := range cs { + if !kbIDPattern.MatchString(c[0]) { + continue + } + _, _ = s.db.ExecContext(ctx, `CREATE INDEX IF NOT EXISTS `+kbTextIndex(c[0])+` ON `+kbPartition(c[1])+ + ` USING gin (tsv) WHERE collection = `+sqlLiteral(c[0])) + } +} + func sqlLiteral(v string) string { return "'" + strings.ReplaceAll(v, "'", "''") + "'" } // kbScope sets, for this transaction, which rows of kb_chunks it may touch. @@ -263,14 +321,22 @@ func (s *DB) CreateCollection(ctx context.Context, c KBCollection) (KBCollection 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.Dims < 1 || c.Dims > MaxHalfvecDims || c.Dims > MaxEmbeddingDims && !s.kbHalfvec: + limit := MaxEmbeddingDims + if s.kbHalfvec { + limit = MaxHalfvecDims + } + return KBCollection{}, fmt.Errorf("%d dimensions: an indexed embedding has 1 to %d", c.Dims, limit) 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_") + c.VecType = "vector" + if s.kbHalfvec { + c.VecType = "halfvec" + } if c.CreatedAt.IsZero() { c.CreatedAt = time.Now() } @@ -292,9 +358,9 @@ func (s *DB) CreateCollection(ctx context.Context, c KBCollection) (KBCollection 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, + (id, org_id, owner_scope, owner_id, name, embed_model, dims, vec_type, created_at, created_by) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + c.ID, c.OrgID, string(c.OwnerScope), c.OwnerID, c.Name, c.EmbedModel, c.Dims, c.VecType, millis(c.CreatedAt), c.CreatedBy); err != nil { return err } @@ -303,12 +369,18 @@ func (s *DB) CreateCollection(ctx context.Context, c KBCollection) (KBCollection ` PARTITION OF kb_chunks FOR VALUES IN (`+sqlLiteral(c.OrgID)+`)`); err != nil { return err } - // A partial index per collection, on its organization's partition + // Partial indexes 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)) + // names one collection at a time, uses exactly these — a text + // search over one collection does not walk its whole + // organization's index. + if _, err := tx.ExecContext(ctx, `CREATE INDEX `+kbVectorIndex(c.ID)+` ON `+part+ + ` USING hnsw ((embedding`+c.cast()+`) `+c.VecType+`_cosine_ops)`+ + ` WHERE collection = `+sqlLiteral(c.ID)); err != nil { + return err + } + _, err = tx.ExecContext(ctx, `CREATE INDEX `+kbTextIndex(c.ID)+` ON `+part+ + ` USING gin (tsv) WHERE collection = `+sqlLiteral(c.ID)) return err }) if err != nil { @@ -317,13 +389,13 @@ func (s *DB) CreateCollection(ctx context.Context, c KBCollection) (KBCollection return c, nil } -const kbCollectionCols = `id, org_id, owner_scope, owner_id, name, embed_model, dims, created_at, created_by` +const kbCollectionCols = `id, org_id, owner_scope, owner_id, name, embed_model, dims, vec_type, 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) + err := sc.Scan(&c.ID, &c.OrgID, &scope, &c.OwnerID, &c.Name, &c.EmbedModel, &c.Dims, &c.VecType, &at, &c.CreatedBy) c.OwnerScope, c.CreatedAt = Scope(scope), fromMillis(at) return c, err } @@ -377,7 +449,10 @@ func (s *DB) DeleteCollection(ctx context.Context, id string) error { if _, err := c.ExecContext(ctx, `DELETE FROM kb_defaults WHERE collections = ''`); err != nil { return err } - _, err = c.ExecContext(ctx, `DROP INDEX IF EXISTS `+kbVectorIndex(id)) + if _, err := c.ExecContext(ctx, `DROP INDEX IF EXISTS `+kbVectorIndex(id)); err != nil { + return err + } + _, err = c.ExecContext(ctx, `DROP INDEX IF EXISTS `+kbTextIndex(id)) return err }) } @@ -729,6 +804,14 @@ func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) { if err := kbScope(ctx, tx, q.OrgID, ids); err != nil { return err } + // pgvector's HNSW search returns at most ef_search rows (40 by + // default), whatever LIMIT asks: more candidates need it raised. + if len(q.Embeddings) > 0 { + ef := min(max(candidates(q.K), 40), 1000) + if _, err := tx.ExecContext(ctx, `SET LOCAL hnsw.ef_search = `+strconv.Itoa(ef)); err != nil { + return err + } + } for _, c := range q.Collections { // The literal collection and cast are what let the planner use // that collection's partial index. @@ -739,7 +822,7 @@ func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) { if len(v) != c.Dims { return fmt.Errorf("the query has %d dimensions; collection %s takes %d", len(v), c.Name, c.Dims) } - cast := `::vector(` + strconv.Itoa(c.Dims) + `)` + cast := c.cast() sim = `1 - (k.embedding` + cast + ` <=> ?` + cast + `)` args = append(args, vectorText(v)) if err := query(tx, false, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body, `+sim+`, 0::float8 @@ -753,15 +836,20 @@ func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) { if tsq != "" { // The text search's finds are measured against the question's // meaning too, so the two are compared on one scale. + // Candidates by the cheap rank first, from the collection's + // own index; the share of words held and the similarity are + // worked out for those alone, not for every chunk that holds + // one word of the question. targs := append([]any{}, args...) for _, t := range terms { targs = append(targs, t) } - targs = append(targs, tsq, q.OrgID, candidates(q.K)) + targs = append(targs, q.OrgID, tsq, tsq, candidates(q.K)) if err := query(tx, true, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body, `+sim+`, `+coverage+` AS coverage - FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id, to_tsquery('simple', ?) tq - WHERE k.org_id = ? AND k.collection = `+lit+` AND k.tsv @@ tq - ORDER BY coverage DESC, ts_rank(k.tsv, tq) DESC LIMIT ?`, targs...); err != nil { + FROM (SELECT * FROM kb_chunks k + WHERE k.org_id = ? AND k.collection = `+lit+` AND k.tsv @@ to_tsquery('simple', ?) + ORDER BY ts_rank(k.tsv, to_tsquery('simple', ?)) DESC LIMIT ?) k + JOIN kb_documents d ON d.id = k.doc_id`, targs...); err != nil { return err } } @@ -887,8 +975,10 @@ func (s *DB) kbForgetScope(ctx context.Context, c conn, scope Scope, id string) return err } if kbIDPattern.MatchString(cid) { - if _, err := c.ExecContext(ctx, `DROP INDEX IF EXISTS `+kbVectorIndex(cid)); err != nil { - return err + for _, idx := range []string{kbVectorIndex(cid), kbTextIndex(cid)} { + if _, err := c.ExecContext(ctx, `DROP INDEX IF EXISTS `+idx); err != nil { + return err + } } } } diff --git a/internal/store/knowledge_test.go b/internal/store/knowledge_test.go index 6f81c61..7b98bc8 100644 --- a/internal/store/knowledge_test.go +++ b/internal/store/knowledge_test.go @@ -3,6 +3,7 @@ package store import ( "context" "errors" + "fmt" "math" "strings" "testing" @@ -320,10 +321,85 @@ func TestPostgresScopesWithCollectionsAreNotDeleted(t *testing.T) { if err := s.DeleteProject(ctx, "proj2"); err != nil { t.Errorf("deleting the project once its collection is gone: %v", err) } + for _, idx := range []string{kbVectorIndex(c.ID), kbTextIndex(c.ID)} { + if indexExists(t, s, idx) { + t.Errorf("the collection's index %s is left behind", idx) + } + } +} + +func indexExists(t *testing.T, s *DB, name string) bool { + t.Helper() var n int - if err := s.raw.QueryRowContext(ctx, `SELECT COUNT(*) FROM pg_indexes WHERE indexname = $1`, - kbVectorIndex(c.ID)).Scan(&n); err != nil || n != 0 { - t.Errorf("the collection's index is left behind (%d, %v)", n, err) + if err := s.raw.QueryRowContext(context.Background(), `SELECT COUNT(*) FROM pg_indexes WHERE indexname = $1`, + name).Scan(&n); err != nil { + t.Fatal(err) + } + return n == 1 +} + +func TestPostgresCollectionsHaveTheirOwnIndexes(t *testing.T) { + ctx := context.Background() + s := openKB(t) + kbTwoOrgs(t, s) + c := mustCollection(t, s, ScopeProject, "proj1", "", "docs") + if !indexExists(t, s, kbVectorIndex(c.ID)) || !indexExists(t, s, kbTextIndex(c.ID)) { + t.Fatal("a new collection lacks its vector or text index") + } + // The test image's pgvector has halfvec; a collection there uses it, and + // is searched as it is indexed. + if !s.kbHalfvec || c.VecType != "halfvec" { + t.Errorf("halfvec %v, the collection's type %q", s.kbHalfvec, c.VecType) + } + if got, err := s.Collection(ctx, c.ID); err != nil || got.VecType != c.VecType { + t.Errorf("read back %q, %v", got.VecType, err) + } + var def string + if err := s.raw.QueryRowContext(ctx, `SELECT indexdef FROM pg_indexes WHERE indexname = $1`, kbVectorIndex(c.ID)).Scan(&def); err != nil || + !strings.Contains(def, "halfvec_cosine_ops") { + t.Errorf("the vector index is %q (%v)", def, err) + } + if _, err := s.CreateCollection(ctx, KBCollection{OwnerScope: ScopeProject, OwnerID: "proj1", + Name: "wide", EmbedModel: "emb", Dims: 3072}); err != nil { + t.Errorf("3072 dimensions, which halfvec indexes: %v", err) + } + if _, err := s.CreateCollection(ctx, KBCollection{OwnerScope: ScopeProject, OwnerID: "proj1", + Name: "too-wide", EmbedModel: "emb", Dims: MaxHalfvecDims + 1}); err == nil { + t.Error("more dimensions than any index takes were accepted") + } + + // A collection made before text indexes were gets one when the store + // opens. + if _, err := s.raw.ExecContext(ctx, `DROP INDEX `+kbTextIndex(c.ID)); err != nil { + t.Fatal(err) + } + s.kbBackfillIndexes(ctx) + if !indexExists(t, s, kbTextIndex(c.ID)) { + t.Error("the text index was not put back") + } + + // Asking for more than HNSW's default 40 is answered in full. (A table + // this small is read without its index: with one, 40 is all HNSW gives + // unless ef_search is raised, as an EXPLAIN on 20000 rows showed.) + var chunks []KBChunk + for i := range 120 { + chunks = append(chunks, KBChunk{Body: fmt.Sprintf("parça %d izin", i), Embedding: unit(i % 3)}) + } + if _, _, err := s.PutDocument(ctx, c, KBDocument{Name: "big.md", Digest: "big"}, chunks); err != nil { + t.Fatal(err) + } + hits, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{c}, + Embeddings: map[string][]float32{"emb": unit(0)}, K: 50}) + if err != nil || len(hits) != 50 { + t.Errorf("K=50 gave %d hits, %v", len(hits), err) + } +} + +func TestVersionAtLeast(t *testing.T) { + for v, want := range map[string]bool{"0.6.0": false, "0.7.0": true, "0.8.7": true, "1.0": true, "0": false, "x.y": false} { + if got := versionAtLeast(v, 0, 7); got != want { + t.Errorf("%s: %v", v, got) + } } } diff --git a/internal/store/migrations/postgres/0039_kb_vec_type.sql b/internal/store/migrations/postgres/0039_kb_vec_type.sql new file mode 100644 index 0000000..5651123 --- /dev/null +++ b/internal/store/migrations/postgres/0039_kb_vec_type.sql @@ -0,0 +1,5 @@ +-- How a collection's vectors are indexed: vector (single precision) or +-- halfvec (half, pgvector 0.7 and later), which halves the index a search +-- has to hold in memory for the same recall. Fixed when the collection is +-- made: its index and its queries name the same type. +ALTER TABLE kb_collections ADD COLUMN vec_type TEXT NOT NULL DEFAULT 'vector'; diff --git a/internal/store/sqlite.go b/internal/store/sqlite.go index 74a6005..9ebad5f 100644 --- a/internal/store/sqlite.go +++ b/internal/store/sqlite.go @@ -52,6 +52,9 @@ type DB struct { // kbErr is why this store cannot hold knowledge bases, nil when it can; // set once at open (knowledge.go). kbErr error + // kbHalfvec is pgvector 0.7 or later here: new collections index their + // vectors as halfvec. + kbHalfvec bool } // OpenSQLite opens (creating if needed) a database and brings it up to date.