Fix OCM usage counting & Update pricing
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user