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) + } +}