diff --git a/relay/channel/claude/message_delta_usage_patch_test.go b/relay/channel/claude/message_delta_usage_patch_test.go new file mode 100644 index 00000000..43312587 --- /dev/null +++ b/relay/channel/claude/message_delta_usage_patch_test.go @@ -0,0 +1,111 @@ +package claude + +import ( + "testing" + + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/setting/model_setting" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestPatchClaudeMessageDeltaUsageDataPreserveUnknownFields(t *testing.T) { + originalData := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":53},"vendor_meta":{"trace_id":"trace_001"}}` + usage := &dto.ClaudeUsage{ + InputTokens: 100, + CacheReadInputTokens: 30, + CacheCreationInputTokens: 50, + } + + patchedData := patchClaudeMessageDeltaUsageData(originalData, usage) + + require.Equal(t, "message_delta", gjson.Get(patchedData, "type").String()) + require.Equal(t, "end_turn", gjson.Get(patchedData, "delta.stop_reason").String()) + require.Equal(t, "trace_001", gjson.Get(patchedData, "vendor_meta.trace_id").String()) + require.EqualValues(t, 53, gjson.Get(patchedData, "usage.output_tokens").Int()) + require.EqualValues(t, 100, gjson.Get(patchedData, "usage.input_tokens").Int()) + require.EqualValues(t, 30, gjson.Get(patchedData, "usage.cache_read_input_tokens").Int()) + require.EqualValues(t, 50, gjson.Get(patchedData, "usage.cache_creation_input_tokens").Int()) +} + +func TestPatchClaudeMessageDeltaUsageDataZeroValueChecks(t *testing.T) { + originalData := `{"type":"message_delta","usage":{"output_tokens":53,"input_tokens":9,"cache_read_input_tokens":0}}` + usage := &dto.ClaudeUsage{ + InputTokens: 100, + CacheReadInputTokens: 30, + CacheCreationInputTokens: 0, + } + + patchedData := patchClaudeMessageDeltaUsageData(originalData, usage) + + require.EqualValues(t, 9, gjson.Get(patchedData, "usage.input_tokens").Int()) + require.EqualValues(t, 30, gjson.Get(patchedData, "usage.cache_read_input_tokens").Int()) + assert.False(t, gjson.Get(patchedData, "usage.cache_creation_input_tokens").Exists()) +} + +func TestShouldSkipClaudeMessageDeltaUsagePatch(t *testing.T) { + originGlobalPassThrough := model_setting.GetGlobalSettings().PassThroughRequestEnabled + t.Cleanup(func() { + model_setting.GetGlobalSettings().PassThroughRequestEnabled = originGlobalPassThrough + }) + + model_setting.GetGlobalSettings().PassThroughRequestEnabled = true + assert.True(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{})) + + model_setting.GetGlobalSettings().PassThroughRequestEnabled = false + assert.True(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ChannelSetting: dto.ChannelSettings{PassThroughBodyEnabled: true}}, + })) + assert.False(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ChannelSetting: dto.ChannelSettings{PassThroughBodyEnabled: false}}, + })) +} + +func TestBuildMessageDeltaPatchUsage(t *testing.T) { + t.Run("merge missing fields from claudeInfo", func(t *testing.T) { + claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{OutputTokens: 53}} + claudeInfo := &ClaudeResponseInfo{ + Usage: &dto.Usage{ + PromptTokens: 100, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 30, + CachedCreationTokens: 50, + }, + ClaudeCacheCreation5mTokens: 10, + ClaudeCacheCreation1hTokens: 20, + }, + } + + usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo) + require.NotNil(t, usage) + require.EqualValues(t, 100, usage.InputTokens) + require.EqualValues(t, 30, usage.CacheReadInputTokens) + require.EqualValues(t, 50, usage.CacheCreationInputTokens) + require.EqualValues(t, 53, usage.OutputTokens) + require.NotNil(t, usage.CacheCreation) + require.EqualValues(t, 10, usage.CacheCreation.Ephemeral5mInputTokens) + require.EqualValues(t, 20, usage.CacheCreation.Ephemeral1hInputTokens) + }) + + t.Run("keep upstream non-zero values", func(t *testing.T) { + claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{ + InputTokens: 9, + CacheReadInputTokens: 7, + CacheCreationInputTokens: 6, + }} + claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{ + PromptTokens: 100, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 30, + CachedCreationTokens: 50, + }, + }} + + usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo) + require.EqualValues(t, 9, usage.InputTokens) + require.EqualValues(t, 7, usage.CacheReadInputTokens) + require.EqualValues(t, 6, usage.CacheCreationInputTokens) + }) +} diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index bdb376ed..069c784c 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -21,6 +21,8 @@ import ( "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" ) const ( @@ -544,6 +546,78 @@ type ClaudeResponseInfo struct { Done bool } +func buildMessageDeltaPatchUsage(claudeResponse *dto.ClaudeResponse, claudeInfo *ClaudeResponseInfo) *dto.ClaudeUsage { + usage := &dto.ClaudeUsage{} + if claudeResponse != nil && claudeResponse.Usage != nil { + *usage = *claudeResponse.Usage + } + + if claudeInfo == nil || claudeInfo.Usage == nil { + return usage + } + + if usage.InputTokens == 0 && claudeInfo.Usage.PromptTokens > 0 { + usage.InputTokens = claudeInfo.Usage.PromptTokens + } + if usage.CacheReadInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedTokens > 0 { + usage.CacheReadInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedTokens + } + if usage.CacheCreationInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens > 0 { + usage.CacheCreationInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens + } + if usage.CacheCreation == nil && (claudeInfo.Usage.ClaudeCacheCreation5mTokens > 0 || claudeInfo.Usage.ClaudeCacheCreation1hTokens > 0) { + usage.CacheCreation = &dto.ClaudeCacheCreationUsage{ + Ephemeral5mInputTokens: claudeInfo.Usage.ClaudeCacheCreation5mTokens, + Ephemeral1hInputTokens: claudeInfo.Usage.ClaudeCacheCreation1hTokens, + } + } + return usage +} + +func shouldSkipClaudeMessageDeltaUsagePatch(info *relaycommon.RelayInfo) bool { + if model_setting.GetGlobalSettings().PassThroughRequestEnabled { + return true + } + if info == nil { + return false + } + return info.ChannelSetting.PassThroughBodyEnabled +} + +func patchClaudeMessageDeltaUsageData(data string, usage *dto.ClaudeUsage) string { + if data == "" || usage == nil { + return data + } + + data = setMessageDeltaUsageInt(data, "usage.input_tokens", usage.InputTokens) + data = setMessageDeltaUsageInt(data, "usage.cache_read_input_tokens", usage.CacheReadInputTokens) + data = setMessageDeltaUsageInt(data, "usage.cache_creation_input_tokens", usage.CacheCreationInputTokens) + + if usage.CacheCreation != nil { + data = setMessageDeltaUsageInt(data, "usage.cache_creation.ephemeral_5m_input_tokens", usage.CacheCreation.Ephemeral5mInputTokens) + data = setMessageDeltaUsageInt(data, "usage.cache_creation.ephemeral_1h_input_tokens", usage.CacheCreation.Ephemeral1hInputTokens) + } + + return data +} + +func setMessageDeltaUsageInt(data string, path string, localValue int) string { + if localValue <= 0 { + return data + } + + upstreamValue := gjson.Get(data, path) + if upstreamValue.Exists() && upstreamValue.Int() > 0 { + return data + } + + patchedData, err := sjson.Set(data, path, localValue) + if err != nil { + return data + } + return patchedData +} + func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *dto.ChatCompletionsStreamResponse, claudeInfo *ClaudeResponseInfo) bool { if claudeInfo == nil { return false @@ -638,6 +712,12 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud if claudeResponse.Message != nil { info.UpstreamModelName = claudeResponse.Message.Model } + } else if claudeResponse.Type == "message_delta" { + // 确保 message_delta 的 usage 包含完整的 input_tokens 和 cache 相关字段 + // 解决 AWS Bedrock 等上游返回的 message_delta 缺少这些字段的问题 + if !shouldSkipClaudeMessageDeltaUsagePatch(info) { + data = patchClaudeMessageDeltaUsageData(data, buildMessageDeltaPatchUsage(&claudeResponse, claudeInfo)) + } } helper.ClaudeChunkData(c, claudeResponse, data) } else if info.RelayFormat == types.RelayFormatOpenAI { diff --git a/relay/channel/claude/relay_claude_test.go b/relay/channel/claude/relay_claude_test.go new file mode 100644 index 00000000..e34c861a --- /dev/null +++ b/relay/channel/claude/relay_claude_test.go @@ -0,0 +1,175 @@ +package claude + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/dto" +) + +func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) { + claudeInfo := &ClaudeResponseInfo{ + Usage: &dto.Usage{}, + } + claudeResponse := &dto.ClaudeResponse{ + Type: "message_start", + Message: &dto.ClaudeMediaMessage{ + Id: "msg_123", + Model: "claude-3-5-sonnet", + Usage: &dto.ClaudeUsage{ + InputTokens: 100, + OutputTokens: 1, + CacheCreationInputTokens: 50, + CacheReadInputTokens: 30, + }, + }, + } + + ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo) + if !ok { + t.Fatal("expected true") + } + if claudeInfo.Usage.PromptTokens != 100 { + t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens) + } + if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 { + t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens) + } + if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 { + t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens) + } + if claudeInfo.ResponseId != "msg_123" { + t.Errorf("ResponseId = %s, want msg_123", claudeInfo.ResponseId) + } + if claudeInfo.Model != "claude-3-5-sonnet" { + t.Errorf("Model = %s, want claude-3-5-sonnet", claudeInfo.Model) + } +} + +func TestFormatClaudeResponseInfo_MessageDelta_FullUsage(t *testing.T) { + // message_start 先积累 usage + claudeInfo := &ClaudeResponseInfo{ + Usage: &dto.Usage{ + PromptTokens: 100, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 30, + CachedCreationTokens: 50, + }, + CompletionTokens: 1, + }, + } + + // message_delta 带完整 usage（原生 Anthropic 场景） + claudeResponse := &dto.ClaudeResponse{ + Type: "message_delta", + Usage: &dto.ClaudeUsage{ + InputTokens: 100, + OutputTokens: 200, + CacheCreationInputTokens: 50, + CacheReadInputTokens: 30, + }, + } + + ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo) + if !ok { + t.Fatal("expected true") + } + if claudeInfo.Usage.PromptTokens != 100 { + t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens) + } + if claudeInfo.Usage.CompletionTokens != 200 { + t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens) + } + if claudeInfo.Usage.TotalTokens != 300 { + t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens) + } + if !claudeInfo.Done { + t.Error("expected Done = true") + } +} + +func TestFormatClaudeResponseInfo_MessageDelta_OnlyOutputTokens(t *testing.T) { + // 模拟 Bedrock: message_start 已积累 usage + claudeInfo := &ClaudeResponseInfo{ + Usage: &dto.Usage{ + PromptTokens: 100, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 30, + CachedCreationTokens: 50, + }, + CompletionTokens: 1, + ClaudeCacheCreation5mTokens: 10, + ClaudeCacheCreation1hTokens: 20, + }, + } + + // Bedrock 的 message_delta 只有 output_tokens，缺少 input_tokens 和 cache 字段 + claudeResponse := &dto.ClaudeResponse{ + Type: "message_delta", + Usage: &dto.ClaudeUsage{ + OutputTokens: 200, + // InputTokens, CacheCreationInputTokens, CacheReadInputTokens 都是 0 + }, + } + + ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo) + if !ok { + t.Fatal("expected true") + } + // PromptTokens 应保持 message_start 的值（因为 message_delta 的 InputTokens=0，不更新） + if claudeInfo.Usage.PromptTokens != 100 { + t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens) + } + if claudeInfo.Usage.CompletionTokens != 200 { + t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens) + } + if claudeInfo.Usage.TotalTokens != 300 { + t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens) + } + // cache 字段应保持 message_start 的值 + if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 { + t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens) + } + if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 { + t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens) + } + if claudeInfo.Usage.ClaudeCacheCreation5mTokens != 10 { + t.Errorf("ClaudeCacheCreation5mTokens = %d, want 10", claudeInfo.Usage.ClaudeCacheCreation5mTokens) + } + if claudeInfo.Usage.ClaudeCacheCreation1hTokens != 20 { + t.Errorf("ClaudeCacheCreation1hTokens = %d, want 20", claudeInfo.Usage.ClaudeCacheCreation1hTokens) + } + if !claudeInfo.Done { + t.Error("expected Done = true") + } +} + +func TestFormatClaudeResponseInfo_NilClaudeInfo(t *testing.T) { + claudeResponse := &dto.ClaudeResponse{Type: "message_start"} + ok := FormatClaudeResponseInfo(claudeResponse, nil, nil) + if ok { + t.Error("expected false for nil claudeInfo") + } +} + +func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) { + text := "hello" + claudeInfo := &ClaudeResponseInfo{ + Usage: &dto.Usage{}, + ResponseText: strings.Builder{}, + } + claudeResponse := &dto.ClaudeResponse{ + Type: "content_block_delta", + Delta: &dto.ClaudeMediaMessage{ + Text: &text, + }, + } + + ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo) + if !ok { + t.Fatal("expected true") + } + if claudeInfo.ResponseText.String() != "hello" { + t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello") + } +}