relay-text.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735
  1. package controller
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "net/http"
  10. "one-api/common"
  11. "one-api/model"
  12. "strings"
  13. "time"
  14. "github.com/gin-gonic/gin"
  15. )
  16. const (
  17. APITypeOpenAI = iota
  18. APITypeClaude
  19. APITypePaLM
  20. APITypeBaidu
  21. APITypeZhipu
  22. APITypeAli
  23. APITypeXunfei
  24. APITypeAIProxyLibrary
  25. APITypeTencent
  26. APITypeGemini
  27. )
  28. var httpClient *http.Client
  29. var impatientHTTPClient *http.Client
  30. func init() {
  31. if common.RelayTimeout == 0 {
  32. httpClient = &http.Client{}
  33. } else {
  34. httpClient = &http.Client{
  35. Timeout: time.Duration(common.RelayTimeout) * time.Second,
  36. }
  37. }
  38. impatientHTTPClient = &http.Client{
  39. Timeout: 5 * time.Second,
  40. }
  41. }
  42. func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
  43. channelType := c.GetInt("channel")
  44. channelId := c.GetInt("channel_id")
  45. tokenId := c.GetInt("token_id")
  46. userId := c.GetInt("id")
  47. group := c.GetString("group")
  48. startTime := time.Now()
  49. var textRequest GeneralOpenAIRequest
  50. err := common.UnmarshalBodyReusable(c, &textRequest)
  51. if err != nil {
  52. return errorWrapper(err, "bind_request_body_failed", http.StatusBadRequest)
  53. }
  54. if relayMode == RelayModeModerations && textRequest.Model == "" {
  55. textRequest.Model = "text-moderation-latest"
  56. }
  57. if relayMode == RelayModeEmbeddings && textRequest.Model == "" {
  58. textRequest.Model = c.Param("model")
  59. }
  60. // request validation
  61. if textRequest.Model == "" {
  62. return errorWrapper(errors.New("model is required"), "required_field_missing", http.StatusBadRequest)
  63. }
  64. switch relayMode {
  65. case RelayModeCompletions:
  66. if textRequest.Prompt == "" {
  67. return errorWrapper(errors.New("field prompt is required"), "required_field_missing", http.StatusBadRequest)
  68. }
  69. case RelayModeChatCompletions:
  70. if textRequest.Messages == nil || len(textRequest.Messages) == 0 {
  71. return errorWrapper(errors.New("field messages is required"), "required_field_missing", http.StatusBadRequest)
  72. }
  73. case RelayModeEmbeddings:
  74. case RelayModeModerations:
  75. if textRequest.Input == "" {
  76. return errorWrapper(errors.New("field input is required"), "required_field_missing", http.StatusBadRequest)
  77. }
  78. case RelayModeEdits:
  79. if textRequest.Instruction == "" {
  80. return errorWrapper(errors.New("field instruction is required"), "required_field_missing", http.StatusBadRequest)
  81. }
  82. }
  83. // map model name
  84. modelMapping := c.GetString("model_mapping")
  85. isModelMapped := false
  86. if modelMapping != "" && modelMapping != "{}" {
  87. modelMap := make(map[string]string)
  88. err := json.Unmarshal([]byte(modelMapping), &modelMap)
  89. if err != nil {
  90. return errorWrapper(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError)
  91. }
  92. if modelMap[textRequest.Model] != "" {
  93. textRequest.Model = modelMap[textRequest.Model]
  94. isModelMapped = true
  95. }
  96. }
  97. apiType := APITypeOpenAI
  98. switch channelType {
  99. case common.ChannelTypeAnthropic:
  100. apiType = APITypeClaude
  101. case common.ChannelTypeBaidu:
  102. apiType = APITypeBaidu
  103. case common.ChannelTypePaLM:
  104. apiType = APITypePaLM
  105. case common.ChannelTypeZhipu:
  106. apiType = APITypeZhipu
  107. case common.ChannelTypeAli:
  108. apiType = APITypeAli
  109. case common.ChannelTypeXunfei:
  110. apiType = APITypeXunfei
  111. case common.ChannelTypeAIProxyLibrary:
  112. apiType = APITypeAIProxyLibrary
  113. case common.ChannelTypeTencent:
  114. apiType = APITypeTencent
  115. case common.ChannelTypeGemini:
  116. apiType = APITypeGemini
  117. }
  118. baseURL := common.ChannelBaseURLs[channelType]
  119. requestURL := c.Request.URL.String()
  120. if c.GetString("base_url") != "" {
  121. baseURL = c.GetString("base_url")
  122. }
  123. fullRequestURL := getFullRequestURL(baseURL, requestURL, channelType)
  124. switch apiType {
  125. case APITypeOpenAI:
  126. if channelType == common.ChannelTypeAzure {
  127. // https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api
  128. query := c.Request.URL.Query()
  129. apiVersion := query.Get("api-version")
  130. if apiVersion == "" {
  131. apiVersion = c.GetString("api_version")
  132. }
  133. requestURL := strings.Split(requestURL, "?")[0]
  134. requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, apiVersion)
  135. baseURL = c.GetString("base_url")
  136. task := strings.TrimPrefix(requestURL, "/v1/")
  137. model_ := textRequest.Model
  138. model_ = strings.Replace(model_, ".", "", -1)
  139. // https://github.com/songquanpeng/one-api/issues/67
  140. model_ = strings.TrimSuffix(model_, "-0301")
  141. model_ = strings.TrimSuffix(model_, "-0314")
  142. model_ = strings.TrimSuffix(model_, "-0613")
  143. fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/%s", baseURL, model_, task)
  144. }
  145. case APITypeClaude:
  146. fullRequestURL = "https://api.anthropic.com/v1/complete"
  147. if baseURL != "" {
  148. fullRequestURL = fmt.Sprintf("%s/v1/complete", baseURL)
  149. }
  150. case APITypeBaidu:
  151. switch textRequest.Model {
  152. case "ERNIE-Bot":
  153. fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions"
  154. case "ERNIE-Bot-turbo":
  155. fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant"
  156. case "ERNIE-Bot-4":
  157. fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro"
  158. case "BLOOMZ-7B":
  159. fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/bloomz_7b1"
  160. case "Embedding-V1":
  161. fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1"
  162. }
  163. apiKey := c.Request.Header.Get("Authorization")
  164. apiKey = strings.TrimPrefix(apiKey, "Bearer ")
  165. var err error
  166. if apiKey, err = getBaiduAccessToken(apiKey); err != nil {
  167. return errorWrapper(err, "invalid_baidu_config", http.StatusInternalServerError)
  168. }
  169. fullRequestURL += "?access_token=" + apiKey
  170. case APITypePaLM:
  171. fullRequestURL = "https://generativelanguage.googleapis.com/v1beta2/models/chat-bison-001:generateMessage"
  172. if baseURL != "" {
  173. fullRequestURL = fmt.Sprintf("%s/v1beta2/models/chat-bison-001:generateMessage", baseURL)
  174. }
  175. apiKey := c.Request.Header.Get("Authorization")
  176. apiKey = strings.TrimPrefix(apiKey, "Bearer ")
  177. fullRequestURL += "?key=" + apiKey
  178. case APITypeGemini:
  179. requestBaseURL := "https://generativelanguage.googleapis.com"
  180. if baseURL != "" {
  181. requestBaseURL = baseURL
  182. }
  183. version := "v1beta"
  184. if c.GetString("api_version") != "" {
  185. version = c.GetString("api_version")
  186. }
  187. action := "generateContent"
  188. if textRequest.Stream {
  189. action = "streamGenerateContent"
  190. }
  191. fullRequestURL = fmt.Sprintf("%s/%s/models/%s:%s", requestBaseURL, version, textRequest.Model, action)
  192. apiKey := c.Request.Header.Get("Authorization")
  193. apiKey = strings.TrimPrefix(apiKey, "Bearer ")
  194. fullRequestURL += "?key=" + apiKey
  195. //log.Println(fullRequestURL)
  196. case APITypeZhipu:
  197. method := "invoke"
  198. if textRequest.Stream {
  199. method = "sse-invoke"
  200. }
  201. fullRequestURL = fmt.Sprintf("https://open.bigmodel.cn/api/paas/v3/model-api/%s/%s", textRequest.Model, method)
  202. case APITypeAli:
  203. fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
  204. if relayMode == RelayModeEmbeddings {
  205. fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
  206. }
  207. case APITypeTencent:
  208. fullRequestURL = "https://hunyuan.cloud.tencent.com/hyllm/v1/chat/completions"
  209. case APITypeAIProxyLibrary:
  210. fullRequestURL = fmt.Sprintf("%s/api/library/ask", baseURL)
  211. }
  212. var promptTokens int
  213. var completionTokens int
  214. switch relayMode {
  215. case RelayModeChatCompletions:
  216. promptTokens, err = countTokenMessages(textRequest.Messages, textRequest.Model)
  217. if err != nil {
  218. return errorWrapper(err, "count_token_messages_failed", http.StatusInternalServerError)
  219. }
  220. case RelayModeCompletions:
  221. promptTokens = countTokenInput(textRequest.Prompt, textRequest.Model)
  222. case RelayModeModerations:
  223. promptTokens = countTokenInput(textRequest.Input, textRequest.Model)
  224. }
  225. modelPrice := common.GetModelPrice(textRequest.Model)
  226. groupRatio := common.GetGroupRatio(group)
  227. var preConsumedQuota int
  228. var ratio float64
  229. var modelRatio float64
  230. if modelPrice == -1 {
  231. preConsumedTokens := common.PreConsumedQuota
  232. if textRequest.MaxTokens != 0 {
  233. preConsumedTokens = promptTokens + int(textRequest.MaxTokens)
  234. }
  235. modelRatio = common.GetModelRatio(textRequest.Model)
  236. ratio = modelRatio * groupRatio
  237. preConsumedQuota = int(float64(preConsumedTokens) * ratio)
  238. } else {
  239. preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatio)
  240. }
  241. userQuota, err := model.CacheGetUserQuota(userId)
  242. if err != nil {
  243. return errorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError)
  244. }
  245. if userQuota < 0 || userQuota-preConsumedQuota < 0 {
  246. return errorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
  247. }
  248. err = model.CacheDecreaseUserQuota(userId, preConsumedQuota)
  249. if err != nil {
  250. return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError)
  251. }
  252. if userQuota > 100*preConsumedQuota {
  253. // in this case, we do not pre-consume quota
  254. // because the user has enough quota
  255. preConsumedQuota = 0
  256. common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
  257. }
  258. if preConsumedQuota > 0 {
  259. userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
  260. if err != nil {
  261. return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
  262. }
  263. }
  264. var requestBody io.Reader
  265. if isModelMapped {
  266. jsonStr, err := json.Marshal(textRequest)
  267. if err != nil {
  268. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  269. }
  270. requestBody = bytes.NewBuffer(jsonStr)
  271. } else {
  272. requestBody = c.Request.Body
  273. }
  274. switch apiType {
  275. case APITypeClaude:
  276. claudeRequest := requestOpenAI2Claude(textRequest)
  277. jsonStr, err := json.Marshal(claudeRequest)
  278. if err != nil {
  279. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  280. }
  281. requestBody = bytes.NewBuffer(jsonStr)
  282. case APITypeBaidu:
  283. var jsonData []byte
  284. var err error
  285. switch relayMode {
  286. case RelayModeEmbeddings:
  287. baiduEmbeddingRequest := embeddingRequestOpenAI2Baidu(textRequest)
  288. jsonData, err = json.Marshal(baiduEmbeddingRequest)
  289. default:
  290. baiduRequest := requestOpenAI2Baidu(textRequest)
  291. jsonData, err = json.Marshal(baiduRequest)
  292. }
  293. if err != nil {
  294. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  295. }
  296. requestBody = bytes.NewBuffer(jsonData)
  297. case APITypePaLM:
  298. palmRequest := requestOpenAI2PaLM(textRequest)
  299. jsonStr, err := json.Marshal(palmRequest)
  300. if err != nil {
  301. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  302. }
  303. requestBody = bytes.NewBuffer(jsonStr)
  304. case APITypeGemini:
  305. geminiChatRequest := requestOpenAI2Gemini(textRequest)
  306. jsonStr, err := json.Marshal(geminiChatRequest)
  307. if err != nil {
  308. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  309. }
  310. requestBody = bytes.NewBuffer(jsonStr)
  311. case APITypeZhipu:
  312. zhipuRequest := requestOpenAI2Zhipu(textRequest)
  313. jsonStr, err := json.Marshal(zhipuRequest)
  314. if err != nil {
  315. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  316. }
  317. requestBody = bytes.NewBuffer(jsonStr)
  318. case APITypeAli:
  319. var jsonStr []byte
  320. var err error
  321. switch relayMode {
  322. case RelayModeEmbeddings:
  323. aliEmbeddingRequest := embeddingRequestOpenAI2Ali(textRequest)
  324. jsonStr, err = json.Marshal(aliEmbeddingRequest)
  325. default:
  326. aliRequest := requestOpenAI2Ali(textRequest)
  327. jsonStr, err = json.Marshal(aliRequest)
  328. }
  329. if err != nil {
  330. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  331. }
  332. requestBody = bytes.NewBuffer(jsonStr)
  333. case APITypeTencent:
  334. apiKey := c.Request.Header.Get("Authorization")
  335. apiKey = strings.TrimPrefix(apiKey, "Bearer ")
  336. appId, secretId, secretKey, err := parseTencentConfig(apiKey)
  337. if err != nil {
  338. return errorWrapper(err, "invalid_tencent_config", http.StatusInternalServerError)
  339. }
  340. tencentRequest := requestOpenAI2Tencent(textRequest)
  341. tencentRequest.AppId = appId
  342. tencentRequest.SecretId = secretId
  343. jsonStr, err := json.Marshal(tencentRequest)
  344. if err != nil {
  345. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  346. }
  347. sign := getTencentSign(*tencentRequest, secretKey)
  348. c.Request.Header.Set("Authorization", sign)
  349. requestBody = bytes.NewBuffer(jsonStr)
  350. case APITypeAIProxyLibrary:
  351. aiProxyLibraryRequest := requestOpenAI2AIProxyLibrary(textRequest)
  352. aiProxyLibraryRequest.LibraryId = c.GetString("library_id")
  353. jsonStr, err := json.Marshal(aiProxyLibraryRequest)
  354. if err != nil {
  355. return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
  356. }
  357. requestBody = bytes.NewBuffer(jsonStr)
  358. }
  359. var req *http.Request
  360. var resp *http.Response
  361. isStream := textRequest.Stream
  362. if apiType != APITypeXunfei { // cause xunfei use websocket
  363. req, err = http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
  364. // 设置GetBody函数,该函数返回一个新的io.ReadCloser,该io.ReadCloser返回与原始请求体相同的数据
  365. req.GetBody = func() (io.ReadCloser, error) {
  366. return io.NopCloser(requestBody), nil
  367. }
  368. if err != nil {
  369. return errorWrapper(err, "new_request_failed", http.StatusInternalServerError)
  370. }
  371. apiKey := c.Request.Header.Get("Authorization")
  372. apiKey = strings.TrimPrefix(apiKey, "Bearer ")
  373. switch apiType {
  374. case APITypeOpenAI:
  375. if channelType == common.ChannelTypeAzure {
  376. req.Header.Set("api-key", apiKey)
  377. } else {
  378. req.Header.Set("Authorization", c.Request.Header.Get("Authorization"))
  379. if c.Request.Header.Get("OpenAI-Organization") != "" {
  380. req.Header.Set("OpenAI-Organization", c.Request.Header.Get("OpenAI-Organization"))
  381. }
  382. if channelType == common.ChannelTypeOpenRouter {
  383. req.Header.Set("HTTP-Referer", "https://github.com/songquanpeng/one-api")
  384. req.Header.Set("X-Title", "One API")
  385. }
  386. }
  387. case APITypeClaude:
  388. req.Header.Set("x-api-key", apiKey)
  389. anthropicVersion := c.Request.Header.Get("anthropic-version")
  390. if anthropicVersion == "" {
  391. anthropicVersion = "2023-06-01"
  392. }
  393. req.Header.Set("anthropic-version", anthropicVersion)
  394. case APITypeZhipu:
  395. token := getZhipuToken(apiKey)
  396. req.Header.Set("Authorization", token)
  397. case APITypeAli:
  398. req.Header.Set("Authorization", "Bearer "+apiKey)
  399. if textRequest.Stream {
  400. req.Header.Set("X-DashScope-SSE", "enable")
  401. }
  402. case APITypeTencent:
  403. req.Header.Set("Authorization", apiKey)
  404. case APITypeGemini:
  405. req.Header.Set("Content-Type", "application/json")
  406. default:
  407. req.Header.Set("Authorization", "Bearer "+apiKey)
  408. }
  409. if apiType != APITypeGemini {
  410. // 设置公共头部...
  411. req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
  412. req.Header.Set("Accept", c.Request.Header.Get("Accept"))
  413. if isStream && c.Request.Header.Get("Accept") == "" {
  414. req.Header.Set("Accept", "text/event-stream")
  415. }
  416. }
  417. //req.HeaderBar.Set("Connection", c.Request.HeaderBar.Get("Connection"))
  418. resp, err = httpClient.Do(req)
  419. if err != nil {
  420. return errorWrapper(err, "do_request_failed", http.StatusInternalServerError)
  421. }
  422. err = req.Body.Close()
  423. if err != nil {
  424. return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
  425. }
  426. err = c.Request.Body.Close()
  427. if err != nil {
  428. return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
  429. }
  430. isStream = isStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
  431. if resp.StatusCode != http.StatusOK {
  432. if preConsumedQuota != 0 {
  433. go func(ctx context.Context) {
  434. // return pre-consumed quota
  435. err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0, false)
  436. if err != nil {
  437. common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
  438. }
  439. }(c.Request.Context())
  440. }
  441. return relayErrorHandler(resp)
  442. }
  443. }
  444. var textResponse TextResponse
  445. tokenName := c.GetString("token_name")
  446. defer func(ctx context.Context) {
  447. // c.Writer.Flush()
  448. go func() {
  449. promptTokens = textResponse.Usage.PromptTokens
  450. completionTokens = textResponse.Usage.CompletionTokens
  451. quota := 0
  452. if modelPrice == -1 {
  453. completionRatio := common.GetCompletionRatio(textRequest.Model)
  454. quota = promptTokens + int(float64(completionTokens)*completionRatio)
  455. quota = int(float64(quota) * ratio)
  456. if ratio != 0 && quota <= 0 {
  457. quota = 1
  458. }
  459. } else {
  460. quota = int(modelPrice * common.QuotaPerUnit * groupRatio)
  461. }
  462. totalTokens := promptTokens + completionTokens
  463. if totalTokens == 0 {
  464. // in this case, must be some error happened
  465. // we cannot just return, because we may have to return the pre-consumed quota
  466. quota = 0
  467. }
  468. quotaDelta := quota - preConsumedQuota
  469. err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota, true)
  470. if err != nil {
  471. common.LogError(ctx, "error consuming token remain quota: "+err.Error())
  472. }
  473. err = model.CacheUpdateUserQuota(userId)
  474. if err != nil {
  475. common.LogError(ctx, "error update user quota cache: "+err.Error())
  476. }
  477. // record all the consume log even if quota is 0
  478. useTimeSeconds := time.Now().Unix() - startTime.Unix()
  479. var logContent string
  480. if modelPrice == -1 {
  481. logContent = fmt.Sprintf("模型倍率 %.2f,分组倍率 %.2f,用时 %d秒", modelRatio, groupRatio, useTimeSeconds)
  482. } else {
  483. logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f,用时 %d秒", modelPrice, groupRatio, useTimeSeconds)
  484. }
  485. logModel := textRequest.Model
  486. if strings.HasPrefix(logModel, "gpt-4-gizmo") {
  487. logModel = "gpt-4-gizmo-*"
  488. logContent += fmt.Sprintf(",模型 %s", textRequest.Model)
  489. }
  490. model.RecordConsumeLog(ctx, userId, channelId, promptTokens, completionTokens, logModel, tokenName, quota, logContent, tokenId, userQuota)
  491. model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
  492. model.UpdateChannelUsedQuota(channelId, quota)
  493. //if quota != 0 {
  494. //
  495. //}
  496. }()
  497. }(c.Request.Context())
  498. switch apiType {
  499. case APITypeOpenAI:
  500. if isStream {
  501. err, responseText := openaiStreamHandler(c, resp, relayMode)
  502. if err != nil {
  503. return err
  504. }
  505. textResponse.Usage.PromptTokens = promptTokens
  506. textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
  507. return nil
  508. } else {
  509. err, usage := openaiHandler(c, resp, promptTokens, textRequest.Model)
  510. if err != nil {
  511. return err
  512. }
  513. if usage != nil {
  514. textResponse.Usage = *usage
  515. }
  516. return nil
  517. }
  518. case APITypeClaude:
  519. if isStream {
  520. err, responseText := claudeStreamHandler(c, resp)
  521. if err != nil {
  522. return err
  523. }
  524. textResponse.Usage.PromptTokens = promptTokens
  525. textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
  526. return nil
  527. } else {
  528. err, usage := claudeHandler(c, resp, promptTokens, textRequest.Model)
  529. if err != nil {
  530. return err
  531. }
  532. if usage != nil {
  533. textResponse.Usage = *usage
  534. }
  535. return nil
  536. }
  537. case APITypeBaidu:
  538. if isStream {
  539. err, usage := baiduStreamHandler(c, resp)
  540. if err != nil {
  541. return err
  542. }
  543. if usage != nil {
  544. textResponse.Usage = *usage
  545. }
  546. return nil
  547. } else {
  548. var err *OpenAIErrorWithStatusCode
  549. var usage *Usage
  550. switch relayMode {
  551. case RelayModeEmbeddings:
  552. err, usage = baiduEmbeddingHandler(c, resp)
  553. default:
  554. err, usage = baiduHandler(c, resp)
  555. }
  556. if err != nil {
  557. return err
  558. }
  559. if usage != nil {
  560. textResponse.Usage = *usage
  561. }
  562. return nil
  563. }
  564. case APITypePaLM:
  565. if textRequest.Stream { // PaLM2 API does not support stream
  566. err, responseText := palmStreamHandler(c, resp)
  567. if err != nil {
  568. return err
  569. }
  570. textResponse.Usage.PromptTokens = promptTokens
  571. textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
  572. return nil
  573. } else {
  574. err, usage := palmHandler(c, resp, promptTokens, textRequest.Model)
  575. if err != nil {
  576. return err
  577. }
  578. if usage != nil {
  579. textResponse.Usage = *usage
  580. }
  581. return nil
  582. }
  583. case APITypeGemini:
  584. if textRequest.Stream {
  585. err, responseText := geminiChatStreamHandler(c, resp)
  586. if err != nil {
  587. return err
  588. }
  589. textResponse.Usage.PromptTokens = promptTokens
  590. textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
  591. return nil
  592. } else {
  593. err, usage := geminiChatHandler(c, resp, promptTokens, textRequest.Model)
  594. if err != nil {
  595. return err
  596. }
  597. if usage != nil {
  598. textResponse.Usage = *usage
  599. }
  600. return nil
  601. }
  602. case APITypeZhipu:
  603. if isStream {
  604. err, usage := zhipuStreamHandler(c, resp)
  605. if err != nil {
  606. return err
  607. }
  608. if usage != nil {
  609. textResponse.Usage = *usage
  610. }
  611. // zhipu's API does not return prompt tokens & completion tokens
  612. textResponse.Usage.PromptTokens = textResponse.Usage.TotalTokens
  613. return nil
  614. } else {
  615. err, usage := zhipuHandler(c, resp)
  616. if err != nil {
  617. return err
  618. }
  619. if usage != nil {
  620. textResponse.Usage = *usage
  621. }
  622. // zhipu's API does not return prompt tokens & completion tokens
  623. textResponse.Usage.PromptTokens = textResponse.Usage.TotalTokens
  624. return nil
  625. }
  626. case APITypeAli:
  627. if isStream {
  628. err, usage := aliStreamHandler(c, resp)
  629. if err != nil {
  630. return err
  631. }
  632. if usage != nil {
  633. textResponse.Usage = *usage
  634. }
  635. return nil
  636. } else {
  637. var err *OpenAIErrorWithStatusCode
  638. var usage *Usage
  639. switch relayMode {
  640. case RelayModeEmbeddings:
  641. err, usage = aliEmbeddingHandler(c, resp)
  642. default:
  643. err, usage = aliHandler(c, resp)
  644. }
  645. if err != nil {
  646. return err
  647. }
  648. if usage != nil {
  649. textResponse.Usage = *usage
  650. }
  651. return nil
  652. }
  653. case APITypeXunfei:
  654. auth := c.Request.Header.Get("Authorization")
  655. auth = strings.TrimPrefix(auth, "Bearer ")
  656. splits := strings.Split(auth, "|")
  657. if len(splits) != 3 {
  658. return errorWrapper(errors.New("invalid auth"), "invalid_auth", http.StatusBadRequest)
  659. }
  660. var err *OpenAIErrorWithStatusCode
  661. var usage *Usage
  662. if isStream {
  663. err, usage = xunfeiStreamHandler(c, textRequest, splits[0], splits[1], splits[2])
  664. } else {
  665. err, usage = xunfeiHandler(c, textRequest, splits[0], splits[1], splits[2])
  666. }
  667. if err != nil {
  668. return err
  669. }
  670. if usage != nil {
  671. textResponse.Usage = *usage
  672. }
  673. return nil
  674. case APITypeAIProxyLibrary:
  675. if isStream {
  676. err, usage := aiProxyLibraryStreamHandler(c, resp)
  677. if err != nil {
  678. return err
  679. }
  680. if usage != nil {
  681. textResponse.Usage = *usage
  682. }
  683. return nil
  684. } else {
  685. err, usage := aiProxyLibraryHandler(c, resp)
  686. if err != nil {
  687. return err
  688. }
  689. if usage != nil {
  690. textResponse.Usage = *usage
  691. }
  692. return nil
  693. }
  694. case APITypeTencent:
  695. if isStream {
  696. err, responseText := tencentStreamHandler(c, resp)
  697. if err != nil {
  698. return err
  699. }
  700. textResponse.Usage.PromptTokens = promptTokens
  701. textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
  702. return nil
  703. } else {
  704. err, usage := tencentHandler(c, resp)
  705. if err != nil {
  706. return err
  707. }
  708. if usage != nil {
  709. textResponse.Usage = *usage
  710. }
  711. return nil
  712. }
  713. default:
  714. return errorWrapper(errors.New("unknown api type"), "unknown_api_type", http.StatusInternalServerError)
  715. }
  716. }