120 lines
2.9 KiB
Go
120 lines
2.9 KiB
Go
package llm
|
|
|
|
import (
|
|
"encoding/json"
|
|
"runtime"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
)
|
|
|
|
var prepareGenerationBenchmarkSink *PreparedGeneration
|
|
|
|
func BenchmarkPrepareGeneration(b *testing.B) {
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
|
|
BaseURL: "https://api.example.test/v1",
|
|
Model: "benchmark-model",
|
|
})
|
|
if err != nil {
|
|
b.Fatalf("construct benchmark client: %v", err)
|
|
}
|
|
|
|
prompt := domain.RenderedPrompt{
|
|
Messages: []domain.RenderedMessage{{Role: "user", Content: "Summarize this request."}},
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
request domain.GenerateRequest
|
|
}{
|
|
{
|
|
name: "no_extras",
|
|
request: domain.GenerateRequest{
|
|
Prompt: prompt,
|
|
Target: domain.ExecutionTarget{Model: "benchmark-model"},
|
|
},
|
|
},
|
|
{
|
|
name: "small_extras",
|
|
request: domain.GenerateRequest{
|
|
Prompt: prompt,
|
|
Target: domain.ExecutionTarget{Model: "benchmark-model", ExtraParams: map[string]any{
|
|
"priority": "low",
|
|
"attempt": 3,
|
|
"enabled": true,
|
|
}},
|
|
},
|
|
},
|
|
{
|
|
name: "1MiB_string_extra",
|
|
request: domain.GenerateRequest{
|
|
Prompt: prompt,
|
|
Target: domain.ExecutionTarget{Model: "benchmark-model", ExtraParams: map[string]any{
|
|
"payload": strings.Repeat("x", 1<<20),
|
|
}},
|
|
},
|
|
},
|
|
{
|
|
name: "nested_json_extras",
|
|
request: domain.GenerateRequest{
|
|
Prompt: prompt,
|
|
Target: domain.ExecutionTarget{Model: "benchmark-model", ExtraParams: map[string]any{
|
|
"options": map[string]any{
|
|
"enabled": true,
|
|
"limits": []any{2, 4, 8},
|
|
"tags": []string{"summary", "offline"},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
{
|
|
name: "structured_output",
|
|
request: domain.GenerateRequest{
|
|
Prompt: prompt,
|
|
Target: domain.ExecutionTarget{Model: "benchmark-model"},
|
|
StructuredOutput: &domain.StructuredOutputSpec{
|
|
Type: domain.StructuredOutputJSONSchema,
|
|
JSONSchema: &domain.StructuredOutputJSONSpec{
|
|
Name: "summary",
|
|
Strict: true,
|
|
Schema: map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"summary": map[string]any{"type": "string"},
|
|
"count": map[string]any{"type": "integer"},
|
|
},
|
|
"required": []string{"summary", "count"},
|
|
"additionalProperties": false,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
b.Run(tc.name, func(b *testing.B) {
|
|
prepared, err := client.PrepareGeneration(tc.request)
|
|
if err != nil {
|
|
b.Fatalf("check prepared request: %v", err)
|
|
}
|
|
if !json.Valid(prepared.payload) {
|
|
b.Fatal("prepared request is not valid JSON")
|
|
}
|
|
|
|
b.ReportAllocs()
|
|
b.SetBytes(int64(len(prepared.payload)))
|
|
b.ResetTimer()
|
|
for range b.N {
|
|
prepared, err := client.PrepareGeneration(tc.request)
|
|
if err != nil {
|
|
b.Fatalf("prepare generation: %v", err)
|
|
}
|
|
prepareGenerationBenchmarkSink = prepared
|
|
}
|
|
b.StopTimer()
|
|
runtime.KeepAlive(prepareGenerationBenchmarkSink)
|
|
})
|
|
}
|
|
}
|