diff --git a/BACKLOG.md b/BACKLOG.md index dc96c28..7acb0c1 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -975,10 +975,8 @@ the model's context; added to the one leading system message (or a new one); `x-llmgw-sources`; before guardInput; never cached; metrics. Default collections built (kb_defaults on org/team/project/key, narrowest wins, 30 s cache, `ctl kb default`, console action; never refuse); the portal chat's Documents picker (collection, default or none; sources listed -under the answer; a header the page sends itself is dropped). Left: -`cache_salt` per organization (not sent: llama.cpp may refuse an unknown -field, and the router does not know the engine before choosing the -replica); returned passages +under the answer; a header the page sends itself is dropped). `cache_salt` per organization on requests given passages, sent to vLLM +replicas only (chosen replica's engine). Left: returned passages are not audited per request (the header and the kept request record them). ### F4. Adapter training (L, after F1–F3) diff --git a/docs/knowledge.md b/docs/knowledge.md index 4b069cd..95bd240 100644 --- a/docs/knowledge.md +++ b/docs/knowledge.md @@ -316,6 +316,10 @@ Behaviour to know about: - **Other APIs.** Requests through `/v1/responses` and `/v1/messages` (Codex, Claude Code) work the same way, since they are answered as chat completions. +- **Separate prefix caches.** On vLLM replicas, a request given passages + carries a `cache_salt` derived from its organization. One organization's + passages then never share cached blocks with another's, so how fast an + answer starts says nothing about what another tenant asked. - **No answer cache.** Answers from knowledge bases are never served from the answer cache, because two keys of one project can read different collections. diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 5339876..c97e689 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -409,6 +409,9 @@ type parsedRequest struct { guardIn guard.Verdict // knowledge is set when passages from knowledge bases were added. knowledge bool + // kbOrg is the organization whose passages were added, for the prefix + // cache's salt. + kbOrg string // noStore is a policy that keeps nothing of this request. noStore bool // inCluster is a project whose requests may not leave the cluster. diff --git a/internal/router/rag.go b/internal/router/rag.go index 3df234e..c5e93a7 100644 --- a/internal/router/rag.go +++ b/internal/router/rag.go @@ -3,6 +3,8 @@ package router import ( "bytes" "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "errors" "fmt" @@ -223,6 +225,7 @@ func (r *Router) retrieve(ctx context.Context, w http.ResponseWriter, req *parse req.UpstreamBody = body req.PromptTokens = len(body) / 4 req.knowledge = true + req.kbOrg = e.OrgID // Set now: the proxy writes the headers when the answer starts. w.Header().Set(sourcesHeader, strings.Join(sources, ",")) result := "found" @@ -425,3 +428,19 @@ func (r *Router) kbDefaults(ctx context.Context, p auth.Principal, project strin r.kbDefMu.Unlock() return names } + +// withCacheSalt sets vLLM's cache_salt to one derived from an organization: +// the same for its requests, so they still share a cache among themselves. +func withCacheSalt(body []byte, org string) []byte { + var top map[string]json.RawMessage + if json.Unmarshal(body, &top) != nil { + return body + } + sum := sha256.Sum256([]byte("llmgw-kb\x00" + org)) + top["cache_salt"], _ = json.Marshal(hex.EncodeToString(sum[:16])) + out, err := json.Marshal(top) + if err != nil { + return body + } + return out +} diff --git a/internal/router/rag_test.go b/internal/router/rag_test.go index 89e9310..30d5e62 100644 --- a/internal/router/rag_test.go +++ b/internal/router/rag_test.go @@ -9,6 +9,7 @@ import ( "strings" "testing" + "gitea.mynodes.xyz/ilker/llmgw/internal/config" "gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync" "gitea.mynodes.xyz/ilker/llmgw/internal/store" ) @@ -51,6 +52,20 @@ func TestPassagesJoinTheLeadingSystemMessage(t *testing.T) { } } +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"}}]}]}` @@ -142,6 +157,11 @@ func TestAChatIsAnsweredFromItsCollections(t *testing.T) { 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, @@ -181,3 +201,55 @@ func mustID(t *testing.T, s *store.DB, name string) string { 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) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 8f87eab..e01c540 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -1079,6 +1079,14 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa if isAdapter && !onDemand && dec.Endpoint.Engine != "vllm" { body = withAdapter(body, model, ad.index, ad.scale) } + // Passages of one organization's documents never share cached + // blocks with another's: a prefix found warm would say what another + // tenant asked about, by how fast its answer started. vLLM keys its + // prefix cache by this salt; other engines are not sent a field + // they may refuse. + if req.kbOrg != "" && dec.Endpoint.Engine == "vllm" && req.ContentType == "" { + body = withCacheSalt(body, req.kbOrg) + } // What computes the answer, said before the proxy writes the headers. // A model can have replicas on cards and on CPUs at once, and the two // differ tenfold in speed: a reply that came slowly is otherwise