BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: support cached input token billing

Browse Source
This commit is contained in:
wangbo
2026-06-23 17:17:57 +08:00
parent 6089aa6085
commit 7f32446466
16 changed files with 570 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/clients/clients_test.go
+109 -1
View File
@@ -631,7 +631,7 @@ func TestOpenAIClientChatStreamContract(t *testing.T) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-stream\",\"object\":\"chat.completion.chunk\",\"model\":\"deepseek-v4-flash\",\"choices\":[{\"delta\":{\"content\":\"hello\"}}],\"usage\":null}\n\n"))
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-stream\",\"object\":\"chat.completion.chunk\",\"model\":\"deepseek-v4-flash\",\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":\"stop\"}],\"usage\":null}\n\n"))
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-stream\",\"object\":\"chat.completion.chunk\",\"model\":\"deepseek-v4-flash\",\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n"))
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-stream\",\"object\":\"chat.completion.chunk\",\"model\":\"deepseek-v4-flash\",\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3,\"prompt_tokens_details\":{\"cached_tokens\":1}}}\n\n"))
_, _ = w.Write([]byte("data: [DONE]\n\n"))
}))
defer server.Close()
@@ -662,6 +662,9 @@ func TestOpenAIClientChatStreamContract(t *testing.T) {
if response.Usage.TotalTokens != 3 {
t.Fatalf("unexpected usage: %+v", response.Usage)
}
if response.Usage.CachedInputTokens != 1 {
t.Fatalf("expected cached input tokens from stream usage, got %+v", response.Usage)
}
choices, _ := response.Result["choices"].([]any)
choice, _ := choices[0].(map[string]any)
message, _ := choice["message"].(map[string]any)
@@ -670,6 +673,111 @@ func TestOpenAIClientChatStreamContract(t *testing.T) {
}
}
func TestOpenAIClientChatStreamExtractsCachedInputTokensFromProviderUsageShapes(t *testing.T) {
cases := []struct {
name string
usage map[string]any
cached int
}{
{
name: "aliyun dashscope openai-compatible",
usage: map[string]any{
"prompt_tokens": 10,
"completion_tokens": 2,
"total_tokens": 12,
"prompt_tokens_details": map[string]any{
"cached_tokens": 7,
},
},
cached: 7,
},
{
name: "volces responses-style input details",
usage: map[string]any{
"prompt_tokens": 9,
"completion_tokens": 3,
"total_tokens": 12,
"input_tokens_details": map[string]any{
"cached_tokens": 6,
},
},
cached: 6,
},
{
name: "deepseek prompt cache hit tokens",
usage: map[string]any{
"prompt_tokens": 8,
"completion_tokens": 4,
"total_tokens": 12,
"prompt_cache_hit_tokens": 5,
"prompt_cache_miss_tokens": 3,
},
cached: 5,
},
{
name: "glm openai-compatible",
usage: map[string]any{
"prompt_tokens": 11,
"completion_tokens": 1,
"total_tokens": 12,
"prompt_tokens_details": map[string]any{
"cached_tokens": 4,
},
},
cached: 4,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
var gotIncludeUsage bool
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode request: %v", err)
}
streamOptions, _ := body["stream_options"].(map[string]any)
gotIncludeUsage, _ = streamOptions["include_usage"].(bool)
usageBytes, _ := json.Marshal(tc.usage)
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-cache\",\"object\":\"chat.completion.chunk\",\"model\":\"cache-test\",\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}],\"usage\":null}\n\n"))
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-cache\",\"object\":\"chat.completion.chunk\",\"model\":\"cache-test\",\"choices\":[],\"usage\":" + string(usageBytes) + "}\n\n"))
_, _ = w.Write([]byte("data: [DONE]\n\n"))
}))
defer server.Close()
response, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "chat.completions",
Model: "Cache-Test",
Body: map[string]any{
"model": "Cache-Test",
"messages": []any{map[string]any{"role": "user", "content": "ping"}},
"stream": true,
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ModelName: "cache-test",
Credentials: map[string]any{"apiKey": "test-key"},
},
})
if err != nil {
t.Fatalf("run openai stream client: %v", err)
}
if !gotIncludeUsage {
t.Fatal("expected upstream stream_options.include_usage=true")
}
if response.Usage.CachedInputTokens != tc.cached {
t.Fatalf("expected cached input tokens %d, got %+v", tc.cached, response.Usage)
}
resultUsage, _ := response.Result["usage"].(map[string]any)
promptDetails, _ := resultUsage["prompt_tokens_details"].(map[string]any)
if intFromAny(promptDetails["cached_tokens"]) != tc.cached {
t.Fatalf("result usage should expose normalized cached_tokens, got %+v", resultUsage)
}
})
}
}
func TestOpenAIClientChatStreamPreservesStructuredDeltas(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
apps/api/internal/clients/gemini.go
+11 -2
View File
@@ -519,7 +519,11 @@ func geminiUsage(raw map[string]any) Usage {
input := intFromAny(usageMap["prompt_tokens"])
output := intFromAny(usageMap["completion_tokens"])
total := intFromAny(usageMap["total_tokens"])
return Usage{InputTokens: input, OutputTokens: output, TotalTokens: total}
cachedInput := cachedInputTokensFromOpenAIUsage(usageMap)
if cachedInput > input && input > 0 {
cachedInput = input
}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, TotalTokens: total}
}
func geminiUsageMap(raw map[string]any) map[string]any {
@@ -527,8 +531,13 @@ func geminiUsageMap(raw map[string]any) map[string]any {
input := intFromAny(meta["promptTokenCount"])
output := intFromAny(meta["candidatesTokenCount"])
total := intFromAny(meta["totalTokenCount"])
cachedInput := intFromAny(firstPresent(meta["cachedContentTokenCount"], meta["cached_content_token_count"]))
if total == 0 {
total = input + output
}
return map[string]any{"prompt_tokens": input, "completion_tokens": output, "total_tokens": total}
usage := map[string]any{"prompt_tokens": input, "completion_tokens": output, "total_tokens": total}
if cachedInput > 0 {
usage["prompt_tokens_details"] = map[string]any{"cached_tokens": cachedInput}
}
return usage
}
apps/api/internal/clients/helpers.go
+43 -2
View File
@@ -185,11 +185,17 @@ func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts
}},
}
if usage.TotalTokens > 0 {
out["usage"] = map[string]any{
usageMap := map[string]any{
"prompt_tokens": usage.InputTokens,
"completion_tokens": usage.OutputTokens,
"total_tokens": usage.TotalTokens,
}
if usage.CachedInputTokens > 0 {
usageMap["prompt_tokens_details"] = map[string]any{
"cached_tokens": usage.CachedInputTokens,
}
}
out["usage"] = usageMap
}
return out
}
@@ -877,7 +883,42 @@ func usageFromOpenAI(result map[string]any) Usage {
if total == 0 {
total = input + output
}
return Usage{InputTokens: input, OutputTokens: output, TotalTokens: total}
cachedInput := cachedInputTokensFromOpenAIUsage(usage)
if cachedInput > input && input > 0 {
cachedInput = input
}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, TotalTokens: total}
}
func cachedInputTokensFromOpenAIUsage(usage map[string]any) int {
if len(usage) == 0 {
return 0
}
promptDetails, _ := firstPresent(usage["prompt_tokens_details"], usage["promptTokensDetails"]).(map[string]any)
inputDetails, _ := firstPresent(usage["input_tokens_details"], usage["inputTokensDetails"]).(map[string]any)
usageMetadata, _ := usage["usageMetadata"].(map[string]any)
return intFromAny(firstPresent(
promptDetails["cached_tokens"],
promptDetails["cachedTokens"],
promptDetails["cache_read_input_tokens"],
promptDetails["cacheReadInputTokens"],
inputDetails["cached_tokens"],
inputDetails["cachedTokens"],
inputDetails["cache_read_input_tokens"],
inputDetails["cacheReadInputTokens"],
usage["prompt_cache_hit_tokens"],
usage["promptCacheHitTokens"],
usage["cached_tokens"],
usage["cachedTokens"],
usage["cached_input_tokens"],
usage["cachedInputTokens"],
usage["cache_read_input_tokens"],
usage["cacheReadInputTokens"],
usage["cached_content_token_count"],
usage["cachedContentTokenCount"],
usageMetadata["cached_content_token_count"],
usageMetadata["cachedContentTokenCount"],
))
}
func requestIDFromHTTPResponse(resp *http.Response) string {
apps/api/internal/clients/types.go
+4 -3
View File
@@ -36,9 +36,10 @@ type Response struct {
}
type Usage struct {
InputTokens int
OutputTokens int
TotalTokens int
InputTokens int
OutputTokens int
CachedInputTokens int
TotalTokens int
}
type Progress struct {