Ilker ManapandClaude Opus 5.5
2c7b682587
The best text match took the whole text weight, however little of the
question it matched: one everyday word ("most") put the remote-work page
above the travel one for "the most I can spend on a hotel". Now a chunk
gains the weight in proportion to the share of the question's words it
holds, counted once whatever their case.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
896 lines
31 KiB
Go
896 lines
31 KiB
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 cosine similarity to the question, plus up to
|
||
// textWeight for a full-text match: higher is better, and only the
|
||
// order is meant to be compared.
|
||
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
|
||
}
|
||
|
||
// textWeight is how much a full-text match adds to a chunk's similarity to
|
||
// the question, at most. The embedding model reads what a question means;
|
||
// full text is there for what it cannot — a code, a name, a number. Fused by
|
||
// rank, as first built, the two were equals and the gap between a close
|
||
// meaning and a distant one was thrown away: "how much leave do fathers
|
||
// get" found the annual-leave page over the paternity one on the PoC,
|
||
// though the paternity page was far nearer (0.25 against 0.44).
|
||
const textWeight = 0.3
|
||
|
||
// candidates is how many each search brings to the fusion for k results:
|
||
// with only k from each, the meaning search's third-best never met the
|
||
// text search's first, ties were broken by id, and the order was luck.
|
||
func candidates(k int) int { return max(4*k, 20) }
|
||
|
||
// SearchChunks finds the best chunks for a query across collections: a
|
||
// nearest-neighbour search and a full-text search in each, put on one scale. 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))
|
||
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)
|
||
}
|
||
type key struct {
|
||
doc string
|
||
ord int
|
||
}
|
||
type found struct {
|
||
hit KBHit
|
||
// sim is the cosine similarity to the question, when the question
|
||
// was embedded for this chunk's collection; text is the share of
|
||
// the question's words it holds, 0 when the text search did not
|
||
// find it.
|
||
sim, text float64
|
||
hasSim bool
|
||
}
|
||
hits := map[key]*found{}
|
||
add := func(h KBHit, sim sql.NullFloat64, text float64) {
|
||
k := key{h.DocID, h.Ord}
|
||
f, ok := hits[k]
|
||
if !ok {
|
||
f = &found{hit: h}
|
||
hits[k] = f
|
||
}
|
||
if sim.Valid {
|
||
f.sim, f.hasSim = sim.Float64, true
|
||
}
|
||
f.text = max(f.text, text)
|
||
}
|
||
query := func(tx conn, text bool, q string, args ...any) error {
|
||
rows, err := tx.QueryContext(ctx, q, args...)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer rows.Close()
|
||
for rows.Next() {
|
||
var h KBHit
|
||
var sim sql.NullFloat64
|
||
var rank float64
|
||
if err := rows.Scan(&h.Collection, &h.DocID, &h.Document, &h.Ord, &h.Body, &sim, &rank); err != nil {
|
||
return err
|
||
}
|
||
if !text {
|
||
rank = 0
|
||
}
|
||
add(h, sim, rank)
|
||
}
|
||
return rows.Err()
|
||
}
|
||
terms := queryTerms(q.Text)
|
||
tsq := strings.Join(terms, " | ")
|
||
// coverage is the share of the question's words a chunk holds: one
|
||
// everyday word in common is a little evidence, all of them a lot.
|
||
var coverage string
|
||
if len(terms) > 0 {
|
||
parts := make([]string, len(terms))
|
||
for i := range terms {
|
||
parts[i] = `(k.tsv @@ to_tsquery('simple', ?))::int`
|
||
}
|
||
coverage = `(` + strings.Join(parts, " + ") + `)::float8 / ` + strconv.Itoa(len(terms))
|
||
}
|
||
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 {
|
||
// The literal collection and cast are what let the planner use
|
||
// that collection's partial index.
|
||
lit := sqlLiteral(c.ID)
|
||
sim, args := "NULL::float8", []any{}
|
||
v, embedded := q.Embeddings[c.EmbedModel]
|
||
if embedded {
|
||
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) + `)`
|
||
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
|
||
FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id
|
||
WHERE k.org_id = ? AND k.collection = `+lit+`
|
||
ORDER BY k.embedding`+cast+` <=> ?`+cast+` LIMIT ?`,
|
||
vectorText(v), q.OrgID, vectorText(v), candidates(q.K)); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if tsq != "" {
|
||
// The text search's finds are measured against the question's
|
||
// meaning too, so the two are compared on one scale.
|
||
targs := append([]any{}, args...)
|
||
for _, t := range terms {
|
||
targs = append(targs, t)
|
||
}
|
||
targs = append(targs, tsq, q.OrgID, 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 {
|
||
return err
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
// One scale: similarity, plus the share of the question's words held,
|
||
// weighted down.
|
||
for _, f := range hits {
|
||
f.hit.Score = f.sim + textWeight*f.text
|
||
}
|
||
out := make([]KBHit, 0, len(hits))
|
||
for _, f := range hits {
|
||
out = append(out, f.hit)
|
||
}
|
||
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
|
||
}
|
||
|
||
// stopWords are words nearly every passage has, in Turkish and English:
|
||
// matching on them only ranks whatever is longest first.
|
||
var stopWords = func() map[string]bool {
|
||
m := map[string]bool{}
|
||
for _, w := range strings.Fields(`
|
||
acaba ama ancak artık aslında az bana bazı belki ben beni benim beri bile bir birçok biri birkaç
|
||
biz bize bizi bizim bu buna bunda bundan bunlar bunları bunların bunu bunun burada çok çünkü da
|
||
daha dahi de defa diye en gibi göre hangi hangisi hem hep hepsi her hiç için ile ise kaç kadar
|
||
ki kim kime kimi mı mi mu mü nasıl ne neden nerde nerede nereye niçin niye o olan olarak oldu
|
||
olduğu olsa olur ona ondan onlar onları onların onu onun orada öyle şey şu şuna şunda şundan şunu
|
||
tüm ve veya ya yani var yok mısın misin musun müsün midir mıdır nedir hangileri lütfen
|
||
a an and are as at be by can do does for from has have how i in is it its me my of on or our
|
||
should that the their there this to was what when where which who why will with you your
|
||
`) {
|
||
m[w] = true
|
||
}
|
||
return m
|
||
}()
|
||
|
||
// 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 { return strings.Join(queryTerms(text), " | ") }
|
||
|
||
// queryTerms are a question's words that say something: not stop words, not
|
||
// single letters, each once whatever its case.
|
||
func queryTerms(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)
|
||
}) {
|
||
low := strings.ToLowerSpecial(unicode.TurkishCase, w)
|
||
if len([]rune(w)) < 2 || seen[low] || stopWords[low] {
|
||
continue
|
||
}
|
||
seen[low] = true
|
||
words = append(words, w)
|
||
if len(words) == 32 {
|
||
break
|
||
}
|
||
}
|
||
return 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
|
||
}
|