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") } }