llmgw/internal/store/knowledge_test.go

package store

import (
	"context"
	"errors"
	"math"
	"strings"
	"testing"
	"time"
)

func TestKnowledgeIsOffOnSQLite(t *testing.T) {
	s := openDB(t)
	tree(t, s)
	err := s.KnowledgeAvailable()
	if !errors.Is(err, ErrKnowledgeOff) || !strings.Contains(err.Error(), "SQLite") {
		t.Fatalf("SQLite store says %v", err)
	}
	if _, err := s.CreateCollection(context.Background(), KBCollection{
		OwnerScope: ScopeProject, OwnerID: "proj1", Name: "x", EmbedModel: "m", Dims: 3,
	}); !errors.Is(err, ErrKnowledgeOff) {
		t.Errorf("creating a collection on SQLite gave %v", err)
	}
	// Deleting from the tree still works without the tables.
	if err := s.DeleteProject(context.Background(), "proj1"); err != nil {
		t.Errorf("deleting a project on SQLite: %v", err)
	}
}

// kbTwoOrgs is org1 (team1/proj1, team2/proj2) and org2 (team3/proj3).
func kbTwoOrgs(t *testing.T, s *DB) {
	t.Helper()
	ctx := context.Background()
	tree(t, s)
	for _, err := range []error{
		s.PutTeam(ctx, Team{ID: "team2", OrgID: "org1", Name: "Finance"}),
		s.PutProject(ctx, Project{ID: "proj2", TeamID: "team2", Name: "Payroll"}),
		s.PutOrg(ctx, Org{ID: "org2", Name: "Other"}),
		s.PutTeam(ctx, Team{ID: "team3", OrgID: "org2", Name: "Research"}),
		s.PutProject(ctx, Project{ID: "proj3", TeamID: "team3", Name: "Chat"}),
	} {
		if err != nil {
			t.Fatal(err)
		}
	}
}

// unit is a 3-dimensional unit vector pointing mostly along axis i.
func unit(i int) []float32 {
	v := []float32{0.1, 0.1, 0.1}
	v[i] = 1
	n := float32(math.Sqrt(float64(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])))
	for j := range v {
		v[j] /= n
	}
	return v
}

func mustCollection(t *testing.T, s *DB, scope Scope, owner, org, name string) KBCollection {
	t.Helper()
	c, err := s.CreateCollection(context.Background(), KBCollection{
		OwnerScope: scope, OwnerID: owner, OrgID: org, Name: name, EmbedModel: "emb", Dims: 3,
	})
	if err != nil {
		t.Fatalf("creating %s: %v", name, err)
	}
	return c
}

func mustDoc(t *testing.T, s *DB, c KBCollection, name string, bodies ...string) {
	t.Helper()
	var chunks []KBChunk
	for i, b := range bodies {
		chunks = append(chunks, KBChunk{Body: b, Embedding: unit(i % 3)})
	}
	if _, _, err := s.PutDocument(context.Background(), c, KBDocument{Name: name, Digest: name + "-1", Bytes: 10}, chunks); err != nil {
		t.Fatalf("adding %s: %v", name, err)
	}
}

func openKB(t *testing.T) *DB {
	t.Helper()
	s := openPG(t)
	if err := s.KnowledgeAvailable(); err != nil {
		t.Skipf("the test database has no pgvector: %v", err)
	}
	return s
}

func names(cs []KBCollection) string {
	var out []string
	for _, c := range cs {
		out = append(out, c.Name)
	}
	return strings.Join(out, ",")
}

func TestPostgresCollectionsReachOnlyTheirChain(t *testing.T) {
	ctx := context.Background()
	s := openKB(t)
	kbTwoOrgs(t, s)
	mustCollection(t, s, ScopeOrg, "org1", "", "handbook")
	mustCollection(t, s, ScopeTeam, "team1", "", "research")
	pay := mustCollection(t, s, ScopeProject, "proj2", "", "payroll")
	mustCollection(t, s, ScopeProject, "proj3", "", "other-org")

	read := func(e Effective) string {
		t.Helper()
		got, err := s.CollectionsFor(ctx, e, false)
		if err != nil {
			t.Fatal(err)
		}
		return names(got)
	}
	inProj1 := Effective{KeyID: "k1", ProjectID: "proj1", TeamID: "team1", OrgID: "org1"}
	if got := read(inProj1); got != "handbook,research" {
		t.Errorf("proj1 reads %q", got)
	}
	if got := read(Effective{KeyID: "k3", ProjectID: "proj3", TeamID: "team3", OrgID: "org2"}); got != "other-org" {
		t.Errorf("org2 reads %q", got)
	}
	// A personal key with no project reads nothing, whatever its person
	// belongs to.
	if got := read(Effective{KeyID: "kp", UserID: "u1"}); got != "" {
		t.Errorf("a caller outside any organization reads %q", got)
	}

	// A grant to a key adds that collection for that key alone.
	_, k := mustKey(t, s, "k1", "proj1")
	if err := s.Grant(ctx, KBGrant{Scope: ScopeKey, ScopeID: k.ID, Collection: pay.ID, Perm: KBRead}); err != nil {
		t.Fatal(err)
	}
	if got := read(inProj1); got != "handbook,payroll,research" {
		t.Errorf("proj1 with a key grant reads %q", got)
	}
	if got := read(Effective{KeyID: "k-other", ProjectID: "proj1", TeamID: "team1", OrgID: "org1"}); got != "handbook,research" {
		t.Errorf("another key of proj1 reads %q", got)
	}
	// Read is not write.
	w, err := s.CollectionsFor(ctx, inProj1, true)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(names(w), "payroll") {
		t.Errorf("a read grant gave write: %s", names(w))
	}
}

func TestPostgresGrantsNeverCrossOrganizations(t *testing.T) {
	ctx := context.Background()
	s := openKB(t)
	kbTwoOrgs(t, s)
	c := mustCollection(t, s, ScopeProject, "proj1", "", "docs")
	for _, g := range []KBGrant{
		{Scope: ScopeOrg, ScopeID: "org2"},
		{Scope: ScopeTeam, ScopeID: "team3"},
		{Scope: ScopeProject, ScopeID: "proj3"},
	} {
		g.Collection, g.Perm = c.ID, KBRead
		if err := s.Grant(ctx, g); err == nil || !strings.Contains(err.Error(), "another organization") {
			t.Errorf("granting to %s %s gave %v", g.Scope, g.ScopeID, err)
		}
	}
	_, k3 := mustKey(t, s, "k3", "proj3")
	if err := s.Grant(ctx, KBGrant{Scope: ScopeKey, ScopeID: k3.ID, Collection: c.ID, Perm: KBRead}); err == nil {
		t.Error("a key of another organization was granted a collection")
	}
	// A person who belongs to nothing in the organization cannot own or be
	// granted one there.
	if err := s.CreateUser(ctx, User{ID: "u1", Email: "u1@example.com", Name: "u", PasswordHash: "x",
		State: UserActive, CreatedAt: time.Now()}); err != nil {
		t.Fatal(err)
	}
	if err := s.Grant(ctx, KBGrant{Scope: ScopeUser, ScopeID: "u1", Collection: c.ID, Perm: KBRead}); err == nil {
		t.Error("a person with no membership in the organization was granted a collection")
	}
	if err := s.AddMembership(ctx, "u1", ScopeProject, "proj2", "member", "t", time.Now()); err != nil {
		t.Fatal(err)
	}
	if err := s.Grant(ctx, KBGrant{Scope: ScopeUser, ScopeID: "u1", Collection: c.ID, Perm: KBRead}); err != nil {
		t.Errorf("a member of a project in the organization: %v", err)
	}
}

func TestPostgresSearchReturnsOnlyAllowedChunks(t *testing.T) {
	ctx := context.Background()
	s := openKB(t)
	kbTwoOrgs(t, s)
	mine := mustCollection(t, s, ScopeProject, "proj1", "", "mine")
	secret := mustCollection(t, s, ScopeProject, "proj2", "", "secret")
	foreign := mustCollection(t, s, ScopeProject, "proj3", "", "foreign")
	mustDoc(t, s, mine, "a.md", "vacation policy is twenty days", "expenses go through finance")
	mustDoc(t, s, secret, "salaries.md", "vacation salaries are confidential")
	mustDoc(t, s, foreign, "theirs.md", "vacation at the other company")

	hits, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{mine},
		Embeddings: map[string][]float32{"emb": unit(0)}, Text: "vacation policy", K: 5})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) == 0 || hits[0].Body != "vacation policy is twenty days" || hits[0].Document != "a.md" {
		t.Fatalf("best hit %+v", hits)
	}
	for _, h := range hits {
		if h.Collection != mine.ID {
			t.Errorf("a search of %s returned a chunk of %s: %q", mine.Name, h.Collection, h.Body)
		}
	}

	// A collection of another organization is refused, not searched.
	if _, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{mine, foreign},
		Embeddings: map[string][]float32{"emb": unit(0)}, Text: "vacation", K: 5}); err == nil {
		t.Error("a search across organizations was run")
	}
}

func TestPostgresRowSecurityHoldsWithoutTheFilter(t *testing.T) {
	// The second line: a query that forgets its WHERE gets nothing outside
	// the scope its transaction set, and nothing at all without one.
	ctx := context.Background()
	s := openKB(t)
	kbTwoOrgs(t, s)
	mine := mustCollection(t, s, ScopeProject, "proj1", "", "mine")
	other := mustCollection(t, s, ScopeProject, "proj2", "", "other")
	mustDoc(t, s, mine, "a.md", "one", "two")
	mustDoc(t, s, other, "b.md", "three")

	count := func(org string, colls ...string) int {
		t.Helper()
		var n int
		err := s.inTx(ctx, func(c conn) error {
			if org != "" {
				if err := kbScope(ctx, c, org, colls); err != nil {
					return err
				}
			}
			return c.QueryRowContext(ctx, `SELECT COUNT(*) FROM kb_chunks`).Scan(&n)
		})
		if err != nil {
			t.Fatal(err)
		}
		return n
	}
	if n := count(""); n != 0 {
		t.Errorf("with no scope set, %d chunks are visible", n)
	}
	if n := count("org1", mine.ID); n != 2 {
		t.Errorf("scoped to one collection, %d chunks are visible", n)
	}
	if n := count("org2", mine.ID); n != 0 {
		t.Errorf("a collection named under the wrong organization shows %d chunks", n)
	}
	// Nor can a write put a chunk outside its scope.
	err := s.inTx(ctx, func(c conn) error {
		if err := kbScope(ctx, c, "org1", []string{mine.ID}); err != nil {
			return err
		}
		docs, _ := s.Documents(ctx, other.ID)
		_, err := c.ExecContext(ctx, `INSERT INTO kb_chunks (org_id, collection, doc_id, ord, body, embedding)
			VALUES ('org1', ?, ?, 9, 'x', '[1,0,0]')`, other.ID, docs[0].ID)
		return err
	})
	if err == nil || !strings.Contains(err.Error(), "row-level security") {
		t.Errorf("writing into another collection gave %v", err)
	}
}

func TestPostgresDocumentsReplaceWholeAndSkipTheSame(t *testing.T) {
	ctx := context.Background()
	s := openKB(t)
	kbTwoOrgs(t, s)
	c := mustCollection(t, s, ScopeProject, "proj1", "", "docs")
	put := func(digest string, bodies ...string) bool {
		t.Helper()
		var chunks []KBChunk
		for _, b := range bodies {
			chunks = append(chunks, KBChunk{Body: b, Embedding: unit(0)})
		}
		_, changed, err := s.PutDocument(ctx, c, KBDocument{Name: "a.md", Digest: digest}, chunks)
		if err != nil {
			t.Fatal(err)
		}
		return changed
	}
	if !put("d1", "old one", "old two") {
		t.Error("a new document was not a change")
	}
	if put("d1", "ignored") {
		t.Error("the same digest again was a change")
	}
	if !put("d2", "new") {
		t.Error("a new digest was not a change")
	}
	hits, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{c},
		Embeddings: map[string][]float32{"emb": unit(0)}, K: 10})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) != 1 || hits[0].Body != "new" {
		t.Errorf("after replacing, the collection holds %+v", hits)
	}
	if _, _, err := s.PutDocument(ctx, c, KBDocument{Name: "b.md", Digest: "x"},
		[]KBChunk{{Body: "x", Embedding: []float32{1, 0}}}); err == nil {
		t.Error("a chunk of the wrong dimension was accepted")
	}
}

func TestPostgresScopesWithCollectionsAreNotDeleted(t *testing.T) {
	ctx := context.Background()
	s := openKB(t)
	kbTwoOrgs(t, s)
	c := mustCollection(t, s, ScopeProject, "proj2", "", "docs")
	mustDoc(t, s, c, "a.md", "x")
	if err := s.DeleteProject(ctx, "proj2"); !errors.Is(err, ErrNotEmpty) || !strings.Contains(err.Error(), "collection") {
		t.Errorf("deleting a project with a collection gave %v", err)
	}
	if err := s.DeleteCollection(ctx, c.ID); err != nil {
		t.Fatal(err)
	}
	if err := s.DeleteProject(ctx, "proj2"); err != nil {
		t.Errorf("deleting the project once its collection is gone: %v", err)
	}
	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)
	}
}

func TestPostgresAPersonsCollectionsGoWithTheAccount(t *testing.T) {
	ctx := context.Background()
	s := openKB(t)
	kbTwoOrgs(t, s)
	if err := s.CreateUser(ctx, User{ID: "u1", Email: "u1@example.com", Name: "u", PasswordHash: "x",
		State: UserActive, CreatedAt: time.Now()}); err != nil {
		t.Fatal(err)
	}
	if _, err := s.CreateCollection(ctx, KBCollection{OwnerScope: ScopeUser, OwnerID: "u1", OrgID: "org1",
		Name: "notes", EmbedModel: "emb", Dims: 3}); err == nil {
		t.Error("a person outside the organization made a collection in it")
	}
	if err := s.AddMembership(ctx, "u1", ScopeTeam, "team1", "member", "t", time.Now()); err != nil {
		t.Fatal(err)
	}
	c := mustCollection(t, s, ScopeUser, "u1", "org1", "notes")
	mustDoc(t, s, c, "n.md", "x")
	if err := s.DeleteUser(ctx, "u1", time.Now()); err != nil {
		t.Fatal(err)
	}
	if _, err := s.Collection(ctx, c.ID); !errors.Is(err, ErrNotFound) {
		t.Errorf("the person's collection outlived the account: %v", err)
	}
}

func TestOrQueryMatchesAnyWord(t *testing.T) {
	if got := orQuery("Yıllık izin kaç gün? İzin & 'x' -- a"); got != "Yıllık | izin | kaç | gün | İzin" {
		t.Errorf("orQuery gave %q", got)
	}
	if got := orQuery("?!"); got != "" {
		t.Errorf("no words gave %q", got)
	}
}