mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-03 06:35:00 +08:00
feat(executor): support token usage parsing for plugin executors
- Add `ParsePluginExecutorResponseUsage` to extract token usage from non-streaming plugin responses across Claude, Gemini, Interactions, Antigravity, and OpenAI/Codex protocols. - Add `ObservePluginExecutorStreamUsage` to observe and aggregate token usage across streaming chunks. Closes: #5340
This commit is contained in:
Luis Pater
220
internal/runtime/executor/helps/plugin_executor_usage.go
Normal file
220
internal/runtime/executor/helps/plugin_executor_usage.go
Normal file
@@ -0,0 +1,220 @@
|
||||
package helps
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// ParsePluginExecutorResponseUsage extracts token usage from a non-streaming plugin executor response.
|
||||
func ParsePluginExecutorResponseUsage(protocol string, payload []byte) usage.Detail {
|
||||
if len(payload) == 0 {
|
||||
return usage.Detail{}
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(protocol)) {
|
||||
case "claude":
|
||||
return parseClaudePayloadUsage(payload)
|
||||
case "gemini":
|
||||
return ParseGeminiUsage(payload)
|
||||
case "interactions", "interactions-response":
|
||||
return ParseInteractionsUsage(payload)
|
||||
case "antigravity":
|
||||
return ParseAntigravityUsage(payload)
|
||||
case "codex", "openai-response":
|
||||
if detail, ok := ParseCodexUsage(payload); ok {
|
||||
return detail
|
||||
}
|
||||
return ParseOpenAIUsage(payload)
|
||||
default:
|
||||
return ParseOpenAIUsage(payload)
|
||||
}
|
||||
}
|
||||
|
||||
// ObservePluginExecutorStreamUsage parses streaming chunks and updates a stream usage buffer.
|
||||
func ObservePluginExecutorStreamUsage(protocol string, payload []byte, buffer *StreamUsageBuffer) {
|
||||
if buffer == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(protocol)) {
|
||||
case "claude":
|
||||
IterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := parseClaudeStreamLine(line); ok {
|
||||
ObserveMergedStreamUsage(buffer, detail)
|
||||
}
|
||||
})
|
||||
case "gemini":
|
||||
IterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := ParseGeminiStreamUsage(line); ok {
|
||||
buffer.Observe(detail, ok)
|
||||
}
|
||||
})
|
||||
case "interactions", "interactions-response":
|
||||
IterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := ParseInteractionsStreamUsage(line); ok {
|
||||
ObserveMergedStreamUsage(buffer, detail)
|
||||
}
|
||||
})
|
||||
case "antigravity":
|
||||
IterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := ParseAntigravityStreamUsage(line); ok {
|
||||
buffer.Observe(detail, ok)
|
||||
}
|
||||
})
|
||||
case "codex", "openai-response":
|
||||
IterateStreamLines(payload, func(line []byte) {
|
||||
if jsonBytes := ExtractStreamJSONPayload(line); len(jsonBytes) > 0 {
|
||||
if detail, ok := ParseCodexUsage(jsonBytes); ok {
|
||||
buffer.Observe(detail, ok)
|
||||
return
|
||||
}
|
||||
}
|
||||
buffer.ObserveOpenAIStream(line)
|
||||
})
|
||||
default:
|
||||
IterateStreamLines(payload, func(line []byte) {
|
||||
buffer.ObserveOpenAIStream(line)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ObservePluginExecutorStreamTTFT inspects a streaming payload and records TTFT on the reporter.
|
||||
func ObservePluginExecutorStreamTTFT(protocol string, reporter *UsageReporter, payload []byte) {
|
||||
if reporter == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
reporter.RecordFirstPacket()
|
||||
switch strings.ToLower(strings.TrimSpace(protocol)) {
|
||||
case "claude":
|
||||
ObserveClaudeTokenEvent(reporter, payload)
|
||||
case "gemini", "antigravity", "interactions", "interactions-response":
|
||||
ObserveGeminiTokenEvent(reporter, payload)
|
||||
case "codex", "openai-response":
|
||||
ObserveResponsesTokenEvent(reporter, payload)
|
||||
default:
|
||||
ObserveChatTokenEvent(reporter, payload)
|
||||
}
|
||||
}
|
||||
|
||||
func parseClaudePayloadUsage(payload []byte) usage.Detail {
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}
|
||||
}
|
||||
usageNode := gjson.GetBytes(payload, "usage")
|
||||
if !usageNode.Exists() {
|
||||
usageNode = gjson.GetBytes(payload, "message.usage")
|
||||
}
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
return ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`))
|
||||
}
|
||||
|
||||
func parseClaudeStreamLine(line []byte) (usage.Detail, bool) {
|
||||
payload := ExtractStreamJSONPayload(line)
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
usageNode := gjson.GetBytes(payload, "usage")
|
||||
if !usageNode.Exists() {
|
||||
usageNode = gjson.GetBytes(payload, "message.usage")
|
||||
}
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`))
|
||||
return detail, true
|
||||
}
|
||||
|
||||
// ObserveMergedStreamUsage updates buffer with merged usage details.
|
||||
func ObserveMergedStreamUsage(buffer *StreamUsageBuffer, update usage.Detail) {
|
||||
if buffer == nil {
|
||||
return
|
||||
}
|
||||
if existing, ok := buffer.Detail(); ok {
|
||||
merged := MergeStreamUsageDetail(existing, update)
|
||||
buffer.Observe(merged, true)
|
||||
return
|
||||
}
|
||||
buffer.Observe(update, true)
|
||||
}
|
||||
|
||||
// MergeStreamUsageDetail merges existing stream usage with a newer update.
|
||||
func MergeStreamUsageDetail(existing, update usage.Detail) usage.Detail {
|
||||
merged := update
|
||||
if merged.InputTokens == 0 && existing.InputTokens > 0 {
|
||||
merged.InputTokens = existing.InputTokens
|
||||
}
|
||||
if merged.CachedTokens == 0 && existing.CachedTokens > 0 {
|
||||
merged.CachedTokens = existing.CachedTokens
|
||||
}
|
||||
if merged.CacheReadTokens == 0 && existing.CacheReadTokens > 0 {
|
||||
merged.CacheReadTokens = existing.CacheReadTokens
|
||||
}
|
||||
if merged.CacheCreationTokens == 0 && existing.CacheCreationTokens > 0 {
|
||||
merged.CacheCreationTokens = existing.CacheCreationTokens
|
||||
}
|
||||
if merged.OutputTokens == 0 && existing.OutputTokens > 0 {
|
||||
merged.OutputTokens = existing.OutputTokens
|
||||
}
|
||||
if merged.ReasoningTokens == 0 && existing.ReasoningTokens > 0 {
|
||||
merged.ReasoningTokens = existing.ReasoningTokens
|
||||
}
|
||||
if merged.ResponseServiceTier == "" {
|
||||
merged.ResponseServiceTier = existing.ResponseServiceTier
|
||||
}
|
||||
cached := merged.CacheReadTokens + merged.CacheCreationTokens
|
||||
if cached == 0 {
|
||||
cached = merged.CachedTokens
|
||||
}
|
||||
calculatedTotal := merged.InputTokens + merged.OutputTokens + cached
|
||||
if merged.TotalTokens == 0 || merged.TotalTokens < calculatedTotal {
|
||||
merged.TotalTokens = calculatedTotal
|
||||
}
|
||||
nonReasoningOutput := merged.OutputTokens - merged.ReasoningTokens
|
||||
if nonReasoningOutput < 0 {
|
||||
nonReasoningOutput = 0
|
||||
}
|
||||
merged.TokenBreakdown = usage.NewIndependentTokenBreakdown(
|
||||
merged.InputTokens,
|
||||
merged.CacheReadTokens,
|
||||
merged.CacheCreationTokens,
|
||||
nonReasoningOutput,
|
||||
merged.ReasoningTokens,
|
||||
merged.TotalTokens,
|
||||
)
|
||||
return merged
|
||||
}
|
||||
|
||||
// IterateStreamLines splits payload by newline and invokes fn for non-empty lines.
|
||||
func IterateStreamLines(payload []byte, fn func(line []byte)) {
|
||||
for _, line := range bytes.Split(payload, []byte("\n")) {
|
||||
trimmed := bytes.TrimSpace(line)
|
||||
if len(trimmed) == 0 {
|
||||
continue
|
||||
}
|
||||
fn(trimmed)
|
||||
}
|
||||
}
|
||||
|
||||
// ExtractStreamJSONPayload extracts SSE data/json payload from a raw line.
|
||||
func ExtractStreamJSONPayload(line []byte) []byte {
|
||||
trimmed := bytes.TrimSpace(line)
|
||||
if len(trimmed) == 0 {
|
||||
return nil
|
||||
}
|
||||
if bytes.Equal(trimmed, []byte("[DONE]")) {
|
||||
return nil
|
||||
}
|
||||
if bytes.HasPrefix(trimmed, []byte("event:")) {
|
||||
return nil
|
||||
}
|
||||
if bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||
trimmed = bytes.TrimSpace(bytes.TrimPrefix(trimmed, []byte("data:")))
|
||||
}
|
||||
if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("[DONE]")) {
|
||||
return nil
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
Reference in New Issue
Block a user