upstream_response_model_bench_test.go 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. package service
  2. import (
  3. "strings"
  4. "testing"
  5. "github.com/tidwall/gjson"
  6. )
  7. var (
  8. upstreamResponseModelBenchmarkSink string
  9. upstreamResponseModelBenchmarkTerminal = []byte(`{"type":"response.completed","response":{"id":"resp_123","model":"gpt-5.5-2026-04-23"},"usage":{"input_tokens":128,"output_tokens":64}}`)
  10. upstreamResponseModelBenchmarkDelta = []byte(`{"type":"response.output_text.delta","delta":"hello"}`)
  11. upstreamResponseModelBenchmarkWrapper = []byte(`{"response":{"response":{"modelVersion":"gemini-3-pro","candidates":[]}}}`)
  12. )
  13. func BenchmarkUpstreamResponseModelOpenAI(b *testing.B) {
  14. tests := []struct {
  15. name string
  16. payload []byte
  17. }{
  18. {name: "terminal_with_model", payload: upstreamResponseModelBenchmarkTerminal},
  19. {name: "delta_without_model", payload: upstreamResponseModelBenchmarkDelta},
  20. }
  21. for _, tt := range tests {
  22. b.Run(tt.name, func(b *testing.B) {
  23. b.Run("legacy", func(b *testing.B) {
  24. b.ReportAllocs()
  25. for b.Loop() {
  26. upstreamResponseModelBenchmarkSink = benchmarkLegacyOpenAIModel(tt.payload)
  27. }
  28. })
  29. b.Run("optimized", func(b *testing.B) {
  30. b.ReportAllocs()
  31. for b.Loop() {
  32. upstreamResponseModelBenchmarkSink = firstValidTrimmedGJSONString(tt.payload, "response.model", "model")
  33. }
  34. })
  35. })
  36. }
  37. }
  38. func BenchmarkUpstreamResponseModelAntigravityWrapper(b *testing.B) {
  39. b.Run("legacy", func(b *testing.B) {
  40. b.ReportAllocs()
  41. for b.Loop() {
  42. upstreamResponseModelBenchmarkSink = benchmarkLegacyWrappedGeminiModel(upstreamResponseModelBenchmarkWrapper)
  43. }
  44. })
  45. b.Run("optimized", func(b *testing.B) {
  46. b.ReportAllocs()
  47. for b.Loop() {
  48. upstreamResponseModelBenchmarkSink = firstValidTrimmedGJSONString(
  49. upstreamResponseModelBenchmarkWrapper,
  50. "modelVersion",
  51. "response.modelVersion",
  52. "response.response.modelVersion",
  53. )
  54. }
  55. })
  56. }
  57. func benchmarkLegacyOpenAIModel(payload []byte) string {
  58. if len(payload) == 0 || !gjson.ValidBytes(payload) {
  59. return ""
  60. }
  61. return benchmarkFirstTrimmedGJSONModel(
  62. gjson.GetBytes(payload, "response.model"),
  63. gjson.GetBytes(payload, "model"),
  64. )
  65. }
  66. func benchmarkLegacyWrappedGeminiModel(payload []byte) string {
  67. if inner := gjson.GetBytes(payload, "response"); inner.Exists() {
  68. payload = []byte(inner.Raw)
  69. }
  70. if len(payload) == 0 || !gjson.ValidBytes(payload) {
  71. return ""
  72. }
  73. return benchmarkFirstTrimmedGJSONModel(
  74. gjson.GetBytes(payload, "modelVersion"),
  75. gjson.GetBytes(payload, "response.modelVersion"),
  76. )
  77. }
  78. func benchmarkFirstTrimmedGJSONModel(values ...gjson.Result) string {
  79. for _, value := range values {
  80. if !value.Exists() || value.Type != gjson.String {
  81. continue
  82. }
  83. if model := strings.TrimSpace(value.String()); model != "" {
  84. return model
  85. }
  86. }
  87. return ""
  88. }