package router import ( "bytes" "context" "crypto/sha256" "encoding/json" "fmt" "io" "net/http" "net/http/httptest" "strings" "sync" "sync/atomic" "testing" "time" "github.com/prometheus/client_golang/prometheus/testutil" "gitea.mynodes.xyz/ilker/llmgw/internal/guard" "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" "gitea.mynodes.xyz/ilker/llmgw/internal/store" ) // cachingRouter is pricedRouter with its project reusing answers for an // hour, reproducible ones only unless anyAnswer. func cachingRouter(t *testing.T, engineURL string, anyAnswer bool) (*Router, string) { t.Helper() rt, secret := pricedRouter(t, engineURL) if err := rt.store.SetAnswerCache(t.Context(), mustProjectOf(t, rt, "key-1"), time.Hour, anyAnswer); err != nil { t.Fatal(err) } return rt, secret } func ask(temperature any, text string) map[string]any { b := map[string]any{"model": "q14", "messages": []any{map[string]any{"role": "user", "content": text}}} if temperature != nil { b["temperature"] = temperature } return b } // usageOf waits for the record of one request. func usageOf(t *testing.T, rt *Router, id string) store.Usage { t.Helper() deadline := time.Now().Add(5 * time.Second) for time.Now().Before(deadline) { if u, ok, err := rt.store.UsageByRequest(t.Context(), id); err == nil && ok { return u } time.Sleep(20 * time.Millisecond) } t.Fatalf("no usage was recorded for %s", id) return store.Usage{} } func TestAProjectThatReusesAnswersAnswersTheSameRequestWithoutTheModel(t *testing.T) { e := newFakeEngine(t, "qwen2.5-14b-instruct") rt, secret := cachingRouter(t, e.srv.URL, false) h := rt.Handler() auth := func(id string) map[string]string { return map[string]string{"Authorization": "Bearer " + secret, "X-Request-Id": id} } first := post(t, h, ask(0, "what is 2+2"), auth("first")) if first.Code != 200 || first.Header().Get(cacheHeader) != "miss" { t.Fatalf("first: %d %q %s", first.Code, first.Header().Get(cacheHeader), first.Body.String()) } // The same request, said differently: keys in another order, another // end user, the temperature written as a float, by the canonical name. again := map[string]any{"messages": []any{map[string]any{"content": "what is 2+2", "role": "user"}}, "temperature": 0.0, "model": "qwen2.5-14b-instruct", "user": "someone-else"} second := post(t, h, again, auth("second")) if second.Code != 200 || second.Header().Get(cacheHeader) != "hit" { t.Fatalf("second: %d %q %s", second.Code, second.Header().Get(cacheHeader), second.Body.String()) } if e.calls.Load() != 1 { t.Errorf("the engine was asked %d times; the second answer should have been reused", e.calls.Load()) } if second.Body.String() != first.Body.String() { t.Errorf("reused %s, first was %s", second.Body.String(), first.Body.String()) } if second.Header().Get("Age") == "" { t.Error("a reused answer did not say its age") } // Recorded, as a request that generated nothing. miss, hit := usageOf(t, rt, "first"), usageOf(t, rt, "second") if miss.CacheHit || miss.CostMicros == 0 || miss.PromptTokens == 0 { t.Errorf("the answer the model gave: %+v", miss) } if !hit.CacheHit || hit.CostMicros != 0 || hit.PromptTokens != 0 || hit.CompletionTokens != 0 || hit.Status != 200 || hit.ServedModel != "qwen2.5-14b-instruct" { t.Errorf("the reused answer: %+v", hit) } // A sampled answer is asked for each time, and says so. for i := range 2 { rec := post(t, h, ask(0.7, "write a poem"), auth(fmt.Sprintf("poem-%d", i))) if rec.Header().Get(cacheHeader) != "bypass" { t.Errorf("a sampled request: %q", rec.Header().Get(cacheHeader)) } } if e.calls.Load() != 3 { t.Errorf("the engine was asked %d times, expected each sampled request to reach it", e.calls.Load()) } // A seed makes it reproducible. seeded := ask(0.7, "write a poem") seeded["seed"] = 42 post(t, h, seeded, auth("seed-1")) if rec := post(t, h, seeded, auth("seed-2")); rec.Header().Get(cacheHeader) != "hit" { t.Errorf("a seeded request: %q", rec.Header().Get(cacheHeader)) } // A client that asks for a fresh answer gets one. hdr := auth("fresh") hdr["Cache-Control"] = "no-cache" if rec := post(t, h, ask(0, "what is 2+2"), hdr); rec.Header().Get(cacheHeader) != "bypass" { t.Errorf("Cache-Control: no-cache: %q", rec.Header().Get(cacheHeader)) } hdr = auth("fresh-2") hdr[cacheHeader] = "no" if rec := post(t, h, ask(0, "what is 2+2"), hdr); rec.Header().Get(cacheHeader) != "bypass" { t.Errorf("x-llmgw-cache: no: %q", rec.Header().Get(cacheHeader)) } if n := testutil.ToFloat64(rt.m.cacheRequests.WithLabelValues("qwen2.5-14b-instruct", "hit")); n != 2 { t.Errorf("%v hits counted", n) } if n := testutil.ToFloat64(rt.m.cacheRequests.WithLabelValues("qwen2.5-14b-instruct", "miss")); n != 2 { t.Errorf("%v misses counted", n) } if n := testutil.ToFloat64(rt.m.cacheBytes); n <= 0 { t.Errorf("the cache holds %v bytes", n) } // Turned off, nothing is reused, and nothing is said about it. _ = rt.store.SetAnswerCache(t.Context(), mustProjectOf(t, rt, "key-1"), 0, false) rt.reqLogs.by = nil calls := e.calls.Load() if rec := post(t, h, ask(0, "what is 2+2"), auth("off")); rec.Header().Get(cacheHeader) != "" || e.calls.Load() != calls+1 { t.Errorf("a project that reuses nothing: %q, %d calls", rec.Header().Get(cacheHeader), e.calls.Load()-calls) } } func TestAProjectThatReusesAnyAnswerReusesSampledOnes(t *testing.T) { e := newFakeEngine(t, "qwen2.5-14b-instruct") rt, secret := cachingRouter(t, e.srv.URL, true) auth := map[string]string{"Authorization": "Bearer " + secret} post(t, rt.Handler(), ask(nil, "write a poem"), auth) if rec := post(t, rt.Handler(), ask(nil, "write a poem"), auth); rec.Header().Get(cacheHeader) != "hit" { t.Errorf("%q", rec.Header().Get(cacheHeader)) } } func TestKeysOfOneProjectNeverSeeAnotherProjectsAnswers(t *testing.T) { e := newFakeEngine(t, "qwen2.5-14b-instruct") rt, secret := cachingRouter(t, e.srv.URL, false) team, err := rt.store.CreateTeam(t.Context(), mustOrgOf(t, rt), "other") if err != nil { t.Fatal(err) } other, err := rt.store.CreateProject(t.Context(), team, "elsewhere") if err != nil { t.Fatal(err) } _ = rt.store.SetAnswerCache(t.Context(), other, time.Hour, true) secret2 := store.KeyPrefix + strings.Repeat("b", 43) sum := sha256.Sum256([]byte(secret2)) if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-2", Hash: sum[:], Prefix: "bbbb", Name: "other", ProjectID: other, CreatedAt: time.Now()}); err != nil { t.Fatal(err) } post(t, rt.Handler(), ask(0, "the secret plan"), map[string]string{"Authorization": "Bearer " + secret}) rec := post(t, rt.Handler(), ask(0, "the secret plan"), map[string]string{"Authorization": "Bearer " + secret2}) if rec.Header().Get(cacheHeader) != "miss" || e.calls.Load() != 2 { t.Errorf("another project's key: %q, the engine asked %d times", rec.Header().Get(cacheHeader), e.calls.Load()) } } func mustOrgOf(t *testing.T, rt *Router) string { t.Helper() orgs, err := rt.store.Tree(t.Context()) if err != nil || len(orgs) == 0 { t.Fatalf("no organization: %v", err) } return orgs[0].ID } func TestAFailedAnswerIsNotReused(t *testing.T) { e := newFakeEngine(t, "qwen2.5-14b-instruct") rt, secret := cachingRouter(t, e.srv.URL, false) auth := map[string]string{"Authorization": "Bearer " + secret} e.status.Store(http.StatusBadRequest) if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code == 200 { t.Fatalf("%d", rec.Code) } e.status.Store(0) if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code != 200 || rec.Header().Get(cacheHeader) != "miss" { t.Errorf("after a failure: %d %q", rec.Code, rec.Header().Get(cacheHeader)) } } // switchGuard blocks every request, or puts every one under no_store, when // told to. type switchGuard struct { block, noStore atomic.Bool // answers, when set, stops answers saying "hi". answers atomic.Bool } func (g *switchGuard) Input(context.Context, guard.Subject, []byte) (guard.Verdict, error) { if g.block.Load() { return guard.Verdict{Action: guard.Block, Policy: "nothing"}, nil } return guard.Verdict{}, nil } func (g *switchGuard) Output(context.Context, guard.Subject) guard.OutputFilter { if !g.answers.Load() { return nil } return &stopHi{} } func (g *switchGuard) Directives(guard.Subject) guard.Directives { return guard.Directives{NoStore: g.noStore.Load()} } func TestTheGuardrailsStillDecideAboutAnAnswerFromTheCache(t *testing.T) { e := newFakeEngine(t, "qwen2.5-14b-instruct") rt, secret := cachingRouter(t, e.srv.URL, false) g := &switchGuard{} rt.SetGuard(g) auth := map[string]string{"Authorization": "Bearer " + secret} if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Header().Get(cacheHeader) != "miss" { t.Fatalf("%q", rec.Header().Get(cacheHeader)) } // A policy that now refuses the request refuses it, answer kept or not. g.block.Store(true) if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code != http.StatusBadRequest || rec.Header().Get(cacheHeader) != "" { t.Errorf("a blocked request: %d %q %s", rec.Code, rec.Header().Get(cacheHeader), rec.Body.String()) } g.block.Store(false) // A policy that keeps nothing neither reads the cache nor fills it. g.noStore.Store(true) if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code != 200 || rec.Header().Get(cacheHeader) != "bypass" { t.Errorf("under no_store: %d %q", rec.Code, rec.Header().Get(cacheHeader)) } if rec := post(t, rt.Handler(), ask(0, "under no_store"), auth); rec.Header().Get(cacheHeader) != "bypass" { t.Errorf("under no_store: %q", rec.Header().Get(cacheHeader)) } g.noStore.Store(false) if rec := post(t, rt.Handler(), ask(0, "under no_store"), auth); rec.Header().Get(cacheHeader) != "miss" { t.Errorf("an answer given under no_store was kept: %q", rec.Header().Get(cacheHeader)) } // An output policy added since judges the reused answer too. g.answers.Store(true) rec := post(t, rt.Handler(), ask(0, "hello"), auth) if rec.Header().Get(cacheHeader) != "hit" || !strings.Contains(rec.Body.String(), "content_filter") || rec.Header().Get(guardHeader) != "blocked" { t.Errorf("an answer the output policy stops: %q %q %s", rec.Header().Get(cacheHeader), rec.Header().Get(guardHeader), rec.Body.String()) } // And an answer it stopped is not kept. if rec := post(t, rt.Handler(), ask(0, "new question"), auth); rec.Header().Get(cacheHeader) != "miss" { t.Fatalf("%q", rec.Header().Get(cacheHeader)) } g.answers.Store(false) if rec := post(t, rt.Handler(), ask(0, "new question"), auth); rec.Header().Get(cacheHeader) != "miss" { t.Errorf("an answer a guard stopped was kept: %q", rec.Header().Get(cacheHeader)) } } // greetingEngine streams "Hello world" to a request that asks for a // stream, with usage when asked, and answers whole otherwise. func greetingEngine(t *testing.T) (*httptest.Server, *atomic.Int64) { t.Helper() var calls atomic.Int64 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/metrics" { return } calls.Add(1) var b struct { Stream bool `json:"stream"` StreamOptions struct { IncludeUsage bool `json:"include_usage"` } `json:"stream_options"` } body, _ := io.ReadAll(r.Body) _ = json.Unmarshal(body, &b) if !b.Stream { w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{"id":"c1","object":"chat.completion","created":1,"model":"qwen2.5-14b-instruct",`+ `"choices":[{"index":0,"message":{"role":"assistant","content":"Hello world"},"finish_reason":"stop"}],`+ `"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`) return } w.Header().Set("Content-Type", "text/event-stream") head := `{"id":"c1","object":"chat.completion.chunk","created":1,"model":"qwen2.5-14b-instruct","choices":` for _, c := range []string{ `[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}`, `[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}`, `[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}`, `[{"index":0,"delta":{},"finish_reason":"stop"}]}`, } { fmt.Fprintf(w, "data: %s%s\n\n", head, c) } if b.StreamOptions.IncludeUsage { fmt.Fprintf(w, "data: %s[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n\n", head) } _, _ = io.WriteString(w, "data: [DONE]\n\n") })) t.Cleanup(srv.Close) return srv, &calls } // streamedText is the content of a streamed chat answer, and whether it // carried usage and ended with [DONE]. func streamedText(t *testing.T, sse string) (text string, usage, done bool) { t.Helper() for _, line := range strings.Split(sse, "\n") { data, ok := strings.CutPrefix(line, "data: ") if !ok { continue } if data == "[DONE]" { done = true continue } var c struct { Object string `json:"object"` Choices []struct { Delta struct { Content string `json:"content"` } `json:"delta"` } `json:"choices"` Usage *struct { CompletionTokens int `json:"completion_tokens"` } `json:"usage"` } if err := json.Unmarshal([]byte(data), &c); err != nil || c.Object != "chat.completion.chunk" { t.Fatalf("not a chunk: %q", data) } for _, ch := range c.Choices { text += ch.Delta.Content } if c.Usage != nil && c.Usage.CompletionTokens == 2 { usage = true } } return text, usage, done } func TestAStreamedAnswerIsReusedForClientsThatStreamAndThatDoNot(t *testing.T) { srv, calls := greetingEngine(t) rt, secret := cachingRouter(t, srv.URL, false) auth := map[string]string{"Authorization": "Bearer " + secret} streamed := func(includeUsage bool) map[string]any { b := ask(0, "greet me") b["stream"] = true if includeUsage { b["stream_options"] = map[string]any{"include_usage": true} } return b } first := post(t, rt.Handler(), streamed(false), auth) if first.Header().Get(cacheHeader) != "miss" { t.Fatalf("%q %s", first.Header().Get(cacheHeader), first.Body.String()) } if text, _, _ := streamedText(t, first.Body.String()); text != "Hello world" { t.Fatalf("streamed %q", text) } // Whole, to a client that does not stream. whole := post(t, rt.Handler(), ask(0, "greet me"), auth) var c struct { Object string `json:"object"` Choices []struct { Message struct { Role, Content string } `json:"message"` FinishReason string `json:"finish_reason"` } `json:"choices"` Usage struct { PromptTokens, CompletionTokens int } `json:"usage"` } if err := json.Unmarshal(whole.Body.Bytes(), &c); err != nil || whole.Header().Get(cacheHeader) != "hit" { t.Fatalf("%q %s %v", whole.Header().Get(cacheHeader), whole.Body.String(), err) } if c.Object != "chat.completion" || len(c.Choices) != 1 || c.Choices[0].Message.Content != "Hello world" || c.Choices[0].Message.Role != "assistant" || c.Choices[0].FinishReason != "stop" { t.Errorf("put back together as %s", whole.Body.String()) } // Streamed again, with the usage only to a client that asks for it. for _, withUsage := range []bool{false, true} { rec := post(t, rt.Handler(), streamed(withUsage), auth) text, usage, done := streamedText(t, rec.Body.String()) if rec.Header().Get(cacheHeader) != "hit" || text != "Hello world" || !done || usage != withUsage || !strings.HasPrefix(rec.Header().Get("Content-Type"), "text/event-stream") { t.Errorf("streamed from the cache (usage %v): %q %q usage=%v done=%v\n%s", withUsage, rec.Header().Get(cacheHeader), text, usage, done, rec.Body.String()) } } if calls.Load() != 1 { t.Errorf("the engine was asked %d times", calls.Load()) } } func TestManyClientsReadingAndFillingTheCacheAtOnce(t *testing.T) { e := newFakeEngine(t, "qwen2.5-14b-instruct") rt, secret := cachingRouter(t, e.srv.URL, false) h := rt.Handler() auth := map[string]string{"Authorization": "Bearer " + secret} var wg sync.WaitGroup for i := range 40 { wg.Go(func() { rec := post(t, h, ask(0, fmt.Sprintf("question %d", i%5)), auth) if rec.Code != 200 { t.Errorf("%d %s", rec.Code, rec.Body.String()) } }) } wg.Wait() for i := range 5 { if rec := post(t, h, ask(0, fmt.Sprintf("question %d", i)), auth); rec.Header().Get(cacheHeader) != "hit" { t.Errorf("question %d: %q", i, rec.Header().Get(cacheHeader)) } } } func TestTheAnswerCacheDropsWhatWasUsedLongestAgoPastItsBudget(t *testing.T) { var held int64 c := newAnswerCache(1000, 400, func(b int64, _ int) { held = b }) now := time.Now() put := func(key string) bool { return c.put(&cachedAnswer{key: key, project: "p", body: bytes.Repeat([]byte("x"), 100), at: now, expires: now.Add(time.Hour)}, now) } for _, k := range []string{"a", "b", "c"} { put(k) } c.get("a", time.Hour, now) // a is now the most recently used put("d") put("e") if c.get("b", time.Hour, now) != nil { t.Error("b, used longest ago, was kept past the budget") } if c.get("a", time.Hour, now) == nil { t.Error("a, used recently, was dropped") } if held > 1000 { t.Errorf("holding %d bytes of 1000", held) } if c.put(&cachedAnswer{key: "big", body: bytes.Repeat([]byte("x"), 500)}, now) { t.Error("an answer larger than an entry may be was kept") } // Past its TTL, or past a TTL shortened since, it is gone. if c.get("a", time.Hour, now.Add(2*time.Hour)) != nil { t.Error("an expired answer was returned") } if c.get("e", time.Minute, now.Add(2*time.Minute)) != nil { t.Error("an answer older than the project's TTL now was returned") } } func TestAnAnswerIsLookedUpByWhatChangesIt(t *testing.T) { key := func(project, model, body string) string { k, ok := answerKey(project, "/v1/chat/completions", model, []byte(body)) if !ok { t.Fatalf("no key for %s", body) } return k } base := key("p", "m", `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0}`) same := []string{ `{"messages":[{"content":"hi","role":"user"}],"temperature":0.0,"model":"alias"}`, `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0,"stream":true,"stream_options":{"include_usage":true}}`, `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0e0,"user":"u","metadata":{"a":"b"}}`, } for _, b := range same { if key("p", "m", b) != base { t.Errorf("%s was looked up apart", b) } } different := []string{ `{"model":"m","messages":[{"role":"user","content":"hi!"}],"temperature":0}`, `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0,"max_tokens":5}`, `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0,"seed":9007199254740993}`, } for _, b := range different { if key("p", "m", b) == base { t.Errorf("%s was looked up as the same", b) } } if key("p", "m", `{"seed":9007199254740993}`) == key("p", "m", `{"seed":9007199254740992}`) { t.Error("two large seeds were taken for one") } if key("other", "m", `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0}`) == base { t.Error("another project's request was looked up as the same") } if key("p", "other", `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0}`) == base { t.Error("another model's request was looked up as the same") } if _, ok := answerKey("p", "/v1/chat/completions", "m", []byte(`not json`)); ok { t.Error("a body that is not JSON was given a key") } } func TestAStreamedAnswerIsPutBackTogetherAndTakenApartAgain(t *testing.T) { sse := strings.Join([]string{ `data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}`, `data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]}`, `data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":"}}]},"finish_reason":null}]}`, `data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"Paris\"}"}}]},"finish_reason":null}]}`, `data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, `data: [DONE]`, }, "\n\n") + "\n\n" body, ok := assembleStream("/v1/chat/completions", []byte(sse), &proxy.Usage{PromptTokens: 3, CompletionTokens: 4}) if !ok { t.Fatal("not put back together") } var c struct { Choices []struct { Message struct { Content *string `json:"content"` ToolCalls []toolCall `json:"tool_calls"` } `json:"message"` FinishReason string `json:"finish_reason"` } `json:"choices"` Usage struct { TotalTokens int `json:"total_tokens"` } `json:"usage"` } if err := json.Unmarshal(body, &c); err != nil { t.Fatal(err) } if len(c.Choices) != 1 || c.Choices[0].Message.Content != nil || len(c.Choices[0].Message.ToolCalls) != 1 || c.Choices[0].Message.ToolCalls[0].Function.Arguments != `{"city":"Paris"}` || c.Choices[0].Message.ToolCalls[0].ID != "call_1" || c.Choices[0].FinishReason != "tool_calls" || c.Usage.TotalTokens != 7 { t.Errorf("put back together as %s", body) } // Taken apart, and put back together again, it is the same answer. events, err := streamFromAnswer("/v1/chat/completions", body, false) if err != nil { t.Fatal(err) } again, ok := assembleStream("/v1/chat/completions", []byte("data: "+strings.Join(append(events, "[DONE]"), "\n\ndata: ")+"\n\n"), &proxy.Usage{PromptTokens: 3, CompletionTokens: 4}) if !ok || !bytes.Equal(again, body) { t.Errorf("round trip:\n%s\n%s", body, again) } // What it does not understand, and a stream cut short, are not kept. for _, bad := range []string{ strings.Replace(sse, `"delta":{}`, `"delta":{"audio":{"id":"a"}}`, 1), strings.Replace(sse, `"finish_reason":null}]}`, `"finish_reason":null,"logprobs":{"content":[]}}]}`, 1), strings.Replace(sse, "data: [DONE]", "", 1), } { if _, ok := assembleStream("/v1/chat/completions", []byte(bad), nil); ok { t.Errorf("kept: %s", bad) } } // A completion, too. text := `data: {"id":"t","object":"text_completion","created":1,"model":"m","choices":[{"index":0,"text":"Hel","finish_reason":null}]}` + "\n\n" + `data: {"id":"t","object":"text_completion","created":1,"model":"m","choices":[{"index":0,"text":"lo","finish_reason":"length"}]}` + "\n\n" + "data: [DONE]\n\n" tb, ok := assembleStream("/v1/completions", []byte(text), nil) if !ok || !strings.Contains(string(tb), `"text":"Hello"`) || !strings.Contains(string(tb), `"object":"text_completion"`) { t.Errorf("a completion: %s", tb) } } func TestAnAnswerFromTheCacheReachesAnthropicClientsInTheirShape(t *testing.T) { srv, calls := greetingEngine(t) rt, secret := cachingRouter(t, srv.URL, false) send := func(stream bool) *httptest.ResponseRecorder { b, _ := json.Marshal(map[string]any{"model": "q14", "max_tokens": 50, "temperature": 0, "stream": stream, "messages": []any{map[string]any{"role": "user", "content": "greet me"}}}) req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(b)) req.Header.Set("Content-Type", "application/json") req.Header.Set("x-api-key", secret) rec := httptest.NewRecorder() rt.Handler().ServeHTTP(rec, req) return rec } first := send(true) if first.Code != 200 || !strings.Contains(first.Body.String(), "Hello") { t.Fatalf("%d %s", first.Code, first.Body.String()) } whole := send(false) if whole.Code != 200 || !strings.Contains(whole.Body.String(), `"text":"Hello world"`) { t.Errorf("whole: %d %s", whole.Code, whole.Body.String()) } streamed := send(true) if !strings.Contains(streamed.Body.String(), "content_block_delta") || !strings.Contains(streamed.Body.String(), "Hello world") || !strings.Contains(streamed.Body.String(), "message_stop") { t.Errorf("streamed: %s", streamed.Body.String()) } if calls.Load() != 1 { t.Errorf("the engine was asked %d times", calls.Load()) } }