llmgw/internal/router/kb_test.go

package router

import (
	"crypto/sha256"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

func TestKnowledgeBasesSayWhyTheyAreOffOnSQLite(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	code, _, body := call(t, rt.Handler(), secret, http.MethodGet, "/v1/kb/collections", nil)
	if code != http.StatusNotImplemented || !strings.Contains(body, "SQLite") {
		t.Errorf("got %d %s", code, body)
	}
}

// kbRouter is a router on the PostgreSQL test database, with pgvector.
func kbRouter(t *testing.T) *Router {
	t.Helper()
	return kbRouterWith(t, newFakeEngine(t, "qwen2.5-14b-instruct").srv.URL)
}

// kbRouterWith is a router on the PostgreSQL test database, with pgvector,
// in front of an engine at engineURL.
func kbRouterWith(t *testing.T, engineURL string) *Router {
	t.Helper()
	dsn := os.Getenv("LLMGW_TEST_POSTGRES")
	if dsn == "" {
		t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run these")
	}
	if u, err := url.Parse(dsn); err != nil || !strings.Contains(u.Path, "test") {
		t.Fatalf("refusing a database whose name does not say test: %s", dsn)
	}
	cfg := baseConfig(config.Endpoint{ID: "a", URL: engineURL, Models: []string{"qwen2.5-14b-instruct"}})
	cfg.Auth.Enabled = true
	cfg.Auth.Store = dsn
	rt := newTestRouter(t, cfg)
	if err := rt.store.KnowledgeAvailable(); err != nil {
		t.Skipf("no pgvector in the test database: %v", err)
	}
	return rt
}

// kbKey makes a key in a project, or a person's key with no project.
func kbKey(t *testing.T, s *store.DB, project, user string) string {
	t.Helper()
	secret := store.KeyPrefix + strings.ToLower(store.NewID("")) + strings.Repeat("k", 27)
	sum := sha256.Sum256([]byte(secret))
	k := store.Key{ID: store.NewID("key-"), Hash: sum[:], Prefix: secret[:8], Name: "t", ProjectID: project, CreatedAt: time.Now()}
	var err error
	if user != "" {
		err = s.CreateUserKey(t.Context(), user, k)
	} else {
		err = s.CreateKey(t.Context(), k)
	}
	if err != nil {
		t.Fatal(err)
	}
	return secret
}

func vec(x ...float32) []float32 { return x }

func doc(name, digest string, embeddings ...[]float32) map[string]any {
	var chunks []map[string]any
	for i, e := range embeddings {
		chunks = append(chunks, map[string]any{"text": name + " parça " + string(rune('a'+i)), "embedding": e})
	}
	return map[string]any{"name": name, "digest": digest, "bytes": 10, "chunks": chunks}
}

func TestKnowledgeBasesKeepTenantsApart(t *testing.T) {
	rt := kbRouter(t)
	h := rt.Handler()
	ctx := t.Context()
	s := rt.store
	// Unique names: the test database is shared and not emptied here.
	sfx := strings.ToLower(store.NewID(""))
	orgA, _ := s.CreateOrg(ctx, "a-"+sfx)
	teamA, _ := s.CreateTeam(ctx, orgA, "t")
	hr, _ := s.CreateProject(ctx, teamA, "hr")
	sales, _ := s.CreateProject(ctx, teamA, "sales")
	orgB, _ := s.CreateOrg(ctx, "b-"+sfx)
	teamB, _ := s.CreateTeam(ctx, orgB, "t")
	other, _ := s.CreateProject(ctx, teamB, "x")
	c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: hr,
		Name: "politikalar", EmbedModel: "emb", Dims: 3})
	if err != nil {
		t.Fatal(err)
	}
	hrKey, salesKey, otherKey := kbKey(t, s, hr, ""), kbKey(t, s, sales, ""), kbKey(t, s, other, "")
	base := "/v1/kb/collections/"

	// The owner's project writes, by name or by id.
	code, out, body := call(t, h, hrKey, http.MethodPut, base+"politikalar/documents", doc("izin.pdf", "d1", vec(1, 0, 0), vec(0, 1, 0)))
	if code != http.StatusOK || out["changed"] != true || out["chunks"] != float64(2) {
		t.Fatalf("owner's put: %d %s", code, body)
	}
	code, out, _ = call(t, h, hrKey, http.MethodPut, base+c.ID+"/documents", doc("izin.pdf", "d1", vec(1, 0, 0)))
	if code != http.StatusOK || out["changed"] != false {
		t.Errorf("the same digest again: %d %v", code, out)
	}
	code, _, body = call(t, h, hrKey, http.MethodGet, base+"politikalar/manifest", nil)
	if code != http.StatusOK || !strings.Contains(body, `"digest":"d1"`) || !strings.Contains(body, `"writable":true`) {
		t.Errorf("manifest: %d %s", code, body)
	}

	// Another organization, and another project without a grant, see no
	// such collection — not a refusal that says it exists.
	for name, key := range map[string]string{"other org": otherKey, "other project": salesKey} {
		for _, m := range []string{http.MethodGet + " manifest", http.MethodPut + " documents"} {
			method, path, _ := strings.Cut(m, " ")
			code, _, body := call(t, h, key, method, base+c.ID+"/"+path, doc("x", "y", vec(1, 0, 0)))
			if code != http.StatusNotFound {
				t.Errorf("%s %s: %d %s", name, m, code, body)
			}
		}
	}

	// A read grant lets sales read, not write.
	if err := s.Grant(ctx, store.KBGrant{Scope: store.ScopeProject, ScopeID: sales, Collection: c.ID, Perm: store.KBRead}); err != nil {
		t.Fatal(err)
	}
	if code, _, body := call(t, h, salesKey, http.MethodGet, base+"politikalar/manifest", nil); code != http.StatusOK ||
		!strings.Contains(body, `"writable":false`) {
		t.Errorf("read grant, manifest: %d %s", code, body)
	}
	if code, _, _ := call(t, h, salesKey, http.MethodPut, base+"politikalar/documents", doc("x", "y", vec(1, 0, 0))); code != http.StatusNotFound {
		t.Errorf("read grant, put: %d", code)
	}
	if code, _, body := call(t, h, salesKey, http.MethodGet, "/v1/kb/collections", nil); code != http.StatusOK ||
		!strings.Contains(body, "politikalar") || !strings.Contains(body, `"owner_name":"a-`+sfx+` / t / hr"`) {
		t.Errorf("read grant, list: %d %s", code, body)
	}
	if _, _, body := call(t, h, otherKey, http.MethodGet, "/v1/kb/collections", nil); strings.Contains(body, "politikalar") {
		t.Errorf("another organization lists it: %s", body)
	}

	// Vectors of the wrong size are refused before anything is written.
	if code, _, body := call(t, h, hrKey, http.MethodPut, base+"politikalar/documents", doc("b.md", "d", vec(1, 0))); code != http.StatusBadRequest ||
		!strings.Contains(body, "dimensions") {
		t.Errorf("wrong dimensions: %d %s", code, body)
	}

	if code, _, body := call(t, h, hrKey, http.MethodDelete, base+"politikalar/documents?name=izin.pdf", nil); code != http.StatusOK {
		t.Errorf("delete: %d %s", code, body)
	}
	if code, _, _ := call(t, h, hrKey, http.MethodDelete, base+"politikalar/documents?name=izin.pdf", nil); code != http.StatusNotFound {
		t.Errorf("deleting again: %d", code)
	}
}

func TestAPersonalKeyReadsOnlyInTheProjectItNames(t *testing.T) {
	rt := kbRouter(t)
	h := rt.Handler()
	ctx := t.Context()
	s := rt.store
	sfx := strings.ToLower(store.NewID(""))
	org, _ := s.CreateOrg(ctx, "p-"+sfx)
	team, _ := s.CreateTeam(ctx, org, "t")
	hr, _ := s.CreateProject(ctx, team, "hr")
	fin, _ := s.CreateProject(ctx, team, "fin")
	uid := store.NewID("u-")
	if err := s.CreateUser(ctx, store.User{ID: uid, Email: uid + "@example.com", Name: "u", PasswordHash: "x",
		State: store.UserActive, CreatedAt: time.Now()}); err != nil {
		t.Fatal(err)
	}
	if err := s.AddMembership(ctx, uid, store.ScopeProject, hr, "member", "t", time.Now()); err != nil {
		t.Fatal(err)
	}
	if _, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: hr,
		Name: "ik", EmbedModel: "emb", Dims: 3}); err != nil {
		t.Fatal(err)
	}
	if _, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: fin,
		Name: "finans", EmbedModel: "emb", Dims: 3}); err != nil {
		t.Fatal(err)
	}
	key := kbKey(t, s, "", uid)
	get := func(project string) (int, string) {
		req, _ := http.NewRequest(http.MethodGet, "/v1/kb/collections", nil)
		req.Header.Set("Authorization", "Bearer "+key)
		if project != "" {
			req.Header.Set(projectHeader, project)
		}
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		return rec.Code, rec.Body.String()
	}
	if code, body := get(""); code != http.StatusOK || strings.Contains(body, "ik") {
		t.Errorf("no project named: %d %s", code, body)
	}
	if code, body := get(hr); code != http.StatusOK || !strings.Contains(body, `"ik"`) || strings.Contains(body, "finans") {
		t.Errorf("in hr: %d %s", code, body)
	}
	if code, _ := get(fin); code != http.StatusForbidden {
		t.Errorf("a project the person is not in: %d", code)
	}
}