openai_images.go 49 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677
  1. package service
  2. import (
  3. "bufio"
  4. "bytes"
  5. "context"
  6. "crypto/sha256"
  7. "encoding/base64"
  8. "encoding/hex"
  9. "encoding/json"
  10. "fmt"
  11. "io"
  12. "mime"
  13. "mime/multipart"
  14. "net/http"
  15. "net/textproto"
  16. "os"
  17. "strconv"
  18. "strings"
  19. "sync/atomic"
  20. "time"
  21. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  22. "github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
  23. "github.com/gin-gonic/gin"
  24. "github.com/imroc/req/v3"
  25. "github.com/tidwall/gjson"
  26. "github.com/tidwall/sjson"
  27. )
  28. const (
  29. openAIImagesGenerationsEndpoint = "/v1/images/generations"
  30. openAIImagesEditsEndpoint = "/v1/images/edits"
  31. openAIImagesGenerationsURL = "https://api.openai.com/v1/images/generations"
  32. openAIImagesEditsURL = "https://api.openai.com/v1/images/edits"
  33. openAIChatGPTStartURL = "https://chatgpt.com/"
  34. openAIChatGPTFilesURL = "https://chatgpt.com/backend-api/files"
  35. openAIImageBackendUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
  36. openAIImageMaxDownloadBytes = 20 << 20 // 20MB per image download
  37. openAIImageMaxUploadPartSize = 20 << 20 // 20MB per multipart upload part
  38. openAIImagesResponsesMainModel = "gpt-5.6-luna"
  39. openAIImagesVerbatimPromptInstructions = "When invoking the image_generation tool, use the user's image prompt verbatim. Do not rewrite, expand, summarize, embellish, translate, normalize punctuation, or add or remove visual details or constraints. Preserve the original language, wording, capitalization, quotes, and punctuation exactly."
  40. )
  41. // openAIImagesResponsesMainModelValue selects the Responses driver independently
  42. // of the image_generation tool model. An environment override lets operators
  43. // recover from upstream model retirement without rebuilding the gateway.
  44. func openAIImagesResponsesMainModelValue() string {
  45. if model := strings.TrimSpace(os.Getenv("SUB2API_IMAGES_MAIN_MODEL")); model != "" {
  46. return model
  47. }
  48. return openAIImagesResponsesMainModel
  49. }
  50. type OpenAIImagesCapability string
  51. const (
  52. OpenAIImagesCapabilityBasic OpenAIImagesCapability = "images-basic"
  53. OpenAIImagesCapabilityNative OpenAIImagesCapability = "images-native"
  54. )
  55. type OpenAIImagesUpload struct {
  56. FieldName string
  57. FileName string
  58. ContentType string
  59. Data []byte
  60. Width int
  61. Height int
  62. }
  63. type OpenAIImagesRequest struct {
  64. Endpoint string
  65. ContentType string
  66. Multipart bool
  67. Model string
  68. ExplicitModel bool
  69. Prompt string
  70. Stream bool
  71. N int
  72. Size string
  73. ExplicitSize bool
  74. SizeTier string
  75. ResponseFormat string
  76. Quality string
  77. Background string
  78. OutputFormat string
  79. Moderation string
  80. InputFidelity string
  81. Style string
  82. OutputCompression *int
  83. PartialImages *int
  84. HasMask bool
  85. HasNativeOptions bool
  86. RequiredCapability OpenAIImagesCapability
  87. InputImageURLs []string
  88. MaskImageURL string
  89. Uploads []OpenAIImagesUpload
  90. MaskUpload *OpenAIImagesUpload
  91. Body []byte
  92. bodyHash string
  93. }
  94. func (r *OpenAIImagesRequest) ModerationBody() []byte {
  95. if r == nil {
  96. return nil
  97. }
  98. payload := map[string]any{}
  99. if prompt := strings.TrimSpace(r.Prompt); prompt != "" {
  100. payload["prompt"] = prompt
  101. }
  102. images := r.moderationImages()
  103. if len(images) > 0 {
  104. payload["images"] = images
  105. }
  106. if len(payload) == 0 {
  107. return nil
  108. }
  109. body, err := json.Marshal(payload)
  110. if err != nil {
  111. return nil
  112. }
  113. return body
  114. }
  115. func (r *OpenAIImagesRequest) moderationImages() []map[string]string {
  116. if r == nil {
  117. return nil
  118. }
  119. images := make([]map[string]string, 0, len(r.InputImageURLs)+len(r.Uploads)+1)
  120. for _, imageURL := range r.InputImageURLs {
  121. imageURL = strings.TrimSpace(imageURL)
  122. if imageURL != "" {
  123. images = append(images, map[string]string{"image_url": imageURL})
  124. }
  125. }
  126. for _, upload := range r.Uploads {
  127. if dataURL := upload.ModerationDataURL(); dataURL != "" {
  128. images = append(images, map[string]string{"image_url": dataURL})
  129. }
  130. }
  131. if maskURL := strings.TrimSpace(r.MaskImageURL); maskURL != "" {
  132. images = append(images, map[string]string{"image_url": maskURL})
  133. }
  134. if r.MaskUpload != nil {
  135. if dataURL := r.MaskUpload.ModerationDataURL(); dataURL != "" {
  136. images = append(images, map[string]string{"image_url": dataURL})
  137. }
  138. }
  139. return images
  140. }
  141. func (u OpenAIImagesUpload) ModerationDataURL() string {
  142. if len(u.Data) == 0 {
  143. return ""
  144. }
  145. contentType := strings.TrimSpace(u.ContentType)
  146. if contentType == "" {
  147. contentType = http.DetectContentType(u.Data)
  148. }
  149. if !strings.HasPrefix(strings.ToLower(contentType), "image/") {
  150. return ""
  151. }
  152. return fmt.Sprintf("data:%s;base64,%s", contentType, base64.StdEncoding.EncodeToString(u.Data))
  153. }
  154. func (r *OpenAIImagesRequest) IsEdits() bool {
  155. return r != nil && r.Endpoint == openAIImagesEditsEndpoint
  156. }
  157. func (r *OpenAIImagesRequest) StickySessionSeed() string {
  158. if r == nil {
  159. return ""
  160. }
  161. parts := []string{
  162. "openai-images",
  163. strings.TrimSpace(r.Endpoint),
  164. strings.TrimSpace(r.Model),
  165. strings.TrimSpace(r.Size),
  166. strings.TrimSpace(r.Prompt),
  167. }
  168. seed := strings.Join(parts, "|")
  169. if strings.TrimSpace(r.Prompt) == "" && r.bodyHash != "" {
  170. seed += "|body=" + r.bodyHash
  171. }
  172. return seed
  173. }
  174. func (s *OpenAIGatewayService) ParseOpenAIImagesRequest(c *gin.Context, body []byte) (*OpenAIImagesRequest, error) {
  175. if c == nil || c.Request == nil {
  176. return nil, fmt.Errorf("missing request context")
  177. }
  178. endpoint := normalizeOpenAIImagesEndpointPath(c.Request.URL.Path)
  179. if endpoint == "" {
  180. return nil, fmt.Errorf("unsupported images endpoint")
  181. }
  182. contentType := strings.TrimSpace(c.GetHeader("Content-Type"))
  183. req := &OpenAIImagesRequest{
  184. Endpoint: endpoint,
  185. ContentType: contentType,
  186. N: 1,
  187. Body: body,
  188. }
  189. if len(body) > 0 {
  190. sum := sha256.Sum256(body)
  191. req.bodyHash = hex.EncodeToString(sum[:8])
  192. }
  193. mediaType, _, err := mime.ParseMediaType(contentType)
  194. if err == nil && strings.EqualFold(mediaType, "multipart/form-data") {
  195. req.Multipart = true
  196. if parseErr := parseOpenAIImagesMultipartRequest(body, contentType, req); parseErr != nil {
  197. return nil, parseErr
  198. }
  199. } else {
  200. if len(body) == 0 {
  201. return nil, fmt.Errorf("request body is empty")
  202. }
  203. if !gjson.ValidBytes(body) {
  204. return nil, fmt.Errorf("failed to parse request body")
  205. }
  206. if parseErr := parseOpenAIImagesJSONRequest(body, req); parseErr != nil {
  207. return nil, parseErr
  208. }
  209. }
  210. applyOpenAIImagesDefaults(req)
  211. if err := validateOpenAIImagesModel(req.Model); err != nil {
  212. return nil, err
  213. }
  214. req.SizeTier = normalizeOpenAIImageSizeTier(req.Size)
  215. req.RequiredCapability = classifyOpenAIImagesCapability(req)
  216. return req, nil
  217. }
  218. func parseOpenAIImagesJSONRequest(body []byte, req *OpenAIImagesRequest) error {
  219. if modelResult := gjson.GetBytes(body, "model"); modelResult.Exists() {
  220. req.Model = strings.TrimSpace(modelResult.String())
  221. req.ExplicitModel = req.Model != ""
  222. }
  223. req.Prompt = strings.TrimSpace(gjson.GetBytes(body, "prompt").String())
  224. if streamResult := gjson.GetBytes(body, "stream"); streamResult.Exists() {
  225. if streamResult.Type != gjson.True && streamResult.Type != gjson.False {
  226. return fmt.Errorf("invalid stream field type")
  227. }
  228. req.Stream = streamResult.Bool()
  229. }
  230. if nResult := gjson.GetBytes(body, "n"); nResult.Exists() {
  231. if nResult.Type != gjson.Number {
  232. return fmt.Errorf("invalid n field type")
  233. }
  234. req.N = int(nResult.Int())
  235. if req.N <= 0 {
  236. return fmt.Errorf("n must be greater than 0")
  237. }
  238. }
  239. if sizeResult := gjson.GetBytes(body, "size"); sizeResult.Exists() {
  240. req.Size = strings.TrimSpace(sizeResult.String())
  241. req.ExplicitSize = req.Size != ""
  242. }
  243. req.ResponseFormat = strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "response_format").String()))
  244. req.Quality = strings.TrimSpace(gjson.GetBytes(body, "quality").String())
  245. req.Background = strings.TrimSpace(gjson.GetBytes(body, "background").String())
  246. req.OutputFormat = strings.TrimSpace(gjson.GetBytes(body, "output_format").String())
  247. req.Moderation = strings.TrimSpace(gjson.GetBytes(body, "moderation").String())
  248. req.InputFidelity = strings.TrimSpace(gjson.GetBytes(body, "input_fidelity").String())
  249. req.Style = strings.TrimSpace(gjson.GetBytes(body, "style").String())
  250. req.HasMask = gjson.GetBytes(body, "mask").Exists()
  251. if outputCompression := gjson.GetBytes(body, "output_compression"); outputCompression.Exists() {
  252. if outputCompression.Type != gjson.Number {
  253. return fmt.Errorf("invalid output_compression field type")
  254. }
  255. v := int(outputCompression.Int())
  256. req.OutputCompression = &v
  257. }
  258. if partialImages := gjson.GetBytes(body, "partial_images"); partialImages.Exists() {
  259. if partialImages.Type != gjson.Number {
  260. return fmt.Errorf("invalid partial_images field type")
  261. }
  262. v := int(partialImages.Int())
  263. req.PartialImages = &v
  264. }
  265. if req.IsEdits() {
  266. images := gjson.GetBytes(body, "images")
  267. if images.Exists() {
  268. if !images.IsArray() {
  269. return fmt.Errorf("invalid images field type")
  270. }
  271. for _, item := range images.Array() {
  272. if imageURL := strings.TrimSpace(item.Get("image_url").String()); imageURL != "" {
  273. req.InputImageURLs = append(req.InputImageURLs, imageURL)
  274. continue
  275. }
  276. if item.Get("file_id").Exists() {
  277. return fmt.Errorf("images[].file_id is not supported (use images[].image_url instead)")
  278. }
  279. }
  280. }
  281. if maskImageURL := strings.TrimSpace(gjson.GetBytes(body, "mask.image_url").String()); maskImageURL != "" {
  282. req.MaskImageURL = maskImageURL
  283. req.HasMask = true
  284. }
  285. if gjson.GetBytes(body, "mask.file_id").Exists() {
  286. return fmt.Errorf("mask.file_id is not supported (use mask.image_url instead)")
  287. }
  288. if len(req.InputImageURLs) == 0 {
  289. return fmt.Errorf("images[].image_url is required")
  290. }
  291. }
  292. req.HasNativeOptions = hasOpenAINativeImageOptions(func(path string) bool {
  293. return gjson.GetBytes(body, path).Exists()
  294. })
  295. return nil
  296. }
  297. func parseOpenAIImagesMultipartRequest(body []byte, contentType string, req *OpenAIImagesRequest) error {
  298. _, params, err := mime.ParseMediaType(contentType)
  299. if err != nil {
  300. return fmt.Errorf("invalid multipart content-type: %w", err)
  301. }
  302. boundary := strings.TrimSpace(params["boundary"])
  303. if boundary == "" {
  304. return fmt.Errorf("multipart boundary is required")
  305. }
  306. reader := multipart.NewReader(bytes.NewReader(body), boundary)
  307. for {
  308. part, err := reader.NextPart()
  309. if err == io.EOF {
  310. break
  311. }
  312. if err != nil {
  313. return fmt.Errorf("read multipart body: %w", err)
  314. }
  315. name := strings.TrimSpace(part.FormName())
  316. if name == "" {
  317. _ = part.Close()
  318. continue
  319. }
  320. data, err := io.ReadAll(io.LimitReader(part, openAIImageMaxUploadPartSize))
  321. _ = part.Close()
  322. if err != nil {
  323. return fmt.Errorf("read multipart field %s: %w", name, err)
  324. }
  325. fileName := strings.TrimSpace(part.FileName())
  326. if fileName != "" {
  327. partContentType := strings.TrimSpace(part.Header.Get("Content-Type"))
  328. if name == "mask" && len(data) > 0 {
  329. req.HasMask = true
  330. width, height := parseOpenAIImageDimensions(part.Header)
  331. maskUpload := OpenAIImagesUpload{
  332. FieldName: name,
  333. FileName: fileName,
  334. ContentType: partContentType,
  335. Data: data,
  336. Width: width,
  337. Height: height,
  338. }
  339. req.MaskUpload = &maskUpload
  340. }
  341. if name == "image" || strings.HasPrefix(name, "image[") {
  342. width, height := parseOpenAIImageDimensions(part.Header)
  343. req.Uploads = append(req.Uploads, OpenAIImagesUpload{
  344. FieldName: name,
  345. FileName: fileName,
  346. ContentType: partContentType,
  347. Data: data,
  348. Width: width,
  349. Height: height,
  350. })
  351. }
  352. continue
  353. }
  354. value := strings.TrimSpace(string(data))
  355. switch name {
  356. case "model":
  357. req.Model = value
  358. req.ExplicitModel = value != ""
  359. case "prompt":
  360. req.Prompt = value
  361. case "size":
  362. req.Size = value
  363. req.ExplicitSize = value != ""
  364. case "response_format":
  365. req.ResponseFormat = strings.ToLower(value)
  366. case "stream":
  367. parsed, err := strconv.ParseBool(value)
  368. if err != nil {
  369. return fmt.Errorf("invalid stream field value")
  370. }
  371. req.Stream = parsed
  372. case "n":
  373. n, err := strconv.Atoi(value)
  374. if err != nil || n <= 0 {
  375. return fmt.Errorf("n must be a positive integer")
  376. }
  377. req.N = n
  378. case "quality":
  379. req.Quality = value
  380. req.HasNativeOptions = true
  381. case "background":
  382. req.Background = value
  383. req.HasNativeOptions = true
  384. case "output_format":
  385. req.OutputFormat = value
  386. req.HasNativeOptions = true
  387. case "moderation":
  388. req.Moderation = value
  389. req.HasNativeOptions = true
  390. case "input_fidelity":
  391. req.InputFidelity = value
  392. req.HasNativeOptions = true
  393. case "style":
  394. req.Style = value
  395. req.HasNativeOptions = true
  396. case "output_compression":
  397. n, err := strconv.Atoi(value)
  398. if err != nil {
  399. return fmt.Errorf("invalid output_compression field value")
  400. }
  401. req.OutputCompression = &n
  402. req.HasNativeOptions = true
  403. case "partial_images":
  404. n, err := strconv.Atoi(value)
  405. if err != nil {
  406. return fmt.Errorf("invalid partial_images field value")
  407. }
  408. req.PartialImages = &n
  409. req.HasNativeOptions = true
  410. default:
  411. if isOpenAINativeImageOption(name) && value != "" {
  412. req.HasNativeOptions = true
  413. }
  414. }
  415. }
  416. if len(req.Uploads) == 0 && req.IsEdits() {
  417. return fmt.Errorf("image file is required")
  418. }
  419. return nil
  420. }
  421. func parseOpenAIImageDimensions(_ textproto.MIMEHeader) (int, int) {
  422. return 0, 0
  423. }
  424. func applyOpenAIImagesDefaults(req *OpenAIImagesRequest) {
  425. if req == nil {
  426. return
  427. }
  428. if req.N <= 0 {
  429. req.N = 1
  430. }
  431. if strings.TrimSpace(req.Model) != "" {
  432. req.Model = strings.TrimSpace(req.Model)
  433. return
  434. }
  435. req.Model = "gpt-image-2"
  436. }
  437. func isOpenAIImageGenerationModel(model string) bool {
  438. return IsGPTImageGenerationModel(model) || isGrokImageGenerationModel(model)
  439. }
  440. // IsGPTImageGenerationModel identifies the GPT native image-generation model family.
  441. func IsGPTImageGenerationModel(model string) bool {
  442. model = strings.ToLower(strings.TrimSpace(model))
  443. return strings.HasPrefix(model, "gpt-image-")
  444. }
  445. func isGrokImageGenerationModel(model string) bool {
  446. model = strings.ToLower(strings.TrimSpace(model))
  447. return model == "grok-imagine" ||
  448. model == "grok-imagine-edit" ||
  449. strings.HasPrefix(model, "grok-imagine-image")
  450. }
  451. func validateOpenAIImagesModel(model string) error {
  452. model = strings.TrimSpace(model)
  453. if isOpenAIImageGenerationModel(model) {
  454. return nil
  455. }
  456. if model == "" {
  457. return fmt.Errorf("images endpoint requires an image model")
  458. }
  459. return fmt.Errorf("images endpoint requires an image model, got %q", model)
  460. }
  461. func normalizeOpenAIImagesEndpointPath(path string) string {
  462. trimmed := strings.TrimSpace(path)
  463. switch {
  464. case strings.Contains(trimmed, "/images/generations"):
  465. return openAIImagesGenerationsEndpoint
  466. case strings.Contains(trimmed, "/images/edits"):
  467. return openAIImagesEditsEndpoint
  468. default:
  469. return ""
  470. }
  471. }
  472. func classifyOpenAIImagesCapability(req *OpenAIImagesRequest) OpenAIImagesCapability {
  473. if req == nil {
  474. return OpenAIImagesCapabilityNative
  475. }
  476. if req.ExplicitModel || req.ExplicitSize {
  477. return OpenAIImagesCapabilityNative
  478. }
  479. model := strings.ToLower(strings.TrimSpace(req.Model))
  480. if !strings.HasPrefix(model, "gpt-image-") {
  481. return OpenAIImagesCapabilityNative
  482. }
  483. if req.Stream || req.N != 1 || req.HasMask || req.HasNativeOptions {
  484. return OpenAIImagesCapabilityNative
  485. }
  486. if req.IsEdits() && !req.Multipart {
  487. return OpenAIImagesCapabilityNative
  488. }
  489. if req.ResponseFormat != "" && req.ResponseFormat != "b64_json" {
  490. return OpenAIImagesCapabilityNative
  491. }
  492. return OpenAIImagesCapabilityBasic
  493. }
  494. func hasOpenAINativeImageOptions(exists func(path string) bool) bool {
  495. for _, path := range []string{
  496. "background",
  497. "quality",
  498. "style",
  499. "output_format",
  500. "output_compression",
  501. "moderation",
  502. "input_fidelity",
  503. "partial_images",
  504. } {
  505. if exists(path) {
  506. return true
  507. }
  508. }
  509. return false
  510. }
  511. func isOpenAINativeImageOption(name string) bool {
  512. switch strings.TrimSpace(strings.ToLower(name)) {
  513. case "background", "quality", "style", "output_format", "output_compression", "moderation", "input_fidelity", "partial_images":
  514. return true
  515. default:
  516. return false
  517. }
  518. }
  519. func normalizeOpenAIImageSizeTier(size string) string {
  520. return NormalizeImageBillingTierOrDefault(size)
  521. }
  522. func (s *OpenAIGatewayService) ForwardImages(
  523. ctx context.Context,
  524. c *gin.Context,
  525. account *Account,
  526. body []byte,
  527. parsed *OpenAIImagesRequest,
  528. channelMappedModel string,
  529. ) (*OpenAIForwardResult, error) {
  530. if parsed == nil {
  531. return nil, fmt.Errorf("parsed images request is required")
  532. }
  533. switch account.Type {
  534. case AccountTypeAPIKey:
  535. return s.forwardOpenAIImagesAPIKey(ctx, c, account, body, parsed, channelMappedModel)
  536. case AccountTypeOAuth, AccountTypeSetupToken:
  537. return s.forwardOpenAIImagesOAuth(ctx, c, account, parsed, channelMappedModel)
  538. default:
  539. return nil, fmt.Errorf("unsupported account type: %s", account.Type)
  540. }
  541. }
  542. func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey(
  543. ctx context.Context,
  544. c *gin.Context,
  545. account *Account,
  546. body []byte,
  547. parsed *OpenAIImagesRequest,
  548. channelMappedModel string,
  549. ) (*OpenAIForwardResult, error) {
  550. startTime := time.Now()
  551. requestModel := strings.TrimSpace(parsed.Model)
  552. if mapped := strings.TrimSpace(channelMappedModel); mapped != "" {
  553. requestModel = mapped
  554. }
  555. if err := validateOpenAIImagesModel(requestModel); err != nil {
  556. return nil, err
  557. }
  558. upstreamModel := account.GetMappedModel(requestModel)
  559. if err := validateOpenAIImagesModel(upstreamModel); err != nil {
  560. return nil, err
  561. }
  562. SetOpsUpstreamModel(c, upstreamModel)
  563. logger.LegacyPrintf(
  564. "service.openai_gateway",
  565. "[OpenAI] Images request routing request_model=%s upstream_model=%s endpoint=%s account_type=%s",
  566. strings.TrimSpace(parsed.Model),
  567. upstreamModel,
  568. parsed.Endpoint,
  569. account.Type,
  570. )
  571. forwardBody, forwardContentType, err := rewriteOpenAIImagesModel(body, parsed.ContentType, upstreamModel)
  572. if err != nil {
  573. return nil, err
  574. }
  575. // 生图是长耗时、上游侧已产生实际成本的操作:客户端中途断开不应连带取消上游请求。
  576. // detachStreamUpstreamContext 在非流式时原样返回请求 context,于是客户端一断开
  577. // 就把已经在出图的上游调用打断成 context canceled,网关记 502、不扣费,而上游那边
  578. // 图已经生成并计费。同一端点的 OAuth 分支 forwardOpenAIImagesOAuth 以及 Grok 媒体
  579. // 路径本来就无条件脱钩,这里对齐;上游侧仍由 ResponseHeaderTimeout 兜底。
  580. upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
  581. defer releaseUpstreamCtx()
  582. token, _, err := s.GetAccessToken(upstreamCtx, account)
  583. if err != nil {
  584. return nil, err
  585. }
  586. upstreamReq, err := s.buildOpenAIImagesRequest(upstreamCtx, c, account, forwardBody, forwardContentType, token, parsed.Endpoint)
  587. if err != nil {
  588. return nil, err
  589. }
  590. proxyURL := ""
  591. if account.ProxyID != nil && account.Proxy != nil {
  592. proxyURL = account.Proxy.URL()
  593. }
  594. upstreamStart := time.Now()
  595. resp, err := s.doOpenAIUpstream(upstreamReq, proxyURL, account)
  596. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  597. if err != nil {
  598. safeErr := sanitizeUpstreamErrorMessage(err.Error())
  599. setOpsUpstreamError(c, 0, safeErr, "")
  600. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  601. ProxyID: opsUpstreamProxyID(account),
  602. ProxyName: opsUpstreamProxyName(account),
  603. Platform: account.Platform,
  604. AccountID: account.ID,
  605. AccountName: account.Name,
  606. UpstreamStatusCode: 0,
  607. UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
  608. Kind: "request_error",
  609. Message: safeErr,
  610. })
  611. return nil, fmt.Errorf("upstream request failed: %s", safeErr)
  612. }
  613. if resp.StatusCode >= 400 {
  614. respBody := s.readUpstreamErrorBody(resp)
  615. _ = resp.Body.Close()
  616. respBody = s.redactAgentIdentitySensitiveBody(upstreamCtx, account, respBody)
  617. resp.Body = io.NopCloser(bytes.NewReader(respBody))
  618. upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
  619. upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
  620. if s.shouldFailoverOpenAIUpstreamResponse(account, resp.StatusCode, upstreamMsg, respBody) {
  621. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  622. ProxyID: opsUpstreamProxyID(account),
  623. ProxyName: opsUpstreamProxyName(account),
  624. Platform: account.Platform,
  625. AccountID: account.ID,
  626. AccountName: account.Name,
  627. UpstreamStatusCode: resp.StatusCode,
  628. UpstreamRequestID: resp.Header.Get("x-request-id"),
  629. UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
  630. Kind: "failover",
  631. Message: upstreamMsg,
  632. })
  633. shouldDisable := s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, upstreamModel)
  634. retryableOnSameAccount := !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode)
  635. if account.IsOpenAIOAuthLike() && resp.StatusCode == http.StatusTooManyRequests {
  636. return nil, s.newOpenAIAccountFailoverError(account, resp.StatusCode, resp.Header, respBody, upstreamMsg, shouldDisable, retryableOnSameAccount)
  637. }
  638. if isOpenAIHTTPUpstreamAccessStateError(resp.StatusCode, upstreamMsg, respBody) {
  639. return nil, newOpenAIUpstreamFailoverError(resp.StatusCode, resp.Header, respBody, upstreamMsg, retryableOnSameAccount)
  640. }
  641. return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, RetryableOnSameAccount: retryableOnSameAccount}
  642. }
  643. return s.handleOpenAIImagesErrorResponse(upstreamCtx, resp, c, account, upstreamModel)
  644. }
  645. defer func() { _ = resp.Body.Close() }()
  646. var usage OpenAIUsage
  647. imageCount := parsed.N
  648. var firstTokenMs *int
  649. if parsed.Stream && isEventStreamResponse(resp.Header) {
  650. streamUsage, streamCount, streamSizes, ttft, err := s.handleOpenAIImagesStreamingResponse(resp, c, startTime)
  651. if err != nil {
  652. if streamCount > 0 {
  653. return &OpenAIForwardResult{
  654. RequestID: resp.Header.Get("x-request-id"),
  655. UpstreamHeaders: resp.Header,
  656. Usage: streamUsage,
  657. Model: requestModel,
  658. UpstreamModel: upstreamModel,
  659. Stream: parsed.Stream,
  660. ResponseHeaders: resp.Header.Clone(),
  661. Duration: time.Since(startTime),
  662. FirstTokenMs: ttft,
  663. ImageCount: streamCount,
  664. ImageSize: parsed.SizeTier,
  665. ImageInputSize: parsed.Size,
  666. ImageOutputSizes: streamSizes,
  667. }, err
  668. }
  669. return nil, err
  670. }
  671. usage = streamUsage
  672. imageCount = streamCount
  673. imageOutputSizes := streamSizes
  674. firstTokenMs = ttft
  675. return &OpenAIForwardResult{
  676. RequestID: resp.Header.Get("x-request-id"),
  677. UpstreamHeaders: resp.Header,
  678. Usage: usage,
  679. Model: requestModel,
  680. UpstreamModel: upstreamModel,
  681. Stream: parsed.Stream,
  682. ResponseHeaders: resp.Header.Clone(),
  683. Duration: time.Since(startTime),
  684. FirstTokenMs: firstTokenMs,
  685. ImageCount: imageCount,
  686. ImageSize: parsed.SizeTier,
  687. ImageInputSize: parsed.Size,
  688. ImageOutputSizes: imageOutputSizes,
  689. }, nil
  690. } else {
  691. nonStreamUsage, nonStreamCount, nonStreamSizes, err := s.handleOpenAIImagesNonStreamingResponse(upstreamCtx, resp, c, account, parsed)
  692. if err != nil {
  693. return nil, err
  694. }
  695. usage = nonStreamUsage
  696. if nonStreamCount > 0 {
  697. imageCount = nonStreamCount
  698. }
  699. return &OpenAIForwardResult{
  700. RequestID: resp.Header.Get("x-request-id"),
  701. UpstreamHeaders: resp.Header,
  702. Usage: usage,
  703. Model: requestModel,
  704. UpstreamModel: upstreamModel,
  705. Stream: parsed.Stream,
  706. ResponseHeaders: resp.Header.Clone(),
  707. Duration: time.Since(startTime),
  708. FirstTokenMs: firstTokenMs,
  709. ImageCount: imageCount,
  710. ImageSize: parsed.SizeTier,
  711. ImageInputSize: parsed.Size,
  712. ImageOutputSizes: nonStreamSizes,
  713. }, nil
  714. }
  715. }
  716. func (s *OpenAIGatewayService) buildOpenAIImagesRequest(
  717. ctx context.Context,
  718. c *gin.Context,
  719. account *Account,
  720. body []byte,
  721. contentType string,
  722. token string,
  723. endpoint string,
  724. ) (*http.Request, error) {
  725. targetURL := openAIImagesGenerationsURL
  726. if endpoint == openAIImagesEditsEndpoint {
  727. targetURL = openAIImagesEditsURL
  728. }
  729. baseURL := account.GetOpenAIBaseURL()
  730. if baseURL != "" {
  731. validatedURL, err := s.validateUpstreamBaseURL(baseURL)
  732. if err != nil {
  733. return nil, err
  734. }
  735. targetURL = buildOpenAIImagesURL(validatedURL, endpoint)
  736. }
  737. req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
  738. if err != nil {
  739. return nil, err
  740. }
  741. req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
  742. authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
  743. if err != nil {
  744. return nil, fmt.Errorf("build openai authentication headers: %w", err)
  745. }
  746. for key, values := range authHeaders {
  747. for _, value := range values {
  748. req.Header.Add(key, value)
  749. }
  750. }
  751. for key, values := range c.Request.Header {
  752. if !openaiPassthroughAllowedHeaders[strings.ToLower(key)] {
  753. continue
  754. }
  755. for _, value := range values {
  756. req.Header.Add(key, value)
  757. }
  758. }
  759. customUA := account.GetOpenAIUserAgent()
  760. if customUA != "" {
  761. req.Header.Set("User-Agent", customUA)
  762. }
  763. if strings.TrimSpace(contentType) != "" {
  764. req.Header.Set("Content-Type", contentType)
  765. }
  766. // 账号级请求头覆写(仅 openai api_key 账号启用时生效;OAuth 路径 no-op)
  767. account.ApplyHeaderOverrides(req.Header)
  768. return req, nil
  769. }
  770. func buildOpenAIImagesURL(base string, endpoint string) string {
  771. return buildOpenAIEndpointURL(base, endpoint)
  772. }
  773. func rewriteOpenAIImagesModel(body []byte, contentType string, model string) ([]byte, string, error) {
  774. model = strings.TrimSpace(model)
  775. if model == "" {
  776. return body, contentType, nil
  777. }
  778. mediaType, _, err := mime.ParseMediaType(contentType)
  779. if err == nil && strings.EqualFold(mediaType, "multipart/form-data") {
  780. rewrittenBody, rewrittenType, rewriteErr := rewriteOpenAIImagesMultipartModel(body, contentType, model)
  781. return rewrittenBody, rewrittenType, rewriteErr
  782. }
  783. rewritten, err := sjson.SetBytes(body, "model", model)
  784. if err != nil {
  785. return nil, "", fmt.Errorf("rewrite image request model: %w", err)
  786. }
  787. return rewritten, contentType, nil
  788. }
  789. func rewriteOpenAIImagesMultipartModel(body []byte, contentType string, model string) ([]byte, string, error) {
  790. _, params, err := mime.ParseMediaType(contentType)
  791. if err != nil {
  792. return nil, "", fmt.Errorf("parse multipart content-type: %w", err)
  793. }
  794. boundary := strings.TrimSpace(params["boundary"])
  795. if boundary == "" {
  796. return nil, "", fmt.Errorf("multipart boundary is required")
  797. }
  798. reader := multipart.NewReader(bytes.NewReader(body), boundary)
  799. var buffer bytes.Buffer
  800. writer := multipart.NewWriter(&buffer)
  801. modelWritten := false
  802. for {
  803. part, err := reader.NextPart()
  804. if err == io.EOF {
  805. break
  806. }
  807. if err != nil {
  808. return nil, "", fmt.Errorf("read multipart body: %w", err)
  809. }
  810. formName := strings.TrimSpace(part.FormName())
  811. partHeader := cloneMultipartHeader(part.Header)
  812. target, err := writer.CreatePart(partHeader)
  813. if err != nil {
  814. _ = part.Close()
  815. return nil, "", fmt.Errorf("create multipart part: %w", err)
  816. }
  817. if formName == "model" && part.FileName() == "" {
  818. if _, err := target.Write([]byte(model)); err != nil {
  819. _ = part.Close()
  820. return nil, "", fmt.Errorf("rewrite multipart model: %w", err)
  821. }
  822. modelWritten = true
  823. _ = part.Close()
  824. continue
  825. }
  826. if _, err := io.Copy(target, part); err != nil {
  827. _ = part.Close()
  828. return nil, "", fmt.Errorf("copy multipart part: %w", err)
  829. }
  830. _ = part.Close()
  831. }
  832. if !modelWritten {
  833. if err := writer.WriteField("model", model); err != nil {
  834. return nil, "", fmt.Errorf("append multipart model field: %w", err)
  835. }
  836. }
  837. if err := writer.Close(); err != nil {
  838. return nil, "", fmt.Errorf("finalize multipart body: %w", err)
  839. }
  840. return buffer.Bytes(), writer.FormDataContentType(), nil
  841. }
  842. func cloneMultipartHeader(src textproto.MIMEHeader) textproto.MIMEHeader {
  843. dst := make(textproto.MIMEHeader, len(src))
  844. for key, values := range src {
  845. copied := make([]string, len(values))
  846. copy(copied, values)
  847. dst[key] = copied
  848. }
  849. return dst
  850. }
  851. func (s *OpenAIGatewayService) handleOpenAIImagesNonStreamingResponse(
  852. ctx context.Context,
  853. resp *http.Response,
  854. c *gin.Context,
  855. account *Account,
  856. parsed *OpenAIImagesRequest,
  857. ) (OpenAIUsage, int, []string, error) {
  858. body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
  859. if err != nil {
  860. return OpenAIUsage{}, 0, nil, err
  861. }
  862. body = s.backfillOpenAIImagesB64JSON(ctx, account, parsed, body)
  863. responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
  864. contentType := "application/json"
  865. if s.cfg != nil && !s.cfg.Security.ResponseHeaders.Enabled {
  866. if upstreamType := resp.Header.Get("Content-Type"); upstreamType != "" {
  867. contentType = upstreamType
  868. }
  869. }
  870. c.Data(resp.StatusCode, contentType, body)
  871. usage, _ := extractOpenAIUsageFromJSONBytes(body)
  872. return usage, extractOpenAIImageCountFromJSONBytes(body), collectOpenAIResponseImageOutputSizesFromJSONBytes(body), nil
  873. }
  874. func (s *OpenAIGatewayService) handleOpenAIImagesStreamingResponse(
  875. resp *http.Response,
  876. c *gin.Context,
  877. startTime time.Time,
  878. ) (OpenAIUsage, int, []string, *int, error) {
  879. responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
  880. contentType := strings.TrimSpace(resp.Header.Get("Content-Type"))
  881. if contentType == "" {
  882. contentType = "text/event-stream"
  883. }
  884. c.Status(resp.StatusCode)
  885. c.Header("Content-Type", contentType)
  886. flusher, ok := c.Writer.(http.Flusher)
  887. if !ok {
  888. return OpenAIUsage{}, 0, nil, nil, fmt.Errorf("streaming is not supported by response writer")
  889. }
  890. usage := OpenAIUsage{}
  891. imageCounter := newOpenAIImageOutputCounter()
  892. var firstTokenMs *int
  893. clientDisconnected := false
  894. lastDownstreamWriteAt := time.Now()
  895. var fallbackBody bytes.Buffer
  896. fallbackBytes := int64(0)
  897. fallbackLimit := resolveUpstreamResponseReadLimit(s.cfg)
  898. seenSSEData := false
  899. fallbackTooLarge := false
  900. var sseData openAISSEDataAccumulator
  901. processSSEData := func(dataBytes []byte) {
  902. seenSSEData = true
  903. fallbackBody.Reset()
  904. fallbackBytes = 0
  905. mergeOpenAIUsage(&usage, dataBytes)
  906. imageCounter.AddSSEData(dataBytes)
  907. }
  908. flushSSEEvent := func() {
  909. sseData.Flush(processSSEData)
  910. }
  911. processLine := func(line []byte) {
  912. if len(line) == 0 {
  913. return
  914. }
  915. if firstTokenMs == nil {
  916. ms := int(time.Since(startTime).Milliseconds())
  917. firstTokenMs = &ms
  918. }
  919. if !clientDisconnected {
  920. if _, writeErr := c.Writer.Write(line); writeErr != nil {
  921. clientDisconnected = true
  922. logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Images stream client disconnected, continue draining upstream for billing")
  923. } else {
  924. flusher.Flush()
  925. lastDownstreamWriteAt = time.Now()
  926. }
  927. }
  928. trimmedLine := strings.TrimRight(string(line), "\r\n")
  929. if _, ok := extractOpenAISSEDataLine(trimmedLine); ok || strings.TrimSpace(trimmedLine) == "" {
  930. sseData.AddLine(trimmedLine, processSSEData)
  931. return
  932. }
  933. if !seenSSEData && !fallbackTooLarge {
  934. fallbackBytes += int64(len(line))
  935. if fallbackBytes <= fallbackLimit {
  936. _, _ = fallbackBody.Write(line)
  937. } else {
  938. fallbackTooLarge = true
  939. fallbackBody.Reset()
  940. }
  941. }
  942. }
  943. finalizeFallbackBody := func() {
  944. if seenSSEData || fallbackBody.Len() == 0 {
  945. return
  946. }
  947. body := bytes.TrimSpace(fallbackBody.Bytes())
  948. if len(body) == 0 {
  949. return
  950. }
  951. mergeOpenAIUsage(&usage, body)
  952. imageCounter.AddJSONResponse(body)
  953. }
  954. streamInterval := s.openAIImageStreamDataInterval()
  955. keepaliveInterval := s.openAIImageStreamKeepaliveInterval()
  956. if streamInterval <= 0 && keepaliveInterval <= 0 {
  957. reader := bufio.NewReader(resp.Body)
  958. for {
  959. line, err := reader.ReadBytes('\n')
  960. processLine(line)
  961. if err == io.EOF {
  962. break
  963. }
  964. if err != nil {
  965. flushSSEEvent()
  966. return usage, imageCounter.Count(), imageCounter.Sizes(), firstTokenMs, err
  967. }
  968. }
  969. flushSSEEvent()
  970. finalizeFallbackBody()
  971. return usage, imageCounter.Count(), imageCounter.Sizes(), firstTokenMs, nil
  972. }
  973. type readEvent struct {
  974. line []byte
  975. err error
  976. }
  977. events := make(chan readEvent, 16)
  978. done := make(chan struct{})
  979. sendEvent := func(ev readEvent) bool {
  980. select {
  981. case events <- ev:
  982. return true
  983. case <-done:
  984. return false
  985. }
  986. }
  987. var lastReadAt int64
  988. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  989. go func() {
  990. defer close(events)
  991. reader := bufio.NewReader(resp.Body)
  992. for {
  993. line, err := reader.ReadBytes('\n')
  994. if len(line) > 0 {
  995. atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
  996. }
  997. if len(line) > 0 && !sendEvent(readEvent{line: line}) {
  998. return
  999. }
  1000. if err == io.EOF {
  1001. return
  1002. }
  1003. if err != nil {
  1004. _ = sendEvent(readEvent{err: err})
  1005. return
  1006. }
  1007. }
  1008. }()
  1009. defer close(done)
  1010. var intervalTicker *time.Ticker
  1011. if streamInterval > 0 {
  1012. intervalTicker = time.NewTicker(streamInterval)
  1013. defer intervalTicker.Stop()
  1014. }
  1015. var intervalCh <-chan time.Time
  1016. if intervalTicker != nil {
  1017. intervalCh = intervalTicker.C
  1018. }
  1019. var keepaliveTicker *time.Ticker
  1020. if keepaliveInterval > 0 {
  1021. keepaliveTicker = time.NewTicker(keepaliveInterval)
  1022. defer keepaliveTicker.Stop()
  1023. }
  1024. var keepaliveCh <-chan time.Time
  1025. if keepaliveTicker != nil {
  1026. keepaliveCh = keepaliveTicker.C
  1027. }
  1028. for {
  1029. select {
  1030. case ev, ok := <-events:
  1031. if !ok {
  1032. flushSSEEvent()
  1033. finalizeFallbackBody()
  1034. return usage, imageCounter.Count(), imageCounter.Sizes(), firstTokenMs, nil
  1035. }
  1036. if ev.err != nil {
  1037. flushSSEEvent()
  1038. return usage, imageCounter.Count(), imageCounter.Sizes(), firstTokenMs, ev.err
  1039. }
  1040. processLine(ev.line)
  1041. case <-intervalCh:
  1042. lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
  1043. if time.Since(lastRead) < streamInterval {
  1044. continue
  1045. }
  1046. if clientDisconnected {
  1047. return usage, imageCounter.Count(), imageCounter.Sizes(), firstTokenMs, fmt.Errorf("image stream incomplete after timeout")
  1048. }
  1049. logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Images stream data interval timeout: interval=%s", streamInterval)
  1050. _ = s.writeOpenAIImagesStreamEvent(c, flusher, "error", buildOpenAIImagesStreamErrorBody(fmt.Sprintf("upstream image stream idle for %s", streamInterval)))
  1051. return usage, imageCounter.Count(), imageCounter.Sizes(), firstTokenMs, fmt.Errorf("image stream data interval timeout")
  1052. case <-keepaliveCh:
  1053. if clientDisconnected || time.Since(lastDownstreamWriteAt) < keepaliveInterval {
  1054. continue
  1055. }
  1056. if _, writeErr := io.WriteString(c.Writer, ":\n\n"); writeErr != nil {
  1057. clientDisconnected = true
  1058. logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Images stream client disconnected during keepalive, continue draining upstream for billing")
  1059. continue
  1060. }
  1061. flusher.Flush()
  1062. lastDownstreamWriteAt = time.Now()
  1063. }
  1064. }
  1065. }
  1066. func (s *OpenAIGatewayService) openAIImageStreamDataInterval() time.Duration {
  1067. if s == nil || s.cfg == nil || s.cfg.Gateway.ImageStreamDataIntervalTimeout <= 0 {
  1068. return 0
  1069. }
  1070. return time.Duration(s.cfg.Gateway.ImageStreamDataIntervalTimeout) * time.Second
  1071. }
  1072. func (s *OpenAIGatewayService) openAIImageStreamKeepaliveInterval() time.Duration {
  1073. if s == nil || s.cfg == nil || s.cfg.Gateway.ImageStreamKeepaliveInterval <= 0 {
  1074. return 0
  1075. }
  1076. return time.Duration(s.cfg.Gateway.ImageStreamKeepaliveInterval) * time.Second
  1077. }
  1078. func extractOpenAIImagesBillableCountFromJSONBytes(body []byte) int {
  1079. if count := extractOpenAIImageCountFromJSONBytes(body); count > 0 {
  1080. return count
  1081. }
  1082. if len(body) == 0 || !gjson.ValidBytes(body) {
  1083. return 0
  1084. }
  1085. if count := int(gjson.GetBytes(body, "usage.images").Int()); count > 0 {
  1086. return count
  1087. }
  1088. if count := int(gjson.GetBytes(body, "tool_usage.image_gen.images").Int()); count > 0 {
  1089. return count
  1090. }
  1091. eventType := strings.TrimSpace(gjson.GetBytes(body, "type").String())
  1092. if eventType == "" || !strings.HasSuffix(eventType, ".completed") {
  1093. return 0
  1094. }
  1095. if gjson.GetBytes(body, "b64_json").Exists() || gjson.GetBytes(body, "url").Exists() {
  1096. return 1
  1097. }
  1098. return 0
  1099. }
  1100. func mergeOpenAIUsage(dst *OpenAIUsage, body []byte) {
  1101. if dst == nil {
  1102. return
  1103. }
  1104. if parsed, ok := extractOpenAIUsageFromJSONBytes(body); ok {
  1105. if parsed.InputTokens > 0 {
  1106. dst.InputTokens = parsed.InputTokens
  1107. }
  1108. if parsed.OutputTokens > 0 {
  1109. dst.OutputTokens = parsed.OutputTokens
  1110. }
  1111. if parsed.CacheReadInputTokens > 0 {
  1112. dst.CacheReadInputTokens = parsed.CacheReadInputTokens
  1113. }
  1114. if parsed.ImageInputTokens > 0 {
  1115. dst.ImageInputTokens = parsed.ImageInputTokens
  1116. }
  1117. if parsed.ImageOutputTokens > 0 {
  1118. dst.ImageOutputTokens = parsed.ImageOutputTokens
  1119. }
  1120. }
  1121. }
  1122. func extractOpenAIImageCountFromJSONBytes(body []byte) int {
  1123. return countOpenAIResponseImageOutputsFromJSONBytes(body)
  1124. }
  1125. type openAIImagePointerInfo struct {
  1126. Pointer string
  1127. DownloadURL string
  1128. B64JSON string
  1129. MimeType string
  1130. Prompt string
  1131. }
  1132. func collectOpenAIImagePointers(body []byte) []openAIImagePointerInfo {
  1133. if len(body) == 0 {
  1134. return nil
  1135. }
  1136. prompt := ""
  1137. for _, path := range []string{
  1138. "message.metadata.dalle.prompt",
  1139. "metadata.dalle.prompt",
  1140. "revised_prompt",
  1141. } {
  1142. if value := strings.TrimSpace(gjson.GetBytes(body, path).String()); value != "" {
  1143. prompt = value
  1144. break
  1145. }
  1146. }
  1147. matches := openAIImagePointerMatches(body)
  1148. out := make([]openAIImagePointerInfo, 0, len(matches))
  1149. for _, pointer := range matches {
  1150. out = append(out, openAIImagePointerInfo{Pointer: pointer, Prompt: prompt})
  1151. }
  1152. return mergeOpenAIImagePointerInfos(out, collectOpenAIImageInlineAssets(body, prompt))
  1153. }
  1154. func openAIImagePointerMatches(body []byte) []string {
  1155. raw := string(body)
  1156. matches := make([]string, 0, 4)
  1157. for _, prefix := range []string{"file-service://", "sediment://"} {
  1158. start := 0
  1159. for {
  1160. idx := strings.Index(raw[start:], prefix)
  1161. if idx < 0 {
  1162. break
  1163. }
  1164. idx += start
  1165. end := idx + len(prefix)
  1166. for end < len(raw) {
  1167. ch := raw[end]
  1168. if ch != '-' && ch != '_' &&
  1169. (ch < '0' || ch > '9') &&
  1170. (ch < 'a' || ch > 'z') &&
  1171. (ch < 'A' || ch > 'Z') {
  1172. break
  1173. }
  1174. end++
  1175. }
  1176. matches = append(matches, raw[idx:end])
  1177. start = end
  1178. }
  1179. }
  1180. return dedupeStrings(matches)
  1181. }
  1182. func mergeOpenAIImagePointerInfos(existing []openAIImagePointerInfo, next []openAIImagePointerInfo) []openAIImagePointerInfo {
  1183. if len(next) == 0 {
  1184. return existing
  1185. }
  1186. seen := make(map[string]openAIImagePointerInfo, len(existing)+len(next))
  1187. out := make([]openAIImagePointerInfo, 0, len(existing)+len(next))
  1188. for _, item := range existing {
  1189. if key := item.identityKey(); key != "" {
  1190. seen[key] = item
  1191. }
  1192. out = append(out, item)
  1193. }
  1194. for _, item := range next {
  1195. key := item.identityKey()
  1196. if key == "" {
  1197. continue
  1198. }
  1199. if existingItem, ok := seen[key]; ok {
  1200. merged := mergeOpenAIImagePointerInfo(existingItem, item)
  1201. if merged != existingItem {
  1202. for i := range out {
  1203. if out[i].identityKey() == key {
  1204. out[i] = merged
  1205. break
  1206. }
  1207. }
  1208. seen[key] = merged
  1209. }
  1210. continue
  1211. }
  1212. seen[key] = item
  1213. out = append(out, item)
  1214. }
  1215. return out
  1216. }
  1217. func (i openAIImagePointerInfo) identityKey() string {
  1218. switch {
  1219. case strings.TrimSpace(i.Pointer) != "":
  1220. return "pointer:" + strings.TrimSpace(i.Pointer)
  1221. case strings.TrimSpace(i.DownloadURL) != "":
  1222. return "download:" + strings.TrimSpace(i.DownloadURL)
  1223. case strings.TrimSpace(i.B64JSON) != "":
  1224. b64 := strings.TrimSpace(i.B64JSON)
  1225. if len(b64) > 64 {
  1226. b64 = b64[:64]
  1227. }
  1228. return "b64:" + b64
  1229. default:
  1230. return ""
  1231. }
  1232. }
  1233. func mergeOpenAIImagePointerInfo(existing, next openAIImagePointerInfo) openAIImagePointerInfo {
  1234. merged := existing
  1235. if strings.TrimSpace(merged.Pointer) == "" {
  1236. merged.Pointer = next.Pointer
  1237. }
  1238. if strings.TrimSpace(merged.DownloadURL) == "" {
  1239. merged.DownloadURL = next.DownloadURL
  1240. }
  1241. if strings.TrimSpace(merged.B64JSON) == "" {
  1242. merged.B64JSON = next.B64JSON
  1243. }
  1244. if strings.TrimSpace(merged.MimeType) == "" {
  1245. merged.MimeType = next.MimeType
  1246. }
  1247. if strings.TrimSpace(merged.Prompt) == "" {
  1248. merged.Prompt = next.Prompt
  1249. }
  1250. return merged
  1251. }
  1252. func resolveOpenAIImageBytes(
  1253. ctx context.Context,
  1254. client *req.Client,
  1255. headers http.Header,
  1256. conversationID string,
  1257. pointer openAIImagePointerInfo,
  1258. errorBodyReadLimit int64,
  1259. ) ([]byte, error) {
  1260. if normalized := normalizeOpenAIImageBase64(pointer.B64JSON); normalized != "" {
  1261. return base64.StdEncoding.DecodeString(normalized)
  1262. }
  1263. if downloadURL := strings.TrimSpace(pointer.DownloadURL); downloadURL != "" {
  1264. return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit)
  1265. }
  1266. if strings.TrimSpace(pointer.Pointer) == "" {
  1267. return nil, fmt.Errorf("image asset is missing pointer, url, and base64 data")
  1268. }
  1269. downloadURL, err := fetchOpenAIImageDownloadURL(ctx, client, headers, conversationID, pointer.Pointer, errorBodyReadLimit)
  1270. if err != nil {
  1271. return nil, err
  1272. }
  1273. return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit)
  1274. }
  1275. func normalizeOpenAIImageBase64(raw string) string {
  1276. raw = strings.TrimSpace(raw)
  1277. if raw == "" {
  1278. return ""
  1279. }
  1280. if strings.HasPrefix(strings.ToLower(raw), "data:") {
  1281. if idx := strings.Index(raw, ","); idx >= 0 && idx+1 < len(raw) {
  1282. raw = raw[idx+1:]
  1283. }
  1284. }
  1285. raw = strings.TrimSpace(raw)
  1286. raw = strings.TrimRight(raw, "=") + strings.Repeat("=", (4-len(raw)%4)%4)
  1287. if raw == "" {
  1288. return ""
  1289. }
  1290. if _, err := base64.StdEncoding.DecodeString(raw); err != nil {
  1291. return ""
  1292. }
  1293. return raw
  1294. }
  1295. func collectOpenAIImageInlineAssets(body []byte, fallbackPrompt string) []openAIImagePointerInfo {
  1296. if len(body) == 0 || !gjson.ValidBytes(body) {
  1297. return nil
  1298. }
  1299. var decoded any
  1300. if err := json.Unmarshal(body, &decoded); err != nil {
  1301. return nil
  1302. }
  1303. var out []openAIImagePointerInfo
  1304. walkOpenAIImageInlineAssets(decoded, strings.TrimSpace(fallbackPrompt), &out)
  1305. return out
  1306. }
  1307. func walkOpenAIImageInlineAssets(node any, prompt string, out *[]openAIImagePointerInfo) {
  1308. switch value := node.(type) {
  1309. case map[string]any:
  1310. localPrompt := prompt
  1311. for _, key := range []string{"revised_prompt", "image_gen_title", "prompt"} {
  1312. if v, ok := value[key].(string); ok && strings.TrimSpace(v) != "" {
  1313. localPrompt = strings.TrimSpace(v)
  1314. break
  1315. }
  1316. }
  1317. item := openAIImagePointerInfo{
  1318. Prompt: localPrompt,
  1319. Pointer: firstNonEmptyString(value["asset_pointer"], value["pointer"]),
  1320. DownloadURL: firstNonEmptyString(value["download_url"], value["url"], value["image_url"]),
  1321. B64JSON: firstNonEmptyString(value["b64_json"], value["base64"], value["image_base64"]),
  1322. MimeType: firstNonEmptyString(value["mime_type"], value["mimeType"], value["content_type"]),
  1323. }
  1324. switch {
  1325. case strings.HasPrefix(strings.TrimSpace(item.Pointer), "file-service://"),
  1326. strings.HasPrefix(strings.TrimSpace(item.Pointer), "sediment://"),
  1327. isLikelyOpenAIImageDownloadURL(item.DownloadURL),
  1328. normalizeOpenAIImageBase64(item.B64JSON) != "":
  1329. *out = append(*out, item)
  1330. }
  1331. for _, child := range value {
  1332. walkOpenAIImageInlineAssets(child, localPrompt, out)
  1333. }
  1334. case []any:
  1335. for _, child := range value {
  1336. walkOpenAIImageInlineAssets(child, prompt, out)
  1337. }
  1338. }
  1339. }
  1340. func firstNonEmptyString(values ...any) string {
  1341. for _, value := range values {
  1342. if s, ok := value.(string); ok && strings.TrimSpace(s) != "" {
  1343. return strings.TrimSpace(s)
  1344. }
  1345. }
  1346. return ""
  1347. }
  1348. func isLikelyOpenAIImageDownloadURL(raw string) bool {
  1349. raw = strings.TrimSpace(raw)
  1350. if raw == "" {
  1351. return false
  1352. }
  1353. if strings.HasPrefix(strings.ToLower(raw), "data:image/") {
  1354. return true
  1355. }
  1356. if !strings.HasPrefix(strings.ToLower(raw), "http://") && !strings.HasPrefix(strings.ToLower(raw), "https://") {
  1357. return false
  1358. }
  1359. lower := strings.ToLower(raw)
  1360. return strings.Contains(lower, "/download") ||
  1361. strings.Contains(lower, ".png") ||
  1362. strings.Contains(lower, ".jpg") ||
  1363. strings.Contains(lower, ".jpeg") ||
  1364. strings.Contains(lower, ".webp")
  1365. }
  1366. func fetchOpenAIImageDownloadURL(
  1367. ctx context.Context,
  1368. client *req.Client,
  1369. headers http.Header,
  1370. conversationID string,
  1371. pointer string,
  1372. errorBodyReadLimit int64,
  1373. ) (string, error) {
  1374. url := ""
  1375. allowConversationRetry := false
  1376. switch {
  1377. case strings.HasPrefix(pointer, "file-service://"):
  1378. fileID := strings.TrimPrefix(pointer, "file-service://")
  1379. url = fmt.Sprintf("%s/%s/download", openAIChatGPTFilesURL, fileID)
  1380. case strings.HasPrefix(pointer, "sediment://"):
  1381. attachmentID := strings.TrimPrefix(pointer, "sediment://")
  1382. url = fmt.Sprintf("https://chatgpt.com/backend-api/conversation/%s/attachment/%s/download", conversationID, attachmentID)
  1383. allowConversationRetry = true
  1384. default:
  1385. return "", fmt.Errorf("unsupported image pointer: %s", pointer)
  1386. }
  1387. var lastErr error
  1388. for attempt := 0; attempt < 8; attempt++ {
  1389. var result struct {
  1390. DownloadURL string `json:"download_url"`
  1391. }
  1392. resp, err := client.R().
  1393. SetContext(ctx).
  1394. SetHeaders(headerToMap(headers)).
  1395. SetSuccessResult(&result).
  1396. Get(url)
  1397. if err != nil {
  1398. lastErr = err
  1399. } else if resp.IsSuccessState() && strings.TrimSpace(result.DownloadURL) != "" {
  1400. return strings.TrimSpace(result.DownloadURL), nil
  1401. } else {
  1402. statusErr := newOpenAIImageStatusError(resp, "fetch image download url failed", errorBodyReadLimit)
  1403. if !allowConversationRetry || !isOpenAIImageTransientConversationNotFoundError(statusErr) {
  1404. return "", statusErr
  1405. }
  1406. lastErr = statusErr
  1407. }
  1408. if attempt == 7 {
  1409. break
  1410. }
  1411. timer := time.NewTimer(750 * time.Millisecond)
  1412. select {
  1413. case <-ctx.Done():
  1414. if !timer.Stop() {
  1415. <-timer.C
  1416. }
  1417. return "", ctx.Err()
  1418. case <-timer.C:
  1419. }
  1420. }
  1421. if lastErr == nil {
  1422. lastErr = fmt.Errorf("fetch image download url failed")
  1423. }
  1424. return "", lastErr
  1425. }
  1426. func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers http.Header, downloadURL string, errorBodyReadLimit int64) ([]byte, error) {
  1427. request := client.R().
  1428. SetContext(ctx).
  1429. DisableAutoReadResponse()
  1430. if strings.HasPrefix(downloadURL, openAIChatGPTStartURL) {
  1431. downloadHeaders := cloneHTTPHeader(headers)
  1432. downloadHeaders.Set("Accept", "image/*,*/*;q=0.8")
  1433. downloadHeaders.Del("Content-Type")
  1434. request.SetHeaders(headerToMap(downloadHeaders))
  1435. } else {
  1436. userAgent := strings.TrimSpace(headers.Get("User-Agent"))
  1437. if userAgent == "" {
  1438. userAgent = openAIImageBackendUserAgent
  1439. }
  1440. request.SetHeader("User-Agent", userAgent)
  1441. }
  1442. resp, err := request.Get(downloadURL)
  1443. if err != nil {
  1444. return nil, err
  1445. }
  1446. defer func() {
  1447. if resp != nil && resp.Body != nil {
  1448. _ = resp.Body.Close()
  1449. }
  1450. }()
  1451. if resp.StatusCode < 200 || resp.StatusCode >= 300 {
  1452. return nil, newOpenAIImageStatusError(resp, "download image bytes failed", errorBodyReadLimit)
  1453. }
  1454. return io.ReadAll(io.LimitReader(resp.Body, openAIImageMaxDownloadBytes))
  1455. }
  1456. type openAIImageStatusError struct {
  1457. StatusCode int
  1458. Message string
  1459. ResponseBody []byte
  1460. ResponseHeaders http.Header
  1461. RequestID string
  1462. URL string
  1463. }
  1464. func (e *openAIImageStatusError) Error() string {
  1465. if e == nil {
  1466. return "openai image backend request failed"
  1467. }
  1468. if e.Message != "" {
  1469. return e.Message
  1470. }
  1471. if e.StatusCode > 0 {
  1472. return fmt.Sprintf("openai image backend request failed: status %d", e.StatusCode)
  1473. }
  1474. return "openai image backend request failed"
  1475. }
  1476. func newOpenAIImageStatusError(resp *req.Response, fallback string, errorBodyReadLimit int64) error {
  1477. if resp == nil {
  1478. if strings.TrimSpace(fallback) == "" {
  1479. fallback = "openai image backend request failed"
  1480. }
  1481. return fmt.Errorf("%s", fallback)
  1482. }
  1483. statusCode := resp.StatusCode
  1484. headers := http.Header(nil)
  1485. requestID := ""
  1486. requestURL := ""
  1487. body := []byte(nil)
  1488. if resp.Response != nil {
  1489. headers = resp.Header.Clone()
  1490. requestID = strings.TrimSpace(resp.Header.Get("x-request-id"))
  1491. if resp.Request != nil && resp.Request.URL != nil {
  1492. requestURL = resp.Request.URL.String()
  1493. }
  1494. if resp.Body != nil {
  1495. if errorBodyReadLimit <= 0 {
  1496. errorBodyReadLimit = openAIUpstreamErrorBodyReadLimit
  1497. }
  1498. body, _ = io.ReadAll(io.LimitReader(resp.Body, errorBodyReadLimit))
  1499. _ = resp.Body.Close()
  1500. }
  1501. }
  1502. message := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(body))
  1503. if message == "" {
  1504. prefix := strings.TrimSpace(fallback)
  1505. if prefix == "" {
  1506. prefix = "openai image backend request failed"
  1507. }
  1508. message = fmt.Sprintf("%s: status %d", prefix, statusCode)
  1509. }
  1510. return &openAIImageStatusError{
  1511. StatusCode: statusCode,
  1512. Message: message,
  1513. ResponseBody: body,
  1514. ResponseHeaders: headers,
  1515. RequestID: requestID,
  1516. URL: requestURL,
  1517. }
  1518. }
  1519. func isOpenAIImageTransientConversationNotFoundError(err error) bool {
  1520. statusErr, ok := err.(*openAIImageStatusError)
  1521. if !ok || statusErr == nil || statusErr.StatusCode != http.StatusNotFound {
  1522. return false
  1523. }
  1524. msg := strings.ToLower(strings.TrimSpace(statusErr.Message))
  1525. if strings.Contains(msg, "conversation_not_found") {
  1526. return true
  1527. }
  1528. if strings.Contains(msg, "conversation") && strings.Contains(msg, "not found") {
  1529. return true
  1530. }
  1531. bodyMsg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(statusErr.ResponseBody)))
  1532. if strings.Contains(bodyMsg, "conversation_not_found") {
  1533. return true
  1534. }
  1535. return strings.Contains(bodyMsg, "conversation") && strings.Contains(bodyMsg, "not found")
  1536. }
  1537. func cloneHTTPHeader(src http.Header) http.Header {
  1538. dst := make(http.Header, len(src))
  1539. for key, values := range src {
  1540. copied := make([]string, len(values))
  1541. copy(copied, values)
  1542. dst[key] = copied
  1543. }
  1544. return dst
  1545. }
  1546. func headerToMap(header http.Header) map[string]string {
  1547. if len(header) == 0 {
  1548. return nil
  1549. }
  1550. result := make(map[string]string, len(header))
  1551. for key, values := range header {
  1552. if len(values) == 0 {
  1553. continue
  1554. }
  1555. result[key] = values[0]
  1556. }
  1557. return result
  1558. }
  1559. func dedupeStrings(values []string) []string {
  1560. if len(values) == 0 {
  1561. return nil
  1562. }
  1563. seen := make(map[string]struct{}, len(values))
  1564. out := make([]string, 0, len(values))
  1565. for _, value := range values {
  1566. if _, ok := seen[value]; ok {
  1567. continue
  1568. }
  1569. seen[value] = struct{}{}
  1570. out = append(out, value)
  1571. }
  1572. return out
  1573. }