feat(usage): add support for detailed token breakdown in usage tracking
- Introduced `CacheReadTokens` and `CacheCreationTokens` to enhance token breakdown. - Refactored `parseClaudeUsageNode` for cleaner and reusable logic. - Adjusted helpers and updated token calculations to align with the new fields.
This commit is contained in:
@@ -49,11 +49,13 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
|
||||
requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))
|
||||
|
||||
tokens := tokenStats{
|
||||
InputTokens: record.Detail.InputTokens,
|
||||
OutputTokens: record.Detail.OutputTokens,
|
||||
ReasoningTokens: record.Detail.ReasoningTokens,
|
||||
CachedTokens: record.Detail.CachedTokens,
|
||||
TotalTokens: record.Detail.TotalTokens,
|
||||
InputTokens: record.Detail.InputTokens,
|
||||
OutputTokens: record.Detail.OutputTokens,
|
||||
ReasoningTokens: record.Detail.ReasoningTokens,
|
||||
CachedTokens: record.Detail.CachedTokens,
|
||||
CacheReadTokens: record.Detail.CacheReadTokens,
|
||||
CacheCreationTokens: record.Detail.CacheCreationTokens,
|
||||
TotalTokens: record.Detail.TotalTokens,
|
||||
}
|
||||
if tokens.TotalTokens == 0 {
|
||||
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
|
||||
@@ -116,11 +118,13 @@ type requestDetail struct {
|
||||
}
|
||||
|
||||
type tokenStats struct {
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
ReasoningTokens int64 `json:"reasoning_tokens"`
|
||||
CachedTokens int64 `json:"cached_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
ReasoningTokens int64 `json:"reasoning_tokens"`
|
||||
CachedTokens int64 `json:"cached_tokens"`
|
||||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||||
CacheCreationTokens int64 `json:"cache_creation_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
}
|
||||
|
||||
type failDetail struct {
|
||||
|
||||
@@ -116,6 +116,8 @@ func hasNonZeroTokenUsage(detail usage.Detail) bool {
|
||||
detail.OutputTokens != 0 ||
|
||||
detail.ReasoningTokens != 0 ||
|
||||
detail.CachedTokens != 0 ||
|
||||
detail.CacheReadTokens != 0 ||
|
||||
detail.CacheCreationTokens != 0 ||
|
||||
detail.TotalTokens != 0
|
||||
}
|
||||
|
||||
@@ -356,17 +358,7 @@ func ParseClaudeUsage(data []byte) usage.Detail {
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: usageNode.Get("input_tokens").Int(),
|
||||
OutputTokens: usageNode.Get("output_tokens").Int(),
|
||||
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
|
||||
}
|
||||
if detail.CachedTokens == 0 {
|
||||
// fall back to creation tokens when read tokens are absent
|
||||
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
|
||||
}
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
|
||||
return detail
|
||||
return parseClaudeUsageNode(usageNode)
|
||||
}
|
||||
|
||||
func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
@@ -378,16 +370,24 @@ func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
return parseClaudeUsageNode(usageNode), true
|
||||
}
|
||||
|
||||
func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail {
|
||||
cacheReadTokens := usageNode.Get("cache_read_input_tokens").Int()
|
||||
cacheCreationTokens := usageNode.Get("cache_creation_input_tokens").Int()
|
||||
detail := usage.Detail{
|
||||
InputTokens: usageNode.Get("input_tokens").Int(),
|
||||
OutputTokens: usageNode.Get("output_tokens").Int(),
|
||||
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
|
||||
InputTokens: usageNode.Get("input_tokens").Int(),
|
||||
OutputTokens: usageNode.Get("output_tokens").Int(),
|
||||
CachedTokens: cacheReadTokens,
|
||||
CacheReadTokens: cacheReadTokens,
|
||||
CacheCreationTokens: cacheCreationTokens,
|
||||
}
|
||||
if detail.CachedTokens == 0 {
|
||||
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
|
||||
detail.CachedTokens = detail.CacheCreationTokens
|
||||
}
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
|
||||
return detail, true
|
||||
return detail
|
||||
}
|
||||
|
||||
func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {
|
||||
|
||||
@@ -34,11 +34,13 @@ type Failure struct {
|
||||
|
||||
// Detail holds the token usage breakdown.
|
||||
type Detail struct {
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
ReasoningTokens int64
|
||||
CachedTokens int64
|
||||
TotalTokens int64
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
ReasoningTokens int64
|
||||
CachedTokens int64
|
||||
CacheReadTokens int64
|
||||
CacheCreationTokens int64
|
||||
TotalTokens int64
|
||||
}
|
||||
|
||||
type requestedModelAliasContextKey struct{}
|
||||
|
||||
Reference in New Issue
Block a user