diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index bdb376ed..ad9d2234 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -544,6 +544,30 @@ type ClaudeResponseInfo struct { Done bool } +// enrichMessageDeltaUsage 补全 message_delta 事件中缺失的 input_tokens 和 cache 相关字段 +// 当上游(如 AWS Bedrock)的 message_delta 不包含这些字段时,从 claudeInfo 中积累的数据补全 +func enrichMessageDeltaUsage(claudeResponse *dto.ClaudeResponse, claudeInfo *ClaudeResponseInfo) { + if claudeResponse.Usage == nil { + claudeResponse.Usage = &dto.ClaudeUsage{} + } + if claudeResponse.Usage.InputTokens == 0 && claudeInfo.Usage.PromptTokens > 0 { + claudeResponse.Usage.InputTokens = claudeInfo.Usage.PromptTokens + } + if claudeResponse.Usage.CacheReadInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedTokens > 0 { + claudeResponse.Usage.CacheReadInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedTokens + } + if claudeResponse.Usage.CacheCreationInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens > 0 { + claudeResponse.Usage.CacheCreationInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens + } + if claudeResponse.Usage.CacheCreation == nil && + (claudeInfo.Usage.ClaudeCacheCreation5mTokens > 0 || claudeInfo.Usage.ClaudeCacheCreation1hTokens > 0) { + claudeResponse.Usage.CacheCreation = &dto.ClaudeCacheCreationUsage{ + Ephemeral5mInputTokens: claudeInfo.Usage.ClaudeCacheCreation5mTokens, + Ephemeral1hInputTokens: claudeInfo.Usage.ClaudeCacheCreation1hTokens, + } + } +} + func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *dto.ChatCompletionsStreamResponse, claudeInfo *ClaudeResponseInfo) bool { if claudeInfo == nil { return false @@ -638,6 +662,13 @@ 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 缺少这些字段的问题 + enrichMessageDeltaUsage(&claudeResponse, claudeInfo) + if newData, err := json.Marshal(claudeResponse); err == nil { + data = string(newData) + } } 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..82c91018 --- /dev/null +++ b/relay/channel/claude/relay-claude_test.go @@ -0,0 +1,305 @@ +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") + } +} + +// TestEnrichMessageDeltaUsage 测试 message_delta 事件的 usage 补全逻辑 +// 这是修复 issue #2881 的核心逻辑:当上游(如 Bedrock)的 message_delta 缺少 +// input_tokens 和 cache 相关字段时,用 claudeInfo 中积累的数据补全 +func TestEnrichMessageDeltaUsage(t *testing.T) { + tests := []struct { + name string + claudeInfo *ClaudeResponseInfo + deltaUsage *dto.ClaudeUsage + wantInput int + wantCacheRead int + wantCacheCreate int + wantOutput int + want5m int + want1h int + }{ + { + name: "Bedrock: delta 只有 output_tokens,从 claudeInfo 补全其他字段", + claudeInfo: &ClaudeResponseInfo{ + Usage: &dto.Usage{ + PromptTokens: 100, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 30, + CachedCreationTokens: 50, + }, + ClaudeCacheCreation5mTokens: 10, + ClaudeCacheCreation1hTokens: 20, + }, + }, + deltaUsage: &dto.ClaudeUsage{OutputTokens: 200}, + wantInput: 100, + wantCacheRead: 30, + wantCacheCreate: 50, + wantOutput: 200, + want5m: 10, + want1h: 20, + }, + { + name: "原生 Anthropic: delta 已包含所有字段,不覆盖", + claudeInfo: &ClaudeResponseInfo{ + Usage: &dto.Usage{ + PromptTokens: 100, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 30, + CachedCreationTokens: 50, + }, + }, + }, + deltaUsage: &dto.ClaudeUsage{ + InputTokens: 100, + OutputTokens: 200, + CacheReadInputTokens: 30, + CacheCreationInputTokens: 50, + }, + wantInput: 100, + wantCacheRead: 30, + wantCacheCreate: 50, + wantOutput: 200, + }, + { + name: "delta usage 为 nil,创建并补全", + claudeInfo: &ClaudeResponseInfo{ + Usage: &dto.Usage{ + PromptTokens: 80, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 20, + CachedCreationTokens: 40, + }, + }, + }, + deltaUsage: nil, + wantInput: 80, + wantCacheRead: 20, + wantCacheCreate: 40, + wantOutput: 0, + }, + { + name: "没有 cache 数据,不补全", + claudeInfo: &ClaudeResponseInfo{ + Usage: &dto.Usage{ + PromptTokens: 100, + }, + }, + deltaUsage: &dto.ClaudeUsage{OutputTokens: 50}, + wantInput: 100, + wantCacheRead: 0, + wantCacheCreate: 0, + wantOutput: 50, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claudeResponse := &dto.ClaudeResponse{ + Type: "message_delta", + Usage: tt.deltaUsage, + } + + // 模拟 HandleStreamResponseData 中 Claude 格式的补全逻辑 + enrichMessageDeltaUsage(claudeResponse, tt.claudeInfo) + + if claudeResponse.Usage == nil { + t.Fatal("Usage should not be nil after enrichment") + } + if claudeResponse.Usage.InputTokens != tt.wantInput { + t.Errorf("InputTokens = %d, want %d", claudeResponse.Usage.InputTokens, tt.wantInput) + } + if claudeResponse.Usage.CacheReadInputTokens != tt.wantCacheRead { + t.Errorf("CacheReadInputTokens = %d, want %d", claudeResponse.Usage.CacheReadInputTokens, tt.wantCacheRead) + } + if claudeResponse.Usage.CacheCreationInputTokens != tt.wantCacheCreate { + t.Errorf("CacheCreationInputTokens = %d, want %d", claudeResponse.Usage.CacheCreationInputTokens, tt.wantCacheCreate) + } + if claudeResponse.Usage.OutputTokens != tt.wantOutput { + t.Errorf("OutputTokens = %d, want %d", claudeResponse.Usage.OutputTokens, tt.wantOutput) + } + if tt.want5m > 0 || tt.want1h > 0 { + if claudeResponse.Usage.CacheCreation == nil { + t.Fatal("CacheCreation should not be nil") + } + if claudeResponse.Usage.CacheCreation.Ephemeral5mInputTokens != tt.want5m { + t.Errorf("Ephemeral5mInputTokens = %d, want %d", claudeResponse.Usage.CacheCreation.Ephemeral5mInputTokens, tt.want5m) + } + if claudeResponse.Usage.CacheCreation.Ephemeral1hInputTokens != tt.want1h { + t.Errorf("Ephemeral1hInputTokens = %d, want %d", claudeResponse.Usage.CacheCreation.Ephemeral1hInputTokens, tt.want1h) + } + } + }) + } +}