Fix OCM usage counting & Update pricing

This commit is contained in:
世界
2026-02-24 19:16:40 +08:00
parent d48236da94
commit 22cf675263
6 changed files with 748 additions and 156 deletions

View File

@@ -406,7 +406,9 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
isChatCompletions := path == "/v1/chat/completions"
mediaType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type"))
isStreaming := err == nil && mediaType == "text/event-stream"
if !isStreaming && !isChatCompletions && response.Header.Get("Content-Type") == "" {
isStreaming = true
}
if !isStreaming {
bodyBytes, err := io.ReadAll(response.Body)
if err != nil {
@@ -414,13 +416,14 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
return
}
var responseModel string
var responseModel, serviceTier string
var inputTokens, outputTokens, cachedTokens int64
if isChatCompletions {
var chatCompletion openai.ChatCompletion
if json.Unmarshal(bodyBytes, &chatCompletion) == nil {
responseModel = chatCompletion.Model
serviceTier = string(chatCompletion.ServiceTier)
inputTokens = chatCompletion.Usage.PromptTokens
outputTokens = chatCompletion.Usage.CompletionTokens
cachedTokens = chatCompletion.Usage.PromptTokensDetails.CachedTokens
@@ -429,6 +432,7 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
var responsesResponse responses.Response
if json.Unmarshal(bodyBytes, &responsesResponse) == nil {
responseModel = string(responsesResponse.Model)
serviceTier = string(responsesResponse.ServiceTier)
inputTokens = responsesResponse.Usage.InputTokens
outputTokens = responsesResponse.Usage.OutputTokens
cachedTokens = responsesResponse.Usage.InputTokensDetails.CachedTokens
@@ -440,7 +444,7 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
responseModel = requestModel
}
if responseModel != "" {
s.usageTracker.AddUsage(responseModel, inputTokens, outputTokens, cachedTokens, username)
s.usageTracker.AddUsage(responseModel, inputTokens, outputTokens, cachedTokens, serviceTier, username)
}
}
@@ -455,7 +459,7 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
}
var inputTokens, outputTokens, cachedTokens int64
var responseModel string
var responseModel, serviceTier string
buffer := make([]byte, buf.BufferSize)
var leftover []byte
@@ -490,6 +494,9 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
if chatChunk.Model != "" {
responseModel = chatChunk.Model
}
if chatChunk.ServiceTier != "" {
serviceTier = string(chatChunk.ServiceTier)
}
if chatChunk.Usage.PromptTokens > 0 {
inputTokens = chatChunk.Usage.PromptTokens
cachedTokens = chatChunk.Usage.PromptTokensDetails.CachedTokens
@@ -506,6 +513,9 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
if string(completedEvent.Response.Model) != "" {
responseModel = string(completedEvent.Response.Model)
}
if completedEvent.Response.ServiceTier != "" {
serviceTier = string(completedEvent.Response.ServiceTier)
}
if completedEvent.Response.Usage.InputTokens > 0 {
inputTokens = completedEvent.Response.Usage.InputTokens
cachedTokens = completedEvent.Response.Usage.InputTokensDetails.CachedTokens
@@ -534,7 +544,7 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
if inputTokens > 0 || outputTokens > 0 {
if responseModel != "" {
s.usageTracker.AddUsage(responseModel, inputTokens, outputTokens, cachedTokens, username)
s.usageTracker.AddUsage(responseModel, inputTokens, outputTokens, cachedTokens, serviceTier, username)
}
}
return