From 496625824be18aaacea5cbcde5abddef0b084ebe Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Thu, 8 Oct 2026 15:27:50 +0200 Subject: [PATCH] Chat requests answered from knowledge bases A request naming collections in X-Llmgw-Knowledge has its last question embedded on the caller's account, the best passages found and added to the system message, numbered, with an instruction to cite them and to say when they do not hold the answer; x-llmgw-sources says which were given. Done before the guardrails judge the request, never served from the answer cache, and a collection the key may not read is answered as one that does not exist. Metrics llmgw_kb_retrievals_total and llmgw_kb_retrieval_seconds. Co-Authored-By: Claude Opus 5.5 (1M context) --- BACKLOG.md | 15 +- CHANGELOG.md | 10 +- docs/knowledge.md | 85 ++++++++- internal/router/cache.go | 5 + internal/router/handlers.go | 7 + internal/router/kb.go | 22 +-- internal/router/kb_test.go | 10 +- internal/router/metrics.go | 19 +- internal/router/rag.go | 359 ++++++++++++++++++++++++++++++++++++ internal/router/rag_test.go | 149 +++++++++++++++ 10 files changed, 653 insertions(+), 28 deletions(-) create mode 100644 internal/router/rag.go create mode 100644 internal/router/rag_test.go diff --git a/BACKLOG.md b/BACKLOG.md index c855c40e..8d93f721 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -946,7 +946,7 @@ PostgreSQL tests of store, router and adminapi share LLMGW_TEST_POSTGRES and the store's TRUNCATE each test; run them with `go test -p 1` or one package at a time, or a router test loses its rows mid-run. -### F3. Retrieval on the request path (M) +### F3. Retrieval on the request path (M) — done 2026-10-08 (branch `rag`), but defaults and cache_salt - `llmgw.knowledge: [collections]` in the body (stripped before the engine), or default collections per project/key. - Allowed set from the caller's chain (§7.9.2), cached like limits; grants @@ -966,6 +966,19 @@ at a time, or a router test loses its rows mid-run. (including via cache), that a personal key without a project is refused, metrics (retrieval latency, chunks returned), docs. +Built (`internal/router/rag.go`): `X-Llmgw-Knowledge` (header, as the +other llmgw extensions are, not a body field), `X-Llmgw-Knowledge-K`; +collections from the key's chain (kbEffective, shared with the kb +endpoints), 404 for one not readable; the question embedded through the +ordinary path on the caller's principal; SearchChunks; passages trimmed to +the model's context; added to the one leading system message (or a new +one); `x-llmgw-sources`; before guardInput; never cached; metrics. Left: +default collections per project or key (none yet: a client must name them); +`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); the chat playground has no collection picker; returned passages +are not audited per request (the header and the kept request record them). + ### F4. Adapter training (L, after F1–F3) A run-to-completion job placed by the control plane like a replica and supervised by the agent like an engine, in a Python environment the agent diff --git a/CHANGELOG.md b/CHANGELOG.md index fe7fe23c..4c1354dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -51,8 +51,14 @@ APIs change between minors without a deprecation period. what that machine sent. The gateway takes them at `/v1/kb/…` with the key's own access; administrators make, grant and search collections with `llmgw ctl kb`, `/api/v1/kb` and the console's Knowledge screen - ([docs/knowledge.md](docs/knowledge.md)). Answering requests from them is - next (F3). + ([docs/knowledge.md](docs/knowledge.md)). +- **Answers from knowledge bases** (BACKLOG F3): a chat request naming + collections in `X-Llmgw-Knowledge` has its last question embedded on the + caller's account, the best passages found — by meaning, plus a little for + the question's words — and added to the system message, numbered, with + `x-llmgw-sources` saying which; before the guardrails judge the request, + never from the answer cache, and as 404 for a collection the key may not + read. Through the Responses and Messages APIs too. ## v0.13.0 — 2026-10-08 diff --git a/docs/knowledge.md b/docs/knowledge.md index 99f84d7e..dcd988ee 100644 --- a/docs/knowledge.md +++ b/docs/knowledge.md @@ -7,12 +7,6 @@ passages ("chunks"), has the gateway embed each passage with the person's own key, and sends the text and the vectors. The files themselves never leave the machine. -> **Status.** Collections, access and filling them work now. Answering -> requests from them (the gateway finding the passages a question needs and -> putting them in front of the model) is the next step and is not in this -> build yet. Until then a collection can be searched by an administrator -> (`llmgw ctl kb search`) to check what it holds. - ## Why retrieval, and not training Retrieval is the tool for *knowledge*: policies, procedures, contracts. The @@ -234,6 +228,85 @@ Deleting the ledger is safe: the next run hashes every file once, finds what the collection already has, and sends nothing that is unchanged. `--ledger` uses another file. +## Asking questions of a collection + +A chat request names the collections to answer from in a header. Any client +that can add a header can do it: the OpenAI and Anthropic SDKs, Continue, or +curl. + +```sh +curl https://llm.example.com/v1/chat/completions \ + -H "Authorization: Bearer $LLMGW_API_KEY" \ + -H "X-Llmgw-Knowledge: politikalar" \ + -d '{"model": "qwen2.5-coder-14b-instruct", + "messages": [{"role": "user", "content": "Kaç gün yıllık iznim var?"}]}' +``` + +```python +client.chat.completions.create( + model="qwen2.5-coder-14b-instruct", + messages=[{"role": "user", "content": "Kaç gün yıllık iznim var?"}], + extra_headers={"X-Llmgw-Knowledge": "politikalar"}, +) +``` + +What the gateway does with such a request: + +1. **It finds the collections.** It looks up the ones named, by name or ID, + among those the key may read. A collection it may not read is answered + 404, the same as one that does not exist. +2. **It embeds the question.** The last user message is embedded by the + collection's model, on the caller's own account: their limits, budget and + guardrails apply. +3. **It searches.** The collections are searched by meaning and by words + (below). +4. **It adds the passages.** The best passages are numbered and added to the + system message, with an instruction to answer from them, cite them as + [1], [2], and say so when they do not hold the answer. The system message + is created if the request has none. Passages that would not fit the + model's context are left out, the weakest first. +5. **It judges and sends the request.** The request, passages included, then + goes through the guardrails and to the model as usual. A document + carrying instructions aimed at its reader is judged like anything else + the model reads. + +How it is shaped: + +| | | +|---|---| +| `X-Llmgw-Knowledge` | Collections, comma separated. | +| `X-Llmgw-Knowledge-K` | How many passages, 1 to 20; 5 by default. | +| `X-Llmgw-Project` | For a personal key: the project it acts in. | +| `x-llmgw-sources` (answer) | The passages given, as `document#chunk`, in their numbered order; empty when nothing matched. | + +Behaviour to know about: + +- **No passages matched.** The model is told that nothing in the documents + answers the question, and asked to say so rather than answer from general + knowledge. +- **Other APIs.** Requests through `/v1/responses` and `/v1/messages` + (Codex, Claude Code) work the same way, since they are answered as chat + completions. +- **No answer cache.** Answers from knowledge bases are never served from + the answer cache, because two keys of one project can read different + collections. + +### How a passage is chosen + +Each passage scores the following, and the highest scores win: + +- its similarity in meaning to the question (cosine, 0 to 1), plus +- up to 0.3 for the share of the question's words it contains. + +Words almost every passage has ("ne", "için", "the", "of") are not counted. +The meaning does most of the work: "how much leave do fathers get" finds the +paternity page, not the annual-leave page that shares more of its words. +The word match helps with what meaning cannot carry, such as a form number, +a product code or a name. + +Metrics: `llmgw_kb_retrievals_total{result}` (found, nothing, refused) and +`llmgw_kb_retrieval_seconds`. + ## How it is kept apart Passages live in PostgreSQL, partitioned by organization, with an index per diff --git a/internal/router/cache.go b/internal/router/cache.go index aab55a5d..e26be49a 100644 --- a/internal/router/cache.go +++ b/internal/router/cache.go @@ -251,6 +251,11 @@ func reusable(req *parsedRequest, anyAnswer bool) bool { if req.noStore || req.ThreadID != "" || req.ContentType != "" || !cacheEndpoints[req.Path] { return false } + // An answer from knowledge bases is the caller's own: two keys of one + // project can read different collections. + if req.Header.Get(knowledgeHeader) != "" { + return false + } if v := strings.ToLower(strings.TrimSpace(req.Header.Get(cacheHeader))); v == "no" || v == "off" || v == "bypass" { return false } diff --git a/internal/router/handlers.go b/internal/router/handlers.go index d75fe935..c59b0ba8 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -575,6 +575,13 @@ func (r *Router) handleInference(path string) http.HandlerFunc { w.Header().Set(trimmedHeader, strconv.Itoa(dropped)) } + // The passages a question asked knowledge bases for, before it is + // judged: they are text the model reads, and a document can carry + // instructions aimed at its reader. + if !r.retrieve(req.Context(), w, parsed, start) { + return + } + // Judged as it will be sent: authorised, within its limits, and with // its stored conversation in it, which is text the model reads too. if r.guard != nil && !r.guardInput(req.Context(), w, parsed, start) { diff --git a/internal/router/kb.go b/internal/router/kb.go index 9873162c..a75efa3d 100644 --- a/internal/router/kb.go +++ b/internal/router/kb.go @@ -106,23 +106,15 @@ func (r *Router) kbCaller(w http.ResponseWriter, req *http.Request) (auth.Princi writeError(w, http.StatusNotImplemented, "invalid_request_error", err.Error()) return p, storepkg.Effective{}, false } - e := storepkg.Effective{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, TeamID: p.TeamID, OrgID: p.OrgID} - if proj := strings.TrimSpace(req.Header.Get(projectHeader)); proj != "" && p.ProjectID == "" && p.UserID != "" { - may, err := r.store.CanUseProject(req.Context(), p.UserID, proj) - if err != nil { + e, err := r.kbEffective(req.Context(), p, req.Header.Get(projectHeader)) + if err != nil { + var ke *kbError + if errors.As(err, &ke) { + writeError(w, ke.status, ke.kind, ke.msg) + } else { writeError(w, http.StatusInternalServerError, "server_error", err.Error()) - return p, e, false } - if !may { - writeError(w, http.StatusForbidden, "permission_error", "you are not a member of project "+proj) - return p, e, false - } - eff, err := r.store.ResolveFor(req.Context(), p.KeyID, p.UserID, proj) - if err != nil { - writeError(w, http.StatusInternalServerError, "server_error", err.Error()) - return p, e, false - } - e.ProjectID, e.TeamID, e.OrgID = eff.ProjectID, eff.TeamID, eff.OrgID + return p, e, false } return p, e, true } diff --git a/internal/router/kb_test.go b/internal/router/kb_test.go index 6814d7f9..b82d40bb 100644 --- a/internal/router/kb_test.go +++ b/internal/router/kb_test.go @@ -25,6 +25,13 @@ func TestKnowledgeBasesSayWhyTheyAreOffOnSQLite(t *testing.T) { // kbRouter is a router on the PostgreSQL test database, with pgvector. func kbRouter(t *testing.T) *Router { + t.Helper() + return kbRouterWith(t, newFakeEngine(t, "qwen2.5-14b-instruct").srv.URL) +} + +// kbRouterWith is a router on the PostgreSQL test database, with pgvector, +// in front of an engine at engineURL. +func kbRouterWith(t *testing.T, engineURL string) *Router { t.Helper() dsn := os.Getenv("LLMGW_TEST_POSTGRES") if dsn == "" { @@ -33,8 +40,7 @@ func kbRouter(t *testing.T) *Router { if u, err := url.Parse(dsn); err != nil || !strings.Contains(u.Path, "test") { t.Fatalf("refusing a database whose name does not say test: %s", dsn) } - e := newFakeEngine(t, "qwen2.5-14b-instruct") - cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"qwen2.5-14b-instruct"}}) + cfg := baseConfig(config.Endpoint{ID: "a", URL: engineURL, Models: []string{"qwen2.5-14b-instruct"}}) cfg.Auth.Enabled = true cfg.Auth.Store = dsn rt := newTestRouter(t, cfg) diff --git a/internal/router/metrics.go b/internal/router/metrics.go index c3b64f61..10df8529 100644 --- a/internal/router/metrics.go +++ b/internal/router/metrics.go @@ -45,8 +45,12 @@ type metrics struct { // cacheRequests counts what the answer cache did, cacheBytes and // cacheEntries what it holds. cacheRequests *prometheus.CounterVec - cacheBytes prometheus.Gauge - cacheEntries prometheus.Gauge + // kbRetrievals counts requests answered from knowledge bases by what + // the search found; kbRetrieval is how long finding it took. + kbRetrievals *prometheus.CounterVec + kbRetrieval prometheus.Histogram + cacheBytes prometheus.Gauge + cacheEntries prometheus.Gauge tenant *tenantMetrics guard *guardMetrics @@ -138,6 +142,16 @@ func newMetrics() *metrics { "hit (answered from it), miss (asked the model, and kept the answer if it could), " + "bypass (not one whose answer is reused).", }, []string{"model", "result"}), + kbRetrievals: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_kb_retrievals_total", + Help: "Chat requests that named knowledge-base collections, by what came of it: " + + "found (passages given to the model), nothing (none matched), refused (not answered).", + }, []string{"result"}), + kbRetrieval: prometheus.NewHistogram(prometheus.HistogramOpts{ + Name: "llmgw_kb_retrieval_seconds", + Help: "Time to find a request's passages: the question embedded, and the collections searched.", + Buckets: []float64{0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10}, + }), cacheBytes: prometheus.NewGauge(prometheus.GaugeOpts{ Name: "llmgw_cache_bytes", Help: "Memory the answers kept for reuse take in this process, roughly.", @@ -161,6 +175,7 @@ func (m *metrics) collectors() []prometheus.Collector { m.adapterAffinity, m.adapterLoads, m.adapterEvictions, m.kvDecisions, m.kvTokenize, m.cacheRequests, m.cacheBytes, m.cacheEntries, + m.kbRetrievals, m.kbRetrieval, }, slices.Concat(m.tenant.collectors(), m.guard.collectors(), m.token.collectors())...) } diff --git a/internal/router/rag.go b/internal/router/rag.go new file mode 100644 index 00000000..89f8774c --- /dev/null +++ b/internal/router/rag.go @@ -0,0 +1,359 @@ +package router + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" + storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// Answering from knowledge bases (BACKLOG F3): a chat request that names +// collections in X-Llmgw-Knowledge has the passages its last question needs +// found in them and put in front of the model, with their sources, before it +// is judged and sent. Which collections it may name is decided by the key's +// place in the tree, exactly as llmgw-kb's own requests are; one it may not +// read is answered as one that does not exist. + +const ( + // knowledgeHeader names the collections, by name or id, comma separated. + knowledgeHeader = "X-Llmgw-Knowledge" + // knowledgeKHeader is how many passages; kbDefaultK when absent. + knowledgeKHeader = "X-Llmgw-Knowledge-K" + // sourcesHeader answers which passages were given: document#chunk, each + // path-escaped, in the order they were numbered. + sourcesHeader = "x-llmgw-sources" + + kbDefaultK = 5 + kbMaxK = 20 +) + +// kbContextIntro is what the model is told about the passages. In English, +// which every model reads; it is asked to answer in the question's language. +const kbContextIntro = "Below are passages from the organisation's own documents, found for the user's last " + + "question. Answer from them, and cite the ones you use as [1], [2] and so on. If they do not hold the " + + "answer, say so rather than guessing. Answer in the language of the question." + +const kbNothingFound = "The user asked a question meant to be answered from the organisation's documents, " + + "and nothing in them matched it. Say so, rather than answering from general knowledge as if it were " + + "the organisation's." + +// kbEffective is a caller's place in the tree for reading collections: a +// personal key has one only in a project it names and belongs to. +func (r *Router) kbEffective(ctx context.Context, p auth.Principal, project string) (storepkg.Effective, error) { + e := storepkg.Effective{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, TeamID: p.TeamID, OrgID: p.OrgID} + project = strings.TrimSpace(project) + if project == "" || p.ProjectID != "" || p.UserID == "" { + return e, nil + } + may, err := r.store.CanUseProject(ctx, p.UserID, project) + if err != nil { + return e, err + } + if !may { + return e, &kbError{status: http.StatusForbidden, kind: "permission_error", msg: "you are not a member of project " + project} + } + eff, err := r.store.ResolveFor(ctx, p.KeyID, p.UserID, project) + if err != nil { + return e, err + } + e.ProjectID, e.TeamID, e.OrgID = eff.ProjectID, eff.TeamID, eff.OrgID + return e, nil +} + +// kbError is a refusal with the status and kind a client is answered with. +type kbError struct { + status int + kind, msg string + retryAfter string +} + +func (e *kbError) Error() string { return e.msg } + +// retrieve finds and adds the passages a request asked for. It answers false +// when it has answered the client itself. +func (r *Router) retrieve(ctx context.Context, w http.ResponseWriter, req *parsedRequest, start time.Time) bool { + names := splitList(req.Header.Get(knowledgeHeader)) + if len(names) == 0 { + return true + } + fail := func(err error) bool { + var ke *kbError + if !errors.As(err, &ke) { + ke = &kbError{status: http.StatusInternalServerError, kind: "server_error", msg: err.Error()} + } + if ke.retryAfter != "" { + w.Header().Set("Retry-After", ke.retryAfter) + } + writeError(w, ke.status, ke.kind, ke.msg) + r.m.kbRetrievals.WithLabelValues("refused").Inc() + r.m.observeRequest(r.metricModel(req.Model), "", ke.status, time.Since(start), nil, 0) + return false + } + if req.Path != "/v1/chat/completions" { + return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error", + msg: knowledgeHeader + " is for chat completions"}) + } + p, ok := auth.FromContext(ctx) + if !ok || r.store == nil { + return fail(&kbError{status: http.StatusNotImplemented, kind: "invalid_request_error", + msg: "knowledge bases need authentication and the store, which this gateway runs without"}) + } + if err := r.store.KnowledgeAvailable(); err != nil { + return fail(&kbError{status: http.StatusNotImplemented, kind: "invalid_request_error", msg: err.Error()}) + } + k := kbDefaultK + if v := req.Header.Get(knowledgeKHeader); v != "" { + n, err := strconv.Atoi(v) + if err != nil || n < 1 || n > kbMaxK { + return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error", + msg: fmt.Sprintf("%s is 1 to %d", knowledgeKHeader, kbMaxK)}) + } + k = n + } + began := time.Now() + e, err := r.kbEffective(ctx, p, req.Header.Get(projectHeader)) + if err != nil { + return fail(err) + } + allowed, err := r.store.CollectionsFor(ctx, e, false) + if err != nil { + return fail(err) + } + var colls []storepkg.KBCollection + for _, n := range names { + c, ok := pickCollection(allowed, n) + if !ok { + msg := "no collection " + n + if e.OrgID == "" { + msg += " (a personal key names its project in " + projectHeader + ")" + } + return fail(&kbError{status: http.StatusNotFound, kind: "not_found_error", msg: msg}) + } + colls = append(colls, c) + } + + question := lastUserText(req.UpstreamBody) + if strings.TrimSpace(question) == "" { + return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error", + msg: knowledgeHeader + " needs a last user message with text to search for"}) + } + embeddings := map[string][]float32{} + for _, c := range colls { + if _, done := embeddings[c.EmbedModel]; done { + continue + } + v, err := r.embedFor(ctx, p, c.EmbedModel, question) + if err != nil { + return fail(err) + } + embeddings[c.EmbedModel] = v + } + hits, err := r.store.SearchChunks(ctx, storepkg.KBSearch{OrgID: e.OrgID, Collections: colls, + Embeddings: embeddings, Text: question, K: k}) + if err != nil { + return fail(err) + } + + // What fits: the context the model has, less what the request already + // holds and the answer it asked room for — the last passages go first. + room := -1 + if _, entry := r.resolve(req.Model); entry.ContextLength > 0 { + room = (entry.ContextLength-req.MaxTokens)*4 - len(req.UpstreamBody) - len(kbContextIntro) - 64 + } + var b strings.Builder + var sources []string + for i, h := range hits { + block := fmt.Sprintf("\n\n[%d] %s\n%s", i+1, h.Document, strings.TrimSpace(h.Body)) + if room >= 0 && b.Len()+len(block) > room { + break + } + b.WriteString(block) + sources = append(sources, url.PathEscape(h.Document)+"#"+strconv.Itoa(h.Ord)) + } + text := kbNothingFound + if len(sources) > 0 { + text = kbContextIntro + b.String() + } + body, err := withSystemText(req.UpstreamBody, text) + if err != nil { + return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error", msg: err.Error()}) + } + req.UpstreamBody = body + req.PromptTokens = len(body) / 4 + // Set now: the proxy writes the headers when the answer starts. + w.Header().Set(sourcesHeader, strings.Join(sources, ",")) + result := "found" + if len(sources) == 0 { + result = "nothing" + } + r.m.kbRetrievals.WithLabelValues(result).Inc() + r.m.kbRetrieval.Observe(time.Since(began).Seconds()) + return true +} + +// pickCollection finds a collection by id, then by name. +func pickCollection(cs []storepkg.KBCollection, want string) (storepkg.KBCollection, bool) { + for _, c := range cs { + if c.ID == want { + return c, true + } + } + for _, c := range cs { + if strings.EqualFold(c.Name, want) { + return c, true + } + } + return storepkg.KBCollection{}, false +} + +func splitList(v string) []string { + var out []string + seen := map[string]bool{} + for _, s := range strings.Split(v, ",") { + if s = strings.TrimSpace(s); s != "" && !seen[s] { + seen[s] = true + out = append(out, s) + } + } + return out +} + +// embedFor has a model embed the question as the caller: through the +// ordinary path, so the caller's limits, budget and guardrails apply to it +// and it is counted to them. +func (r *Router) embedFor(ctx context.Context, p auth.Principal, model, text string) ([]float32, error) { + body, _ := json.Marshal(map[string]any{"model": model, "input": []string{text}}) + hreq, err := http.NewRequestWithContext(auth.WithPrincipal(ctx, p), http.MethodPost, "/v1/embeddings", bytes.NewReader(body)) + if err != nil { + return nil, err + } + hreq.Header.Set("Content-Type", "application/json") + w := &bufferWriter{header: http.Header{}} + r.handleInference("/v1/embeddings")(w, hreq) + status := w.status + if status == 0 { + status = http.StatusOK + } + if status != http.StatusOK { + msg := errorMessageOf(w.buf.Bytes()) + ke := &kbError{status: status, kind: "server_error", + msg: fmt.Sprintf("the knowledge base's model %s could not embed the question: %s", model, msg), + retryAfter: w.header.Get("Retry-After")} + if status < 500 { + ke.kind = "invalid_request_error" + } + return nil, ke + } + var res struct { + Data []struct { + Embedding []float32 `json:"embedding"` + } `json:"data"` + } + if err := json.Unmarshal(w.buf.Bytes(), &res); err != nil || len(res.Data) != 1 { + return nil, fmt.Errorf("%s answered no embedding", model) + } + return res.Data[0].Embedding, nil +} + +func errorMessageOf(b []byte) string { + var e struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + if json.Unmarshal(b, &e) == nil && e.Error.Message != "" { + return e.Error.Message + } + return strings.TrimSpace(string(b)) +} + +// lastUserText is the text of a chat request's last user message: its +// content, or the text parts of it. +func lastUserText(body []byte) string { + var in struct { + Messages []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"messages"` + } + if json.Unmarshal(body, &in) != nil { + return "" + } + for i := len(in.Messages) - 1; i >= 0; i-- { + m := in.Messages[i] + if m.Role != "user" { + continue + } + var s string + if json.Unmarshal(m.Content, &s) == nil { + return s + } + var parts []struct { + Type string `json:"type"` + Text string `json:"text"` + } + if json.Unmarshal(m.Content, &parts) == nil { + var b strings.Builder + for _, p := range parts { + if p.Type == "text" { + b.WriteString(p.Text + "\n") + } + } + return b.String() + } + return "" + } + return "" +} + +// withSystemText adds text to a chat request's system message: appended to +// the first one when it leads, or as a new one at the start. One system +// message, first, because some models' templates refuse anything else. +func withSystemText(body []byte, text string) ([]byte, error) { + var top map[string]json.RawMessage + if err := json.Unmarshal(body, &top); err != nil { + return nil, fmt.Errorf("the body is not a JSON object: %w", err) + } + var msgs []map[string]json.RawMessage + if err := json.Unmarshal(top["messages"], &msgs); err != nil { + return nil, errors.New("the request has no messages") + } + added := false + if len(msgs) > 0 { + var role string + _ = json.Unmarshal(msgs[0]["role"], &role) + if role == "system" || role == "developer" { + var s string + if json.Unmarshal(msgs[0]["content"], &s) == nil { + msgs[0]["content"], _ = json.Marshal(strings.TrimRight(s, "\n") + "\n\n" + text) + added = true + } else { + var parts []json.RawMessage + if json.Unmarshal(msgs[0]["content"], &parts) == nil { + part, _ := json.Marshal(map[string]string{"type": "text", "text": "\n\n" + text}) + msgs[0]["content"], _ = json.Marshal(append(parts, part)) + added = true + } + } + } + } + if !added { + content, _ := json.Marshal(text) + role, _ := json.Marshal("system") + msgs = append([]map[string]json.RawMessage{{"role": role, "content": content}}, msgs...) + } + var err error + if top["messages"], err = json.Marshal(msgs); err != nil { + return nil, err + } + return json.Marshal(top) +} diff --git a/internal/router/rag_test.go b/internal/router/rag_test.go new file mode 100644 index 00000000..31924dc7 --- /dev/null +++ b/internal/router/rag_test.go @@ -0,0 +1,149 @@ +package router + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "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 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) + 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) + } + + // 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) + } +}