chat_handler.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669
  1. package relay
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "log"
  8. "net/http"
  9. "one-api/common"
  10. "one-api/constant"
  11. "one-api/dto"
  12. "one-api/logger"
  13. "one-api/model"
  14. relaycommon "one-api/relay/common"
  15. relayconstant "one-api/relay/constant"
  16. "one-api/relay/helper"
  17. "one-api/service"
  18. "one-api/setting"
  19. "strconv"
  20. "strings"
  21. "time"
  22. "github.com/gin-gonic/gin"
  23. )
  24. func RelayMidjourneyImage(c *gin.Context) {
  25. taskId := c.Param("id")
  26. midjourneyTask := model.GetByOnlyMJId(taskId)
  27. if midjourneyTask == nil {
  28. c.JSON(400, gin.H{
  29. "error": "midjourney_task_not_found",
  30. })
  31. return
  32. }
  33. var httpClient *http.Client
  34. if channel, err := model.CacheGetChannel(midjourneyTask.ChannelId); err == nil {
  35. proxy := channel.GetSetting().Proxy
  36. if proxy != "" {
  37. if httpClient, err = service.NewProxyHttpClient(proxy); err != nil {
  38. c.JSON(400, gin.H{
  39. "error": "proxy_url_invalid",
  40. })
  41. return
  42. }
  43. }
  44. }
  45. if httpClient == nil {
  46. httpClient = service.GetHttpClient()
  47. }
  48. resp, err := httpClient.Get(midjourneyTask.ImageUrl)
  49. if err != nil {
  50. c.JSON(http.StatusInternalServerError, gin.H{
  51. "error": "http_get_image_failed",
  52. })
  53. return
  54. }
  55. defer resp.Body.Close()
  56. if resp.StatusCode != http.StatusOK {
  57. responseBody, _ := io.ReadAll(resp.Body)
  58. c.JSON(resp.StatusCode, gin.H{
  59. "error": string(responseBody),
  60. })
  61. return
  62. }
  63. // 从Content-Type头获取MIME类型
  64. contentType := resp.Header.Get("Content-Type")
  65. if contentType == "" {
  66. // 如果无法确定内容类型,则默认为jpeg
  67. contentType = "image/jpeg"
  68. }
  69. // 设置响应的内容类型
  70. c.Writer.Header().Set("Content-Type", contentType)
  71. // 将图片流式传输到响应体
  72. _, err = io.Copy(c.Writer, resp.Body)
  73. if err != nil {
  74. log.Println("Failed to stream image:", err)
  75. }
  76. return
  77. }
  78. func RelayMidjourneyNotify(c *gin.Context) *dto.MidjourneyResponse {
  79. var midjRequest dto.MidjourneyDto
  80. err := common.UnmarshalBodyReusable(c, &midjRequest)
  81. if err != nil {
  82. return &dto.MidjourneyResponse{
  83. Code: 4,
  84. Description: "bind_request_body_failed",
  85. Properties: nil,
  86. Result: "",
  87. }
  88. }
  89. midjourneyTask := model.GetByOnlyMJId(midjRequest.MjId)
  90. if midjourneyTask == nil {
  91. return &dto.MidjourneyResponse{
  92. Code: 4,
  93. Description: "midjourney_task_not_found",
  94. Properties: nil,
  95. Result: "",
  96. }
  97. }
  98. midjourneyTask.Progress = midjRequest.Progress
  99. midjourneyTask.PromptEn = midjRequest.PromptEn
  100. midjourneyTask.State = midjRequest.State
  101. midjourneyTask.SubmitTime = midjRequest.SubmitTime
  102. midjourneyTask.StartTime = midjRequest.StartTime
  103. midjourneyTask.FinishTime = midjRequest.FinishTime
  104. midjourneyTask.ImageUrl = midjRequest.ImageUrl
  105. midjourneyTask.VideoUrl = midjRequest.VideoUrl
  106. videoUrlsStr, _ := json.Marshal(midjRequest.VideoUrls)
  107. midjourneyTask.VideoUrls = string(videoUrlsStr)
  108. midjourneyTask.Status = midjRequest.Status
  109. midjourneyTask.FailReason = midjRequest.FailReason
  110. err = midjourneyTask.Update()
  111. if err != nil {
  112. return &dto.MidjourneyResponse{
  113. Code: 4,
  114. Description: "update_midjourney_task_failed",
  115. }
  116. }
  117. return nil
  118. }
  119. func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjourneyTask dto.MidjourneyDto) {
  120. midjourneyTask.MjId = originTask.MjId
  121. midjourneyTask.Progress = originTask.Progress
  122. midjourneyTask.PromptEn = originTask.PromptEn
  123. midjourneyTask.State = originTask.State
  124. midjourneyTask.SubmitTime = originTask.SubmitTime
  125. midjourneyTask.StartTime = originTask.StartTime
  126. midjourneyTask.FinishTime = originTask.FinishTime
  127. midjourneyTask.ImageUrl = ""
  128. if originTask.ImageUrl != "" && setting.MjForwardUrlEnabled {
  129. midjourneyTask.ImageUrl = setting.ServerAddress + "/mj/image/" + originTask.MjId
  130. if originTask.Status != "SUCCESS" {
  131. midjourneyTask.ImageUrl += "?rand=" + strconv.FormatInt(time.Now().UnixNano(), 10)
  132. }
  133. } else {
  134. midjourneyTask.ImageUrl = originTask.ImageUrl
  135. }
  136. if originTask.VideoUrl != "" {
  137. midjourneyTask.VideoUrl = originTask.VideoUrl
  138. }
  139. midjourneyTask.Status = originTask.Status
  140. midjourneyTask.FailReason = originTask.FailReason
  141. midjourneyTask.Action = originTask.Action
  142. midjourneyTask.Description = originTask.Description
  143. midjourneyTask.Prompt = originTask.Prompt
  144. if originTask.Buttons != "" {
  145. var buttons []dto.ActionButton
  146. err := json.Unmarshal([]byte(originTask.Buttons), &buttons)
  147. if err == nil {
  148. midjourneyTask.Buttons = buttons
  149. }
  150. }
  151. if originTask.VideoUrls != "" {
  152. var videoUrls []dto.ImgUrls
  153. err := json.Unmarshal([]byte(originTask.VideoUrls), &videoUrls)
  154. if err == nil {
  155. midjourneyTask.VideoUrls = videoUrls
  156. }
  157. }
  158. if originTask.Properties != "" {
  159. var properties dto.Properties
  160. err := json.Unmarshal([]byte(originTask.Properties), &properties)
  161. if err == nil {
  162. midjourneyTask.Properties = &properties
  163. }
  164. }
  165. return
  166. }
  167. func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse {
  168. startTime := time.Now().UnixNano() / int64(time.Millisecond)
  169. tokenId := c.GetInt("token_id")
  170. userId := c.GetInt("id")
  171. //group := c.GetString("group")
  172. channelId := c.GetInt("channel_id")
  173. relayInfo := relaycommon.GenRelayInfo(c)
  174. var swapFaceRequest dto.SwapFaceRequest
  175. err := common.UnmarshalBodyReusable(c, &swapFaceRequest)
  176. if err != nil {
  177. return service.MidjourneyErrorWrapper(constant.MjRequestError, "bind_request_body_failed")
  178. }
  179. if swapFaceRequest.SourceBase64 == "" || swapFaceRequest.TargetBase64 == "" {
  180. return service.MidjourneyErrorWrapper(constant.MjRequestError, "sour_base64_and_target_base64_is_required")
  181. }
  182. modelName := service.CoverActionToModelName(constant.MjActionSwapFace)
  183. priceData := helper.ModelPriceHelperPerCall(c, relayInfo)
  184. userQuota, err := model.GetUserQuota(userId, false)
  185. if err != nil {
  186. return &dto.MidjourneyResponse{
  187. Code: 4,
  188. Description: err.Error(),
  189. }
  190. }
  191. if userQuota-priceData.Quota < 0 {
  192. return &dto.MidjourneyResponse{
  193. Code: 4,
  194. Description: "quota_not_enough",
  195. }
  196. }
  197. requestURL := getMjRequestPath(c.Request.URL.String())
  198. baseURL := c.GetString("base_url")
  199. fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
  200. mjResp, _, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
  201. if err != nil {
  202. return &mjResp.Response
  203. }
  204. defer func() {
  205. if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 {
  206. err := service.PostConsumeQuota(relayInfo, priceData.Quota, 0, true)
  207. if err != nil {
  208. logger.SysError("error consuming token remain quota: " + err.Error())
  209. }
  210. tokenName := c.GetString("token_name")
  211. logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace)
  212. other := service.GenerateMjOtherInfo(priceData)
  213. model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{
  214. ChannelId: channelId,
  215. ModelName: modelName,
  216. TokenName: tokenName,
  217. Quota: priceData.Quota,
  218. Content: logContent,
  219. TokenId: tokenId,
  220. UserQuota: userQuota,
  221. Group: relayInfo.UsingGroup,
  222. Other: other,
  223. })
  224. model.UpdateUserUsedQuotaAndRequestCount(userId, priceData.Quota)
  225. model.UpdateChannelUsedQuota(channelId, priceData.Quota)
  226. }
  227. }()
  228. midjResponse := &mjResp.Response
  229. midjourneyTask := &model.Midjourney{
  230. UserId: userId,
  231. Code: midjResponse.Code,
  232. Action: constant.MjActionSwapFace,
  233. MjId: midjResponse.Result,
  234. Prompt: "InsightFace",
  235. PromptEn: "",
  236. Description: midjResponse.Description,
  237. State: "",
  238. SubmitTime: startTime,
  239. StartTime: time.Now().UnixNano() / int64(time.Millisecond),
  240. FinishTime: 0,
  241. ImageUrl: "",
  242. Status: "",
  243. Progress: "0%",
  244. FailReason: "",
  245. ChannelId: c.GetInt("channel_id"),
  246. Quota: priceData.Quota,
  247. }
  248. err = midjourneyTask.Insert()
  249. if err != nil {
  250. return service.MidjourneyErrorWrapper(constant.MjRequestError, "insert_midjourney_task_failed")
  251. }
  252. c.Writer.WriteHeader(mjResp.StatusCode)
  253. respBody, err := json.Marshal(midjResponse)
  254. if err != nil {
  255. return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
  256. }
  257. _, err = io.Copy(c.Writer, bytes.NewBuffer(respBody))
  258. if err != nil {
  259. return service.MidjourneyErrorWrapper(constant.MjRequestError, "copy_response_body_failed")
  260. }
  261. return nil
  262. }
  263. func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse {
  264. taskId := c.Param("id")
  265. userId := c.GetInt("id")
  266. originTask := model.GetByMJId(userId, taskId)
  267. if originTask == nil {
  268. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_no_found")
  269. }
  270. channel, err := model.GetChannelById(originTask.ChannelId, true)
  271. if err != nil {
  272. return service.MidjourneyErrorWrapper(constant.MjRequestError, "get_channel_info_failed")
  273. }
  274. if channel.Status != common.ChannelStatusEnabled {
  275. return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用")
  276. }
  277. c.Set("channel_id", originTask.ChannelId)
  278. c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
  279. requestURL := getMjRequestPath(c.Request.URL.String())
  280. fullRequestURL := fmt.Sprintf("%s%s", channel.GetBaseURL(), requestURL)
  281. midjResponseWithStatus, _, err := service.DoMidjourneyHttpRequest(c, time.Second*30, fullRequestURL)
  282. if err != nil {
  283. return &midjResponseWithStatus.Response
  284. }
  285. midjResponse := &midjResponseWithStatus.Response
  286. c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
  287. respBody, err := json.Marshal(midjResponse)
  288. if err != nil {
  289. return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
  290. }
  291. service.IOCopyBytesGracefully(c, nil, respBody)
  292. return nil
  293. }
  294. func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse {
  295. userId := c.GetInt("id")
  296. var err error
  297. var respBody []byte
  298. switch relayMode {
  299. case relayconstant.RelayModeMidjourneyTaskFetch:
  300. taskId := c.Param("id")
  301. originTask := model.GetByMJId(userId, taskId)
  302. if originTask == nil {
  303. return &dto.MidjourneyResponse{
  304. Code: 4,
  305. Description: "task_no_found",
  306. }
  307. }
  308. midjourneyTask := coverMidjourneyTaskDto(c, originTask)
  309. respBody, err = json.Marshal(midjourneyTask)
  310. if err != nil {
  311. return &dto.MidjourneyResponse{
  312. Code: 4,
  313. Description: "unmarshal_response_body_failed",
  314. }
  315. }
  316. case relayconstant.RelayModeMidjourneyTaskFetchByCondition:
  317. var condition = struct {
  318. IDs []string `json:"ids"`
  319. }{}
  320. err = c.BindJSON(&condition)
  321. if err != nil {
  322. return &dto.MidjourneyResponse{
  323. Code: 4,
  324. Description: "do_request_failed",
  325. }
  326. }
  327. var tasks []dto.MidjourneyDto
  328. if len(condition.IDs) != 0 {
  329. originTasks := model.GetByMJIds(userId, condition.IDs)
  330. for _, originTask := range originTasks {
  331. midjourneyTask := coverMidjourneyTaskDto(c, originTask)
  332. tasks = append(tasks, midjourneyTask)
  333. }
  334. }
  335. if tasks == nil {
  336. tasks = make([]dto.MidjourneyDto, 0)
  337. }
  338. respBody, err = json.Marshal(tasks)
  339. if err != nil {
  340. return &dto.MidjourneyResponse{
  341. Code: 4,
  342. Description: "unmarshal_response_body_failed",
  343. }
  344. }
  345. }
  346. c.Writer.Header().Set("Content-Type", "application/json")
  347. _, err = io.Copy(c.Writer, bytes.NewBuffer(respBody))
  348. if err != nil {
  349. return &dto.MidjourneyResponse{
  350. Code: 4,
  351. Description: "copy_response_body_failed",
  352. }
  353. }
  354. return nil
  355. }
  356. func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyResponse {
  357. //tokenId := c.GetInt("token_id")
  358. //channelType := c.GetInt("channel")
  359. userId := c.GetInt("id")
  360. group := c.GetString("group")
  361. channelId := c.GetInt("channel_id")
  362. relayInfo := relaycommon.GenRelayInfo(c)
  363. consumeQuota := true
  364. var midjRequest dto.MidjourneyRequest
  365. err := common.UnmarshalBodyReusable(c, &midjRequest)
  366. if err != nil {
  367. return service.MidjourneyErrorWrapper(constant.MjRequestError, "bind_request_body_failed")
  368. }
  369. if relayMode == relayconstant.RelayModeMidjourneyAction { // midjourney plus,需要从customId中获取任务信息
  370. mjErr := service.CoverPlusActionToNormalAction(&midjRequest)
  371. if mjErr != nil {
  372. return mjErr
  373. }
  374. relayMode = relayconstant.RelayModeMidjourneyChange
  375. }
  376. if relayMode == relayconstant.RelayModeMidjourneyVideo {
  377. midjRequest.Action = constant.MjActionVideo
  378. }
  379. if relayMode == relayconstant.RelayModeMidjourneyImagine { //绘画任务,此类任务可重复
  380. if midjRequest.Prompt == "" {
  381. return service.MidjourneyErrorWrapper(constant.MjRequestError, "prompt_is_required")
  382. }
  383. midjRequest.Action = constant.MjActionImagine
  384. } else if relayMode == relayconstant.RelayModeMidjourneyDescribe { //按图生文任务,此类任务可重复
  385. midjRequest.Action = constant.MjActionDescribe
  386. } else if relayMode == relayconstant.RelayModeMidjourneyEdits { //编辑任务,此类任务可重复
  387. midjRequest.Action = constant.MjActionEdits
  388. } else if relayMode == relayconstant.RelayModeMidjourneyShorten { //缩短任务,此类任务可重复,plus only
  389. midjRequest.Action = constant.MjActionShorten
  390. } else if relayMode == relayconstant.RelayModeMidjourneyBlend { //绘画任务,此类任务可重复
  391. midjRequest.Action = constant.MjActionBlend
  392. } else if relayMode == relayconstant.RelayModeMidjourneyUpload { //绘画任务,此类任务可重复
  393. midjRequest.Action = constant.MjActionUpload
  394. } else if midjRequest.TaskId != "" { //放大、变换任务,此类任务,如果重复且已有结果,远端api会直接返回最终结果
  395. mjId := ""
  396. if relayMode == relayconstant.RelayModeMidjourneyChange {
  397. if midjRequest.TaskId == "" {
  398. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_id_is_required")
  399. } else if midjRequest.Action == "" {
  400. return service.MidjourneyErrorWrapper(constant.MjRequestError, "action_is_required")
  401. } else if midjRequest.Index == 0 {
  402. return service.MidjourneyErrorWrapper(constant.MjRequestError, "index_is_required")
  403. }
  404. //action = midjRequest.Action
  405. mjId = midjRequest.TaskId
  406. } else if relayMode == relayconstant.RelayModeMidjourneySimpleChange {
  407. if midjRequest.Content == "" {
  408. return service.MidjourneyErrorWrapper(constant.MjRequestError, "content_is_required")
  409. }
  410. params := service.ConvertSimpleChangeParams(midjRequest.Content)
  411. if params == nil {
  412. return service.MidjourneyErrorWrapper(constant.MjRequestError, "content_parse_failed")
  413. }
  414. mjId = params.TaskId
  415. midjRequest.Action = params.Action
  416. } else if relayMode == relayconstant.RelayModeMidjourneyModal {
  417. //if midjRequest.MaskBase64 == "" {
  418. // return service.MidjourneyErrorWrapper(constant.MjRequestError, "mask_base64_is_required")
  419. //}
  420. mjId = midjRequest.TaskId
  421. midjRequest.Action = constant.MjActionModal
  422. } else if relayMode == relayconstant.RelayModeMidjourneyVideo {
  423. midjRequest.Action = constant.MjActionVideo
  424. if midjRequest.TaskId == "" {
  425. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_id_is_required")
  426. } else if midjRequest.Action == "" {
  427. return service.MidjourneyErrorWrapper(constant.MjRequestError, "action_is_required")
  428. }
  429. mjId = midjRequest.TaskId
  430. }
  431. originTask := model.GetByMJId(userId, mjId)
  432. if originTask == nil {
  433. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_not_found")
  434. } else { //原任务的Status=SUCCESS,则可以做放大UPSCALE、变换VARIATION等动作,此时必须使用原来的请求地址才能正确处理
  435. if setting.MjActionCheckSuccessEnabled {
  436. if originTask.Status != "SUCCESS" && relayMode != relayconstant.RelayModeMidjourneyModal {
  437. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_status_not_success")
  438. }
  439. }
  440. channel, err := model.GetChannelById(originTask.ChannelId, true)
  441. if err != nil {
  442. return service.MidjourneyErrorWrapper(constant.MjRequestError, "get_channel_info_failed")
  443. }
  444. if channel.Status != common.ChannelStatusEnabled {
  445. return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用")
  446. }
  447. c.Set("base_url", channel.GetBaseURL())
  448. c.Set("channel_id", originTask.ChannelId)
  449. c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
  450. log.Printf("检测到此操作为放大、变换、重绘,获取原channel信息: %s,%s", strconv.Itoa(originTask.ChannelId), channel.GetBaseURL())
  451. }
  452. midjRequest.Prompt = originTask.Prompt
  453. //if channelType == common.ChannelTypeMidjourneyPlus {
  454. // // plus
  455. //} else {
  456. // // 普通版渠道
  457. //
  458. //}
  459. }
  460. if midjRequest.Action == constant.MjActionInPaint || midjRequest.Action == constant.MjActionCustomZoom {
  461. consumeQuota = false
  462. }
  463. //baseURL := common.ChannelBaseURLs[channelType]
  464. requestURL := getMjRequestPath(c.Request.URL.String())
  465. baseURL := c.GetString("base_url")
  466. //midjRequest.NotifyHook = "http://127.0.0.1:3000/mj/notify"
  467. fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
  468. modelName := service.CoverActionToModelName(midjRequest.Action)
  469. priceData := helper.ModelPriceHelperPerCall(c, relayInfo)
  470. userQuota, err := model.GetUserQuota(userId, false)
  471. if err != nil {
  472. return &dto.MidjourneyResponse{
  473. Code: 4,
  474. Description: err.Error(),
  475. }
  476. }
  477. if consumeQuota && userQuota-priceData.Quota < 0 {
  478. return &dto.MidjourneyResponse{
  479. Code: 4,
  480. Description: "quota_not_enough",
  481. }
  482. }
  483. midjResponseWithStatus, responseBody, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
  484. if err != nil {
  485. return &midjResponseWithStatus.Response
  486. }
  487. midjResponse := &midjResponseWithStatus.Response
  488. defer func() {
  489. if consumeQuota && midjResponseWithStatus.StatusCode == 200 {
  490. err := service.PostConsumeQuota(relayInfo, priceData.Quota, 0, true)
  491. if err != nil {
  492. logger.SysError("error consuming token remain quota: " + err.Error())
  493. }
  494. tokenName := c.GetString("token_name")
  495. logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s,ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result)
  496. other := service.GenerateMjOtherInfo(priceData)
  497. model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{
  498. ChannelId: channelId,
  499. ModelName: modelName,
  500. TokenName: tokenName,
  501. Quota: priceData.Quota,
  502. Content: logContent,
  503. TokenId: relayInfo.TokenId,
  504. UserQuota: userQuota,
  505. Group: group,
  506. Other: other,
  507. })
  508. model.UpdateUserUsedQuotaAndRequestCount(userId, priceData.Quota)
  509. model.UpdateChannelUsedQuota(channelId, priceData.Quota)
  510. }
  511. }()
  512. // 文档:https://github.com/novicezk/midjourney-proxy/blob/main/docs/api.md
  513. //1-提交成功
  514. // 21-任务已存在(处理中或者有结果了) {"code":21,"description":"任务已存在","result":"0741798445574458","properties":{"status":"SUCCESS","imageUrl":"https://xxxx"}}
  515. // 22-排队中 {"code":22,"description":"排队中,前面还有1个任务","result":"0741798445574458","properties":{"numberOfQueues":1,"discordInstanceId":"1118138338562560102"}}
  516. // 23-队列已满,请稍后再试 {"code":23,"description":"队列已满,请稍后尝试","result":"14001929738841620","properties":{"discordInstanceId":"1118138338562560102"}}
  517. // 24-prompt包含敏感词 {"code":24,"description":"可能包含敏感词","properties":{"promptEn":"nude body","bannedWord":"nude"}}
  518. // other: 提交错误,description为错误描述
  519. midjourneyTask := &model.Midjourney{
  520. UserId: userId,
  521. Code: midjResponse.Code,
  522. Action: midjRequest.Action,
  523. MjId: midjResponse.Result,
  524. Prompt: midjRequest.Prompt,
  525. PromptEn: "",
  526. Description: midjResponse.Description,
  527. State: "",
  528. SubmitTime: time.Now().UnixNano() / int64(time.Millisecond),
  529. StartTime: 0,
  530. FinishTime: 0,
  531. ImageUrl: "",
  532. Status: "",
  533. Progress: "0%",
  534. FailReason: "",
  535. ChannelId: c.GetInt("channel_id"),
  536. Quota: priceData.Quota,
  537. }
  538. if midjResponse.Code == 3 {
  539. //无实例账号自动禁用渠道(No available account instance)
  540. channel, err := model.GetChannelById(midjourneyTask.ChannelId, true)
  541. if err != nil {
  542. logger.SysError("get_channel_null: " + err.Error())
  543. }
  544. if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled {
  545. model.UpdateChannelStatus(midjourneyTask.ChannelId, "", 2, "No available account instance")
  546. }
  547. }
  548. if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 {
  549. //非1-提交成功,21-任务已存在和22-排队中,则记录错误原因
  550. midjourneyTask.FailReason = midjResponse.Description
  551. consumeQuota = false
  552. }
  553. if midjResponse.Code == 21 { //21-任务已存在(处理中或者有结果了)
  554. // 将 properties 转换为一个 map
  555. properties, ok := midjResponse.Properties.(map[string]interface{})
  556. if ok {
  557. imageUrl, ok1 := properties["imageUrl"].(string)
  558. status, ok2 := properties["status"].(string)
  559. if ok1 && ok2 {
  560. midjourneyTask.ImageUrl = imageUrl
  561. midjourneyTask.Status = status
  562. if status == "SUCCESS" {
  563. midjourneyTask.Progress = "100%"
  564. midjourneyTask.StartTime = time.Now().UnixNano() / int64(time.Millisecond)
  565. midjourneyTask.FinishTime = time.Now().UnixNano() / int64(time.Millisecond)
  566. midjResponse.Code = 1
  567. }
  568. }
  569. }
  570. //修改返回值
  571. if midjRequest.Action != constant.MjActionInPaint && midjRequest.Action != constant.MjActionCustomZoom {
  572. newBody := strings.Replace(string(responseBody), `"code":21`, `"code":1`, -1)
  573. responseBody = []byte(newBody)
  574. }
  575. }
  576. if midjResponse.Code == 1 && midjRequest.Action == "UPLOAD" {
  577. midjourneyTask.Progress = "100%"
  578. midjourneyTask.Status = "SUCCESS"
  579. }
  580. err = midjourneyTask.Insert()
  581. if err != nil {
  582. return &dto.MidjourneyResponse{
  583. Code: 4,
  584. Description: "insert_midjourney_task_failed",
  585. }
  586. }
  587. if midjResponse.Code == 22 { //22-排队中,说明任务已存在
  588. //修改返回值
  589. newBody := strings.Replace(string(responseBody), `"code":22`, `"code":1`, -1)
  590. responseBody = []byte(newBody)
  591. }
  592. //resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
  593. bodyReader := io.NopCloser(bytes.NewBuffer(responseBody))
  594. //for k, v := range resp.Header {
  595. // c.Writer.Header().Set(k, v[0])
  596. //}
  597. c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
  598. _, err = io.Copy(c.Writer, bodyReader)
  599. if err != nil {
  600. return &dto.MidjourneyResponse{
  601. Code: 4,
  602. Description: "copy_response_body_failed",
  603. }
  604. }
  605. err = bodyReader.Close()
  606. if err != nil {
  607. return &dto.MidjourneyResponse{
  608. Code: 4,
  609. Description: "close_response_body_failed",
  610. }
  611. }
  612. return nil
  613. }
  614. type taskChangeParams struct {
  615. ID string
  616. Action string
  617. Index int
  618. }
  619. func getMjRequestPath(path string) string {
  620. requestURL := path
  621. if strings.Contains(requestURL, "/mj-") {
  622. urls := strings.Split(requestURL, "/mj/")
  623. if len(urls) < 2 {
  624. return requestURL
  625. }
  626. requestURL = "/mj/" + urls[1]
  627. }
  628. return requestURL
  629. }