feat: support cached input token billing

2026-06-23 17:17:57 +08:00
parent 6089aa6085
commit 7f32446466
16 changed files with 570 additions and 45 deletions
@@ -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")