From f5482317d7e4cf487e200a9b6baa1672f09ed91f Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Tue, 29 Sep 2026 12:07:50 +0100 Subject: [PATCH] test(bench): add benchmarks for translators, SSE and unified route --- bench_test.go | 226 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 226 insertions(+) create mode 100644 bench_test.go diff --git a/bench_test.go b/bench_test.go new file mode 100644 index 0000000..ec2ccf6 --- /dev/null +++ b/bench_test.go @@ -0,0 +1,226 @@ +package traefikllmgateway + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" +) + +const benchStreamDeltas = 500 + +func benchOpenAIRequest() map[string]any { + return map[string]any{ + "model": "bench-model", + "max_tokens": float64(1024), + "temperature": 0.7, + "messages": []any{ + map[string]any{"role": "system", "content": "You are a helpful assistant."}, + map[string]any{"role": "user", "content": "What is the weather in Paris and Berlin?"}, + map[string]any{"role": "assistant", "content": "", "tool_calls": []any{ + map[string]any{"id": "call_1", "type": "function", "function": map[string]any{"name": "get_weather", "arguments": `{"city":"Paris"}`}}, + map[string]any{"id": "call_2", "type": "function", "function": map[string]any{"name": "get_weather", "arguments": `{"city":"Berlin"}`}}, + }}, + map[string]any{"role": "tool", "tool_call_id": "call_1", "content": "18C, cloudy"}, + map[string]any{"role": "tool", "tool_call_id": "call_2", "content": "15C, rain"}, + map[string]any{"role": "user", "content": "Summarise both in one sentence."}, + }, + "tools": []any{ + map[string]any{"type": "function", "function": map[string]any{ + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": map[string]any{ + "type": "object", + "properties": map[string]any{"city": map[string]any{"type": "string"}}, + "required": []any{"city"}, + }, + }}, + }, + "tool_choice": "auto", + } +} + +func benchAnthropicStream() []byte { + var sb strings.Builder + sb.WriteString("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_01B\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-opus-5\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":1}}}\n\n") + sb.WriteString("event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n") + for i := 0; i < benchStreamDeltas; i++ { + fmt.Fprintf(&sb, "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\" token%d\"}}\n\n", i) + } + sb.WriteString("event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n") + sb.WriteString("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":500}}\n\n") + sb.WriteString("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + return []byte(sb.String()) +} + +func benchGeminiStream() []byte { + var sb strings.Builder + for i := 0; i < benchStreamDeltas; i++ { + fmt.Fprintf(&sb, "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\" token%d\"}]}}],\"responseId\":\"resp_B\"}\n\n", i) + } + sb.WriteString("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":5,\"candidatesTokenCount\":500}}\n\n") + return []byte(sb.String()) +} + +func BenchmarkAnthropicRequestFromOpenAI(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + if _, err := anthropicRequestFromOpenAI(benchOpenAIRequest()); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGeminiRequestFromOpenAI(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + if _, err := geminiRequestFromOpenAI(benchOpenAIRequest()); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkOpenAIResponseFromAnthropic(b *testing.B) { + body := []byte(`{"id":"msg_01XYZ","type":"message","role":"assistant","content":[{"type":"text","text":"The capital of France is Paris."},{"type":"tool_use","id":"tu_1","name":"get_weather","input":{"city":"Paris"}}],"model":"claude-opus-5","stop_reason":"tool_use","stop_sequence":null,"usage":{"input_tokens":12,"output_tokens":8}}`) + b.ReportAllocs() + b.SetBytes(int64(len(body))) + for i := 0; i < b.N; i++ { + if _, _, err := openAIResponseFromAnthropic(body, "gw-model", 1734000000); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkOpenAIResponseFromGemini(b *testing.B) { + body := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"The capital of France is Paris."},{"functionCall":{"name":"get_weather","args":{"city":"Paris"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":12,"candidatesTokenCount":8},"responseId":"resp_1"}`) + b.ReportAllocs() + b.SetBytes(int64(len(body))) + for i := 0; i < b.N; i++ { + if _, _, err := openAIResponseFromGemini(body, "gw-model", 1734000000); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkReadSSE(b *testing.B) { + stream := benchAnthropicStream() + b.ReportAllocs() + b.SetBytes(int64(len(stream))) + for i := 0; i < b.N; i++ { + if err := readSSE(bytes.NewReader(stream), func(sseEvent) error { return nil }); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkAnthropicStreamTranslate(b *testing.B) { + stream := benchAnthropicStream() + b.ReportAllocs() + b.SetBytes(int64(len(stream))) + for i := 0; i < b.N; i++ { + st := newAnthropicStreamState("gw-model", 1734000000) + err := readSSE(bytes.NewReader(stream), func(ev sseEvent) error { + _, terr := st.translate(ev) + return terr + }) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGeminiStreamTranslate(b *testing.B) { + stream := benchGeminiStream() + b.ReportAllocs() + b.SetBytes(int64(len(stream))) + for i := 0; i < b.N; i++ { + st := newGeminiStreamState("gw-model", 1734000000) + err := readSSE(bytes.NewReader(stream), func(ev sseEvent) error { + _, terr := st.translate(ev) + return terr + }) + if err != nil { + b.Fatal(err) + } + } +} + +func newBenchGateway(b *testing.B, provider *ProviderConfig) http.Handler { + b.Helper() + devNull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + b.Fatalf("open %s: %v", os.DevNull, err) + } + origStderr := os.Stderr + os.Stderr = devNull + b.Cleanup(func() { + os.Stderr = origStderr + _ = devNull.Close() + }) + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"bench": provider} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}} + next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + b.Fatalf("New: %v", err) + } + return h +} + +func benchChatRequest(payload []byte) *http.Request { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(payload)) + req.Header.Set("Authorization", "Bearer sk-alice") + req.Header.Set("Content-Type", "application/json") + return req +} + +func BenchmarkUnifiedChatOpenAI(b *testing.B) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"x","object":"chat.completion","model":"bench-model","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`)) + })) + b.Cleanup(upstream.Close) + h := newBenchGateway(b, &ProviderConfig{Type: "openai", BaseURL: upstream.URL, APIKey: "k", Models: []string{"bench-model"}}) + payload := []byte(`{"model":"bench-model","messages":[{"role":"user","content":"hi"}]}`) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + rec := httptest.NewRecorder() + h.ServeHTTP(rec, benchChatRequest(payload)) + if rec.Code != http.StatusOK { + b.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String()) + } + } +} + +func BenchmarkUnifiedChatAnthropicStream(b *testing.B) { + stream := benchAnthropicStream() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write(stream) + })) + b.Cleanup(upstream.Close) + h := newBenchGateway(b, &ProviderConfig{Type: "anthropic", BaseURL: upstream.URL, APIKey: "k", Models: []string{"bench-model"}}) + payload := []byte(`{"model":"bench-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + + b.ReportAllocs() + b.SetBytes(int64(len(stream))) + b.ResetTimer() + for i := 0; i < b.N; i++ { + rec := httptest.NewRecorder() + h.ServeHTTP(rec, benchChatRequest(payload)) + if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), "[DONE]") { + b.Fatalf("status = %d, body tail = %.200s", rec.Code, rec.Body.String()) + } + } +}