diff --git a/litellm-rust/loadtest/cmd/s3stub/main.go b/litellm-rust/loadtest/cmd/s3stub/main.go index 8c4a5efa8cc..8bf9817b7c3 100644 --- a/litellm-rust/loadtest/cmd/s3stub/main.go +++ b/litellm-rust/loadtest/cmd/s3stub/main.go @@ -3,10 +3,16 @@ // credentials, so gateways can sign their uploads as they would for AWS and verify can read them // back without any. Objects live until the process exits. // -// s3stub -addr :9100 +// -max-string N shortens every JSON string longer than N bytes in a stored object to its first N +// bytes, so request logs keep their fields and numbers but not whole prompts, and a long run's logs +// fit in memory. Objects that are not JSON are stored as sent. +// +// s3stub -addr :9100 -max-string 256 package main import ( + "bytes" + "encoding/json" "encoding/xml" "flag" "io" @@ -19,8 +25,9 @@ import ( ) type store struct { - mu sync.RWMutex - objects map[string][]byte // "/" + mu sync.RWMutex + objects map[string][]byte // "/" + maxString int } type listResult struct { @@ -42,8 +49,9 @@ const page = 1000 func main() { addr := flag.String("addr", ":9100", "listen address") + maxString := flag.Int("max-string", 0, "shorten JSON strings longer than this many bytes (0 keeps objects as sent)") flag.Parse() - s := &store{objects: map[string][]byte{}} + s := &store{objects: map[string][]byte{}, maxString: *maxString} http.HandleFunc("GET /health", func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "ok") }) http.HandleFunc("/", s.serve) log.Printf("s3stub listening on %s", *addr) @@ -64,6 +72,9 @@ func (s *store) serve(w http.ResponseWriter, r *http.Request) { http.Error(w, err.Error(), http.StatusBadRequest) return } + if s.maxString > 0 { + body = shorten(body, s.maxString) + } s.mu.Lock() s.objects[bucket+"/"+key] = body s.mu.Unlock() @@ -116,3 +127,36 @@ func (s *store) list(w http.ResponseWriter, r *http.Request, bucket string) { io.WriteString(w, xml.Header) xml.NewEncoder(w).Encode(res) } + +// shorten re-encodes a JSON object with every string cut to max bytes; numbers keep their text. +func shorten(body []byte, max int) []byte { + dec := json.NewDecoder(bytes.NewReader(body)) + dec.UseNumber() + var v any + if err := dec.Decode(&v); err != nil || dec.More() { + return body + } + out, err := json.Marshal(cut(v, max)) + if err != nil { + return body + } + return out +} + +func cut(v any, max int) any { + switch v := v.(type) { + case string: + if len(v) > max { + return v[:max] + } + case []any: + for i := range v { + v[i] = cut(v[i], max) + } + case map[string]any: + for k := range v { + v[k] = cut(v[k], max) + } + } + return v +} diff --git a/litellm-rust/loadtest/cmd/s3stub/main_test.go b/litellm-rust/loadtest/cmd/s3stub/main_test.go new file mode 100644 index 00000000000..0f1644b3009 --- /dev/null +++ b/litellm-rust/loadtest/cmd/s3stub/main_test.go @@ -0,0 +1,15 @@ +package main + +import "testing" + +func TestShortenKeepsFieldsAndNumbers(t *testing.T) { + in := []byte(`{"messages":[{"content":"abcdefgh"}],"response_cost":0.10262000000000001,"n":3,"ok":true}`) + got := string(shorten(in, 3)) + want := `{"messages":[{"content":"abc"}],"n":3,"ok":true,"response_cost":0.10262000000000001}` + if got != want { + t.Fatalf("got %s", got) + } + if got := string(shorten([]byte("not json"), 3)); got != "not json" { + t.Fatalf("non-JSON changed: %s", got) + } +}