| package service | |
| import ( | |
| "github.com/QuantumNous/new-api/common" | |
| "github.com/QuantumNous/new-api/constant" | |
| "github.com/QuantumNous/new-api/dto" | |
| "github.com/gin-gonic/gin" | |
| ) | |
| //func GetPromptTokens(textRequest dto.GeneralOpenAIRequest, relayMode int) (int, error) { | |
| // switch relayMode { | |
| // case constant.RelayModeChatCompletions: | |
| // return CountTokenMessages(textRequest.Messages, textRequest.Model) | |
| // case constant.RelayModeCompletions: | |
| // return CountTokenInput(textRequest.Prompt, textRequest.Model), nil | |
| // case constant.RelayModeModerations: | |
| // return CountTokenInput(textRequest.Input, textRequest.Model), nil | |
| // } | |
| // return 0, errors.New("unknown relay mode") | |
| //} | |
| func ResponseText2Usage(c *gin.Context, responseText string, modeName string, promptTokens int) *dto.Usage { | |
| common.SetContextKey(c, constant.ContextKeyLocalCountTokens, true) | |
| usage := &dto.Usage{} | |
| usage.PromptTokens = promptTokens | |
| ctkm := CountTextToken(responseText, modeName) | |
| usage.CompletionTokens = ctkm | |
| usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens | |
| return usage | |
| } | |
| func ValidUsage(usage *dto.Usage) bool { | |
| return usage != nil && (usage.PromptTokens != 0 || usage.CompletionTokens != 0) | |
| } | |