Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.
 
 
 

292 rindas
8.7 KiB

  1. package xunfei
  2. import (
  3. "crypto/hmac"
  4. "crypto/sha256"
  5. "encoding/base64"
  6. "encoding/json"
  7. "fmt"
  8. "io"
  9. "net/url"
  10. "strings"
  11. "time"
  12. "github.com/QuantumNous/new-api/common"
  13. "github.com/QuantumNous/new-api/constant"
  14. "github.com/QuantumNous/new-api/dto"
  15. "github.com/QuantumNous/new-api/relay/helper"
  16. "github.com/QuantumNous/new-api/types"
  17. "github.com/gin-gonic/gin"
  18. "github.com/gorilla/websocket"
  19. )
  20. // https://console.xfyun.cn/services/cbm
  21. // https://www.xfyun.cn/doc/spark/Web.html
  22. func requestOpenAI2Xunfei(request dto.GeneralOpenAIRequest, xunfeiAppId string, domain string) *XunfeiChatRequest {
  23. messages := make([]XunfeiMessage, 0, len(request.Messages))
  24. shouldCovertSystemMessage := !strings.HasSuffix(request.Model, "3.5")
  25. for _, message := range request.Messages {
  26. if message.Role == "system" && shouldCovertSystemMessage {
  27. messages = append(messages, XunfeiMessage{
  28. Role: "user",
  29. Content: message.StringContent(),
  30. })
  31. messages = append(messages, XunfeiMessage{
  32. Role: "assistant",
  33. Content: "Okay",
  34. })
  35. } else {
  36. messages = append(messages, XunfeiMessage{
  37. Role: message.Role,
  38. Content: message.StringContent(),
  39. })
  40. }
  41. }
  42. xunfeiRequest := XunfeiChatRequest{}
  43. xunfeiRequest.Header.AppId = xunfeiAppId
  44. xunfeiRequest.Parameter.Chat.Domain = domain
  45. xunfeiRequest.Parameter.Chat.Temperature = request.Temperature
  46. xunfeiRequest.Parameter.Chat.TopK = request.N
  47. xunfeiRequest.Parameter.Chat.MaxTokens = request.GetMaxTokens()
  48. xunfeiRequest.Payload.Message.Text = messages
  49. return &xunfeiRequest
  50. }
  51. func responseXunfei2OpenAI(response *XunfeiChatResponse) *dto.OpenAITextResponse {
  52. if len(response.Payload.Choices.Text) == 0 {
  53. response.Payload.Choices.Text = []XunfeiChatResponseTextItem{
  54. {
  55. Content: "",
  56. },
  57. }
  58. }
  59. choice := dto.OpenAITextResponseChoice{
  60. Index: 0,
  61. Message: dto.Message{
  62. Role: "assistant",
  63. Content: response.Payload.Choices.Text[0].Content,
  64. },
  65. FinishReason: constant.FinishReasonStop,
  66. }
  67. fullTextResponse := dto.OpenAITextResponse{
  68. Object: "chat.completion",
  69. Created: common.GetTimestamp(),
  70. Choices: []dto.OpenAITextResponseChoice{choice},
  71. Usage: response.Payload.Usage.Text,
  72. }
  73. return &fullTextResponse
  74. }
  75. func streamResponseXunfei2OpenAI(xunfeiResponse *XunfeiChatResponse) *dto.ChatCompletionsStreamResponse {
  76. if len(xunfeiResponse.Payload.Choices.Text) == 0 {
  77. xunfeiResponse.Payload.Choices.Text = []XunfeiChatResponseTextItem{
  78. {
  79. Content: "",
  80. },
  81. }
  82. }
  83. var choice dto.ChatCompletionsStreamResponseChoice
  84. choice.Delta.SetContentString(xunfeiResponse.Payload.Choices.Text[0].Content)
  85. if xunfeiResponse.Payload.Choices.Status == 2 {
  86. choice.FinishReason = &constant.FinishReasonStop
  87. }
  88. response := dto.ChatCompletionsStreamResponse{
  89. Object: "chat.completion.chunk",
  90. Created: common.GetTimestamp(),
  91. Model: "SparkDesk",
  92. Choices: []dto.ChatCompletionsStreamResponseChoice{choice},
  93. }
  94. return &response
  95. }
  96. func buildXunfeiAuthUrl(hostUrl string, apiKey, apiSecret string) string {
  97. HmacWithShaToBase64 := func(algorithm, data, key string) string {
  98. mac := hmac.New(sha256.New, []byte(key))
  99. mac.Write([]byte(data))
  100. encodeData := mac.Sum(nil)
  101. return base64.StdEncoding.EncodeToString(encodeData)
  102. }
  103. ul, err := url.Parse(hostUrl)
  104. if err != nil {
  105. fmt.Println(err)
  106. }
  107. date := time.Now().UTC().Format(time.RFC1123)
  108. signString := []string{"host: " + ul.Host, "date: " + date, "GET " + ul.Path + " HTTP/1.1"}
  109. sign := strings.Join(signString, "\n")
  110. sha := HmacWithShaToBase64("hmac-sha256", sign, apiSecret)
  111. authUrl := fmt.Sprintf("hmac username=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey,
  112. "hmac-sha256", "host date request-line", sha)
  113. authorization := base64.StdEncoding.EncodeToString([]byte(authUrl))
  114. v := url.Values{}
  115. v.Add("host", ul.Host)
  116. v.Add("date", date)
  117. v.Add("authorization", authorization)
  118. callUrl := hostUrl + "?" + v.Encode()
  119. return callUrl
  120. }
  121. func xunfeiStreamHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.NewAPIError) {
  122. domain, authUrl := getXunfeiAuthUrl(c, apiKey, apiSecret, textRequest.Model)
  123. dataChan, stopChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId)
  124. if err != nil {
  125. return nil, types.NewError(err, types.ErrorCodeDoRequestFailed)
  126. }
  127. helper.SetEventStreamHeaders(c)
  128. var usage dto.Usage
  129. c.Stream(func(w io.Writer) bool {
  130. select {
  131. case xunfeiResponse := <-dataChan:
  132. usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens
  133. usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens
  134. usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens
  135. response := streamResponseXunfei2OpenAI(&xunfeiResponse)
  136. jsonResponse, err := json.Marshal(response)
  137. if err != nil {
  138. common.SysLog("error marshalling stream response: " + err.Error())
  139. return true
  140. }
  141. c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
  142. return true
  143. case <-stopChan:
  144. c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
  145. return false
  146. }
  147. })
  148. return &usage, nil
  149. }
  150. func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.NewAPIError) {
  151. domain, authUrl := getXunfeiAuthUrl(c, apiKey, apiSecret, textRequest.Model)
  152. dataChan, stopChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId)
  153. if err != nil {
  154. return nil, types.NewError(err, types.ErrorCodeDoRequestFailed)
  155. }
  156. var usage dto.Usage
  157. var content string
  158. var xunfeiResponse XunfeiChatResponse
  159. stop := false
  160. for !stop {
  161. select {
  162. case xunfeiResponse = <-dataChan:
  163. if len(xunfeiResponse.Payload.Choices.Text) == 0 {
  164. continue
  165. }
  166. content += xunfeiResponse.Payload.Choices.Text[0].Content
  167. usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens
  168. usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens
  169. usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens
  170. case stop = <-stopChan:
  171. }
  172. }
  173. if len(xunfeiResponse.Payload.Choices.Text) == 0 {
  174. xunfeiResponse.Payload.Choices.Text = []XunfeiChatResponseTextItem{
  175. {
  176. Content: "",
  177. },
  178. }
  179. }
  180. xunfeiResponse.Payload.Choices.Text[0].Content = content
  181. response := responseXunfei2OpenAI(&xunfeiResponse)
  182. jsonResponse, err := json.Marshal(response)
  183. if err != nil {
  184. return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
  185. }
  186. c.Writer.Header().Set("Content-Type", "application/json")
  187. _, _ = c.Writer.Write(jsonResponse)
  188. return &usage, nil
  189. }
  190. func xunfeiMakeRequest(textRequest dto.GeneralOpenAIRequest, domain, authUrl, appId string) (chan XunfeiChatResponse, chan bool, error) {
  191. d := websocket.Dialer{
  192. HandshakeTimeout: 5 * time.Second,
  193. }
  194. conn, resp, err := d.Dial(authUrl, nil)
  195. if err != nil || resp.StatusCode != 101 {
  196. return nil, nil, err
  197. }
  198. data := requestOpenAI2Xunfei(textRequest, appId, domain)
  199. err = conn.WriteJSON(data)
  200. if err != nil {
  201. return nil, nil, err
  202. }
  203. dataChan := make(chan XunfeiChatResponse)
  204. stopChan := make(chan bool)
  205. go func() {
  206. defer func() {
  207. conn.Close()
  208. }()
  209. for {
  210. _, msg, err := conn.ReadMessage()
  211. if err != nil {
  212. common.SysLog("error reading stream response: " + err.Error())
  213. break
  214. }
  215. var response XunfeiChatResponse
  216. err = json.Unmarshal(msg, &response)
  217. if err != nil {
  218. common.SysLog("error unmarshalling stream response: " + err.Error())
  219. break
  220. }
  221. dataChan <- response
  222. if response.Payload.Choices.Status == 2 {
  223. if err != nil {
  224. common.SysLog("error closing websocket connection: " + err.Error())
  225. }
  226. break
  227. }
  228. }
  229. stopChan <- true
  230. }()
  231. return dataChan, stopChan, nil
  232. }
  233. func apiVersion2domain(apiVersion string) string {
  234. switch apiVersion {
  235. case "v1.1":
  236. return "lite"
  237. case "v2.1":
  238. return "generalv2"
  239. case "v3.1":
  240. return "generalv3"
  241. case "v3.5":
  242. return "generalv3.5"
  243. case "v4.0":
  244. return "4.0Ultra"
  245. }
  246. return "general" + apiVersion
  247. }
  248. func getXunfeiAuthUrl(c *gin.Context, apiKey string, apiSecret string, modelName string) (string, string) {
  249. apiVersion := getAPIVersion(c, modelName)
  250. domain := apiVersion2domain(apiVersion)
  251. authUrl := buildXunfeiAuthUrl(fmt.Sprintf("wss://spark-api.xf-yun.com/%s/chat", apiVersion), apiKey, apiSecret)
  252. return domain, authUrl
  253. }
  254. func getAPIVersion(c *gin.Context, modelName string) string {
  255. query := c.Request.URL.Query()
  256. apiVersion := query.Get("api-version")
  257. if apiVersion != "" {
  258. return apiVersion
  259. }
  260. parts := strings.Split(modelName, "-")
  261. if len(parts) == 2 {
  262. apiVersion = parts[1]
  263. return apiVersion
  264. }
  265. apiVersion = c.GetString("api_version")
  266. if apiVersion != "" {
  267. return apiVersion
  268. }
  269. apiVersion = "v1.1"
  270. common.SysLog("api_version not found, using default: " + apiVersion)
  271. return apiVersion
  272. }