llmgw/internal/control/adminapi/train_test.go

package adminapi

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/train"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

type noAgents struct{}

func (noAgents) Send(string, *llmgwv1.ControlMessage) error { return nil }
func (noAgents) Connected(string) bool                      { return false }

func TestATrainingJobIsCheckedBeforeItIsQueued(t *testing.T) {
	models := append(testModels(), catalog.Model{Name: "small", Source: "/s",
		Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}})
	a, _ := setup(t, models, gpuNode("cuda01", 1))
	s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	a.SetTrain(&TrainAdmin{Store: s, Service: train.New(s, noAgents{}, nil)})
	mux := http.NewServeMux()
	a.Register(mux)
	post := func(body string) (int, string) {
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, httptest.NewRequest("POST", "/api/v1/train/jobs", strings.NewReader(body)))
		return rec.Code, rec.Body.String()
	}
	good := map[string]any{"name": "tone", "model": "small", "base_repo": "Qwen/Qwen2.5-0.5B-Instruct", "node": "cuda01",
		"gpu_memory_mib": 2300, "data": "{\"q\":\"a\",\"a\":\"b\"}\n{\"q\":\"c\",\"a\":\"d\"}\n", "params": map[string]any{"epochs": 3},
		"public": true}
	with := func(k string, v any) string {
		m := map[string]any{}
		for kk, vv := range good {
			m[kk] = vv
		}
		m[k] = v
		b, _ := json.Marshal(m)
		return string(b)
	}
	for _, c := range []struct {
		field string
		value any
		says  string
	}{
		{"name", "bad name", "letters, digits"},
		{"model", "nope", "no model nope"},
		{"name", "small", "already a model"},
		{"base_repo", "Qwen", "owner/name"},
		{"node", "gone", "no node gone"},
		{"data", "{\"q\":\"a\"}\n", "no q or no a"},
		{"data", "", "no pairs"},
		{"params", map[string]any{"epochs": 0}, "params.epochs"},
		{"params", map[string]any{"momentum": 0.9}, "params.momentum"},
		{"public", false, "name the adapter's owner"},
		{"data", "{\"q\":\"a\",\"a\":\"b\",\"collection\":\"kb_aaaaaaaaaaaaaaaa\"}\n", "from a knowledge base"},
	} {
		if code, body := post(with(c.field, c.value)); code != http.StatusBadRequest || !strings.Contains(body, c.says) {
			t.Errorf("%s=%v: %d %s", c.field, c.value, code, body)
		}
	}
	code, body := post(with("name", "tone"))
	var v TrainJobView
	_ = json.Unmarshal([]byte(body), &v)
	if code != http.StatusCreated || v.Pairs != 2 || !v.GGUF || v.State != "queued" || v.Params != `{"epochs":3}` {
		t.Fatalf("%d %s", code, body)
	}
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, httptest.NewRequest("DELETE", "/api/v1/train/jobs/"+v.ID, nil))
	if rec.Code != http.StatusConflict {
		t.Errorf("deleting a queued job: %d", rec.Code)
	}
	rec = httptest.NewRecorder()
	mux.ServeHTTP(rec, httptest.NewRequest("GET", "/api/v1/train/jobs/"+v.ID+"/artifacts/peft", nil))
	if rec.Code != http.StatusNotFound {
		t.Errorf("an artifact not made: %d", rec.Code)
	}
}

func TestAnAdapterTrainedOnAKnowledgeBaseBelongsToWhoMayReadIt(t *testing.T) {
	dsn := os.Getenv("LLMGW_TEST_POSTGRES")
	if dsn == "" {
		t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
	}
	ctx := context.Background()
	s, err := store.OpenPostgres(ctx, dsn)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	if s.KnowledgeAvailable() != nil {
		t.Skip("no pgvector")
	}
	models := append(testModels(), catalog.Model{Name: "small", Source: "/s",
		Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}})
	a, _ := setup(t, models, gpuNode("cuda01", 1))
	a.SetQuotas(s)
	a.SetKB(&KBAdmin{Store: s})
	a.SetTrain(&TrainAdmin{Store: s, Service: train.New(s, noAgents{}, nil)})
	sfx := strings.ToLower(store.NewID(""))
	org, _ := s.CreateOrg(ctx, "Tr"+sfx)
	hr, _ := s.CreateTeam(ctx, org, "hr")
	_, _ = s.CreateProject(ctx, hr, "docs")
	sales, _ := s.CreateTeam(ctx, org, "sales")
	_, _ = s.CreateProject(ctx, sales, "crm")
	c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeTeam, OwnerID: hr, Name: "p", EmbedModel: "bge-m3", Dims: 3})
	if err != nil {
		t.Fatal(err)
	}
	mux := http.NewServeMux()
	a.Register(mux)
	post := func(extra string) (int, TrainJobView, string) {
		body := `{"name":"tone-` + store.NewID("") + `","model":"small","base_repo":"Qwen/Q","node":"cuda01","gpu_memory_mib":2000,` +
			`"data":"{\"q\":\"a\",\"a\":\"b\",\"collection\":\"` + c.ID + `\"}\n"` + extra + `}`
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, httptest.NewRequest("POST", "/api/v1/train/jobs", strings.NewReader(body)))
		var v TrainJobView
		_ = json.Unmarshal(rec.Body.Bytes(), &v)
		return rec.Code, v, rec.Body.String()
	}
	if code, v, body := post(""); code != http.StatusCreated || v.Owner != "team:"+hr {
		t.Errorf("by default: %d %s", code, body)
	}
	if code, v, body := post(`,"owner":"project:Tr` + sfx + `/hr/docs"`); code != http.StatusCreated || !strings.HasPrefix(v.Owner, "project:") {
		t.Errorf("a project of the team: %d %s", code, body)
	}
	if code, _, body := post(`,"owner":"project:Tr` + sfx + `/sales/crm"`); code != http.StatusBadRequest || !strings.Contains(body, "may not read") {
		t.Errorf("a project that cannot read it: %d %s", code, body)
	}
	if code, _, body := post(`,"public":true`); code != http.StatusBadRequest || !strings.Contains(body, "not to everyone") {
		t.Errorf("public: %d %s", code, body)
	}
}