Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

248 linhas
7.0 KiB

  1. package zhipu
  2. import (
  3. "bufio"
  4. "encoding/json"
  5. "io"
  6. "net/http"
  7. "strings"
  8. "sync"
  9. "time"
  10. "github.com/QuantumNous/new-api/common"
  11. "github.com/QuantumNous/new-api/constant"
  12. "github.com/QuantumNous/new-api/dto"
  13. relaycommon "github.com/QuantumNous/new-api/relay/common"
  14. "github.com/QuantumNous/new-api/relay/helper"
  15. "github.com/QuantumNous/new-api/service"
  16. "github.com/QuantumNous/new-api/types"
  17. "github.com/gin-gonic/gin"
  18. "github.com/golang-jwt/jwt/v5"
  19. )
  20. // https://open.bigmodel.cn/doc/api#chatglm_std
  21. // chatglm_std, chatglm_lite
  22. // https://open.bigmodel.cn/api/paas/v3/model-api/chatglm_std/invoke
  23. // https://open.bigmodel.cn/api/paas/v3/model-api/chatglm_std/sse-invoke
  24. var zhipuTokens sync.Map
  25. var expSeconds int64 = 24 * 3600
  26. func getZhipuToken(apikey string) string {
  27. data, ok := zhipuTokens.Load(apikey)
  28. if ok {
  29. tokenData := data.(zhipuTokenData)
  30. if time.Now().Before(tokenData.ExpiryTime) {
  31. return tokenData.Token
  32. }
  33. }
  34. split := strings.Split(apikey, ".")
  35. if len(split) != 2 {
  36. common.SysLog("invalid zhipu key: " + apikey)
  37. return ""
  38. }
  39. id := split[0]
  40. secret := split[1]
  41. expMillis := time.Now().Add(time.Duration(expSeconds)*time.Second).UnixNano() / 1e6
  42. expiryTime := time.Now().Add(time.Duration(expSeconds) * time.Second)
  43. timestamp := time.Now().UnixNano() / 1e6
  44. payload := jwt.MapClaims{
  45. "api_key": id,
  46. "exp": expMillis,
  47. "timestamp": timestamp,
  48. }
  49. token := jwt.NewWithClaims(jwt.SigningMethodHS256, payload)
  50. token.Header["alg"] = "HS256"
  51. token.Header["sign_type"] = "SIGN"
  52. tokenString, err := token.SignedString([]byte(secret))
  53. if err != nil {
  54. return ""
  55. }
  56. zhipuTokens.Store(apikey, zhipuTokenData{
  57. Token: tokenString,
  58. ExpiryTime: expiryTime,
  59. })
  60. return tokenString
  61. }
  62. func requestOpenAI2Zhipu(request dto.GeneralOpenAIRequest) *ZhipuRequest {
  63. messages := make([]ZhipuMessage, 0, len(request.Messages))
  64. for _, message := range request.Messages {
  65. if message.Role == "system" {
  66. messages = append(messages, ZhipuMessage{
  67. Role: "system",
  68. Content: message.StringContent(),
  69. })
  70. messages = append(messages, ZhipuMessage{
  71. Role: "user",
  72. Content: "Okay",
  73. })
  74. } else {
  75. messages = append(messages, ZhipuMessage{
  76. Role: message.Role,
  77. Content: message.StringContent(),
  78. })
  79. }
  80. }
  81. return &ZhipuRequest{
  82. Prompt: messages,
  83. Temperature: request.Temperature,
  84. TopP: request.TopP,
  85. Incremental: false,
  86. }
  87. }
  88. func responseZhipu2OpenAI(response *ZhipuResponse) *dto.OpenAITextResponse {
  89. fullTextResponse := dto.OpenAITextResponse{
  90. Id: response.Data.TaskId,
  91. Object: "chat.completion",
  92. Created: common.GetTimestamp(),
  93. Choices: make([]dto.OpenAITextResponseChoice, 0, len(response.Data.Choices)),
  94. Usage: response.Data.Usage,
  95. }
  96. for i, choice := range response.Data.Choices {
  97. openaiChoice := dto.OpenAITextResponseChoice{
  98. Index: i,
  99. Message: dto.Message{
  100. Role: choice.Role,
  101. Content: strings.Trim(choice.Content, "\""),
  102. },
  103. FinishReason: "",
  104. }
  105. if i == len(response.Data.Choices)-1 {
  106. openaiChoice.FinishReason = "stop"
  107. }
  108. fullTextResponse.Choices = append(fullTextResponse.Choices, openaiChoice)
  109. }
  110. return &fullTextResponse
  111. }
  112. func streamResponseZhipu2OpenAI(zhipuResponse string) *dto.ChatCompletionsStreamResponse {
  113. var choice dto.ChatCompletionsStreamResponseChoice
  114. choice.Delta.SetContentString(zhipuResponse)
  115. response := dto.ChatCompletionsStreamResponse{
  116. Object: "chat.completion.chunk",
  117. Created: common.GetTimestamp(),
  118. Model: "chatglm",
  119. Choices: []dto.ChatCompletionsStreamResponseChoice{choice},
  120. }
  121. return &response
  122. }
  123. func streamMetaResponseZhipu2OpenAI(zhipuResponse *ZhipuStreamMetaResponse) (*dto.ChatCompletionsStreamResponse, *dto.Usage) {
  124. var choice dto.ChatCompletionsStreamResponseChoice
  125. choice.Delta.SetContentString("")
  126. choice.FinishReason = &constant.FinishReasonStop
  127. response := dto.ChatCompletionsStreamResponse{
  128. Id: zhipuResponse.RequestId,
  129. Object: "chat.completion.chunk",
  130. Created: common.GetTimestamp(),
  131. Model: "chatglm",
  132. Choices: []dto.ChatCompletionsStreamResponseChoice{choice},
  133. }
  134. return &response, &zhipuResponse.Usage
  135. }
  136. func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
  137. var usage *dto.Usage
  138. scanner := bufio.NewScanner(resp.Body)
  139. scanner.Split(bufio.ScanLines)
  140. dataChan := make(chan string)
  141. metaChan := make(chan string)
  142. stopChan := make(chan bool)
  143. go func() {
  144. for scanner.Scan() {
  145. data := scanner.Text()
  146. lines := strings.Split(data, "\n")
  147. for i, line := range lines {
  148. if len(line) < 5 {
  149. continue
  150. }
  151. if line[:5] == "data:" {
  152. dataChan <- line[5:]
  153. if i != len(lines)-1 {
  154. dataChan <- "\n"
  155. }
  156. } else if line[:5] == "meta:" {
  157. metaChan <- line[5:]
  158. }
  159. }
  160. }
  161. stopChan <- true
  162. }()
  163. helper.SetEventStreamHeaders(c)
  164. c.Stream(func(w io.Writer) bool {
  165. select {
  166. case data := <-dataChan:
  167. response := streamResponseZhipu2OpenAI(data)
  168. jsonResponse, err := json.Marshal(response)
  169. if err != nil {
  170. common.SysLog("error marshalling stream response: " + err.Error())
  171. return true
  172. }
  173. c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
  174. return true
  175. case data := <-metaChan:
  176. var zhipuResponse ZhipuStreamMetaResponse
  177. err := json.Unmarshal([]byte(data), &zhipuResponse)
  178. if err != nil {
  179. common.SysLog("error unmarshalling stream response: " + err.Error())
  180. return true
  181. }
  182. response, zhipuUsage := streamMetaResponseZhipu2OpenAI(&zhipuResponse)
  183. jsonResponse, err := json.Marshal(response)
  184. if err != nil {
  185. common.SysLog("error marshalling stream response: " + err.Error())
  186. return true
  187. }
  188. usage = zhipuUsage
  189. c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
  190. return true
  191. case <-stopChan:
  192. c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
  193. return false
  194. }
  195. })
  196. service.CloseResponseBodyGracefully(resp)
  197. return usage, nil
  198. }
  199. func zhipuHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
  200. var zhipuResponse ZhipuResponse
  201. responseBody, err := io.ReadAll(resp.Body)
  202. if err != nil {
  203. return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError)
  204. }
  205. service.CloseResponseBodyGracefully(resp)
  206. err = json.Unmarshal(responseBody, &zhipuResponse)
  207. if err != nil {
  208. return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
  209. }
  210. if !zhipuResponse.Success {
  211. return nil, types.WithOpenAIError(types.OpenAIError{
  212. Message: zhipuResponse.Msg,
  213. Code: zhipuResponse.Code,
  214. }, resp.StatusCode)
  215. }
  216. fullTextResponse := responseZhipu2OpenAI(&zhipuResponse)
  217. jsonResponse, err := json.Marshal(fullTextResponse)
  218. if err != nil {
  219. return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
  220. }
  221. c.Writer.Header().Set("Content-Type", "application/json")
  222. c.Writer.WriteHeader(resp.StatusCode)
  223. _, err = c.Writer.Write(jsonResponse)
  224. return &fullTextResponse.Usage, nil
  225. }