package adminapi import ( "context" "encoding/json" "net/http" "net/url" "os" "path/filepath" "strings" "testing" "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" "gitea.mynodes.xyz/ilker/llmgw/internal/store" ) func kbModels() []catalog.Model { return append(testModels(), catalog.Model{Name: "bge-m3", Source: "/e", Capabilities: catalog.Capabilities{Embeddings: true}, Profiles: []catalog.Profile{{Name: "cpu", GPUClass: "h100-80g", TP: 1}}}) } func TestKnowledgeBasesListSaysWhyThereAreNone(t *testing.T) { a, _ := setup(t, kbModels()) s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "s.db")) if err != nil { t.Fatal(err) } defer s.Close() a.SetKB(&KBAdmin{Store: s}) var v KBView if code := call(t, a, "GET", "/api/v1/kb", "", &v); code != http.StatusOK || v.Available || !strings.Contains(v.Reason, "SQLite") { t.Errorf("got %d %+v", code, v) } if code := call(t, a, "POST", "/api/v1/kb/collections", `{"owner":"org:x","name":"n","model":"bge-m3"}`, nil); code != http.StatusNotImplemented { t.Errorf("creating on SQLite: %d", code) } } func TestKnowledgeBasesAdministered(t *testing.T) { dsn := os.Getenv("LLMGW_TEST_POSTGRES") if dsn == "" { t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this") } if u, err := url.Parse(dsn); err != nil || !strings.Contains(u.Path, "test") { t.Fatalf("refusing a database whose name does not say test") } ctx := context.Background() s, err := store.OpenPostgres(ctx, dsn) if err != nil { t.Fatal(err) } defer s.Close() if err := s.KnowledgeAvailable(); err != nil { t.Skip(err) } a, _ := setup(t, kbModels()) a.SetQuotas(s) probes := 0 a.SetKB(&KBAdmin{Store: s, Embed: func(_ context.Context, model, _ string) ([]float32, error) { probes++ return make([]float32, 4), nil }}) sfx := strings.ToLower(store.NewID("")) org, _ := s.CreateOrg(ctx, "Acme"+sfx) team, _ := s.CreateTeam(ctx, org, "hr") _, _ = s.CreateProject(ctx, team, "docs") other, _ := s.CreateOrg(ctx, "Other"+sfx) // Owner by path; the dimension asked of the model. var c KBCollectionView rec := func(method, path, body string, out any) int { return call(t, a, method, path, body, out) } body := `{"owner":"team:Acme` + sfx + `/hr","name":"politikalar","model":"bge-m3"}` if code := callCreated(t, a, body, &c); code != http.StatusCreated || c.Dims != 4 || probes != 1 || c.Org != org { t.Fatalf("create: %d %+v (probes %d)", code, c, probes) } // A model that does not embed is refused before anything is written. if code := rec("POST", "/api/v1/kb/collections", `{"owner":"org:`+org+`","name":"x","model":"llama"}`, nil); code != http.StatusBadRequest { t.Errorf("a chat model as the embedder: %d", code) } // Grants by path, never across organizations. if code := rec("POST", "/api/v1/kb/grants", `{"collection":"Acme`+sfx+`/politikalar","target":"project:Acme`+sfx+`/hr/docs","perm":"write"}`, nil); code != http.StatusOK { t.Errorf("grant: %d", code) } if code := rec("POST", "/api/v1/kb/grants", `{"collection":"`+c.ID+`","target":"org:`+other+`"}`, nil); code != http.StatusBadRequest { t.Errorf("a grant to another organization: %d", code) } var v KBView rec("GET", "/api/v1/kb", "", &v) var got *KBCollectionView for i := range v.Collections { if v.Collections[i].ID == c.ID { got = &v.Collections[i] } } if got == nil || len(got.Grants) != 1 || got.Grants[0].Perm != "write" || !strings.Contains(got.OwnerName, "hr") { t.Fatalf("list: %+v", got) } // By ORG/NAME too. if code := rec("DELETE", "/api/v1/kb/grants", `{"collection":"Acme`+sfx+`/politikalar","target":"project:Acme`+sfx+`/hr/docs"}`, nil); code != http.StatusOK { t.Errorf("revoke: %d", code) } if code := rec("DELETE", "/api/v1/kb/collections/"+c.ID, "", nil); code != http.StatusOK { t.Errorf("delete: %d", code) } } func callCreated(t *testing.T, a *API, body string, out any) int { t.Helper() mux := http.NewServeMux() a.Register(mux) r, _ := http.NewRequest("POST", "/api/v1/kb/collections", strings.NewReader(body)) r.Header.Set("Content-Type", "application/json") w := &recorder{h: http.Header{}} mux.ServeHTTP(w, r) if w.code == http.StatusCreated { _ = json.Unmarshal(w.body, out) } else { t.Logf("create answered %d %s", w.code, w.body) } return w.code } type recorder struct { h http.Header code int body []byte } func (r *recorder) Header() http.Header { return r.h } func (r *recorder) Write(b []byte) (int, error) { if r.code == 0 { r.code = http.StatusOK } r.body = append(r.body, b...) return len(b), nil } func (r *recorder) WriteHeader(c int) { r.code = c }