batch_image_processor.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682
  1. package service
  2. import (
  3. "bufio"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "net/http"
  10. "sort"
  11. "strconv"
  12. "strings"
  13. "time"
  14. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  15. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  16. "go.uber.org/zap"
  17. )
  18. const (
  19. BatchImageParsedStatusSucceeded = "succeeded"
  20. BatchImageParsedStatusFailed = "failed"
  21. defaultBatchImageProcessorRequeue = 30 * time.Second
  22. batchImageProviderErrorRequeue = time.Minute
  23. batchImageMaxErrorMessageLength = 1000
  24. )
  25. type BatchImageAccountResolver interface {
  26. ResolveBatchImageAccount(ctx context.Context, accountID int64) (*Account, error)
  27. }
  28. type BatchImageAccountLookup interface {
  29. GetByID(ctx context.Context, id int64) (*Account, error)
  30. }
  31. type BatchImageAccountRepositoryResolver struct {
  32. Repo BatchImageAccountLookup
  33. }
  34. func (r *BatchImageAccountRepositoryResolver) ResolveBatchImageAccount(ctx context.Context, accountID int64) (*Account, error) {
  35. if r == nil || r.Repo == nil {
  36. return nil, ErrAccountNotFound
  37. }
  38. return r.Repo.GetByID(ctx, accountID)
  39. }
  40. type BatchImageProviderProcessor struct {
  41. Repo BatchImageRepository
  42. ProviderRegistry *BatchImageProviderRegistry
  43. AccountResolver BatchImageAccountResolver
  44. Indexer *BatchImageResultIndexer
  45. BillingRepo UsageBillingRepository
  46. AuthCache APIKeyAuthCacheInvalidator
  47. DefaultRequeue time.Duration
  48. }
  49. func (p *BatchImageProviderProcessor) Process(ctx context.Context, batchID string) (BatchImageProcessResult, error) {
  50. if p == nil || p.Repo == nil || p.ProviderRegistry == nil || p.AccountResolver == nil {
  51. return BatchImageProcessResult{}, infraerrors.New(http.StatusInternalServerError, "BATCH_IMAGE_PROCESSOR_NOT_CONFIGURED", "batch image processor is not configured")
  52. }
  53. job, err := p.Repo.GetBatchImageJobByBatchID(ctx, batchID)
  54. if err != nil {
  55. return BatchImageProcessResult{}, err
  56. }
  57. if isBatchImageProcessorDoneStatus(job.Status) {
  58. if err := p.releaseTerminalHold(ctx, job); err != nil {
  59. return BatchImageProcessResult{}, err
  60. }
  61. return BatchImageProcessResult{Terminal: true}, nil
  62. }
  63. provider, ok := p.ProviderRegistry.Get(job.Provider)
  64. if !ok || provider == nil {
  65. return BatchImageProcessResult{}, ErrBatchImageUnsupportedProvider
  66. }
  67. if job.AccountID == nil || *job.AccountID <= 0 {
  68. return BatchImageProcessResult{}, ErrBatchImageMissingAccountID
  69. }
  70. account, err := p.AccountResolver.ResolveBatchImageAccount(ctx, *job.AccountID)
  71. if err != nil {
  72. return BatchImageProcessResult{}, err
  73. }
  74. if !provider.SupportsAccount(account) {
  75. return BatchImageProcessResult{}, ErrBatchImageProviderUnsupportedAccount
  76. }
  77. if strings.TrimSpace(batchImageDerefString(job.ProviderJobName)) == "" {
  78. return BatchImageProcessResult{}, ErrBatchImageMissingProviderJobName
  79. }
  80. if job.Status == BatchImageJobStatusIndexing {
  81. return p.indexAndSettle(ctx, job, provider, account)
  82. }
  83. status, err := provider.Get(ctx, job, account)
  84. if err != nil {
  85. logger.L().Warn("batch_image.provider_status_check_failed",
  86. zap.String("batch_id", job.BatchID),
  87. zap.String("provider", job.Provider),
  88. zap.String("provider_job_name", batchImageDerefString(job.ProviderJobName)),
  89. zap.Error(err),
  90. )
  91. return BatchImageProcessResult{RequeueAfter: batchImageProviderErrorRequeue}, nil
  92. }
  93. if status == nil {
  94. return BatchImageProcessResult{RequeueAfter: p.requeueDelay(0)}, nil
  95. }
  96. if err := p.persistProviderOutputRef(ctx, job, status.ProviderOutputRef); err != nil {
  97. return BatchImageProcessResult{}, err
  98. }
  99. switch status.InternalState {
  100. case BatchProviderStateQueued:
  101. return BatchImageProcessResult{RequeueAfter: p.requeueDelay(status.SuggestedRequeueAfter)}, nil
  102. case BatchProviderStateRunning:
  103. if job.Status != BatchImageJobStatusRunning {
  104. if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusRunning, BatchImageTransitionOptions{
  105. EventType: "provider_status_checked",
  106. EventPayload: map[string]any{"provider_state": status.RawState},
  107. }); err != nil {
  108. return BatchImageProcessResult{}, err
  109. }
  110. job.Status = BatchImageJobStatusRunning
  111. }
  112. return BatchImageProcessResult{RequeueAfter: p.requeueDelay(status.SuggestedRequeueAfter)}, nil
  113. case BatchProviderStateSucceeded:
  114. if job.Status != BatchImageJobStatusIndexing {
  115. if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusIndexing, BatchImageTransitionOptions{
  116. EventType: "indexing_started",
  117. EventPayload: map[string]any{"provider_state": status.RawState},
  118. }); err != nil {
  119. return BatchImageProcessResult{}, err
  120. }
  121. job.Status = BatchImageJobStatusIndexing
  122. }
  123. return p.indexAndSettle(ctx, job, provider, account)
  124. case BatchProviderStateFailed, BatchProviderStateExpired:
  125. code := strings.TrimSpace(status.ErrorCode)
  126. if code == "" && status.InternalState == BatchProviderStateExpired {
  127. code = "PROVIDER_BATCH_EXPIRED"
  128. }
  129. if code == "" {
  130. code = "PROVIDER_BATCH_FAILED"
  131. }
  132. msg := truncateBatchImageMessage(status.ErrorMessage, batchImageMaxErrorMessageLength)
  133. if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusFailed, BatchImageTransitionOptions{
  134. EventType: "job_failed",
  135. EventPayload: map[string]any{"provider_state": status.RawState, "error_code": code},
  136. ErrorCode: batchImageStringPtr(code),
  137. ErrorMessage: batchImageOptionalStringPtr(msg),
  138. }); err != nil {
  139. return BatchImageProcessResult{}, err
  140. }
  141. job.Status = BatchImageJobStatusFailed
  142. if err := p.releaseTerminalHold(ctx, job); err != nil {
  143. return BatchImageProcessResult{}, err
  144. }
  145. return BatchImageProcessResult{Terminal: true}, nil
  146. case BatchProviderStateCancelled:
  147. if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusCancelled, BatchImageTransitionOptions{
  148. EventType: "job_failed",
  149. EventPayload: map[string]any{"provider_state": status.RawState, "error_code": "PROVIDER_BATCH_CANCELLED"},
  150. }); err != nil {
  151. return BatchImageProcessResult{}, err
  152. }
  153. job.Status = BatchImageJobStatusCancelled
  154. if err := p.releaseTerminalHold(ctx, job); err != nil {
  155. return BatchImageProcessResult{}, err
  156. }
  157. return BatchImageProcessResult{Terminal: true}, nil
  158. default:
  159. return BatchImageProcessResult{RequeueAfter: p.requeueDelay(status.SuggestedRequeueAfter)}, nil
  160. }
  161. }
  162. func (p *BatchImageProviderProcessor) indexAndSettle(ctx context.Context, job *BatchImageJob, provider BatchImageProvider, account *Account) (BatchImageProcessResult, error) {
  163. indexer := p.Indexer
  164. if indexer == nil {
  165. indexer = &BatchImageResultIndexer{Repo: p.Repo}
  166. }
  167. if indexer.Repo == nil {
  168. indexer.Repo = p.Repo
  169. }
  170. result, err := indexer.Index(ctx, job, provider, account)
  171. if err != nil {
  172. if errors.Is(err, ErrBatchImageIndexOutputMissing) {
  173. return BatchImageProcessResult{}, err
  174. }
  175. // job 状态已被并发方推进(如已进入 settling/终态):不是索引数据问题,
  176. // 短延迟 requeue 让下一轮按最新状态处理,不能误转 failed。
  177. if errors.Is(err, ErrBatchImageIndexStateConflict) {
  178. return BatchImageProcessResult{RequeueAfter: time.Millisecond}, nil
  179. }
  180. code := "INDEX_PARSE_FAILED"
  181. if errors.Is(err, ErrBatchImageDuplicateCustomID) {
  182. code = "DUPLICATE_CUSTOM_ID_IN_OUTPUT"
  183. }
  184. msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength)
  185. transitionErr := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusFailed, BatchImageTransitionOptions{
  186. EventType: "indexing_failed",
  187. EventPayload: map[string]any{"error_code": code},
  188. ErrorCode: batchImageStringPtr(code),
  189. ErrorMessage: batchImageOptionalStringPtr(msg),
  190. })
  191. if transitionErr != nil {
  192. return BatchImageProcessResult{}, transitionErr
  193. }
  194. job.Status = BatchImageJobStatusFailed
  195. if err := p.releaseTerminalHold(ctx, job); err != nil {
  196. return BatchImageProcessResult{}, err
  197. }
  198. return BatchImageProcessResult{Terminal: true}, nil
  199. }
  200. if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusSettling, BatchImageTransitionOptions{
  201. EventType: "indexing_completed",
  202. EventPayload: map[string]any{
  203. "success_count": result.SuccessCount,
  204. "fail_count": result.FailCount,
  205. "total_count": result.TotalCount,
  206. },
  207. }); err != nil {
  208. return BatchImageProcessResult{}, err
  209. }
  210. return BatchImageProcessResult{RequeueAfter: time.Millisecond}, nil
  211. }
  212. func (p *BatchImageProviderProcessor) releaseTerminalHold(ctx context.Context, job *BatchImageJob) error {
  213. if p == nil || job == nil {
  214. return nil
  215. }
  216. if job.Status != BatchImageJobStatusFailed && job.Status != BatchImageJobStatusCancelled {
  217. return nil
  218. }
  219. if err := releaseBatchImageBalanceHold(ctx, p.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil {
  220. return err
  221. }
  222. if p.AuthCache != nil && job.UserID > 0 {
  223. p.AuthCache.InvalidateAuthCacheByUserID(ctx, job.UserID)
  224. }
  225. return nil
  226. }
  227. func (p *BatchImageProviderProcessor) persistProviderOutputRef(ctx context.Context, job *BatchImageJob, ref string) error {
  228. ref = strings.TrimSpace(ref)
  229. if ref == "" || job == nil || batchImageDerefString(job.ProviderOutputRef) == ref {
  230. return nil
  231. }
  232. if err := p.Repo.UpdateBatchImageJobProviderOutputRef(ctx, job.BatchID, ref); err != nil {
  233. return err
  234. }
  235. job.ProviderOutputRef = &ref
  236. return nil
  237. }
  238. func (p *BatchImageProviderProcessor) requeueDelay(suggested time.Duration) time.Duration {
  239. if suggested > 0 {
  240. return suggested
  241. }
  242. if p != nil && p.DefaultRequeue > 0 {
  243. return p.DefaultRequeue
  244. }
  245. return defaultBatchImageProcessorRequeue
  246. }
  247. func isBatchImageProcessorDoneStatus(status string) bool {
  248. if status == BatchImageJobStatusSettling {
  249. return true
  250. }
  251. return IsTerminalBatchImageJobStatus(status)
  252. }
  253. type BatchImageIndexResult struct {
  254. SuccessCount int
  255. FailCount int
  256. TotalCount int
  257. }
  258. type BatchImageResultIndexer struct {
  259. Repo BatchImageRepository
  260. }
  261. func (i *BatchImageResultIndexer) Index(ctx context.Context, job *BatchImageJob, provider BatchImageProvider, account *Account) (*BatchImageIndexResult, error) {
  262. if i == nil || i.Repo == nil || job == nil || provider == nil {
  263. return nil, ErrBatchImageIndexOutputMissing
  264. }
  265. expected, err := i.listExpectedCustomIDs(ctx, job.BatchID)
  266. if err != nil {
  267. return nil, err
  268. }
  269. r, _, err := provider.OpenResult(ctx, job, account)
  270. if err != nil {
  271. return nil, ErrBatchImageIndexOutputMissing.WithCause(err)
  272. }
  273. defer func() { _ = r.Close() }()
  274. scanner := bufio.NewScanner(r)
  275. scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
  276. seen := make(map[string]int)
  277. unknownCount := 0
  278. var items []CreateBatchImageItemParams
  279. result := &BatchImageIndexResult{}
  280. lineNumber := 0
  281. now := time.Now()
  282. sourceObject := batchImageDerefString(job.ProviderOutputRef)
  283. if sourceObject == "" {
  284. sourceObject = batchImageDerefString(job.ProviderJobName)
  285. }
  286. for scanner.Scan() {
  287. lineNumber++
  288. line := strings.TrimSpace(scanner.Text())
  289. if line == "" {
  290. continue
  291. }
  292. parsed, err := ParseBatchImageResultLine([]byte(line), lineNumber)
  293. if err != nil {
  294. return nil, err
  295. }
  296. // 与提交时的 custom_id 集对账:provider 输出中未知/多余的行不能进入 item 表,
  297. // 否则 success+fail > item_count 会让结算永远校验失败。
  298. if len(expected) > 0 {
  299. if _, ok := expected[parsed.CustomID]; !ok {
  300. unknownCount++
  301. continue
  302. }
  303. }
  304. if firstLine, ok := seen[parsed.CustomID]; ok {
  305. return nil, ErrBatchImageDuplicateCustomID.WithCause(fmt.Errorf("custom id %q duplicated at lines %d and %d", parsed.CustomID, firstLine, lineNumber))
  306. }
  307. seen[parsed.CustomID] = lineNumber
  308. lineNo := parsed.SourceLineNumber
  309. item := CreateBatchImageItemParams{
  310. JobID: job.BatchID,
  311. CustomID: parsed.CustomID,
  312. Status: BatchImageItemStatusFailed,
  313. ProviderSourceObject: batchImageOptionalStringPtr(sourceObject),
  314. SourceLineNumber: &lineNo,
  315. ImageCount: parsed.ImageCount,
  316. IndexedAt: &now,
  317. }
  318. if parsed.Status == BatchImageParsedStatusSucceeded {
  319. item.Status = BatchImageItemStatusSuccess
  320. item.MimeType = batchImageOptionalStringPtr(parsed.MimeType)
  321. item.FileExtension = batchImageOptionalStringPtr(parsed.FileExtension)
  322. result.SuccessCount++
  323. } else {
  324. item.ErrorCode = batchImageOptionalStringPtr(parsed.ErrorCode)
  325. item.ErrorMessage = batchImageOptionalStringPtr(parsed.ErrorMessage)
  326. result.FailCount++
  327. }
  328. items = append(items, item)
  329. result.TotalCount++
  330. }
  331. if err := scanner.Err(); err != nil {
  332. if errors.Is(err, io.ErrUnexpectedEOF) {
  333. return nil, ErrBatchImageIndexParseFailed.WithCause(err)
  334. }
  335. return nil, err
  336. }
  337. // 输出中漏掉的已提交项必须补失败记录,而不是静默消失:
  338. // 否则用户看不到该项,且只按成功数计费会掩盖 provider 的丢单。
  339. missingCount := 0
  340. if len(expected) > 0 {
  341. missingIDs := make([]string, 0)
  342. for customID := range expected {
  343. if _, ok := seen[customID]; !ok {
  344. missingIDs = append(missingIDs, customID)
  345. }
  346. }
  347. sort.Strings(missingIDs)
  348. for _, customID := range missingIDs {
  349. items = append(items, CreateBatchImageItemParams{
  350. JobID: job.BatchID,
  351. CustomID: customID,
  352. Status: BatchImageItemStatusFailed,
  353. ProviderSourceObject: batchImageOptionalStringPtr(sourceObject),
  354. ErrorCode: batchImageStringPtr("PROVIDER_RESULT_MISSING"),
  355. ErrorMessage: batchImageStringPtr("provider output did not include a result for this item"),
  356. IndexedAt: &now,
  357. })
  358. result.FailCount++
  359. result.TotalCount++
  360. }
  361. missingCount = len(missingIDs)
  362. }
  363. if result.TotalCount == 0 {
  364. return nil, ErrBatchImageIndexNoResultLines
  365. }
  366. if unknownCount > 0 || missingCount > 0 {
  367. logger.L().Warn("batch_image.index_reconciled",
  368. zap.String("batch_id", job.BatchID),
  369. zap.Int("unknown_custom_ids", unknownCount),
  370. zap.Int("missing_custom_ids", missingCount),
  371. )
  372. if err := i.Repo.AppendBatchImageEvent(ctx, job.BatchID, "index_reconciled", map[string]any{
  373. "batch_id": job.BatchID,
  374. "unknown_custom_ids": unknownCount,
  375. "missing_custom_ids": missingCount,
  376. }); err != nil {
  377. logger.L().Warn("batch_image.index_reconcile_event_failed",
  378. zap.String("batch_id", job.BatchID),
  379. zap.Error(err),
  380. )
  381. }
  382. }
  383. if err := i.Repo.ReplaceBatchImageItemsForJob(ctx, job.BatchID, items, BatchImageCounts{
  384. SuccessCount: result.SuccessCount,
  385. FailCount: result.FailCount,
  386. }); err != nil {
  387. return nil, err
  388. }
  389. return result, nil
  390. }
  391. // listExpectedCustomIDs 返回该 job 当前 item 表中的全部 custom_id 集合,
  392. // 即提交时预创建(或上一轮索引重建)的完整条目清单,用于与 provider 输出对账。
  393. func (i *BatchImageResultIndexer) listExpectedCustomIDs(ctx context.Context, batchID string) (map[string]struct{}, error) {
  394. const pageSize = 500
  395. expected := make(map[string]struct{})
  396. offset := 0
  397. for {
  398. page, err := i.Repo.ListBatchImageItems(ctx, batchID, BatchImageItemFilter{Limit: pageSize, Offset: offset})
  399. if err != nil {
  400. return nil, err
  401. }
  402. for _, item := range page {
  403. if item != nil {
  404. expected[item.CustomID] = struct{}{}
  405. }
  406. }
  407. if len(page) < pageSize {
  408. return expected, nil
  409. }
  410. offset += len(page)
  411. }
  412. }
  413. type ParsedBatchImageResult struct {
  414. CustomID string
  415. Status string
  416. MimeType string
  417. FileExtension string
  418. ImageCount int
  419. ErrorCode string
  420. ErrorMessage string
  421. SourceLineNumber int
  422. }
  423. func ParseBatchImageResultLine(line []byte, lineNumber int) (*ParsedBatchImageResult, error) {
  424. var obj map[string]any
  425. if err := json.Unmarshal(line, &obj); err != nil {
  426. return nil, ErrBatchImageIndexParseFailed.WithCause(fmt.Errorf("line %d: %w", lineNumber, err))
  427. }
  428. customID := batchImageFirstNonEmptyString(
  429. batchImageMapString(obj, "key"),
  430. batchImageMapString(obj, "custom_id"),
  431. batchImageMapString(obj, "customId"),
  432. batchImageNestedString(obj, "request", "key"),
  433. )
  434. if customID == "" {
  435. return nil, ErrBatchImageIndexParseFailed.WithCause(fmt.Errorf("line %d: missing custom id", lineNumber))
  436. }
  437. parsed := &ParsedBatchImageResult{
  438. CustomID: customID,
  439. SourceLineNumber: lineNumber,
  440. }
  441. imageCount, mimeType := batchImageFindImageParts(obj)
  442. if imageCount > 0 {
  443. parsed.Status = BatchImageParsedStatusSucceeded
  444. parsed.ImageCount = imageCount
  445. parsed.MimeType = mimeType
  446. parsed.FileExtension = batchImageFileExtension(mimeType)
  447. return parsed, nil
  448. }
  449. if code, message, ok := batchImageFailureFromProviderFields(obj); ok {
  450. parsed.Status = BatchImageParsedStatusFailed
  451. parsed.ErrorCode = code
  452. parsed.ErrorMessage = truncateBatchImageMessage(message, batchImageMaxErrorMessageLength)
  453. return parsed, nil
  454. }
  455. if _, hasResponse := obj["response"]; hasResponse || batchImageHasCandidates(obj) {
  456. parsed.Status = BatchImageParsedStatusFailed
  457. parsed.ErrorCode = "EMPTY_IMAGE_OUTPUT"
  458. parsed.ErrorMessage = "provider response contained no image output"
  459. return parsed, nil
  460. }
  461. parsed.Status = BatchImageParsedStatusFailed
  462. parsed.ErrorCode = "PROVIDER_ITEM_FAILED"
  463. parsed.ErrorMessage = "provider result line contained no image output"
  464. return parsed, nil
  465. }
  466. func batchImageFindImageParts(obj map[string]any) (int, string) {
  467. count, mimeType := batchImageFindImagePartsInCandidates(batchImageNestedAny(obj, "response", "candidates"))
  468. if count > 0 {
  469. return count, mimeType
  470. }
  471. return batchImageFindImagePartsInCandidates(obj["candidates"])
  472. }
  473. func batchImageFindImagePartsInCandidates(raw any) (int, string) {
  474. candidates, ok := raw.([]any)
  475. if !ok {
  476. return 0, ""
  477. }
  478. count := 0
  479. firstMime := ""
  480. for _, candidateRaw := range candidates {
  481. candidate, ok := candidateRaw.(map[string]any)
  482. if !ok {
  483. continue
  484. }
  485. partsRaw := batchImageNestedAny(candidate, "content", "parts")
  486. parts, ok := partsRaw.([]any)
  487. if !ok {
  488. continue
  489. }
  490. for _, partRaw := range parts {
  491. part, ok := partRaw.(map[string]any)
  492. if !ok {
  493. continue
  494. }
  495. inline, ok := firstMap(part["inlineData"], part["inline_data"])
  496. if !ok {
  497. continue
  498. }
  499. data := strings.TrimSpace(batchImageMapString(inline, "data"))
  500. mime := batchImageFirstNonEmptyString(batchImageMapString(inline, "mimeType"), batchImageMapString(inline, "mime_type"))
  501. if data == "" || !strings.HasPrefix(strings.ToLower(strings.TrimSpace(mime)), "image/") {
  502. continue
  503. }
  504. count++
  505. if firstMime == "" {
  506. firstMime = strings.TrimSpace(mime)
  507. }
  508. }
  509. }
  510. return count, firstMime
  511. }
  512. func batchImageFailureFromProviderFields(obj map[string]any) (string, string, bool) {
  513. if status, ok := obj["status"].(map[string]any); ok {
  514. message := batchImageFirstNonEmptyString(batchImageMapString(status, "message"), batchImageMapString(status, "details"))
  515. code := batchImageFirstNonEmptyString(batchImageMapString(status, "code"), batchImageMapString(status, "status"))
  516. return batchImageMapFailureCode(code, message), message, true
  517. }
  518. if errObj, ok := obj["error"].(map[string]any); ok {
  519. message := batchImageFirstNonEmptyString(batchImageMapString(errObj, "message"), batchImageMapString(errObj, "details"))
  520. code := batchImageFirstNonEmptyString(batchImageMapString(errObj, "code"), batchImageMapString(errObj, "status"))
  521. return batchImageMapFailureCode(code, message), message, true
  522. }
  523. return "", "", false
  524. }
  525. func batchImageMapFailureCode(code, message string) string {
  526. text := strings.ToLower(strings.TrimSpace(code + " " + message))
  527. switch {
  528. case strings.Contains(text, "safety"), strings.Contains(text, "policy"), strings.Contains(text, "blocked"), strings.Contains(text, "prohibited"):
  529. return "SAFETY_BLOCKED"
  530. case strings.Contains(text, "invalid_argument"), strings.Contains(text, "invalid argument"), strings.Contains(text, "bad request"):
  531. return "INVALID_ARGUMENT"
  532. case strings.Contains(text, "quota"), strings.Contains(text, "rate"), strings.Contains(text, "resource_exhausted"), strings.Contains(text, "too many requests"):
  533. return "PROVIDER_RATE_LIMITED"
  534. default:
  535. return "PROVIDER_ITEM_FAILED"
  536. }
  537. }
  538. func batchImageFileExtension(mimeType string) string {
  539. switch strings.ToLower(strings.TrimSpace(mimeType)) {
  540. case "image/png":
  541. return "png"
  542. case "image/jpeg", "image/jpg":
  543. return "jpg"
  544. case "image/webp":
  545. return "webp"
  546. default:
  547. return ""
  548. }
  549. }
  550. func batchImageHasCandidates(obj map[string]any) bool {
  551. if _, ok := obj["candidates"]; ok {
  552. return true
  553. }
  554. _, ok := batchImageNestedAny(obj, "response", "candidates").([]any)
  555. return ok
  556. }
  557. func batchImageMapString(m map[string]any, key string) string {
  558. if m == nil {
  559. return ""
  560. }
  561. switch v := m[key].(type) {
  562. case string:
  563. return strings.TrimSpace(v)
  564. case json.Number:
  565. return v.String()
  566. case float64:
  567. return strconv.FormatInt(int64(v), 10)
  568. default:
  569. return ""
  570. }
  571. }
  572. func batchImageNestedString(m map[string]any, keys ...string) string {
  573. if nested, ok := batchImageNestedAny(m, keys...).(string); ok {
  574. return strings.TrimSpace(nested)
  575. }
  576. return ""
  577. }
  578. func batchImageNestedAny(m map[string]any, keys ...string) any {
  579. var current any = m
  580. for _, key := range keys {
  581. cm, ok := current.(map[string]any)
  582. if !ok {
  583. return nil
  584. }
  585. current = cm[key]
  586. }
  587. return current
  588. }
  589. func firstMap(values ...any) (map[string]any, bool) {
  590. for _, value := range values {
  591. if m, ok := value.(map[string]any); ok {
  592. return m, true
  593. }
  594. }
  595. return nil, false
  596. }
  597. func batchImageFirstNonEmptyString(values ...string) string {
  598. for _, value := range values {
  599. if strings.TrimSpace(value) != "" {
  600. return strings.TrimSpace(value)
  601. }
  602. }
  603. return ""
  604. }
  605. func batchImageDerefString(v *string) string {
  606. if v == nil {
  607. return ""
  608. }
  609. return strings.TrimSpace(*v)
  610. }
  611. func batchImageStringPtr(v string) *string {
  612. return &v
  613. }
  614. func batchImageOptionalStringPtr(v string) *string {
  615. v = strings.TrimSpace(v)
  616. if v == "" {
  617. return nil
  618. }
  619. return &v
  620. }
  621. func truncateBatchImageMessage(message string, limit int) string {
  622. message = strings.TrimSpace(message)
  623. if limit <= 0 || len(message) <= limit {
  624. return message
  625. }
  626. return message[:limit]
  627. }