refactor: extract parseGeminiFamilyUsageDetail helper to reduce duplication

This commit is contained in:
NguyenSiTrung
2025-12-24 10:22:31 +07:00
parent 872339bceb
commit 969c1a5b72

View File

@@ -275,15 +275,7 @@ func parseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
return detail, true return detail, true
} }
func parseGeminiCLIUsage(data []byte) usage.Detail { func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {
usageNode := gjson.ParseBytes(data)
node := usageNode.Get("response.usageMetadata")
if !node.Exists() {
node = usageNode.Get("response.usage_metadata")
}
if !node.Exists() {
return usage.Detail{}
}
detail := usage.Detail{ detail := usage.Detail{
InputTokens: node.Get("promptTokenCount").Int(), InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(), OutputTokens: node.Get("candidatesTokenCount").Int(),
@@ -297,6 +289,18 @@ func parseGeminiCLIUsage(data []byte) usage.Detail {
return detail return detail
} }
func parseGeminiCLIUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data)
node := usageNode.Get("response.usageMetadata")
if !node.Exists() {
node = usageNode.Get("response.usage_metadata")
}
if !node.Exists() {
return usage.Detail{}
}
return parseGeminiFamilyUsageDetail(node)
}
func parseGeminiUsage(data []byte) usage.Detail { func parseGeminiUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data) usageNode := gjson.ParseBytes(data)
node := usageNode.Get("usageMetadata") node := usageNode.Get("usageMetadata")
@@ -306,17 +310,7 @@ func parseGeminiUsage(data []byte) usage.Detail {
if !node.Exists() { if !node.Exists() {
return usage.Detail{} return usage.Detail{}
} }
detail := usage.Detail{ return parseGeminiFamilyUsageDetail(node)
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
CachedTokens: node.Get("cachedContentTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail
} }
func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) { func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) {
@@ -331,17 +325,7 @@ func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) {
if !node.Exists() { if !node.Exists() {
return usage.Detail{}, false return usage.Detail{}, false
} }
detail := usage.Detail{ return parseGeminiFamilyUsageDetail(node), true
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
CachedTokens: node.Get("cachedContentTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail, true
} }
func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) { func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {
@@ -356,17 +340,7 @@ func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {
if !node.Exists() { if !node.Exists() {
return usage.Detail{}, false return usage.Detail{}, false
} }
detail := usage.Detail{ return parseGeminiFamilyUsageDetail(node), true
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
CachedTokens: node.Get("cachedContentTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail, true
} }
func parseAntigravityUsage(data []byte) usage.Detail { func parseAntigravityUsage(data []byte) usage.Detail {
@@ -381,17 +355,7 @@ func parseAntigravityUsage(data []byte) usage.Detail {
if !node.Exists() { if !node.Exists() {
return usage.Detail{} return usage.Detail{}
} }
detail := usage.Detail{ return parseGeminiFamilyUsageDetail(node)
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
CachedTokens: node.Get("cachedContentTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail
} }
func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) { func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) {
@@ -409,17 +373,7 @@ func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) {
if !node.Exists() { if !node.Exists() {
return usage.Detail{}, false return usage.Detail{}, false
} }
detail := usage.Detail{ return parseGeminiFamilyUsageDetail(node), true
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
CachedTokens: node.Get("cachedContentTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail, true
} }
var stopChunkWithoutUsage sync.Map var stopChunkWithoutUsage sync.Map