Skip to content

phsl / new-api

Go to a project
Commit 81a94c30 by Seefs
Options

fix: patch message_delta usage via gjson/sjson and skip on passthrough

parent 63f246d4
Showing with 177 additions and 147 deletions
relay/channel/claude/message_delta_usage_patch_test.go 0 → 100644
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)
})
}
relay/channel/claude/relay-claude.go
...@@ -21,6 +21,8 @@ import ( ...@@ -21,6 +21,8 @@ import (
"github.com/QuantumNous/new-api/types" "github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
) )
const ( const (
...@@ -544,28 +546,76 @@ type ClaudeResponseInfo struct { ...@@ -544,28 +546,76 @@ type ClaudeResponseInfo struct {
Done bool Done bool
} }
// enrichMessageDeltaUsage 补全 message_delta 事件中缺失的 input_tokens 和 cache 相关字段 func buildMessageDeltaPatchUsage(claudeResponse *dto.ClaudeResponse, claudeInfo *ClaudeResponseInfo) *dto.ClaudeUsage {
// 当上游（如 AWS Bedrock）的 message_delta 不包含这些字段时，从 claudeInfo 中积累的数据补全 usage := &dto.ClaudeUsage{}
func enrichMessageDeltaUsage(claudeResponse *dto.ClaudeResponse, claudeInfo *ClaudeResponseInfo) { if claudeResponse != nil && claudeResponse.Usage != nil {
if claudeResponse.Usage == nil { *usage = *claudeResponse.Usage
claudeResponse.Usage = &dto.ClaudeUsage{}
} }
if claudeResponse.Usage.InputTokens == 0 && claudeInfo.Usage.PromptTokens > 0 {
claudeResponse.Usage.InputTokens = claudeInfo.Usage.PromptTokens if claudeInfo == nil || claudeInfo.Usage == nil {
return usage
}
if usage.InputTokens == 0 && claudeInfo.Usage.PromptTokens > 0 {
usage.InputTokens = claudeInfo.Usage.PromptTokens
} }
if claudeResponse.Usage.CacheReadInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedTokens > 0 { if usage.CacheReadInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedTokens > 0 {
claudeResponse.Usage.CacheReadInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedTokens usage.CacheReadInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedTokens
} }
if claudeResponse.Usage.CacheCreationInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens > 0 { if usage.CacheCreationInputTokens == 0 && claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens > 0 {
claudeResponse.Usage.CacheCreationInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens usage.CacheCreationInputTokens = claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens
} }
if claudeResponse.Usage.CacheCreation == nil && if usage.CacheCreation == nil && (claudeInfo.Usage.ClaudeCacheCreation5mTokens > 0 || claudeInfo.Usage.ClaudeCacheCreation1hTokens > 0) {
(claudeInfo.Usage.ClaudeCacheCreation5mTokens > 0 || claudeInfo.Usage.ClaudeCacheCreation1hTokens > 0) { usage.CacheCreation = &dto.ClaudeCacheCreationUsage{
claudeResponse.Usage.CacheCreation = &dto.ClaudeCacheCreationUsage{
Ephemeral5mInputTokens: claudeInfo.Usage.ClaudeCacheCreation5mTokens, Ephemeral5mInputTokens: claudeInfo.Usage.ClaudeCacheCreation5mTokens,
Ephemeral1hInputTokens: claudeInfo.Usage.ClaudeCacheCreation1hTokens, 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 { func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *dto.ChatCompletionsStreamResponse, claudeInfo *ClaudeResponseInfo) bool {
...@@ -665,9 +715,8 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud ...@@ -665,9 +715,8 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
} else if claudeResponse.Type == "message_delta" { } else if claudeResponse.Type == "message_delta" {
// 确保 message_delta 的 usage 包含完整的 input_tokens 和 cache 相关字段 // 确保 message_delta 的 usage 包含完整的 input_tokens 和 cache 相关字段
// 解决 AWS Bedrock 等上游返回的 message_delta 缺少这些字段的问题 // 解决 AWS Bedrock 等上游返回的 message_delta 缺少这些字段的问题
enrichMessageDeltaUsage(&claudeResponse, claudeInfo) if !shouldSkipClaudeMessageDeltaUsagePatch(info) {
if newData, err := json.Marshal(claudeResponse); err == nil { data = patchClaudeMessageDeltaUsageData(data, buildMessageDeltaPatchUsage(&claudeResponse, claudeInfo))
data = string(newData)
} }
} }
helper.ClaudeChunkData(c, claudeResponse, data) helper.ClaudeChunkData(c, claudeResponse, data)
......
relay/channel/claude/relay-claude_test.go → relay/channel/claude/relay_claude_test.go
...@@ -173,133 +173,3 @@ func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) { ...@@ -173,133 +173,3 @@ func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) {
t.Errorf("ResponseText = %q, want %q", 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)
}
}
})
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment