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/httpapi/chat_completions_mode_test.go
+81
View File
@@ -165,6 +165,87 @@ func TestWriteCompatibleTaskResponseReturnsSSEWhenStreamIsTrue(t *testing.T) {
}
}
func TestWriteCompatibleTaskResponseStreamsCachedUsageProviderShapes(t *testing.T) {
cases := []struct {
name string
usage map[string]any
fragment string
}{
{
name: "aliyun dashscope",
usage: map[string]any{
"prompt_tokens": 10,
"completion_tokens": 2,
"total_tokens": 12,
"prompt_tokens_details": map[string]any{
"cached_tokens": 7,
},
},
fragment: `"prompt_tokens_details":{"cached_tokens":7}`,
},
{
name: "volces responses",
usage: map[string]any{
"prompt_tokens": 9,
"completion_tokens": 3,
"total_tokens": 12,
"input_tokens_details": map[string]any{
"cached_tokens": 6,
},
},
fragment: `"input_tokens_details":{"cached_tokens":6}`,
},
{
name: "deepseek",
usage: map[string]any{
"prompt_tokens": 8,
"completion_tokens": 4,
"total_tokens": 12,
"prompt_cache_hit_tokens": 5,
"prompt_cache_miss_tokens": 3,
},
fragment: `"prompt_cache_hit_tokens":5`,
},
{
name: "glm",
usage: map[string]any{
"prompt_tokens": 11,
"completion_tokens": 1,
"total_tokens": 12,
"prompt_tokens_details": map[string]any{
"cached_tokens": 4,
},
},
fragment: `"prompt_tokens_details":{"cached_tokens":4}`,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
executor := &fakeTaskExecutor{
deltas: []clients.StreamDeltaEvent{{Text: "ok"}},
output: map[string]any{
"id": "chatcmpl-cache",
"object": "chat.completion",
"usage": tc.usage,
},
}
req := httptest.NewRequest(http.MethodPost, "/api/v1/chat/completions", nil)
recorder := httptest.NewRecorder()
writeCompatibleTaskResponse(context.Background(), recorder, req, executor, "chat.completions", "cache-test", store.GatewayTask{ID: "task-test"}, &auth.User{}, true, true)
body := recorder.Body.String()
if !strings.Contains(body, tc.fragment) {
t.Fatalf("SSE body missing cached usage fragment %s: %s", tc.fragment, body)
}
if !strings.Contains(body, `"choices":[]`) || !strings.Contains(body, "data: [DONE]") {
t.Fatalf("SSE body should include final usage chunk and done marker: %s", body)
}
})
}
}
func TestWriteCompatibleTaskResponseStreamsStructuredToolAndReasoningDeltas(t *testing.T) {
executor := &fakeTaskExecutor{
deltas: []clients.StreamDeltaEvent{
apps/api/internal/httpapi/model_catalog.go
+18 -1
View File
@@ -759,7 +759,7 @@ func billingConfigLines(config map[string]any) []string {
func textPricingResource(config map[string]any) bool {
switch strings.ToLower(stringValue(config["resourceType"])) {
case "text", "text_total", "text_input", "text_output":
case "text", "text_total", "text_input", "text_cached_input", "text_output":
return true
default:
return false
@@ -771,6 +771,7 @@ func textPricingLines(config map[string]any) []string {
return nil
}
inputPrice, hasInput := textPriceFromConfig(config, []string{"textInputPer1k", "textInput", "text_input", "inputTokenPrice", "inputPrice"})
cachedInputPrice, hasCachedInput := textPriceFromConfig(config, []string{"textCachedInputPer1k", "textCachedInput", "text_cached_input", "cachedInputTokenPrice", "inputCacheHitTokenPrice", "cachedInputPrice"})
outputPrice, hasOutput := textPriceFromConfig(config, []string{"textOutputPer1k", "textOutput", "text_output", "outputTokenPrice", "outputPrice"})
for _, key := range []string{"formulaConfig", "formula_config"} {
@@ -778,6 +779,9 @@ func textPricingLines(config map[string]any) []string {
if !hasInput {
inputPrice, hasInput = textPriceFromConfig(formulaConfig, []string{"inputTokenPrice", "textInputPer1k", "textInput", "text_input", "inputPrice"})
}
if !hasCachedInput {
cachedInputPrice, hasCachedInput = textPriceFromConfig(formulaConfig, []string{"cachedInputTokenPrice", "textCachedInputPer1k", "textCachedInput", "text_cached_input", "inputCacheHitTokenPrice", "cachedInputPrice"})
}
if !hasOutput {
outputPrice, hasOutput = textPriceFromConfig(formulaConfig, []string{"outputTokenPrice", "textOutputPer1k", "textOutput", "text_output", "outputPrice"})
}
@@ -791,6 +795,9 @@ func textPricingLines(config map[string]any) []string {
if !hasInput {
inputPrice, hasInput = textPriceFromConfig(nested, []string{"inputTokenPrice", "textInputPer1k", "textInput", "text_input", "inputPrice"})
}
if !hasCachedInput {
cachedInputPrice, hasCachedInput = textPriceFromConfig(nested, []string{"cachedInputTokenPrice", "textCachedInputPer1k", "textCachedInput", "text_cached_input", "inputCacheHitTokenPrice", "cachedInputPrice"})
}
if !hasOutput {
outputPrice, hasOutput = textPriceFromConfig(nested, []string{"outputTokenPrice", "textOutputPer1k", "textOutput", "text_output", "outputPrice"})
}
@@ -799,6 +806,9 @@ func textPricingLines(config map[string]any) []string {
if !hasInput {
inputPrice, hasInput = textPriceFromConfig(formulaConfig, []string{"inputTokenPrice", "textInputPer1k", "textInput", "text_input", "inputPrice"})
}
if !hasCachedInput {
cachedInputPrice, hasCachedInput = textPriceFromConfig(formulaConfig, []string{"cachedInputTokenPrice", "textCachedInputPer1k", "textCachedInput", "text_cached_input", "inputCacheHitTokenPrice", "cachedInputPrice"})
}
if !hasOutput {
outputPrice, hasOutput = textPriceFromConfig(formulaConfig, []string{"outputTokenPrice", "textOutputPer1k", "textOutput", "text_output", "outputPrice"})
}
@@ -824,6 +834,10 @@ func textPricingLines(config map[string]any) []string {
if !hasInput {
inputPrice, hasInput = basePrice, true
}
case "text_cached_input":
if !hasCachedInput {
cachedInputPrice, hasCachedInput = basePrice, true
}
case "text_output":
if !hasOutput {
outputPrice, hasOutput = basePrice, true
@@ -835,6 +849,9 @@ func textPricingLines(config map[string]any) []string {
if hasInput {
lines = append(lines, "输入 "+formatNumber(inputPrice)+"/k tokens")
}
if hasCachedInput {
lines = append(lines, "缓存输入 "+formatNumber(cachedInputPrice)+"/k tokens")
}
if hasOutput {
lines = append(lines, "输出 "+formatNumber(outputPrice)+"/k tokens")
}
apps/api/internal/httpapi/openapi_models.go
+13 -3
View File
@@ -294,9 +294,19 @@ type ChatCompletionChoiceMessage struct {
}
type ChatCompletionUsage struct {
PromptTokens int `json:"prompt_tokens,omitempty" example:"12"`
CompletionTokens int `json:"completion_tokens,omitempty" example:"8"`
TotalTokens int `json:"total_tokens,omitempty" example:"20"`
PromptTokens int `json:"prompt_tokens,omitempty" example:"12"`
CompletionTokens int `json:"completion_tokens,omitempty" example:"8"`
TotalTokens int `json:"total_tokens,omitempty" example:"20"`
PromptTokensDetails *ChatPromptTokensDetails `json:"prompt_tokens_details,omitempty"`
InputTokensDetails *ChatInputTokensDetails `json:"input_tokens_details,omitempty"`
}
type ChatPromptTokensDetails struct {
CachedTokens int `json:"cached_tokens,omitempty" example:"4"`
}
type ChatInputTokensDetails struct {
CachedTokens int `json:"cached_tokens,omitempty" example:"4"`
}
type NetworkProxyConfigResponse struct {