Let a conversation continue on a different model than it started on
The request path half. A caller sets x-llmgw-thread and sends only the new turn; llmgw puts the stored history in front of it, sends that to whichever model is serving, and stores the exchange afterwards. A header rather than a body field because some SDKs drop fields they do not know and all of them let a caller set a header. Capturing the reply is opt-in and bounded. This is the hottest path in the system, so a request without a thread accumulates nothing at all, and one with a thread stops accumulating at a cap rather than growing to it — a gateway that kept every reply in memory would be one long answer away from trouble. A reply is truncated on a rune boundary rather than dropped: a shortened turn replays as a shortened turn, while nothing at all replays as the assistant having said nothing, and invalid UTF-8 would break the next turn's body for good rather than for one turn. The streamed path sets the captured reply on every exit, not only the one that looks like the end. A stream can also finish with the client gone or the engine stalled, and what the assistant said up to that point is still what it said. A thread on a project that keeps no conversations is refused. The e2e test asserted that and got a cheerful 200: the request was served, nothing was stored, and the caller would have discovered it only when the model had no idea what they were referring to — which reads as the model being broken. It now asks whether storage is on before doing anything else. History is spliced after the model is finally settled, because a budget fallback may have changed it and trimming is measured against the context of the model that will actually serve. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
cebf9906d2
commit
9eb321ae3e
7 files changed
+595
-3
No files matched your search
@@ -164,6 +164,14 @@ type parsedRequest struct {
|
||||
PromptTokens int
|
||||
MaxTokens int
|
||||
|
||||
// ThreadID is the conversation this request continues, empty for the
|
||||
// stateless path every request took before threads existed.
|
||||
ThreadID string
|
||||
// ThreadTurns are the turns this request itself carried, kept so the
|
||||
// exchange can be stored after a reply without re-reading a body that has
|
||||
// since had the history spliced into it.
|
||||
ThreadTurns []store.Message
|
||||
|
||||
// reservation is carried from the limit check to the point where the
|
||||
// request is settled, so the record counts against the same budget window
|
||||
// the check read.
|
||||
@@ -255,8 +263,26 @@ func (r *Router) handleInference(path string) http.HandlerFunc {
|
||||
w.Header().Set("x-llmgw-queue-ms", strconv.FormatInt(queued.Milliseconds(), 10))
|
||||
}
|
||||
|
||||
// After the model is finally settled — a budget fallback may have
|
||||
// changed it — because trimming is measured against the context of the
|
||||
// model that will actually serve, not the one that was asked for.
|
||||
if dropped, err := r.loadThread(req.Context(), parsed); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
r.m.observeRequest(parsed.Model, "", http.StatusBadRequest, time.Since(start), nil, 0)
|
||||
return
|
||||
} else if dropped > 0 {
|
||||
// Dropping turns is the right answer and must not be a silent one.
|
||||
w.Header().Set(trimmedHeader, strconv.Itoa(dropped))
|
||||
}
|
||||
|
||||
res := r.dispatch(req.Context(), w, parsed)
|
||||
r.finish(w, parsed, res, start)
|
||||
// After the client has been served. A conversation that failed to
|
||||
// record is a worse outcome than a slow reply, but only just, and the
|
||||
// caller has their answer either way.
|
||||
if res.err == nil && res.outcome.Status == http.StatusOK {
|
||||
r.saveThread(context.WithoutCancel(req.Context()), parsed, res.outcome.Reply, parsed.Model)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -390,6 +416,7 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) {
|
||||
ProfileOverride: req.Header.Get("x-llmgw-profile"),
|
||||
MaxTokens: max(b.MaxTokens, b.MaxCompletionTokens),
|
||||
PromptPrefix: promptPrefix(&b),
|
||||
ThreadID: req.Header.Get(threadHeader),
|
||||
}
|
||||
// With authentication on, the key decides the band and the routing profile,
|
||||
// and the headers stop being believed. Trusting the header was only ever
|
||||
|
||||
@@ -154,6 +154,14 @@ type Request struct {
|
||||
// still recording it. Set when the router asked for usage the client did
|
||||
// not request.
|
||||
SuppressUsageChunk bool
|
||||
// CaptureReply asks for the assistant's text back in the Outcome, up to
|
||||
// this many bytes. Zero means capture nothing, which is every request that
|
||||
// is not continuing a stored conversation.
|
||||
//
|
||||
// Opt-in and bounded on purpose. This is the hottest path in the system,
|
||||
// and a gateway that accumulated every reply in memory would be one long
|
||||
// answer away from trouble; a request that stores nothing must pay nothing.
|
||||
CaptureReply int
|
||||
}
|
||||
|
||||
// Usage holds the token counts reported by the engine.
|
||||
@@ -178,6 +186,14 @@ type Outcome struct {
|
||||
Usage *Usage
|
||||
// Bytes is how many bytes were relayed to the client.
|
||||
Bytes int64
|
||||
// Reply is the assistant's text, when CaptureReply asked for it. Truncated
|
||||
// at the cap rather than dropped: a shortened turn replays as a shortened
|
||||
// turn, while nothing at all would replay as the assistant having said
|
||||
// nothing, which is a different and wrong conversation.
|
||||
Reply string
|
||||
// ReplyTruncated says the cap was reached, so nothing downstream records a
|
||||
// half answer as a whole one.
|
||||
ReplyTruncated bool
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -232,11 +248,11 @@ func (p *Proxy) Do(ctx context.Context, w http.ResponseWriter, r Request) (Outco
|
||||
if r.Stream && isEventStream(resp.Header) {
|
||||
return p.relayStream(ctx, cancel, w, resp, r)
|
||||
}
|
||||
return p.relayBuffered(w, resp)
|
||||
return p.relayBuffered(w, resp, r)
|
||||
}
|
||||
|
||||
// relayBuffered copies a non-streaming response through, capturing usage.
|
||||
func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response) (Outcome, error) {
|
||||
func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response, r Request) (Outcome, error) {
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, p.cfg.MaxResponseBytes))
|
||||
if err != nil {
|
||||
return Outcome{Status: resp.StatusCode, Retryable: true}, err
|
||||
@@ -252,6 +268,9 @@ func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response) (Outco
|
||||
Bytes: int64(n),
|
||||
Usage: usageFromJSON(body),
|
||||
}
|
||||
if r.CaptureReply > 0 && resp.StatusCode == http.StatusOK {
|
||||
out.Reply, out.ReplyTruncated = capReply(replyFromJSON(body), r.CaptureReply)
|
||||
}
|
||||
if werr != nil {
|
||||
return out, ErrClientGone
|
||||
}
|
||||
@@ -262,7 +281,7 @@ func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response) (Outco
|
||||
//
|
||||
// abort must cancel the context the upstream request was built with: the idle
|
||||
// watchdog can only stop a silent engine by aborting that read.
|
||||
func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w http.ResponseWriter, resp *http.Response, r Request) (Outcome, error) {
|
||||
func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w http.ResponseWriter, resp *http.Response, r Request) (result Outcome, rerr error) {
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
return Outcome{Status: resp.StatusCode, Retryable: false},
|
||||
@@ -275,6 +294,15 @@ func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w htt
|
||||
|
||||
out := Outcome{Wrote: true, Status: resp.StatusCode}
|
||||
start := time.Now()
|
||||
// Nothing is accumulated unless a thread asked for it, so a request that
|
||||
// stores nothing pays nothing on the hottest path in the system.
|
||||
reply := replyBuilder{cap: r.CaptureReply}
|
||||
// On every exit rather than at the one that looks like the end: a stream
|
||||
// can also end with the client gone or the engine stalled, and what the
|
||||
// assistant said up to that point is still what it said.
|
||||
defer func() {
|
||||
result.Reply, result.ReplyTruncated = reply.result()
|
||||
}()
|
||||
|
||||
// One watchdog covers both phases: until the first event it enforces
|
||||
// FirstTokenTimeout, afterwards IdleTimeout.
|
||||
@@ -295,6 +323,9 @@ func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w htt
|
||||
out.TTFT = time.Since(start)
|
||||
wd.setInterval(p.cfg.IdleTimeout)
|
||||
}
|
||||
if r.CaptureReply > 0 {
|
||||
reply.add(replyDelta(ev.data.String()))
|
||||
}
|
||||
if u := ev.usage(); u != nil {
|
||||
out.Usage = u
|
||||
if r.SuppressUsageChunk {
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// replyFromJSON pulls the assistant's text out of a completed response.
|
||||
//
|
||||
// Only the first choice. A request asking for several is not one continuing a
|
||||
// conversation, and storing one of them as "what the assistant said" would
|
||||
// invent a conversation that never happened.
|
||||
func replyFromJSON(body []byte) string {
|
||||
var b struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
// The completions endpoint, which has no message wrapper.
|
||||
Text string `json:"text"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &b); err != nil || len(b.Choices) == 0 {
|
||||
return ""
|
||||
}
|
||||
if c := b.Choices[0].Message.Content; c != "" {
|
||||
return c
|
||||
}
|
||||
return b.Choices[0].Text
|
||||
}
|
||||
|
||||
// replyDelta pulls the incremental text out of one streamed event.
|
||||
func replyDelta(event string) string {
|
||||
if event == "" || event == "[DONE]" {
|
||||
return ""
|
||||
}
|
||||
var b struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
Text string `json:"text"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(event), &b); err != nil || len(b.Choices) == 0 {
|
||||
return ""
|
||||
}
|
||||
if c := b.Choices[0].Delta.Content; c != "" {
|
||||
return c
|
||||
}
|
||||
return b.Choices[0].Text
|
||||
}
|
||||
|
||||
// capReply truncates at the cap and says whether it had to.
|
||||
//
|
||||
// Truncated rather than dropped, because a shortened turn replays as a
|
||||
// shortened turn while nothing at all replays as the assistant having said
|
||||
// nothing — a different conversation, and a confusing one to be handed back.
|
||||
func capReply(s string, cap int) (string, bool) {
|
||||
if cap <= 0 || len(s) <= cap {
|
||||
return s, false
|
||||
}
|
||||
// On a rune boundary, so a stored reply is never invalid UTF-8 — it goes
|
||||
// back into a JSON body, and a split multi-byte character would make that
|
||||
// body unserialisable.
|
||||
cut := cap
|
||||
for cut > 0 && !isRuneStart(s[cut]) {
|
||||
cut--
|
||||
}
|
||||
return s[:cut], true
|
||||
}
|
||||
|
||||
func isRuneStart(b byte) bool { return b&0xC0 != 0x80 }
|
||||
|
||||
// replyBuilder accumulates a streamed reply up to a cap.
|
||||
//
|
||||
// It stops appending once full rather than growing and truncating at the end,
|
||||
// so a long answer cannot cost more memory than the cap allows however much
|
||||
// the engine sends.
|
||||
type replyBuilder struct {
|
||||
cap int
|
||||
sb strings.Builder
|
||||
truncated bool
|
||||
}
|
||||
|
||||
func (r *replyBuilder) add(s string) {
|
||||
if r.cap <= 0 || s == "" || r.truncated {
|
||||
return
|
||||
}
|
||||
if room := r.cap - r.sb.Len(); len(s) > room {
|
||||
part, _ := capReply(s, room)
|
||||
r.sb.WriteString(part)
|
||||
r.truncated = true
|
||||
return
|
||||
}
|
||||
r.sb.WriteString(s)
|
||||
}
|
||||
|
||||
func (r *replyBuilder) result() (string, bool) { return r.sb.String(), r.truncated }
|
||||
@@ -0,0 +1,95 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNothingIsCapturedUnlessAThreadAsked(t *testing.T) {
|
||||
// This is the hottest path in the system. A gateway that accumulated every
|
||||
// reply in memory would be one long answer away from trouble, so a request
|
||||
// that stores nothing must pay nothing.
|
||||
var r replyBuilder // cap zero, which is every ordinary request
|
||||
r.add("a long assistant reply that nobody asked to keep")
|
||||
if got, truncated := r.result(); got != "" || truncated {
|
||||
t.Errorf("a request with no thread accumulated %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestALongReplyCostsNoMoreThanTheCap(t *testing.T) {
|
||||
// Stopping at the cap rather than growing and truncating afterwards is
|
||||
// what makes the cap a memory bound instead of a formatting rule.
|
||||
r := replyBuilder{cap: 10}
|
||||
for range 1000 {
|
||||
r.add(strings.Repeat("x", 100))
|
||||
}
|
||||
got, truncated := r.result()
|
||||
if len(got) > 10 {
|
||||
t.Errorf("accumulated %d bytes against a cap of 10", len(got))
|
||||
}
|
||||
if !truncated {
|
||||
t.Error("truncation was not reported")
|
||||
}
|
||||
}
|
||||
|
||||
func TestATruncatedReplyIsStillValidUTF8(t *testing.T) {
|
||||
// It goes back into a JSON body on the next turn. A split multi-byte
|
||||
// character would make that body unserialisable, and the thread would
|
||||
// break for good rather than for one turn.
|
||||
r := replyBuilder{cap: 4}
|
||||
r.add("üçüncü") // two-byte runes either side of the cut
|
||||
got, _ := r.result()
|
||||
for i, ch := range got {
|
||||
if ch == '�' {
|
||||
t.Fatalf("invalid UTF-8 at byte %d of %q", i, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheAssistantTextIsFoundInBothResponseShapes(t *testing.T) {
|
||||
chat := `{"choices":[{"message":{"role":"assistant","content":"hello"}}]}`
|
||||
if got := replyFromJSON([]byte(chat)); got != "hello" {
|
||||
t.Errorf("chat completion reply = %q", got)
|
||||
}
|
||||
// The completions endpoint, which has no message wrapper.
|
||||
legacy := `{"choices":[{"text":"hello"}]}`
|
||||
if got := replyFromJSON([]byte(legacy)); got != "hello" {
|
||||
t.Errorf("completion reply = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyTheFirstChoiceIsKept(t *testing.T) {
|
||||
// A request asking for several is not one continuing a conversation, and
|
||||
// storing one of them as "what the assistant said" would invent a
|
||||
// conversation that never happened.
|
||||
body := `{"choices":[{"message":{"content":"one"}},{"message":{"content":"two"}}]}`
|
||||
if got := replyFromJSON([]byte(body)); got != "one" {
|
||||
t.Errorf("reply = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamedDeltasAssembleIntoTheReply(t *testing.T) {
|
||||
r := replyBuilder{cap: 1024}
|
||||
for _, ev := range []string{
|
||||
`{"choices":[{"delta":{"role":"assistant"}}]}`,
|
||||
`{"choices":[{"delta":{"content":"Hel"}}]}`,
|
||||
`{"choices":[{"delta":{"content":"lo"}}]}`,
|
||||
`[DONE]`,
|
||||
} {
|
||||
r.add(replyDelta(ev))
|
||||
}
|
||||
if got, _ := r.result(); got != "Hello" {
|
||||
t.Errorf("assembled %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMalformedEventContributesNothing(t *testing.T) {
|
||||
// An engine sending something unexpected must not corrupt a stored
|
||||
// conversation; the worst case is a turn that is shorter than it was.
|
||||
if got := replyDelta(`not json at all`); got != "" {
|
||||
t.Errorf("garbage produced %q", got)
|
||||
}
|
||||
if got := replyDelta(``); got != "" {
|
||||
t.Errorf("an empty event produced %q", got)
|
||||
}
|
||||
}
|
||||
@@ -560,6 +560,7 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
|
||||
Header: req.Header,
|
||||
Stream: req.Stream,
|
||||
SuppressUsageChunk: req.SuppressUsageChunk,
|
||||
CaptureReply: req.captureBytes(),
|
||||
})
|
||||
// A retryable failure is the endpoint's fault; a client disconnect is not.
|
||||
release(perr == nil || errors.Is(perr, proxy.ErrClientGone))
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
@@ -178,3 +180,115 @@ func ownMessages(body []byte) []store.Message {
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// maxStoredReply bounds what one turn may add to a conversation.
|
||||
//
|
||||
// A cap rather than no cap because this is memory held on the request path for
|
||||
// the length of a reply. Generous enough that ordinary answers are stored
|
||||
// whole, small enough that a runaway generation cannot cost the gateway more
|
||||
// than this per concurrent request.
|
||||
const maxStoredReply = 128 << 10
|
||||
|
||||
// captureBytes says how much of the reply the proxy should hand back.
|
||||
//
|
||||
// Zero for every request not continuing a conversation, which is what keeps
|
||||
// the cost of this feature off the path of requests that do not use it.
|
||||
func (p *parsedRequest) captureBytes() int {
|
||||
if p.ThreadID == "" {
|
||||
return 0
|
||||
}
|
||||
return maxStoredReply
|
||||
}
|
||||
|
||||
// loadThread splices a stored conversation into the request.
|
||||
//
|
||||
// Returns how many turns were dropped to make it fit, or an error to refuse
|
||||
// the request with. A thread on a project that keeps no conversations is
|
||||
// refused rather than ignored: ignoring it would send the model a conversation
|
||||
// with no history, and it would answer confidently about something it was
|
||||
// never told — which reads as the model being wrong rather than as a setting
|
||||
// being off.
|
||||
func (r *Router) loadThread(ctx context.Context, req *parsedRequest) (int, error) {
|
||||
if req.ThreadID == "" || r.store == nil {
|
||||
return 0, nil
|
||||
}
|
||||
if req.ProjectID == "" {
|
||||
// No identity, so no project to scope the thread to. Storing
|
||||
// conversations for an unauthenticated gateway would put every
|
||||
// caller's history in one pool, readable by anyone who guesses an id.
|
||||
return 0, errThreadsNeedAuth
|
||||
}
|
||||
|
||||
// Asked before anything else, because a project that keeps no
|
||||
// conversations must be told so rather than served. Without this the
|
||||
// request succeeds, nothing is stored, and the caller finds out only when
|
||||
// the model has no idea what they are referring to — which reads as the
|
||||
// model being broken. The e2e test that caught this asserted the refusal
|
||||
// and got a cheerful 200.
|
||||
retention, err := r.store.ThreadRetention(ctx, req.ProjectID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if retention == 0 {
|
||||
return 0, store.ErrThreadsNotEnabled
|
||||
}
|
||||
|
||||
// What the caller sent, before the history is spliced in, so the exchange
|
||||
// can be stored afterwards without re-reading a body that has changed.
|
||||
req.ThreadTurns = ownMessages(req.UpstreamBody)
|
||||
|
||||
history, err := r.store.Thread(ctx, req.ProjectID, req.ThreadID)
|
||||
if err != nil && !errors.Is(err, store.ErrNoThread) {
|
||||
return 0, err
|
||||
}
|
||||
// A thread that does not exist yet is the first turn of a new one, which is
|
||||
// the ordinary way to start.
|
||||
if len(history) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
_, entry := r.resolve(req.Model)
|
||||
budget := 0
|
||||
if entry.ContextLength > 0 {
|
||||
// The same four-bytes-per-token estimate the rest of the router uses,
|
||||
// less the output the caller asked for, so trimming and the fit check
|
||||
// cannot disagree about what fits.
|
||||
budget = (entry.ContextLength - req.MaxTokens) * 4
|
||||
if budget < 0 {
|
||||
budget = 0
|
||||
}
|
||||
}
|
||||
|
||||
body, dropped, err := spliceThread(req.UpstreamBody, history, budget)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
req.UpstreamBody = body
|
||||
req.PromptTokens = len(body) / 4
|
||||
return dropped, nil
|
||||
}
|
||||
|
||||
// saveThread stores the exchange once a reply has been produced.
|
||||
//
|
||||
// Best effort, and deliberately after the client has been served: a
|
||||
// conversation that failed to record is a worse outcome than a reply that
|
||||
// never arrived, but only just, and the caller has their answer either way.
|
||||
// The failure is logged rather than swallowed, because a thread silently not
|
||||
// growing is the kind of fault someone discovers days later.
|
||||
func (r *Router) saveThread(ctx context.Context, req *parsedRequest, reply string, model string) {
|
||||
if req.ThreadID == "" || r.store == nil || reply == "" {
|
||||
return
|
||||
}
|
||||
turns := append([]store.Message(nil), req.ThreadTurns...)
|
||||
turns = append(turns, store.Message{Role: "assistant", Content: reply, Model: model})
|
||||
|
||||
if err := r.store.AppendThread(ctx, req.ProjectID, req.ThreadID, turns); err != nil {
|
||||
r.log.Warn("could not store a conversation turn",
|
||||
"thread", req.ThreadID, "project", req.ProjectID, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// errThreadsNeedAuth is refused rather than ignored, for the same reason a
|
||||
// disabled project is.
|
||||
var errThreadsNeedAuth = errors.New(
|
||||
"a conversation needs an API key: without one there is no project to keep it in")
|
||||
@@ -0,0 +1,225 @@
|
||||
package e2e
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// say sends one turn, optionally as part of a thread, and returns what came
|
||||
// back along with the headers that explain what llmgw did.
|
||||
func (h *harness) say(t *testing.T, key, model, thread, text string) (int, http.Header, string) {
|
||||
t.Helper()
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"model": model,
|
||||
"messages": []map[string]string{{"role": "user", "content": text}},
|
||||
"max_tokens": 16,
|
||||
})
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
|
||||
h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if thread != "" {
|
||||
req.Header.Set("x-llmgw-thread", thread)
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
msg, _ := io.ReadAll(resp.Body)
|
||||
return resp.StatusCode, resp.Header, string(msg)
|
||||
}
|
||||
|
||||
// authConfig turns on authentication against a provisioned store. The e2e
|
||||
// tests set this inline elsewhere; threads need it in five places.
|
||||
func authConfig(path string) func(*config.Server) {
|
||||
return func(c *config.Server) {
|
||||
c.Auth = config.Auth{
|
||||
Enabled: true, Store: path,
|
||||
CacheTTL: 50 * time.Millisecond,
|
||||
FlushInterval: 20 * time.Millisecond,
|
||||
BatchSize: 8,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func enableThreads(t *testing.T, path string, hours int) {
|
||||
t.Helper()
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
if err := s.SetThreadRetention(context.Background(), "proj", hours); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func storedTurns(t *testing.T, path, thread string) []store.Message {
|
||||
t.Helper()
|
||||
var out []store.Message
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
msgs, err := s.Thread(context.Background(), "proj", thread)
|
||||
if err != nil {
|
||||
t.Fatalf("reading thread %q: %v", thread, err)
|
||||
}
|
||||
out = msgs
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
func TestAProjectThatNeverAskedStoresNothingAndSaysSo(t *testing.T) {
|
||||
// The property being protected: llmgw keeps no conversations unless a
|
||||
// project turns it on. Refused rather than ignored, because ignoring the
|
||||
// header would send the model a conversation with no history and it would
|
||||
// answer confidently about something it was never told.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
h := start(t, 2, authConfig(path))
|
||||
|
||||
code, _, msg := h.say(t, key, testModel, "th1", "hello")
|
||||
if code != http.StatusBadRequest {
|
||||
t.Fatalf("a thread on a project that keeps none returned %d: %s", code, msg)
|
||||
}
|
||||
if !strings.Contains(msg, "does not keep conversations") {
|
||||
t.Errorf("the refusal does not say why: %s", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAThreadRemembersTheEarlierTurns(t *testing.T) {
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
enableThreads(t, path, 24)
|
||||
h := start(t, 2, authConfig(path))
|
||||
|
||||
if code, _, msg := h.say(t, key, testModel, "th1", "first question"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
if code, _, msg := h.say(t, key, testModel, "th1", "second question"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
turns := storedTurns(t, path, "th1")
|
||||
if len(turns) != 4 {
|
||||
t.Fatalf("stored %d turns, expected two questions and two answers: %+v", len(turns), turns)
|
||||
}
|
||||
if turns[0].Content != "first question" || turns[2].Content != "second question" {
|
||||
t.Errorf("the questions were not stored in order: %+v", turns)
|
||||
}
|
||||
if turns[1].Role != "assistant" || turns[1].Content == "" {
|
||||
t.Errorf("the first answer was not stored: %+v", turns[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAThreadContinuesOnADifferentModel(t *testing.T) {
|
||||
// The whole point. A conversation started on one model is carried to
|
||||
// another, and the stored turns say which model produced each answer —
|
||||
// without that, a thread that moved reads as one voice, and when the
|
||||
// answers disagree that is the first thing anyone asks.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
enableThreads(t, path, 24)
|
||||
h := start(t, 2, func(c *config.Server) {
|
||||
authConfig(path)(c)
|
||||
c.Models = budgetFallbackModels()
|
||||
for i := range c.Endpoints {
|
||||
c.Endpoints[i].Models = []string{testModel, cheapModel}
|
||||
}
|
||||
})
|
||||
|
||||
if code, _, msg := h.say(t, key, testModel, "th1", "started here"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
if code, _, msg := h.say(t, key, cheapModel, "th1", "continued there"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
turns := storedTurns(t, path, "th1")
|
||||
models := map[string]bool{}
|
||||
for _, m := range turns {
|
||||
if m.Role == "assistant" {
|
||||
models[m.Model] = true
|
||||
}
|
||||
}
|
||||
if !models[testModel] || !models[cheapModel] {
|
||||
t.Errorf("the thread does not record both models that answered: %+v", turns)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneProjectCannotReadAnothersThread(t *testing.T) {
|
||||
// Thread ids are opaque, but an unguessable id is a secret and secrets end
|
||||
// up in logs and bug reports. The project scope is what makes a leaked id
|
||||
// useless outside it.
|
||||
path, mint := provision(t)
|
||||
mine := mint(t, "mine")
|
||||
enableThreads(t, path, 24)
|
||||
|
||||
// A key in a second project.
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
if err := s.PutProject(context.Background(),
|
||||
store.Project{ID: "other", TeamID: "team", Name: "Other"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetThreadRetention(context.Background(), "other", 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
theirs := mint(t, "theirs", func(k *store.Key) { k.ProjectID = "other" })
|
||||
|
||||
h := start(t, 2, authConfig(path))
|
||||
if code, _, msg := h.say(t, mine, testModel, "shared-id", "a private thing"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
// The other project using the same id must not see it, and must not be
|
||||
// able to append to it either.
|
||||
code, _, msg := h.say(t, theirs, testModel, "shared-id", "what was said?")
|
||||
if code == http.StatusOK {
|
||||
if strings.Contains(msg, "a private thing") {
|
||||
t.Fatal("one project read another's conversation")
|
||||
}
|
||||
// Served, but from an empty history of its own — check the store
|
||||
// agrees that nothing crossed over.
|
||||
if turns := storedTurns(t, path, "shared-id"); len(turns) > 2 {
|
||||
for _, m := range turns {
|
||||
if m.Content == "a private thing" {
|
||||
t.Fatal("a second project appended into another's thread")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithoutAThreadNothingIsStoredAtAll(t *testing.T) {
|
||||
// Every request that does not opt in must take exactly the path it took
|
||||
// before threads existed.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
enableThreads(t, path, 24)
|
||||
h := start(t, 2, authConfig(path))
|
||||
|
||||
if code, _, msg := h.say(t, key, testModel, "", "no thread here"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
// Pruning with a cutoff far past every retention window deletes every
|
||||
// thread there is, so a count of zero means none were created.
|
||||
n, err := s.PruneThreads(context.Background(), time.Now().Add(1000*time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 0 {
|
||||
t.Errorf("%d threads were created by a request that used none", n)
|
||||
}
|
||||
})
|
||||
}
|
||||
Reference in new issue
Block a user