openai_image_dimensions.go 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. package service
  2. import (
  3. "bufio"
  4. "encoding/base64"
  5. "encoding/binary"
  6. "fmt"
  7. "image"
  8. _ "image/jpeg"
  9. _ "image/png"
  10. "io"
  11. "strings"
  12. )
  13. const maxOpenAIImageDimensionProbeBytes int64 = 1 << 20
  14. func detectOpenAIImageResultSize(encoded string) string {
  15. payload := strings.TrimSpace(encoded)
  16. if strings.HasPrefix(strings.ToLower(payload), "data:") {
  17. comma := strings.IndexByte(payload, ',')
  18. if comma < 0 || comma+1 >= len(payload) {
  19. return ""
  20. }
  21. payload = strings.TrimSpace(payload[comma+1:])
  22. }
  23. if payload == "" {
  24. return ""
  25. }
  26. for _, encoding := range []*base64.Encoding{base64.StdEncoding, base64.RawStdEncoding} {
  27. decoded := base64.NewDecoder(encoding, strings.NewReader(payload))
  28. buffered := bufio.NewReader(io.LimitReader(decoded, maxOpenAIImageDimensionProbeBytes))
  29. prefix, _ := buffered.Peek(30)
  30. if width, height, ok := detectOpenAIWebPDimensions(prefix); ok {
  31. return fmt.Sprintf("%dx%d", width, height)
  32. }
  33. cfg, _, err := image.DecodeConfig(buffered)
  34. if err != nil || cfg.Width <= 0 || cfg.Height <= 0 {
  35. continue
  36. }
  37. return fmt.Sprintf("%dx%d", cfg.Width, cfg.Height)
  38. }
  39. return ""
  40. }
  41. func detectOpenAIWebPDimensions(header []byte) (int, int, bool) {
  42. if len(header) < 16 || string(header[:4]) != "RIFF" || string(header[8:12]) != "WEBP" {
  43. return 0, 0, false
  44. }
  45. switch string(header[12:16]) {
  46. case "VP8X":
  47. if len(header) < 30 {
  48. return 0, 0, false
  49. }
  50. width := 1 + int(header[24]) + int(header[25])<<8 + int(header[26])<<16
  51. height := 1 + int(header[27]) + int(header[28])<<8 + int(header[29])<<16
  52. return width, height, width > 0 && height > 0
  53. case "VP8 ":
  54. if len(header) < 30 || string(header[23:26]) != "\x9d\x01\x2a" {
  55. return 0, 0, false
  56. }
  57. width := int(binary.LittleEndian.Uint16(header[26:28]) & 0x3fff)
  58. height := int(binary.LittleEndian.Uint16(header[28:30]) & 0x3fff)
  59. return width, height, width > 0 && height > 0
  60. case "VP8L":
  61. if len(header) < 25 || header[20] != 0x2f {
  62. return 0, 0, false
  63. }
  64. width := 1 + int(header[21]) + int(header[22]&0x3f)<<8
  65. height := 1 + int(header[22]>>6) + int(header[23])<<2 + int(header[24]&0x0f)<<10
  66. return width, height, width > 0 && height > 0
  67. default:
  68. return 0, 0, false
  69. }
  70. }
  71. func reconcileOpenAIResponsesImageResultSizes(results []openAIResponsesImageResult, firstMeta *openAIResponsesImageResult) {
  72. for i := range results {
  73. // ChatGPT OAuth can normalize requested controls to "auto". The final
  74. // image bytes are authoritative for response metadata and tier billing.
  75. if actualSize := detectOpenAIImageResultSize(results[i].Result); actualSize != "" {
  76. results[i].Size = actualSize
  77. }
  78. }
  79. if firstMeta == nil || len(results) == 0 {
  80. return
  81. }
  82. if size := strings.TrimSpace(results[0].Size); size != "" {
  83. firstMeta.Size = size
  84. }
  85. }