mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-06-09 15:53:18 +08:00
feat(api): implement protocol multiplexer and Redis queue for usage integration
- Added `protocol_multiplexer.go`, enabling support for both HTTP and Redis protocols on a single listener. - Introduced `redis_queue_protocol.go` to handle Redis-compatible RESP commands for queue management. - Integrated `redisqueue` package, supporting in-memory queuing with expiration pruning. - Updated server initialization to manage a shared listener and multiplex connections. - Adjusted `Handler` to adopt `AuthenticateManagementKey` for modular key validation, supporting both HTTP and Redis flows.
This commit is contained in:
145
internal/redisqueue/plugin.go
Normal file
145
internal/redisqueue/plugin.go
Normal file
@@ -0,0 +1,145 @@
|
||||
package redisqueue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
|
||||
)
|
||||
|
||||
func init() {
|
||||
coreusage.RegisterPlugin(&usageQueuePlugin{})
|
||||
}
|
||||
|
||||
type usageQueuePlugin struct{}
|
||||
|
||||
func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Record) {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
if !Enabled() || !internalusage.StatisticsEnabled() {
|
||||
return
|
||||
}
|
||||
|
||||
timestamp := record.RequestedAt
|
||||
if timestamp.IsZero() {
|
||||
timestamp = time.Now()
|
||||
}
|
||||
|
||||
modelName := strings.TrimSpace(record.Model)
|
||||
if modelName == "" {
|
||||
modelName = "unknown"
|
||||
}
|
||||
provider := strings.TrimSpace(record.Provider)
|
||||
if provider == "" {
|
||||
provider = "unknown"
|
||||
}
|
||||
authType := strings.TrimSpace(record.AuthType)
|
||||
if authType == "" {
|
||||
authType = "unknown"
|
||||
}
|
||||
apiKey := strings.TrimSpace(record.APIKey)
|
||||
requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))
|
||||
if requestID == "" {
|
||||
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
|
||||
requestID = strings.TrimSpace(internallogging.GetGinRequestID(ginCtx))
|
||||
}
|
||||
}
|
||||
|
||||
tokens := internalusage.TokenStats{
|
||||
InputTokens: record.Detail.InputTokens,
|
||||
OutputTokens: record.Detail.OutputTokens,
|
||||
ReasoningTokens: record.Detail.ReasoningTokens,
|
||||
CachedTokens: record.Detail.CachedTokens,
|
||||
TotalTokens: record.Detail.TotalTokens,
|
||||
}
|
||||
if tokens.TotalTokens == 0 {
|
||||
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
|
||||
}
|
||||
if tokens.TotalTokens == 0 {
|
||||
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens + tokens.CachedTokens
|
||||
}
|
||||
|
||||
failed := record.Failed
|
||||
if !failed {
|
||||
failed = !resolveSuccess(ctx)
|
||||
}
|
||||
|
||||
detail := internalusage.RequestDetail{
|
||||
Timestamp: timestamp,
|
||||
LatencyMs: record.Latency.Milliseconds(),
|
||||
Source: record.Source,
|
||||
AuthIndex: record.AuthIndex,
|
||||
Tokens: tokens,
|
||||
Failed: failed,
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(queuedUsageDetail{
|
||||
RequestDetail: detail,
|
||||
Provider: provider,
|
||||
Model: modelName,
|
||||
Endpoint: resolveEndpoint(ctx),
|
||||
AuthType: authType,
|
||||
APIKey: apiKey,
|
||||
RequestID: requestID,
|
||||
})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
Enqueue(payload)
|
||||
}
|
||||
|
||||
type queuedUsageDetail struct {
|
||||
internalusage.RequestDetail
|
||||
Provider string `json:"provider"`
|
||||
Model string `json:"model"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
AuthType string `json:"auth_type"`
|
||||
APIKey string `json:"api_key"`
|
||||
RequestID string `json:"request_id"`
|
||||
}
|
||||
|
||||
func resolveSuccess(ctx context.Context) bool {
|
||||
if ctx == nil {
|
||||
return true
|
||||
}
|
||||
ginCtx, ok := ctx.Value("gin").(*gin.Context)
|
||||
if !ok || ginCtx == nil {
|
||||
return true
|
||||
}
|
||||
status := ginCtx.Writer.Status()
|
||||
if status == 0 {
|
||||
return true
|
||||
}
|
||||
return status < http.StatusBadRequest
|
||||
}
|
||||
|
||||
func resolveEndpoint(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
ginCtx, ok := ctx.Value("gin").(*gin.Context)
|
||||
if !ok || ginCtx == nil || ginCtx.Request == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
path := strings.TrimSpace(ginCtx.FullPath())
|
||||
if path == "" && ginCtx.Request.URL != nil {
|
||||
path = strings.TrimSpace(ginCtx.Request.URL.Path)
|
||||
}
|
||||
if path == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
method := strings.TrimSpace(ginCtx.Request.Method)
|
||||
if method == "" {
|
||||
return path
|
||||
}
|
||||
return method + " " + path
|
||||
}
|
||||
Reference in New Issue
Block a user