Ilker ManapandClaude Opus 5.5
26630274f4
A request given passages carries a cache_salt derived from its organization when the replica chosen is vLLM, so how fast an answer starts cannot say what another tenant's documents held. Other engines are not sent the field. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
256 lines
9.4 KiB
Go
256 lines
9.4 KiB
Go
package router
|
||
|
||
import (
|
||
"encoding/json"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"testing"
|
||
|
||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||
"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 TestTheCacheSaltIsAnOrganizations(t *testing.T) {
|
||
body := []byte(`{"model":"m","messages":[],"temperature":0}`)
|
||
a1, a2, b := withCacheSalt(body, "org-a"), withCacheSalt(body, "org-a"), withCacheSalt(body, "org-b")
|
||
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 TestAVLLMReplicaGetsTheOrganizationsSaltOnlyWithPassages(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)
|
||
}
|
||
}
|