feat: support cached input token billing
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
@@ -29,6 +30,44 @@ func TestTokenUsageAmountsFallsBackToInputOutputTotal(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingsSplitsCachedInputTokens(t *testing.T) {
|
||||
service := &Service{}
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
ModelName: "test-model",
|
||||
BillingConfig: map[string]any{
|
||||
"textInputPer1k": 1.0,
|
||||
"textCachedInputPer1k": 0.25,
|
||||
"textOutputPer1k": 2.0,
|
||||
},
|
||||
}
|
||||
response := clients.Response{
|
||||
Usage: clients.Usage{
|
||||
InputTokens: 1000,
|
||||
CachedInputTokens: 400,
|
||||
OutputTokens: 500,
|
||||
TotalTokens: 1500,
|
||||
},
|
||||
}
|
||||
|
||||
lines := service.billings(context.Background(), nil, "chat.completions", nil, candidate, response, false)
|
||||
|
||||
if len(lines) != 3 {
|
||||
t.Fatalf("expected uncached input, cached input, and output lines, got %+v", lines)
|
||||
}
|
||||
uncached, _ := lines[0].(map[string]any)
|
||||
cached, _ := lines[1].(map[string]any)
|
||||
output, _ := lines[2].(map[string]any)
|
||||
if uncached["resourceType"] != "text_input" || uncached["quantity"] != 600 || uncached["amount"] != 0.6 {
|
||||
t.Fatalf("unexpected uncached input billing line: %+v", uncached)
|
||||
}
|
||||
if cached["resourceType"] != "text_cached_input" || cached["quantity"] != 400 || cached["amount"] != 0.1 {
|
||||
t.Fatalf("unexpected cached input billing line: %+v", cached)
|
||||
}
|
||||
if output["resourceType"] != "text_output" || output["quantity"] != 500 || output["amount"] != 1.0 {
|
||||
t.Fatalf("unexpected output billing line: %+v", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEffectiveRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.T) {
|
||||
policy := effectiveRateLimitPolicy(store.RuntimeModelCandidate{
|
||||
PlatformRateLimitPolicy: map[string]any{"rules": []any{
|
||||
|
||||
Reference in New Issue
Block a user