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) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed 2026-10-08 15:27:50 +02:00
1 parent 65116064d2
commit 496625824b
10 files changed
+653 -28

No files matched your search

+14 -1
View File
@@ -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
+8 -2
View File
@@ -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
+79 -6
View File
@@ -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
+5
View File
@@ -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
}
+7
View File
@@ -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) {
+7 -15
View File
@@ -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
}
+8 -2
View File
@@ -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)
+17 -2
View File
@@ -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())...)
}
+359
View File
@@ -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)
}
+149
View File
@@ -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)
}
}