llmgw/internal/router/rag_test.go

package router

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

func TestPassagesJoinTheLeadingSystemMessage(t *testing.T) {
	type msg struct {
		Role    string `json:"role"`
		Content any    `json:"content"`
	}
	read := func(b []byte) []msg {
		var in struct {
			Messages []msg `json:"messages"`
			Model    string
		}
		if err := json.Unmarshal(b, &in); err != nil {
			t.Fatal(err)
		}
		return in.Messages
	}
	// None: one is put first.
	out, err := withSystemText([]byte(`{"model":"m","messages":[{"role":"user","content":"q"}],"temperature":0.1}`), "CTX")
	if err != nil {
		t.Fatal(err)
	}
	if m := read(out); len(m) != 2 || m[0].Role != "system" || m[0].Content != "CTX" || !strings.Contains(string(out), `"temperature":0.1`) {
		t.Errorf("no system message: %s", out)
	}
	// One leading: added to it, never a second.
	out, _ = withSystemText([]byte(`{"messages":[{"role":"system","content":"Be brief."},{"role":"user","content":"q"}]}`), "CTX")
	if m := read(out); len(m) != 2 || m[0].Content != "Be brief.\n\nCTX" {
		t.Errorf("leading system message: %s", out)
	}
	// Parts: a text part added.
	out, _ = withSystemText([]byte(`{"messages":[{"role":"system","content":[{"type":"text","text":"Be brief."}]},{"role":"user","content":"q"}]}`), "CTX")
	if !strings.Contains(string(out), `{"text":"\n\nCTX","type":"text"}`) {
		t.Errorf("parts: %s", out)
	}
	if _, err := withSystemText([]byte(`{"prompt":"x"}`), "CTX"); err == nil {
		t.Error("a request without messages was changed")
	}
}

func TestTheCacheSaltFollowsTheCollections(t *testing.T) {
	body := []byte(`{"model":"m","messages":[],"temperature":0}`)
	a1, a2, b := withCacheSalt(body, "kb_a,kb_b"), withCacheSalt(body, "kb_a,kb_b"), withCacheSalt(body, "kb_a")
	var x, y struct {
		Salt string  `json:"cache_salt"`
		T    float64 `json:"temperature"`
	}
	_ = json.Unmarshal(a1, &x)
	_ = json.Unmarshal(b, &y)
	if string(a1) != string(a2) || x.Salt == "" || x.Salt == y.Salt || len(x.Salt) != 32 || !strings.Contains(string(a1), `"temperature":0`) {
		t.Errorf("salts %s / %s", a1, b)
	}
}

func TestTheLastUserMessageIsTheQuestion(t *testing.T) {
	body := `{"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"a"},
		{"role":"user","content":[{"type":"text","text":"yıllık izin"},{"type":"image_url","image_url":{"url":"x"}}]}]}`
	if got := lastUserText([]byte(body)); strings.TrimSpace(got) != "yıllık izin" {
		t.Errorf("question %q", got)
	}
}

func TestKnowledgeOnSQLiteSaysWhy(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions",
		strings.NewReader(`{"model":"q14","messages":[{"role":"user","content":"hi"}]}`))
	req.Header.Set("Authorization", "Bearer "+secret)
	req.Header.Set(knowledgeHeader, "docs")
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)
	if rec.Code != http.StatusNotImplemented || !strings.Contains(rec.Body.String(), "SQLite") || e.calls.Load() != 0 {
		t.Errorf("got %d %s, engine asked %d times", rec.Code, rec.Body.String(), e.calls.Load())
	}
}

func TestAChatIsAnsweredFromItsCollections(t *testing.T) {
	dsn := os.Getenv("LLMGW_TEST_POSTGRES")
	if dsn == "" {
		t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
	}
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt := kbRouterWith(t, e.srv.URL)
	ctx := t.Context()
	s := rt.store
	sfx := strings.ToLower(store.NewID(""))
	org, _ := s.CreateOrg(ctx, "rag-"+sfx)
	team, _ := s.CreateTeam(ctx, org, "t")
	proj, _ := s.CreateProject(ctx, team, "p")
	other, _ := s.CreateOrg(ctx, "other-"+sfx)
	oteam, _ := s.CreateTeam(ctx, other, "t")
	oproj, _ := s.CreateProject(ctx, oteam, "p")
	if _, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: proj,
		Name: "politika-" + sfx, EmbedModel: "qwen2.5-14b-instruct", Dims: 3}); err != nil {
		t.Fatal(err)
	}
	key, otherKey := kbKey(t, s, proj, ""), kbKey(t, s, oproj, "")
	srv := httptest.NewServer(rt.Handler())
	t.Cleanup(srv.Close)

	dir := t.TempDir()
	_ = os.WriteFile(filepath.Join(dir, "izin.md"), []byte("# Yıllık izin\n\nÇalışanlar yılda yirmi gün izin kullanır.\n"), 0o644)
	_ = os.WriteFile(filepath.Join(dir, "masraf.md"), []byte("# Masraf\n\nFişler ayın beşine kadar verilir.\n"), 0o644)
	h, _ := kbsync.NewHTTP("")
	l, _ := kbsync.OpenLedger(ctx, filepath.Join(t.TempDir(), "l.db"))
	defer l.Close()
	if rep, err := kbsync.Sync(ctx, &kbsync.Client{Base: srv.URL, Key: key, HTTP: h}, l,
		kbsync.Options{Collection: "politika-" + sfx, Roots: []string{dir}}, nil); err != nil || rep.Sent != 2 {
		t.Fatalf("sync %+v %v", rep, err)
	}

	chat := func(key, collections string) *http.Response {
		t.Helper()
		req, _ := http.NewRequest(http.MethodPost, srv.URL+"/v1/chat/completions", strings.NewReader(
			`{"model":"qwen2.5-14b-instruct","messages":[{"role":"system","content":"Kısa yaz."},{"role":"user","content":"kaç gün yıllık izin var"}]}`))
		req.Header.Set("Authorization", "Bearer "+key)
		if collections != "" {
			req.Header.Set(knowledgeHeader, collections)
		}
		req.Header.Set(knowledgeKHeader, "1")
		res, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		res.Body.Close()
		return res
	}
	res := chat(key, "politika-"+sfx)
	if res.StatusCode != http.StatusOK || res.Header.Get(sourcesHeader) != "izin.md#0" {
		t.Fatalf("answered %d, sources %q", res.StatusCode, res.Header.Get(sourcesHeader))
	}
	sent, _ := e.last.Load().([]byte)
	var body struct {
		Messages []struct {
			Role, Content string
		}
	}
	_ = json.Unmarshal(sent, &body)
	if len(body.Messages) != 2 || body.Messages[0].Role != "system" ||
		!strings.HasPrefix(body.Messages[0].Content, "Kısa yaz.\n\nBelow are passages") ||
		!strings.Contains(body.Messages[0].Content, "[1] izin.md\nYıllık izin\n\nÇalışanlar yılda yirmi gün") ||
		strings.Contains(body.Messages[0].Content, "Fişler") {
		t.Errorf("the model was sent:\n%s", sent)
	}

	// The engine here is no vLLM: no salt is sent to it.
	if strings.Contains(string(sent), "cache_salt") {
		t.Errorf("a salt was sent to an engine that is not vLLM:\n%s", sent)
	}

	// A project default answers a request that names nothing, and none
	// turns it off.
	if err := s.SetKBDefault(ctx, store.KBDefault{Scope: store.ScopeProject, ScopeID: proj,
		Collections: []string{mustID(t, s, "politika-"+sfx)}}); err != nil {
		t.Fatal(err)
	}
	if res := chat(key, ""); res.Header.Get(sourcesHeader) != "izin.md#0" {
		t.Errorf("by default: sources %q", res.Header.Get(sourcesHeader))
	}
	if res := chat(key, "none"); res.Header.Get(sourcesHeader) != "" {
		t.Errorf("with none: sources %q", res.Header.Get(sourcesHeader))
	}
	sent, _ = e.last.Load().([]byte)
	if strings.Contains(string(sent), "Below are passages") {
		t.Errorf("none still added passages:\n%s", sent)
	}

	// Another organization's key names it: as if it did not exist, and the
	// model is not asked.
	calls := e.calls.Load()
	if res := chat(otherKey, "politika-"+sfx); res.StatusCode != http.StatusNotFound || e.calls.Load() != calls {
		t.Errorf("another organization: %d, engine asked %d more times", res.StatusCode, e.calls.Load()-calls)
	}
}

func mustID(t *testing.T, s *store.DB, name string) string {
	t.Helper()
	cs, err := s.Collections(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	for _, c := range cs {
		if c.Name == name {
			return c.ID
		}
	}
	t.Fatalf("no collection %s", name)
	return ""
}

func TestAVLLMReplicaGetsASaltOnlyWithPassages(t *testing.T) {
	dsn := os.Getenv("LLMGW_TEST_POSTGRES")
	if dsn == "" {
		t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
	}
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Engine: "vllm", 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.Skip(err)
	}
	ctx := t.Context()
	s := rt.store
	sfx := strings.ToLower(store.NewID(""))
	org, _ := s.CreateOrg(ctx, "salt-"+sfx)
	team, _ := s.CreateTeam(ctx, org, "t")
	proj, _ := s.CreateProject(ctx, team, "p")
	c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: proj,
		Name: "c-" + sfx, EmbedModel: "qwen2.5-14b-instruct", Dims: 3})
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := s.PutDocument(ctx, c, store.KBDocument{Name: "a.md", Digest: "d"},
		[]store.KBChunk{{Body: "yıllık izin yirmi gün", Embedding: []float32{1, 1, 1}}}); err != nil {
		t.Fatal(err)
	}
	key := kbKey(t, s, proj, "")
	ask := func(knowledge string) string {
		req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(
			`{"model":"qwen2.5-14b-instruct","messages":[{"role":"user","content":"yıllık izin"}]}`))
		req.Header.Set("Authorization", "Bearer "+key)
		if knowledge != "" {
			req.Header.Set(knowledgeHeader, knowledge)
		}
		rec := httptest.NewRecorder()
		rt.Handler().ServeHTTP(rec, req)
		if rec.Code != http.StatusOK {
			t.Fatalf("%d %s", rec.Code, rec.Body.String())
		}
		b, _ := e.last.Load().([]byte)
		return string(b)
	}
	if got := ask(c.Name); !strings.Contains(got, `"cache_salt":"`) {
		t.Errorf("no salt with passages:\n%s", got)
	}
	if got := ask(""); strings.Contains(got, "cache_salt") {
		t.Errorf("a salt without passages:\n%s", got)
	}
}

func TestAQuestionIsEmbeddedWithTheModelsQueryPrefix(t *testing.T) {
	// Qwen3-Embedding and its like are trained with an instruction before
	// the query, never before the documents: the documents come through
	// /v1/embeddings as they are; the question gets the prefix here.
	e := newFakeEngine(t, "emb")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"emb"}})
	m := catModel("emb", nil)
	m.Capabilities.Embeddings = true
	m.QueryPrefix = "Instruct: retrieve passages\nQuery: "
	cfg.Models = []catalog.Model{m}
	rt := newTestRouter(t, cfg)
	// The test engine answers with the input's length in bytes and words.
	v, err := rt.embedFor(t.Context(), auth.Principal{}, "emb", "kaç gün izin var")
	if err != nil {
		t.Fatal(err)
	}
	want := "Instruct: retrieve passages\nQuery: kaç gün izin var"
	if int(v[0]) != len(want) || int(v[1]) != len(strings.Fields(want)) {
		t.Errorf("embedded %v, not %q", v, want)
	}
}