usage_helpr.go 982 B

123456789101112131415161718192021222324252627282930
  1. package service
  2. import (
  3. "one-api/dto"
  4. )
  5. //func GetPromptTokens(textRequest dto.GeneralOpenAIRequest, relayMode int) (int, error) {
  6. // switch relayMode {
  7. // case constant.RelayModeChatCompletions:
  8. // return CountTokenMessages(textRequest.Messages, textRequest.Model)
  9. // case constant.RelayModeCompletions:
  10. // return CountTokenInput(textRequest.Prompt, textRequest.Model), nil
  11. // case constant.RelayModeModerations:
  12. // return CountTokenInput(textRequest.Input, textRequest.Model), nil
  13. // }
  14. // return 0, errors.New("unknown relay mode")
  15. //}
  16. func ResponseText2Usage(responseText string, modeName string, promptTokens int) *dto.Usage {
  17. usage := &dto.Usage{}
  18. usage.PromptTokens = promptTokens
  19. ctkm := CountTextToken(responseText, modeName)
  20. usage.CompletionTokens = ctkm
  21. usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
  22. return usage
  23. }
  24. func ValidUsage(usage *dto.Usage) bool {
  25. return usage != nil && (usage.PromptTokens != 0 || usage.CompletionTokens != 0)
  26. }