You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

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