mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-03 13:00:52 +08:00
feat(usage): implement usage tracking infrastructure across executors
- Added `LoggerPlugin` to log usage metrics for observability. - Introduced a new `Manager` to handle usage record queuing and plugin registration. - Integrated new usage reporter and detailed metrics parsing into executors, covering providers like OpenAI, Codex, Claude, and Gemini. - Improved token usage breakdown across streaming and non-streaming responses.
This commit is contained in:
Luis Pater
292
internal/runtime/executor/usage_helpers.go
Normal file
292
internal/runtime/executor/usage_helpers.go
Normal file
@@ -0,0 +1,292 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type usageReporter struct {
|
||||
provider string
|
||||
model string
|
||||
authID string
|
||||
apiKey string
|
||||
requestedAt time.Time
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func newUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *usageReporter {
|
||||
reporter := &usageReporter{
|
||||
provider: provider,
|
||||
model: model,
|
||||
requestedAt: time.Now(),
|
||||
}
|
||||
if auth != nil {
|
||||
reporter.authID = auth.ID
|
||||
}
|
||||
reporter.apiKey = apiKeyFromContext(ctx)
|
||||
return reporter
|
||||
}
|
||||
|
||||
func (r *usageReporter) publish(ctx context.Context, detail usage.Detail) {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
if total > 0 {
|
||||
detail.TotalTokens = total
|
||||
}
|
||||
}
|
||||
if detail.InputTokens == 0 && detail.OutputTokens == 0 && detail.ReasoningTokens == 0 && detail.CachedTokens == 0 && detail.TotalTokens == 0 {
|
||||
return
|
||||
}
|
||||
r.once.Do(func() {
|
||||
usage.PublishRecord(ctx, usage.Record{
|
||||
Provider: r.provider,
|
||||
Model: r.model,
|
||||
APIKey: r.apiKey,
|
||||
AuthID: r.authID,
|
||||
RequestedAt: r.requestedAt,
|
||||
Detail: detail,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func apiKeyFromContext(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
ginCtx, ok := ctx.Value("gin").(*gin.Context)
|
||||
if !ok || ginCtx == nil {
|
||||
return ""
|
||||
}
|
||||
if v, exists := ginCtx.Get("apiKey"); exists {
|
||||
switch value := v.(type) {
|
||||
case string:
|
||||
return value
|
||||
case fmt.Stringer:
|
||||
return value.String()
|
||||
default:
|
||||
return fmt.Sprintf("%v", value)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func parseCodexUsage(data []byte) (usage.Detail, bool) {
|
||||
usageNode := gjson.ParseBytes(data).Get("response.usage")
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: usageNode.Get("input_tokens").Int(),
|
||||
OutputTokens: usageNode.Get("output_tokens").Int(),
|
||||
TotalTokens: usageNode.Get("total_tokens").Int(),
|
||||
}
|
||||
if cached := usageNode.Get("input_tokens_details.cached_tokens"); cached.Exists() {
|
||||
detail.CachedTokens = cached.Int()
|
||||
}
|
||||
if reasoning := usageNode.Get("output_tokens_details.reasoning_tokens"); reasoning.Exists() {
|
||||
detail.ReasoningTokens = reasoning.Int()
|
||||
}
|
||||
return detail, true
|
||||
}
|
||||
|
||||
func parseOpenAIUsage(data []byte) usage.Detail {
|
||||
usageNode := gjson.ParseBytes(data).Get("usage")
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: usageNode.Get("prompt_tokens").Int(),
|
||||
OutputTokens: usageNode.Get("completion_tokens").Int(),
|
||||
TotalTokens: usageNode.Get("total_tokens").Int(),
|
||||
}
|
||||
if cached := usageNode.Get("prompt_tokens_details.cached_tokens"); cached.Exists() {
|
||||
detail.CachedTokens = cached.Int()
|
||||
}
|
||||
if reasoning := usageNode.Get("completion_tokens_details.reasoning_tokens"); reasoning.Exists() {
|
||||
detail.ReasoningTokens = reasoning.Int()
|
||||
}
|
||||
return detail
|
||||
}
|
||||
|
||||
func parseOpenAIStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
payload := jsonPayload(line)
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
usageNode := gjson.GetBytes(payload, "usage")
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: usageNode.Get("prompt_tokens").Int(),
|
||||
OutputTokens: usageNode.Get("completion_tokens").Int(),
|
||||
TotalTokens: usageNode.Get("total_tokens").Int(),
|
||||
}
|
||||
if cached := usageNode.Get("prompt_tokens_details.cached_tokens"); cached.Exists() {
|
||||
detail.CachedTokens = cached.Int()
|
||||
}
|
||||
if reasoning := usageNode.Get("completion_tokens_details.reasoning_tokens"); reasoning.Exists() {
|
||||
detail.ReasoningTokens = reasoning.Int()
|
||||
}
|
||||
return detail, true
|
||||
}
|
||||
|
||||
func parseClaudeUsage(data []byte) usage.Detail {
|
||||
usageNode := gjson.ParseBytes(data).Get("usage")
|
||||
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
|
||||
}
|
||||
|
||||
func parseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
payload := jsonPayload(line)
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
usageNode := gjson.GetBytes(payload, "usage")
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
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 {
|
||||
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
|
||||
}
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
|
||||
return detail, true
|
||||
}
|
||||
|
||||
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{}
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail
|
||||
}
|
||||
|
||||
func parseGeminiUsage(data []byte) usage.Detail {
|
||||
usageNode := gjson.ParseBytes(data)
|
||||
node := usageNode.Get("usageMetadata")
|
||||
if !node.Exists() {
|
||||
node = usageNode.Get("usage_metadata")
|
||||
}
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail
|
||||
}
|
||||
|
||||
func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
payload := jsonPayload(line)
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
node := gjson.GetBytes(payload, "usageMetadata")
|
||||
if !node.Exists() {
|
||||
node = gjson.GetBytes(payload, "usage_metadata")
|
||||
}
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail, true
|
||||
}
|
||||
|
||||
func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
payload := jsonPayload(line)
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
node := gjson.GetBytes(payload, "response.usageMetadata")
|
||||
if !node.Exists() {
|
||||
node = gjson.GetBytes(payload, "usage_metadata")
|
||||
}
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail, true
|
||||
}
|
||||
|
||||
func jsonPayload(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(trimmed[len("data:"):])
|
||||
}
|
||||
if len(trimmed) == 0 || trimmed[0] != '{' {
|
||||
return nil
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
Reference in New Issue
Block a user