feat: support cached input token billing

This commit is contained in:
2026-06-23 17:17:57 +08:00
parent 6089aa6085
commit 7f32446466
16 changed files with 570 additions and 45 deletions
+39
View File
@@ -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{