Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

158 wiersze
5.5 KiB

  1. package openai
  2. import (
  3. "fmt"
  4. "io"
  5. "net/http"
  6. "strings"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/dto"
  9. "github.com/QuantumNous/new-api/logger"
  10. relaycommon "github.com/QuantumNous/new-api/relay/common"
  11. "github.com/QuantumNous/new-api/relay/helper"
  12. "github.com/QuantumNous/new-api/service"
  13. "github.com/QuantumNous/new-api/types"
  14. "github.com/gin-gonic/gin"
  15. )
  16. func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
  17. defer service.CloseResponseBodyGracefully(resp)
  18. // read response body
  19. var responsesResponse dto.OpenAIResponsesResponse
  20. responseBody, err := io.ReadAll(resp.Body)
  21. if err != nil {
  22. return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError)
  23. }
  24. err = common.Unmarshal(responseBody, &responsesResponse)
  25. if err != nil {
  26. return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
  27. }
  28. if oaiError := responsesResponse.GetOpenAIError(); oaiError != nil && oaiError.Type != "" {
  29. apiErr := types.WithOpenAIError(*oaiError, resp.StatusCode)
  30. apiErr.UpstreamBody = service.TruncateBody(string(responseBody))
  31. return nil, apiErr
  32. }
  33. if responsesResponse.HasImageGenerationCall() {
  34. c.Set("image_generation_call", true)
  35. c.Set("image_generation_call_quality", responsesResponse.GetQuality())
  36. c.Set("image_generation_call_size", responsesResponse.GetSize())
  37. }
  38. // 写入新的 response body
  39. service.IOCopyBytesGracefully(c, resp, responseBody)
  40. // compute usage
  41. usage := dto.Usage{}
  42. if responsesResponse.Usage != nil {
  43. usage.PromptTokens = responsesResponse.Usage.InputTokens
  44. usage.CompletionTokens = responsesResponse.Usage.OutputTokens
  45. usage.TotalTokens = responsesResponse.Usage.TotalTokens
  46. if responsesResponse.Usage.InputTokensDetails != nil {
  47. usage.PromptTokensDetails.CachedTokens = responsesResponse.Usage.InputTokensDetails.CachedTokens
  48. }
  49. }
  50. if info == nil || info.ResponsesUsageInfo == nil || info.ResponsesUsageInfo.BuiltInTools == nil {
  51. return &usage, nil
  52. }
  53. // 解析 Tools 用量
  54. for _, tool := range responsesResponse.Tools {
  55. buildToolinfo, ok := info.ResponsesUsageInfo.BuiltInTools[common.Interface2String(tool["type"])]
  56. if !ok || buildToolinfo == nil {
  57. logger.LogError(c, fmt.Sprintf("BuiltInTools not found for tool type: %v", tool["type"]))
  58. continue
  59. }
  60. buildToolinfo.CallCount++
  61. }
  62. return &usage, nil
  63. }
  64. func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
  65. if resp == nil || resp.Body == nil {
  66. logger.LogError(c, "invalid response or response body")
  67. return nil, types.NewError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse)
  68. }
  69. defer service.CloseResponseBodyGracefully(resp)
  70. var usage = &dto.Usage{}
  71. var responseTextBuilder strings.Builder
  72. var streamErr *types.NewAPIError
  73. helper.StreamScannerHandler(c, resp, info, func(data string) bool {
  74. var streamResponse dto.ResponsesStreamResponse
  75. if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil {
  76. sendResponsesStreamData(c, streamResponse, data)
  77. switch streamResponse.Type {
  78. case "response.completed":
  79. if streamResponse.Response != nil {
  80. if streamResponse.Response.Usage != nil {
  81. if streamResponse.Response.Usage.InputTokens != 0 {
  82. usage.PromptTokens = streamResponse.Response.Usage.InputTokens
  83. }
  84. if streamResponse.Response.Usage.OutputTokens != 0 {
  85. usage.CompletionTokens = streamResponse.Response.Usage.OutputTokens
  86. }
  87. if streamResponse.Response.Usage.TotalTokens != 0 {
  88. usage.TotalTokens = streamResponse.Response.Usage.TotalTokens
  89. }
  90. if streamResponse.Response.Usage.InputTokensDetails != nil {
  91. usage.PromptTokensDetails.CachedTokens = streamResponse.Response.Usage.InputTokensDetails.CachedTokens
  92. }
  93. }
  94. if streamResponse.Response.HasImageGenerationCall() {
  95. c.Set("image_generation_call", true)
  96. c.Set("image_generation_call_quality", streamResponse.Response.GetQuality())
  97. c.Set("image_generation_call_size", streamResponse.Response.GetSize())
  98. }
  99. }
  100. case "response.output_text.delta":
  101. responseTextBuilder.WriteString(streamResponse.Delta)
  102. case dto.ResponsesOutputTypeItemDone:
  103. if streamResponse.Item != nil {
  104. switch streamResponse.Item.Type {
  105. case dto.BuildInCallWebSearchCall:
  106. if info != nil && info.ResponsesUsageInfo != nil && info.ResponsesUsageInfo.BuiltInTools != nil {
  107. if webSearchTool, exists := info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; exists && webSearchTool != nil {
  108. webSearchTool.CallCount++
  109. }
  110. }
  111. }
  112. }
  113. case "response.error", "response.failed", "error":
  114. streamErr = handleResponsesStreamError(streamResponse, data)
  115. return false
  116. }
  117. } else {
  118. logger.LogError(c, "failed to unmarshal stream response: "+err.Error())
  119. }
  120. return true
  121. })
  122. if streamErr != nil {
  123. return nil, streamErr
  124. }
  125. if usage.CompletionTokens == 0 {
  126. // 计算输出文本的 token 数量
  127. tempStr := responseTextBuilder.String()
  128. if len(tempStr) > 0 {
  129. // 非正常结束,使用输出文本的 token 数量
  130. completionTokens := service.CountTextToken(tempStr, info.UpstreamModelName)
  131. usage.CompletionTokens = completionTokens
  132. }
  133. }
  134. if usage.PromptTokens == 0 && usage.CompletionTokens != 0 {
  135. usage.PromptTokens = info.GetEstimatePromptTokens()
  136. }
  137. usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
  138. return usage, nil
  139. }