74 lines
2.4 KiB
Go
74 lines
2.4 KiB
Go
package translator
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
|
|
translatorapi "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/translator"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
var benchmarkResponseTranslationOutput []byte
|
|
|
|
func BenchmarkResponseTranslationLargePayload(b *testing.B) {
|
|
payload := strings.Repeat("x", 8<<20)
|
|
cases := []struct {
|
|
name string
|
|
from string
|
|
to string
|
|
rawJSON []byte
|
|
}{
|
|
{
|
|
name: "gemini_to_openai",
|
|
from: "gemini",
|
|
to: "openai",
|
|
rawJSON: []byte(`{"modelVersion":"gemini-test","candidates":[{"index":0,"content":{"parts":[{"text":"` + payload + `"}]},"finishReason":"STOP"}]}`),
|
|
},
|
|
{
|
|
name: "codex_to_openai",
|
|
from: "codex",
|
|
to: "openai",
|
|
rawJSON: []byte(`{"type":"response.completed","response":{"id":"resp_1","created_at":1700000000,"model":"gpt-test","status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"` + payload + `"}]}]}}`),
|
|
},
|
|
{
|
|
name: "claude_to_openai",
|
|
from: "claude",
|
|
to: "openai",
|
|
rawJSON: claudeLargeTextResponse(payload),
|
|
},
|
|
{
|
|
name: "claude_to_openai-response",
|
|
from: "claude",
|
|
to: "openai-response",
|
|
rawJSON: claudeLargeTextResponse(payload),
|
|
},
|
|
}
|
|
|
|
for _, testCase := range cases {
|
|
b.Run(testCase.name, func(b *testing.B) {
|
|
output := translatorapi.ResponseNonStream(testCase.from, testCase.to, context.Background(), "benchmark-model", nil, nil, testCase.rawJSON, nil)
|
|
if !gjson.ValidBytes(output) {
|
|
b.Fatalf("translator generated invalid JSON: %s", output)
|
|
}
|
|
if !bytes.Contains(output, []byte(payload)) {
|
|
b.Fatal("translator dropped the benchmark payload")
|
|
}
|
|
b.ReportAllocs()
|
|
b.SetBytes(int64(len(testCase.rawJSON)))
|
|
b.ResetTimer()
|
|
|
|
for b.Loop() {
|
|
benchmarkResponseTranslationOutput = translatorapi.ResponseNonStream(testCase.from, testCase.to, context.Background(), "benchmark-model", nil, nil, testCase.rawJSON, nil)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func claudeLargeTextResponse(payload string) []byte {
|
|
return []byte("data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"model\":\"claude-test\"}}\n" +
|
|
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\"}}\n" +
|
|
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"" + payload + "\"}}\n" +
|
|
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n")
|
|
}
|