batch_image_public.go 49 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440
  1. package service
  2. import (
  3. "context"
  4. "crypto/sha256"
  5. "encoding/hex"
  6. "encoding/json"
  7. "errors"
  8. "fmt"
  9. "sort"
  10. "strconv"
  11. "strings"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/config"
  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. defaultBatchImageMaxItems = 200
  20. defaultBatchImageMaxOutputImages = 200
  21. defaultBatchImageMaxOutputCount = 4
  22. defaultBatchImageMaxPromptChars = 8000
  23. defaultBatchImageResponseMime = "image/png"
  24. defaultBatchImageImageSize = "1K"
  25. defaultBatchImageDiscountMultiplier = 0.5
  26. defaultBatchImageHoldMultiplier = 0.6
  27. maxBatchImagePublicErrorChars = 500
  28. maxBatchImageReferenceImageBytes = 10 * 1024 * 1024
  29. defaultBatchImageMaxReferenceImages = 1000
  30. defaultBatchImageMaxReferenceBytes = 128 * 1024 * 1024
  31. )
  32. type BatchImageAccountSelectionRepository interface {
  33. GetByID(ctx context.Context, id int64) (*Account, error)
  34. ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error)
  35. ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error)
  36. }
  37. type BatchImageGroupPricingRepository interface {
  38. GetByIDLite(ctx context.Context, id int64) (*Group, error)
  39. }
  40. type BatchImageUserGroupRateRepository interface {
  41. GetByUserAndGroup(ctx context.Context, userID, groupID int64) (*float64, error)
  42. }
  43. type BatchImageSubmitRequest struct {
  44. Model string `json:"model"`
  45. TaskName string `json:"task_name"`
  46. ParentBatchID string `json:"parent_batch_id"`
  47. Provider string `json:"provider"`
  48. Items []BatchImageSubmitItem `json:"items"`
  49. ResponseMimeType string `json:"response_mime_type"`
  50. AspectRatio string `json:"aspect_ratio"`
  51. ImageSize string `json:"image_size"`
  52. Metadata map[string]string `json:"metadata"`
  53. SessionID *string `json:"-"`
  54. }
  55. type BatchImageSubmitItem struct {
  56. CustomID string `json:"custom_id"`
  57. Prompt string `json:"prompt"`
  58. OutputCount int `json:"output_count,omitempty"`
  59. ReferenceImages []BatchImageReferenceInput `json:"reference_images,omitempty"`
  60. }
  61. type BatchImageReferenceInput struct {
  62. ID string `json:"id,omitempty"`
  63. Type string `json:"type,omitempty"`
  64. MimeType string `json:"mime_type"`
  65. Data []byte `json:"data,omitempty"`
  66. FileURI string `json:"file_uri,omitempty"`
  67. }
  68. type BatchImageOwner struct {
  69. UserID int64
  70. APIKeyID int64
  71. GroupID *int64
  72. }
  73. type BatchImagePublicService struct {
  74. Repo BatchImageRepository
  75. AccountRepo BatchImageAccountSelectionRepository
  76. GroupRepo BatchImageGroupPricingRepository
  77. UserGroupRateRepo BatchImageUserGroupRateRepository
  78. Queue BatchImageQueue
  79. ProviderRegistry *BatchImageProviderRegistry
  80. Pricing BatchImagePricingResolver
  81. BillingRepo UsageBillingRepository
  82. AuthCache APIKeyAuthCacheInvalidator
  83. Config *config.Config
  84. }
  85. type BatchImagePricingSnapshot struct {
  86. BaseUnitPrice float64
  87. GroupRateMultiplier float64
  88. AccountRateMultiplier float64
  89. BatchDiscountMultiplier float64
  90. HoldMultiplier float64
  91. BillableUnitPrice float64
  92. HoldUnitPrice float64
  93. EstimatedCost float64
  94. HoldAmount float64
  95. }
  96. type BatchImagePublicBatch struct {
  97. ID string `json:"id"`
  98. Object string `json:"object"`
  99. TaskName string `json:"task_name"`
  100. ParentBatchID *string `json:"parent_batch_id,omitempty"`
  101. Status string `json:"status"`
  102. Model string `json:"model"`
  103. Provider string `json:"provider"`
  104. ItemCount int `json:"item_count"`
  105. SuccessCount int `json:"success_count"`
  106. FailCount int `json:"fail_count"`
  107. EstimatedCost float64 `json:"estimated_cost"`
  108. HoldAmount float64 `json:"hold_amount"`
  109. ActualCost *float64 `json:"actual_cost"`
  110. CreatedAt int64 `json:"created_at"`
  111. SubmittedAt *int64 `json:"submitted_at"`
  112. SettledAt *int64 `json:"settled_at"`
  113. DownloadedAt *int64 `json:"downloaded_at,omitempty"`
  114. OutputDeletedAt *int64 `json:"output_deleted_at,omitempty"`
  115. }
  116. type BatchImagePublicItem struct {
  117. CustomID string `json:"custom_id"`
  118. Status string `json:"status"`
  119. PromptPreview *string `json:"prompt_preview,omitempty"`
  120. MimeType *string `json:"mime_type"`
  121. FileExtension *string `json:"file_extension"`
  122. ImageCount int `json:"image_count"`
  123. Error *BatchImagePublicError `json:"error"`
  124. }
  125. type BatchImagePublicError struct {
  126. Code string `json:"code"`
  127. Message string `json:"message"`
  128. Source string `json:"source,omitempty"`
  129. }
  130. type BatchImagePublicItemsResponse struct {
  131. Object string `json:"object"`
  132. Data []BatchImagePublicItem `json:"data"`
  133. HasMore bool `json:"has_more"`
  134. }
  135. type BatchImagePublicListResponse struct {
  136. Object string `json:"object"`
  137. Data []*BatchImagePublicBatch `json:"data"`
  138. HasMore bool `json:"has_more"`
  139. }
  140. type BatchImagePublicModel struct {
  141. ID string `json:"id"`
  142. Object string `json:"object"`
  143. Provider string `json:"provider"`
  144. }
  145. type BatchImagePublicModelsResponse struct {
  146. Object string `json:"object"`
  147. Data []BatchImagePublicModel `json:"data"`
  148. }
  149. type BatchImageJobsQuery struct {
  150. Status string
  151. TaskName string
  152. Downloaded string
  153. From string
  154. To string
  155. Limit int
  156. Cursor string
  157. }
  158. type BatchImageItemsQuery struct {
  159. Status string
  160. Limit int
  161. Cursor string
  162. }
  163. func NewBatchImagePublicService(repo BatchImageRepository, accountRepo AccountRepository, groupRepo GroupRepository, userGroupRateRepo UserGroupRateRepository, queue BatchImageQueue, pricing *BatchImageModelPricingResolver, billingRepo UsageBillingRepository, authCache APIKeyAuthCacheInvalidator, cfg *config.Config) *BatchImagePublicService {
  164. return &BatchImagePublicService{
  165. Repo: repo,
  166. AccountRepo: accountRepo,
  167. GroupRepo: groupRepo,
  168. UserGroupRateRepo: userGroupRateRepo,
  169. Queue: queue,
  170. ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg),
  171. Pricing: pricing,
  172. BillingRepo: billingRepo,
  173. AuthCache: authCache,
  174. Config: cfg,
  175. }
  176. }
  177. func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOwner, req BatchImageSubmitRequest, idempotencyKey string) (*BatchImagePublicBatch, error) {
  178. if !s.enabled() {
  179. return nil, ErrBatchImageDisabled
  180. }
  181. normalized, err := s.validateSubmitRequest(req)
  182. if err != nil {
  183. return nil, err
  184. }
  185. // 与 ListModels 使用同一鉴权谓词(AllowBatchImageGeneration + Platform==Gemini),
  186. // 避免两个入口校验口径不一致留下防御纵深缺口。
  187. if err := s.ensureGroupAllowsBatchImage(ctx, owner.GroupID); err != nil {
  188. return nil, err
  189. }
  190. requestHash := HashBatchImageSubmitRequest(normalized)
  191. idempotencyKey = strings.TrimSpace(idempotencyKey)
  192. if idempotencyKey != "" {
  193. existing, err := s.Repo.GetBatchImageJobByIdempotencyKey(ctx, owner.UserID, owner.APIKeyID, idempotencyKey)
  194. if err == nil {
  195. if batchImageDerefString(existing.RequestHash) != requestHash {
  196. return nil, ErrBatchImageIdempotencyConflict
  197. }
  198. if existing.Status == BatchImageJobStatusSubmitted && s.Queue != nil {
  199. if enqueueErr := s.Queue.Enqueue(ctx, existing.BatchID); enqueueErr != nil && !errors.Is(enqueueErr, ErrBatchImageAlreadyQueued) {
  200. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, existing.BatchID, "QUEUE_FAILED", sanitizeBatchImagePublicMessage(enqueueErr.Error()), false)
  201. return nil, ErrBatchImageQueueFailed
  202. }
  203. }
  204. return BatchImageJobToPublic(existing), nil
  205. }
  206. if !errors.Is(err, ErrBatchImageJobNotFound) {
  207. return nil, err
  208. }
  209. }
  210. provider, account, err := s.selectProviderAndAccount(ctx, owner, normalized.Provider, normalized.Model)
  211. if err != nil {
  212. return nil, err
  213. }
  214. pricingSnapshot, err := s.resolvePricingSnapshot(ctx, owner, normalized, provider.Name(), account)
  215. if err != nil {
  216. return nil, err
  217. }
  218. parentBatchID := batchImageOptionalStringPtr(normalized.ParentBatchID)
  219. if parentBatchID != nil {
  220. parent, parentErr := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, *parentBatchID)
  221. if parentErr != nil {
  222. return nil, parentErr
  223. }
  224. if parent.ParentBatchID != nil && strings.TrimSpace(*parent.ParentBatchID) != "" {
  225. parentBatchID = batchImageOptionalStringPtr(*parent.ParentBatchID)
  226. }
  227. }
  228. batchID, err := NewBatchImageID()
  229. if err != nil {
  230. return nil, err
  231. }
  232. apiKeyID := owner.APIKeyID
  233. accountID := account.ID
  234. holdID := BatchImageHoldRequestID(batchID)
  235. holdAmount := pricingSnapshot.HoldAmount
  236. job, err := s.Repo.CreateBatchImageJob(ctx, CreateBatchImageJobParams{
  237. BatchID: batchID,
  238. UserID: owner.UserID,
  239. APIKeyID: &apiKeyID,
  240. AccountID: &accountID,
  241. Provider: provider.Name(),
  242. Model: normalized.Model,
  243. TaskName: normalized.TaskName,
  244. ParentBatchID: parentBatchID,
  245. Status: BatchImageJobStatusCreated,
  246. ItemCount: len(normalized.Items),
  247. EstimatedCost: pricingSnapshot.EstimatedCost,
  248. HoldAmount: &holdAmount,
  249. BaseUnitPrice: pricingSnapshot.BaseUnitPrice,
  250. GroupRateMultiplier: pricingSnapshot.GroupRateMultiplier,
  251. AccountRateMultiplier: pricingSnapshot.AccountRateMultiplier,
  252. BatchDiscountMultiplier: pricingSnapshot.BatchDiscountMultiplier,
  253. HoldMultiplier: pricingSnapshot.HoldMultiplier,
  254. BillableUnitPrice: pricingSnapshot.BillableUnitPrice,
  255. HoldUnitPrice: pricingSnapshot.HoldUnitPrice,
  256. PricingSnapshotVersion: 1,
  257. Currency: "USD",
  258. HoldID: &holdID,
  259. IdempotencyKey: batchImageOptionalStringPtr(idempotencyKey),
  260. RequestHash: batchImageStringPtr(requestHash),
  261. SessionID: normalized.SessionID,
  262. })
  263. if err != nil {
  264. return nil, err
  265. }
  266. if err := reserveBatchImageBalanceHold(ctx, s.BillingRepo, job, requestHash); err != nil {
  267. code := "BILLING_HOLD_FAILED"
  268. if errors.Is(err, ErrBatchImageInsufficientBalance) {
  269. code = "INSUFFICIENT_BALANCE"
  270. }
  271. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, code, sanitizeBatchImagePublicMessage(err.Error()), true)
  272. s.hidePreUpstreamSubmitFailure(ctx, owner, job)
  273. return nil, err
  274. }
  275. s.invalidateAuthCache(ctx, owner.UserID)
  276. if err := s.createPendingItems(ctx, job.BatchID, requestHash, normalized.Items); err != nil {
  277. if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
  278. return nil, releaseErr
  279. }
  280. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "ITEM_CREATE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true)
  281. s.hidePreUpstreamSubmitFailure(ctx, owner, job)
  282. return nil, ErrBatchImageQueueFailed
  283. }
  284. input := BatchImageInput{
  285. BatchID: job.BatchID,
  286. Model: normalized.Model,
  287. DisplayName: job.BatchID,
  288. ResponseMimeType: normalized.ResponseMimeType,
  289. AspectRatio: normalized.AspectRatio,
  290. ImageSize: normalized.ImageSize,
  291. Metadata: normalized.Metadata,
  292. Items: make([]BatchImageInputItem, 0, len(normalized.Items)),
  293. }
  294. for _, item := range normalized.Items {
  295. refs := make([]BatchImageReference, 0, len(item.ReferenceImages))
  296. for _, ref := range item.ReferenceImages {
  297. refs = append(refs, BatchImageReference(ref))
  298. }
  299. input.Items = append(input.Items, BatchImageInputItem{
  300. CustomID: item.CustomID,
  301. Prompt: item.Prompt,
  302. ReferenceImages: refs,
  303. })
  304. }
  305. // 上游提交(上传参考图 + 创建批任务)可能长达数分钟且不刷新 updated_at,
  306. // 会被 stale 恢复扫描误判为滞留并退款。提交前转入 uploading 刷新时间戳,
  307. // 提交期间用心跳持续续期。
  308. if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusUploading, BatchImageTransitionOptions{
  309. EventType: "upload_started",
  310. EventPayload: map[string]any{"batch_id": job.BatchID},
  311. }); err != nil {
  312. if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
  313. return nil, releaseErr
  314. }
  315. // 并发 Cancel 等导致的非法转换:job 已处于终态,不再覆盖其状态。
  316. if !errors.Is(err, ErrBatchImageInvalidTransition) {
  317. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "UPLOAD_TRANSITION_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true)
  318. s.hidePreUpstreamSubmitFailure(ctx, owner, job)
  319. }
  320. return nil, err
  321. }
  322. job.Status = BatchImageJobStatusUploading
  323. hbCtx, hbCancel := context.WithCancel(ctx)
  324. hbDone := make(chan struct{})
  325. go s.runSubmitHeartbeat(hbCtx, job.BatchID, hbDone)
  326. providerJob, err := provider.Submit(ctx, job, account, input)
  327. hbCancel()
  328. <-hbDone
  329. if err != nil {
  330. if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
  331. return nil, releaseErr
  332. }
  333. publicErr := batchImageProviderSubmitPublicError(err)
  334. reason := batchImageProviderSubmitRecordCode(publicErr)
  335. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, reason, sanitizeBatchImagePublicMessage(err.Error()), true)
  336. s.hidePreUpstreamSubmitFailure(ctx, owner, job)
  337. return nil, publicErr
  338. }
  339. if providerJob == nil || strings.TrimSpace(providerJob.ProviderJobName) == "" {
  340. if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
  341. return nil, releaseErr
  342. }
  343. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "PROVIDER_SUBMIT_FAILED", "provider job name missing", true)
  344. s.hidePreUpstreamSubmitFailure(ctx, owner, job)
  345. return nil, ErrBatchImageProviderSubmitFailed
  346. }
  347. if err := s.Repo.UpdateBatchImageJobProviderSubmit(ctx, UpdateBatchImageJobProviderSubmitParams{
  348. BatchID: job.BatchID,
  349. ProviderJobName: providerJob.ProviderJobName,
  350. ProviderInputRef: providerJob.ProviderInputRef,
  351. ProviderOutputRef: providerJob.ProviderOutputRef,
  352. GCSInputURI: batchImageGCSRef(provider.Name(), providerJob.ProviderInputRef),
  353. GCSOutputURI: batchImageGCSRef(provider.Name(), providerJob.ProviderOutputRef),
  354. EventPayload: map[string]any{"provider": provider.Name()},
  355. }); err != nil {
  356. // job 可能已被恢复扫描转 failed 并退款:上游批任务已创建成功,
  357. // 必须尽力取消并清理输入,否则上游照常产生成本(孤儿任务)。
  358. s.abortOrphanProviderJob(ctx, provider, job, account, providerJob)
  359. return nil, err
  360. }
  361. if s.Queue != nil {
  362. if err := s.Queue.Enqueue(ctx, job.BatchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
  363. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "QUEUE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), false)
  364. return nil, ErrBatchImageQueueFailed
  365. }
  366. }
  367. created, err := s.Repo.GetBatchImageJobByBatchID(ctx, job.BatchID)
  368. if err != nil {
  369. return nil, err
  370. }
  371. return BatchImageJobToPublic(created), nil
  372. }
  373. func (s *BatchImagePublicService) releaseFailedSubmitHold(ctx context.Context, job *BatchImageJob, requestHash string) error {
  374. if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, requestHash); err != nil {
  375. _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "BILLING_RELEASE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true)
  376. s.enqueueBillingRetry(ctx, job.BatchID)
  377. return ErrBatchImageBillingHoldFailed
  378. }
  379. s.invalidateAuthCache(ctx, job.UserID)
  380. return nil
  381. }
  382. // runSubmitHeartbeat 在 provider.Submit 期间周期性刷新 job 的 updated_at,
  383. // 使 stale 恢复扫描能区分"仍在慢提交"与"进程死亡后的滞留"。
  384. func (s *BatchImagePublicService) runSubmitHeartbeat(ctx context.Context, batchID string, done chan<- struct{}) {
  385. defer close(done)
  386. interval := s.submitHeartbeatInterval()
  387. ticker := time.NewTicker(interval)
  388. defer ticker.Stop()
  389. for {
  390. select {
  391. case <-ctx.Done():
  392. return
  393. case <-ticker.C:
  394. if err := s.Repo.TouchBatchImageJobSubmitting(ctx, batchID); err != nil && ctx.Err() == nil {
  395. logger.L().Warn("batch_image.submit_heartbeat_failed",
  396. zap.String("batch_id", batchID),
  397. zap.Error(err),
  398. )
  399. }
  400. }
  401. }
  402. }
  403. func (s *BatchImagePublicService) submitHeartbeatInterval() time.Duration {
  404. staleAfter := 10 * time.Minute
  405. if s != nil && s.Config != nil && s.Config.BatchImage.StaleActiveAfterSeconds > 0 {
  406. staleAfter = time.Duration(s.Config.BatchImage.StaleActiveAfterSeconds) * time.Second
  407. }
  408. interval := staleAfter / 3
  409. if interval < 15*time.Second {
  410. interval = 15 * time.Second
  411. }
  412. return interval
  413. }
  414. // abortOrphanProviderJob 在上游任务创建成功但本地状态推进失败时,
  415. // 尽力取消上游批任务并清理已上传的输入文件,避免孤儿任务持续产生成本。
  416. func (s *BatchImagePublicService) abortOrphanProviderJob(ctx context.Context, provider BatchImageProvider, job *BatchImageJob, account *Account, providerJob *BatchProviderJob) {
  417. if s == nil || provider == nil || job == nil || providerJob == nil {
  418. return
  419. }
  420. orphan := *job
  421. orphan.ProviderJobName = batchImageOptionalStringPtr(providerJob.ProviderJobName)
  422. orphan.ProviderInputRef = batchImageOptionalStringPtr(providerJob.ProviderInputRef)
  423. orphan.GCSInputURI = batchImageOptionalStringPtr(batchImageGCSRef(provider.Name(), providerJob.ProviderInputRef))
  424. if err := provider.Cancel(ctx, &orphan, account); err != nil {
  425. logger.L().Warn("batch_image.orphan_provider_job_cancel_failed",
  426. zap.String("batch_id", job.BatchID),
  427. zap.String("provider", provider.Name()),
  428. zap.Error(err),
  429. )
  430. }
  431. if err := provider.Cleanup(ctx, &orphan, account, CleanupTargetInput); err != nil {
  432. logger.L().Warn("batch_image.orphan_provider_job_cleanup_failed",
  433. zap.String("batch_id", job.BatchID),
  434. zap.String("provider", provider.Name()),
  435. zap.Error(err),
  436. )
  437. }
  438. if err := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "provider_job_aborted_after_submit", map[string]any{
  439. "batch_id": job.BatchID,
  440. "provider": provider.Name(),
  441. }); err != nil {
  442. logger.L().Warn("batch_image.orphan_provider_job_event_failed",
  443. zap.String("batch_id", job.BatchID),
  444. zap.Error(err),
  445. )
  446. }
  447. }
  448. func (s *BatchImagePublicService) createPendingItems(ctx context.Context, batchID, requestHash string, items []BatchImageSubmitItem) error {
  449. if s == nil || s.Repo == nil || len(items) == 0 {
  450. return nil
  451. }
  452. params := make([]CreateBatchImageItemParams, 0, len(items))
  453. for _, item := range items {
  454. preview := truncateBatchImageMessage(item.Prompt, s.maxPromptChars())
  455. params = append(params, CreateBatchImageItemParams{
  456. JobID: batchID,
  457. CustomID: item.CustomID,
  458. Status: BatchImageItemStatusPending,
  459. RequestHash: batchImageStringPtr(requestHash),
  460. PromptPreview: batchImageStringPtr(preview),
  461. ImageCount: 0,
  462. })
  463. }
  464. return s.Repo.BulkCreateBatchImageItems(ctx, params)
  465. }
  466. func (s *BatchImagePublicService) enqueueBillingRetry(ctx context.Context, batchID string) {
  467. if s == nil || s.Queue == nil {
  468. return
  469. }
  470. if err := s.Queue.Enqueue(ctx, batchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
  471. logger.L().Warn("batch_image.billing_retry_enqueue_failed",
  472. zap.String("batch_id", batchID),
  473. zap.Error(err),
  474. )
  475. if eventErr := s.Repo.AppendBatchImageEvent(ctx, batchID, "billing_retry_enqueue_failed", map[string]any{
  476. "batch_id": batchID,
  477. "error": sanitizeBatchImagePublicMessage(err.Error()),
  478. }); eventErr != nil {
  479. logger.L().Warn("batch_image.billing_retry_event_failed",
  480. zap.String("batch_id", batchID),
  481. zap.Error(eventErr),
  482. )
  483. }
  484. }
  485. }
  486. func (s *BatchImagePublicService) hidePreUpstreamSubmitFailure(ctx context.Context, owner BatchImageOwner, job *BatchImageJob) {
  487. if s == nil || s.Repo == nil || job == nil || job.ProviderJobName != nil {
  488. return
  489. }
  490. if err := s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now()); err != nil {
  491. logger.L().Warn("batch_image.hide_pre_upstream_failure_failed",
  492. zap.String("batch_id", job.BatchID),
  493. zap.Error(err),
  494. )
  495. }
  496. }
  497. func (s *BatchImagePublicService) Get(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) {
  498. job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  499. if err != nil {
  500. return nil, err
  501. }
  502. return BatchImageJobToPublic(job), nil
  503. }
  504. func (s *BatchImagePublicService) List(ctx context.Context, owner BatchImageOwner, query BatchImageJobsQuery) (*BatchImagePublicListResponse, error) {
  505. filter := BatchImageJobFilter{Limit: query.Limit, Offset: parseBatchImageCursor(query.Cursor), ExcludeDeleted: true}
  506. filter.TaskNameLike = strings.TrimSpace(query.TaskName)
  507. switch strings.TrimSpace(query.Status) {
  508. case "", "all":
  509. case "queued":
  510. filter.Status = BatchImageJobStatusSubmitted
  511. case "processing_results":
  512. filter.Status = BatchImageJobStatusIndexing
  513. case "completed":
  514. filter.Status = BatchImageJobStatusCompleted
  515. case "failed":
  516. filter.Status = BatchImageJobStatusFailed
  517. case "cancelled":
  518. filter.Status = BatchImageJobStatusCancelled
  519. case "output_deleted":
  520. filter.Status = BatchImageJobStatusOutputDeleted
  521. default:
  522. filter.Status = strings.TrimSpace(query.Status)
  523. }
  524. switch strings.TrimSpace(strings.ToLower(query.Downloaded)) {
  525. case "", "all":
  526. case "true", "1", "yes", "downloaded":
  527. downloaded := true
  528. filter.Downloaded = &downloaded
  529. case "false", "0", "no", "not_downloaded":
  530. downloaded := false
  531. filter.Downloaded = &downloaded
  532. default:
  533. return nil, ErrBatchImageInvalidItems
  534. }
  535. if from := parseBatchImageListTime(query.From); from != nil {
  536. filter.CreatedAfter = from
  537. }
  538. if to := parseBatchImageListTime(query.To); to != nil {
  539. filter.CreatedBefore = to
  540. }
  541. if filter.Limit <= 0 || filter.Limit > 100 {
  542. filter.Limit = 20
  543. }
  544. jobs, err := s.Repo.ListBatchImageJobsForOwner(ctx, owner.UserID, owner.APIKeyID, filter)
  545. if err != nil {
  546. return nil, err
  547. }
  548. data := make([]*BatchImagePublicBatch, 0, len(jobs))
  549. for _, job := range jobs {
  550. data = append(data, BatchImageJobToPublic(job))
  551. }
  552. return &BatchImagePublicListResponse{
  553. Object: "list",
  554. Data: data,
  555. HasMore: len(data) == filter.Limit,
  556. }, nil
  557. }
  558. func (s *BatchImagePublicService) MarkDownloaded(ctx context.Context, owner BatchImageOwner, batchID string) error {
  559. job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  560. if err != nil {
  561. return err
  562. }
  563. return s.Repo.MarkBatchImageDownloaded(ctx, job.BatchID, time.Now())
  564. }
  565. func (s *BatchImagePublicService) DeleteRecord(ctx context.Context, owner BatchImageOwner, batchID string) error {
  566. job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  567. if err != nil {
  568. return err
  569. }
  570. if !isBatchImageProcessorDoneStatus(job.Status) {
  571. return ErrBatchImageRecordDeleteNotReady
  572. }
  573. return s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now())
  574. }
  575. func (s *BatchImagePublicService) ListModels(ctx context.Context, owner BatchImageOwner) (*BatchImagePublicModelsResponse, error) {
  576. if !s.enabled() {
  577. return nil, ErrBatchImageDisabled
  578. }
  579. if s.Pricing == nil {
  580. return nil, ErrBatchImageSettlementPricingMissing
  581. }
  582. if err := s.ensureGroupAllowsBatchImage(ctx, owner.GroupID); err != nil {
  583. return nil, err
  584. }
  585. modelsByProvider := make(map[string]map[string]struct{})
  586. for _, providerName := range batchImageProviderSelectionOrder("") {
  587. provider, ok := s.ProviderRegistry.Get(providerName)
  588. if !ok || provider == nil {
  589. continue
  590. }
  591. accounts, err := s.listCandidateAccounts(ctx, owner.GroupID, batchImageProviderPlatform(providerName))
  592. if err != nil {
  593. return nil, err
  594. }
  595. for i := range accounts {
  596. account := accounts[i]
  597. if !account.IsSchedulable() || !provider.SupportsAccount(&account) {
  598. continue
  599. }
  600. for _, model := range batchImageModelsFromAccountMapping(&account) {
  601. if _, err := s.Pricing.BatchImageUnitPrice(ctx, &BatchImageJob{Provider: providerName, Model: model}); err != nil {
  602. continue
  603. }
  604. if !account.IsModelSupported(model) {
  605. continue
  606. }
  607. if modelsByProvider[providerName] == nil {
  608. modelsByProvider[providerName] = make(map[string]struct{})
  609. }
  610. modelsByProvider[providerName][model] = struct{}{}
  611. }
  612. }
  613. }
  614. out := make([]BatchImagePublicModel, 0)
  615. for _, providerName := range batchImageProviderSelectionOrder("") {
  616. models := make([]string, 0, len(modelsByProvider[providerName]))
  617. for model := range modelsByProvider[providerName] {
  618. models = append(models, model)
  619. }
  620. sort.Strings(models)
  621. for _, model := range models {
  622. out = append(out, BatchImagePublicModel{
  623. ID: model,
  624. Object: "image.batch.model",
  625. Provider: providerName,
  626. })
  627. }
  628. }
  629. return &BatchImagePublicModelsResponse{Object: "list", Data: out}, nil
  630. }
  631. func (s *BatchImagePublicService) ListItems(ctx context.Context, owner BatchImageOwner, batchID string, query BatchImageItemsQuery) (*BatchImagePublicItemsResponse, error) {
  632. filter := BatchImageItemFilter{Limit: query.Limit, Offset: parseBatchImageCursor(query.Cursor)}
  633. switch strings.TrimSpace(query.Status) {
  634. case "", "all":
  635. case "succeeded", "success":
  636. filter.Status = BatchImageItemStatusSuccess
  637. case "pending":
  638. filter.Status = BatchImageItemStatusPending
  639. case "failed":
  640. filter.Status = BatchImageItemStatusFailed
  641. default:
  642. return nil, ErrBatchImageInvalidItems
  643. }
  644. if filter.Limit <= 0 || filter.Limit > 500 {
  645. filter.Limit = 100
  646. }
  647. items, err := s.Repo.ListBatchImageItemsForOwner(ctx, owner.UserID, owner.APIKeyID, batchID, filter)
  648. if err != nil {
  649. return nil, err
  650. }
  651. data := make([]BatchImagePublicItem, 0, len(items))
  652. for _, item := range items {
  653. data = append(data, BatchImageItemToPublic(item))
  654. }
  655. return &BatchImagePublicItemsResponse{
  656. Object: "list",
  657. Data: data,
  658. HasMore: len(data) == filter.Limit,
  659. }, nil
  660. }
  661. func (s *BatchImagePublicService) Cancel(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) {
  662. job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  663. if err != nil {
  664. return nil, err
  665. }
  666. if isBatchImageProcessorDoneStatus(job.Status) {
  667. if job.Status == BatchImageJobStatusFailed || job.Status == BatchImageJobStatusCancelled {
  668. if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil {
  669. s.enqueueBillingRetry(ctx, job.BatchID)
  670. return nil, ErrBatchImageCancelFailed
  671. }
  672. s.invalidateAuthCache(ctx, owner.UserID)
  673. }
  674. return BatchImageJobToPublic(job), nil
  675. }
  676. if job.ProviderJobName != nil && strings.TrimSpace(*job.ProviderJobName) != "" {
  677. provider, ok := s.ProviderRegistry.Get(job.Provider)
  678. if !ok || provider == nil {
  679. return nil, ErrBatchImageUnsupportedProvider
  680. }
  681. if job.AccountID == nil {
  682. return nil, ErrBatchImageCancelFailed
  683. }
  684. account, err := s.AccountRepo.GetByID(ctx, *job.AccountID)
  685. if err != nil {
  686. return nil, ErrBatchImageCancelFailed
  687. }
  688. if err := provider.Cancel(ctx, job, account); err != nil {
  689. return nil, ErrBatchImageCancelFailed
  690. }
  691. if eventErr := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "job_cancel_requested", map[string]any{"batch_id": job.BatchID}); eventErr != nil {
  692. logger.L().Warn("batch_image.cancel_event_failed",
  693. zap.String("batch_id", job.BatchID),
  694. zap.Error(eventErr),
  695. )
  696. }
  697. if s.Queue != nil {
  698. if err := s.Queue.Enqueue(ctx, job.BatchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
  699. return nil, ErrBatchImageCancelFailed
  700. }
  701. }
  702. updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  703. if err != nil {
  704. return nil, err
  705. }
  706. return BatchImageJobToPublic(updated), nil
  707. }
  708. if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusCancelled, BatchImageTransitionOptions{
  709. EventType: "job_cancelled",
  710. EventPayload: map[string]any{"batch_id": job.BatchID},
  711. }); err != nil {
  712. return nil, err
  713. }
  714. if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil {
  715. s.enqueueBillingRetry(ctx, job.BatchID)
  716. return nil, ErrBatchImageCancelFailed
  717. }
  718. s.invalidateAuthCache(ctx, owner.UserID)
  719. updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  720. if err != nil {
  721. return nil, err
  722. }
  723. return BatchImageJobToPublic(updated), nil
  724. }
  725. func (s *BatchImagePublicService) validateSubmitRequest(req BatchImageSubmitRequest) (BatchImageSubmitRequest, error) {
  726. req.Model = strings.TrimSpace(req.Model)
  727. req.TaskName = strings.TrimSpace(req.TaskName)
  728. req.ParentBatchID = strings.TrimSpace(req.ParentBatchID)
  729. req.Provider = strings.TrimSpace(req.Provider)
  730. req.ResponseMimeType = strings.TrimSpace(req.ResponseMimeType)
  731. req.AspectRatio = strings.TrimSpace(req.AspectRatio)
  732. req.ImageSize = strings.TrimSpace(req.ImageSize)
  733. if req.Model == "" {
  734. return req, ErrBatchImageInvalidModel
  735. }
  736. if req.TaskName == "" {
  737. req.TaskName = defaultBatchImageTaskName(time.Now())
  738. }
  739. if len(req.TaskName) > 255 {
  740. req.TaskName = truncateBatchImageMessage(req.TaskName, 255)
  741. }
  742. if req.Provider != "" && !IsSupportedBatchImageProvider(req.Provider) {
  743. return req, ErrBatchImageUnsupportedProvider
  744. }
  745. if len(req.Items) == 0 {
  746. return req, ErrBatchImageInvalidItems
  747. }
  748. maxItems := s.maxItems()
  749. if len(req.Items) > maxItems {
  750. return req, ErrBatchImageInvalidItems
  751. }
  752. if req.ResponseMimeType == "" {
  753. req.ResponseMimeType = s.defaultResponseMimeType()
  754. }
  755. if req.ImageSize == "" {
  756. req.ImageSize = s.defaultImageSize()
  757. }
  758. if !strings.EqualFold(req.ImageSize, defaultBatchImageImageSize) {
  759. return req, ErrBatchImageInvalidItems
  760. }
  761. req.ImageSize = defaultBatchImageImageSize
  762. req.Metadata = sanitizeBatchImageMetadata(req.Metadata)
  763. seen := make(map[string]struct{}, len(req.Items))
  764. totalReferenceImages := 0
  765. totalInlineReferenceBytes := 0
  766. totalOutputImages := 0
  767. expandedItems := make([]BatchImageSubmitItem, 0, len(req.Items))
  768. for i := range req.Items {
  769. req.Items[i].CustomID = strings.TrimSpace(req.Items[i].CustomID)
  770. if req.Items[i].CustomID == "" {
  771. req.Items[i].CustomID = fmt.Sprintf("item_%06d", i+1)
  772. }
  773. outputCount := req.Items[i].OutputCount
  774. if outputCount == 0 {
  775. outputCount = 1
  776. }
  777. if outputCount < 1 || outputCount > s.maxOutputImagesPerItem() {
  778. return req, ErrBatchImageInvalidItems
  779. }
  780. totalOutputImages += outputCount
  781. if totalOutputImages > s.maxOutputImagesPerJob() {
  782. return req, ErrBatchImageTooManyOutputImages
  783. }
  784. req.Items[i].Prompt = strings.TrimSpace(req.Items[i].Prompt)
  785. if req.Items[i].Prompt == "" {
  786. return req, ErrBatchImageInvalidItems
  787. }
  788. if len(req.Items[i].Prompt) > s.maxPromptChars() {
  789. return req, ErrBatchImagePromptTooLong
  790. }
  791. referenceCount, inlineReferenceBytes, err := normalizeBatchImageReferenceInputs(req.Model, &req.Items[i])
  792. if err != nil {
  793. return req, err
  794. }
  795. totalReferenceImages += referenceCount * outputCount
  796. if totalReferenceImages > s.maxReferenceImagesPerJob() {
  797. return req, ErrBatchImageTooManyReferenceImages
  798. }
  799. totalInlineReferenceBytes += inlineReferenceBytes * outputCount
  800. if totalInlineReferenceBytes > s.maxReferenceInlineBytesPerJob() {
  801. return req, ErrBatchImageReferenceImagesTooLarge
  802. }
  803. for repeatIndex := 1; repeatIndex <= outputCount; repeatIndex++ {
  804. expanded := req.Items[i]
  805. expanded.OutputCount = 0
  806. if outputCount > 1 {
  807. expanded.CustomID = fmt.Sprintf("%s_%0*d", req.Items[i].CustomID, batchImageRepeatSuffixWidth(outputCount), repeatIndex)
  808. }
  809. if _, ok := seen[expanded.CustomID]; ok {
  810. return req, ErrBatchImageDuplicateCustomIDInRequest
  811. }
  812. seen[expanded.CustomID] = struct{}{}
  813. expandedItems = append(expandedItems, expanded)
  814. }
  815. }
  816. req.Items = expandedItems
  817. return req, nil
  818. }
  819. func normalizeBatchImageReferenceInputs(model string, item *BatchImageSubmitItem) (int, int, error) {
  820. if item == nil || len(item.ReferenceImages) == 0 {
  821. return 0, 0, nil
  822. }
  823. maxRefs := maxBatchImageReferenceImagesForModel(model)
  824. if maxRefs <= 0 || len(item.ReferenceImages) > maxRefs {
  825. return 0, 0, ErrBatchImageTooManyReferenceImages
  826. }
  827. out := make([]BatchImageReferenceInput, 0, len(item.ReferenceImages))
  828. inlineBytes := 0
  829. for _, ref := range item.ReferenceImages {
  830. ref.ID = truncateBatchImageMessage(strings.TrimSpace(ref.ID), 80)
  831. ref.Type = truncateBatchImageMessage(strings.TrimSpace(ref.Type), 40)
  832. ref.MimeType = normalizeBatchImageReferenceMimeType(ref.MimeType)
  833. ref.FileURI = strings.TrimSpace(ref.FileURI)
  834. if ref.MimeType == "" {
  835. return 0, 0, ErrBatchImageInvalidReferenceImage
  836. }
  837. if len(ref.Data) == 0 && ref.FileURI == "" {
  838. return 0, 0, ErrBatchImageInvalidReferenceImage
  839. }
  840. if len(ref.Data) > 0 && ref.FileURI != "" {
  841. return 0, 0, ErrBatchImageInvalidReferenceImage
  842. }
  843. if len(ref.Data) > maxBatchImageReferenceImageBytes {
  844. return 0, 0, ErrBatchImageInvalidReferenceImage
  845. }
  846. if ref.FileURI != "" && !strings.HasPrefix(ref.FileURI, "gs://") {
  847. return 0, 0, ErrBatchImageInvalidReferenceImage
  848. }
  849. inlineBytes += len(ref.Data)
  850. out = append(out, ref)
  851. }
  852. item.ReferenceImages = out
  853. return len(out), inlineBytes, nil
  854. }
  855. func normalizeBatchImageReferenceMimeType(v string) string {
  856. switch strings.ToLower(strings.TrimSpace(v)) {
  857. case "image/jpeg", "image/jpg":
  858. return "image/jpeg"
  859. case "image/png":
  860. return "image/png"
  861. case "image/webp":
  862. return "image/webp"
  863. default:
  864. return ""
  865. }
  866. }
  867. func batchImageRepeatSuffixWidth(count int) int {
  868. if count < 10 {
  869. return 2
  870. }
  871. return len(strconv.Itoa(count))
  872. }
  873. func maxBatchImageReferenceImagesForModel(model string) int {
  874. model = strings.ToLower(strings.TrimSpace(model))
  875. if strings.Contains(model, "pro-image") {
  876. return 14
  877. }
  878. if strings.Contains(model, "flash-image") {
  879. return 3
  880. }
  881. return 0
  882. }
  883. func (s *BatchImagePublicService) selectProviderAndAccount(ctx context.Context, owner BatchImageOwner, requestedProvider, model string) (BatchImageProvider, *Account, error) {
  884. providers := batchImageProviderSelectionOrder(requestedProvider)
  885. for _, providerName := range providers {
  886. provider, ok := s.ProviderRegistry.Get(providerName)
  887. if !ok || provider == nil {
  888. continue
  889. }
  890. accounts, err := s.listCandidateAccounts(ctx, owner.GroupID, batchImageProviderPlatform(providerName))
  891. if err != nil {
  892. return nil, nil, err
  893. }
  894. sort.SliceStable(accounts, func(i, j int) bool {
  895. if accounts[i].Priority != accounts[j].Priority {
  896. return accounts[i].Priority > accounts[j].Priority
  897. }
  898. return accounts[i].ID < accounts[j].ID
  899. })
  900. for i := range accounts {
  901. account := accounts[i]
  902. if !account.IsSchedulable() || !account.IsModelSupported(model) {
  903. continue
  904. }
  905. if provider.SupportsAccount(&account) {
  906. return provider, &account, nil
  907. }
  908. }
  909. }
  910. if requestedProvider != "" {
  911. return nil, nil, ErrBatchImageNoAccountAvailable
  912. }
  913. return nil, nil, ErrBatchImageNoAccountAvailable
  914. }
  915. func (s *BatchImagePublicService) listCandidateAccounts(ctx context.Context, groupID *int64, platform string) ([]Account, error) {
  916. if s.AccountRepo == nil {
  917. return nil, ErrBatchImageNoAccountAvailable
  918. }
  919. if groupID != nil && *groupID > 0 {
  920. return s.AccountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, platform)
  921. }
  922. return s.AccountRepo.ListSchedulableByPlatform(ctx, platform)
  923. }
  924. func (s *BatchImagePublicService) ensureGroupAllowsBatchImage(ctx context.Context, groupID *int64) error {
  925. if groupID == nil || *groupID <= 0 {
  926. return nil
  927. }
  928. if s.GroupRepo == nil {
  929. return ErrBatchImageSettlementPricingMissing
  930. }
  931. group, err := s.GroupRepo.GetByIDLite(ctx, *groupID)
  932. if err != nil || group == nil {
  933. return ErrBatchImageSettlementPricingMissing
  934. }
  935. if !group.AllowBatchImageGeneration {
  936. return ErrBatchImageGroupDisabled
  937. }
  938. if group.Platform != PlatformGemini {
  939. return ErrBatchImageGroupDisabled
  940. }
  941. return nil
  942. }
  943. func (s *BatchImagePublicService) resolvePricingSnapshot(ctx context.Context, owner BatchImageOwner, req BatchImageSubmitRequest, provider string, account *Account) (*BatchImagePricingSnapshot, error) {
  944. unit := -1.0
  945. groupMultiplier := 1.0
  946. discountMultiplier := defaultBatchImageDiscountMultiplier
  947. holdMultiplier := defaultBatchImageHoldMultiplier
  948. if owner.GroupID != nil && *owner.GroupID > 0 {
  949. if s.GroupRepo == nil {
  950. return nil, ErrBatchImageSettlementPricingMissing
  951. }
  952. group, err := s.GroupRepo.GetByIDLite(ctx, *owner.GroupID)
  953. if err != nil || group == nil {
  954. return nil, ErrBatchImageSettlementPricingMissing
  955. }
  956. if !group.AllowBatchImageGeneration {
  957. return nil, ErrBatchImageGroupDisabled
  958. }
  959. groupDefaultMultiplier := group.RateMultiplier
  960. if groupDefaultMultiplier < 0 {
  961. groupDefaultMultiplier = 0
  962. }
  963. effectiveGroupMultiplier := groupDefaultMultiplier
  964. if s.UserGroupRateRepo != nil {
  965. userRate, rateErr := s.UserGroupRateRepo.GetByUserAndGroup(ctx, owner.UserID, group.ID)
  966. if rateErr != nil {
  967. return nil, ErrBatchImageSettlementPricingMissing
  968. }
  969. if userRate != nil {
  970. effectiveGroupMultiplier = *userRate
  971. }
  972. }
  973. groupMultiplier = effectiveGroupMultiplier
  974. if group.ImageRateIndependent {
  975. groupMultiplier = group.ImageRateMultiplier
  976. }
  977. if groupMultiplier < 0 {
  978. groupMultiplier = 0
  979. }
  980. discountMultiplier = group.BatchImageDiscountMultiplier
  981. if discountMultiplier < 0 {
  982. discountMultiplier = 0
  983. }
  984. if group.BatchImageHoldMultiplier >= 0 {
  985. holdMultiplier = group.BatchImageHoldMultiplier
  986. }
  987. if configuredUnit := group.GetImagePrice(req.ImageSize); configuredUnit != nil && *configuredUnit >= 0 {
  988. unit = *configuredUnit
  989. }
  990. }
  991. if unit < 0 {
  992. if s.Pricing == nil {
  993. return nil, ErrBatchImageSettlementPricingMissing
  994. }
  995. resolvedUnit, err := s.Pricing.BatchImageUnitPrice(ctx, &BatchImageJob{Provider: provider, Model: req.Model})
  996. if err != nil || resolvedUnit < 0 {
  997. return nil, ErrBatchImageSettlementPricingMissing
  998. }
  999. unit = resolvedUnit
  1000. }
  1001. // 定价不变式:hold 比例不得低于 discount 比例,否则成功率足够高时
  1002. // actualCost > holdAmount,结算永远失败、冻结余额无法解冻。
  1003. // 管理端已校验新配置,此处兜底钳制存量脏数据。
  1004. if holdMultiplier < discountMultiplier {
  1005. logger.L().Warn("batch_image.hold_multiplier_below_discount_clamped",
  1006. zap.Float64("hold_multiplier", holdMultiplier),
  1007. zap.Float64("discount_multiplier", discountMultiplier),
  1008. )
  1009. holdMultiplier = discountMultiplier
  1010. }
  1011. accountMultiplier := 1.0
  1012. if account != nil {
  1013. accountMultiplier = account.BillingRateMultiplier()
  1014. }
  1015. if accountMultiplier < 0 {
  1016. accountMultiplier = 0
  1017. }
  1018. standardUnitPrice := unit * groupMultiplier * accountMultiplier
  1019. billableUnitPrice := standardUnitPrice * discountMultiplier
  1020. holdUnitPrice := standardUnitPrice * holdMultiplier
  1021. return &BatchImagePricingSnapshot{
  1022. BaseUnitPrice: unit,
  1023. GroupRateMultiplier: groupMultiplier,
  1024. AccountRateMultiplier: accountMultiplier,
  1025. BatchDiscountMultiplier: discountMultiplier,
  1026. HoldMultiplier: holdMultiplier,
  1027. BillableUnitPrice: billableUnitPrice,
  1028. HoldUnitPrice: holdUnitPrice,
  1029. EstimatedCost: billableUnitPrice * float64(len(req.Items)),
  1030. HoldAmount: holdUnitPrice * float64(len(req.Items)),
  1031. }, nil
  1032. }
  1033. func (s *BatchImagePublicService) enabled() bool {
  1034. return s != nil && s.Repo != nil && s.AccountRepo != nil && s.Config != nil && s.Config.BatchImage.Enabled
  1035. }
  1036. func (s *BatchImagePublicService) invalidateAuthCache(ctx context.Context, userID int64) {
  1037. if s != nil && s.AuthCache != nil && userID > 0 {
  1038. s.AuthCache.InvalidateAuthCacheByUserID(ctx, userID)
  1039. }
  1040. }
  1041. func (s *BatchImagePublicService) maxItems() int {
  1042. if s != nil && s.Config != nil && s.Config.BatchImage.MaxItemsPerJobDefault > 0 {
  1043. return s.Config.BatchImage.MaxItemsPerJobDefault
  1044. }
  1045. return defaultBatchImageMaxItems
  1046. }
  1047. func (s *BatchImagePublicService) maxOutputImagesPerJob() int {
  1048. if s != nil && s.Config != nil && s.Config.BatchImage.MaxOutputImagesPerJob > 0 {
  1049. return s.Config.BatchImage.MaxOutputImagesPerJob
  1050. }
  1051. return defaultBatchImageMaxOutputImages
  1052. }
  1053. func (s *BatchImagePublicService) maxOutputImagesPerItem() int {
  1054. if s != nil && s.Config != nil && s.Config.BatchImage.MaxOutputImagesPerItem > 0 {
  1055. return s.Config.BatchImage.MaxOutputImagesPerItem
  1056. }
  1057. return defaultBatchImageMaxOutputCount
  1058. }
  1059. func (s *BatchImagePublicService) maxPromptChars() int {
  1060. if s != nil && s.Config != nil && s.Config.BatchImage.MaxPromptCharsPerItem > 0 {
  1061. return s.Config.BatchImage.MaxPromptCharsPerItem
  1062. }
  1063. return defaultBatchImageMaxPromptChars
  1064. }
  1065. func (s *BatchImagePublicService) maxReferenceImagesPerJob() int {
  1066. if s != nil && s.Config != nil && s.Config.BatchImage.MaxReferenceImagesPerJob > 0 {
  1067. return s.Config.BatchImage.MaxReferenceImagesPerJob
  1068. }
  1069. return defaultBatchImageMaxReferenceImages
  1070. }
  1071. func (s *BatchImagePublicService) maxReferenceInlineBytesPerJob() int {
  1072. if s != nil && s.Config != nil && s.Config.BatchImage.MaxReferenceInlineBytesPerJob > 0 {
  1073. return s.Config.BatchImage.MaxReferenceInlineBytesPerJob
  1074. }
  1075. return defaultBatchImageMaxReferenceBytes
  1076. }
  1077. func (s *BatchImagePublicService) defaultResponseMimeType() string {
  1078. if s != nil && s.Config != nil && strings.TrimSpace(s.Config.BatchImage.DefaultResponseMimeType) != "" {
  1079. return strings.TrimSpace(s.Config.BatchImage.DefaultResponseMimeType)
  1080. }
  1081. return defaultBatchImageResponseMime
  1082. }
  1083. func (s *BatchImagePublicService) defaultImageSize() string {
  1084. if s != nil && s.Config != nil && strings.TrimSpace(s.Config.BatchImage.DefaultImageSize) != "" {
  1085. return strings.TrimSpace(s.Config.BatchImage.DefaultImageSize)
  1086. }
  1087. return defaultBatchImageImageSize
  1088. }
  1089. func BatchImageJobToPublic(job *BatchImageJob) *BatchImagePublicBatch {
  1090. if job == nil {
  1091. return nil
  1092. }
  1093. holdAmount := job.EstimatedCost
  1094. if job.HoldAmount != nil {
  1095. holdAmount = *job.HoldAmount
  1096. }
  1097. return &BatchImagePublicBatch{
  1098. ID: job.BatchID,
  1099. Object: "image.batch",
  1100. TaskName: batchImagePublicTaskName(job),
  1101. ParentBatchID: job.ParentBatchID,
  1102. Status: PublicBatchImageStatus(job.Status),
  1103. Model: job.Model,
  1104. Provider: job.Provider,
  1105. ItemCount: job.ItemCount,
  1106. SuccessCount: job.SuccessCount,
  1107. FailCount: job.FailCount,
  1108. EstimatedCost: job.EstimatedCost,
  1109. HoldAmount: holdAmount,
  1110. ActualCost: job.ActualCost,
  1111. CreatedAt: job.CreatedAt.Unix(),
  1112. SubmittedAt: batchImageUnixPtr(job.SubmittedAt),
  1113. SettledAt: batchImageUnixPtr(job.SettledAt),
  1114. DownloadedAt: batchImageUnixPtr(job.DownloadedAt),
  1115. OutputDeletedAt: batchImageUnixPtr(job.OutputDeletedAt),
  1116. }
  1117. }
  1118. func BatchImageItemToPublic(item *BatchImageItem) BatchImagePublicItem {
  1119. out := BatchImagePublicItem{
  1120. CustomID: item.CustomID,
  1121. Status: "failed",
  1122. PromptPreview: item.PromptPreview,
  1123. MimeType: item.MimeType,
  1124. FileExtension: item.FileExtension,
  1125. ImageCount: item.ImageCount,
  1126. }
  1127. if item.Status == BatchImageItemStatusPending {
  1128. out.Status = "pending"
  1129. return out
  1130. }
  1131. if item.Status == BatchImageItemStatusSuccess {
  1132. out.Status = "succeeded"
  1133. return out
  1134. }
  1135. out.Error = &BatchImagePublicError{
  1136. Code: batchImageDerefString(item.ErrorCode),
  1137. Message: sanitizeBatchImagePublicMessage(batchImageDerefString(item.ErrorMessage)),
  1138. Source: batchImageItemErrorSource(item),
  1139. }
  1140. return out
  1141. }
  1142. func batchImageItemErrorSource(item *BatchImageItem) string {
  1143. if item == nil || item.ErrorCode == nil {
  1144. return ""
  1145. }
  1146. code := strings.TrimSpace(*item.ErrorCode)
  1147. if batchImageDerefString(item.ProviderSourceObject) != "" {
  1148. return "provider"
  1149. }
  1150. switch code {
  1151. case "EMPTY_IMAGE_OUTPUT", "PROVIDER_ITEM_FAILED":
  1152. return "provider"
  1153. case "INDEX_OUTPUT_MISSING", "INDEX_PARSE_FAILED", "DUPLICATE_CUSTOM_ID_IN_OUTPUT":
  1154. return "system"
  1155. default:
  1156. return ""
  1157. }
  1158. }
  1159. func PublicBatchImageStatus(status string) string {
  1160. switch status {
  1161. case BatchImageJobStatusCreated, BatchImageJobStatusUploading, BatchImageJobStatusSubmitted:
  1162. return "queued"
  1163. case BatchImageJobStatusRunning:
  1164. return "running"
  1165. case BatchImageJobStatusIndexing:
  1166. return "processing_results"
  1167. case BatchImageJobStatusSettling:
  1168. return "settling"
  1169. case BatchImageJobStatusCompleted:
  1170. return "completed"
  1171. case BatchImageJobStatusFailed:
  1172. return "failed"
  1173. case BatchImageJobStatusCancelled:
  1174. return "cancelled"
  1175. case BatchImageJobStatusOutputDeleted:
  1176. return "output_deleted"
  1177. default:
  1178. return status
  1179. }
  1180. }
  1181. func HashBatchImageSubmitRequest(req BatchImageSubmitRequest) string {
  1182. req.Metadata = sanitizeBatchImageMetadata(req.Metadata)
  1183. b, _ := json.Marshal(req)
  1184. sum := sha256.Sum256(b)
  1185. return hex.EncodeToString(sum[:])
  1186. }
  1187. func batchImageProviderPlatform(provider string) string {
  1188. switch provider {
  1189. case BatchImageProviderGeminiAPI, BatchImageProviderVertex:
  1190. return PlatformGemini
  1191. default:
  1192. return PlatformGemini
  1193. }
  1194. }
  1195. func batchImageProviderSelectionOrder(requestedProvider string) []string {
  1196. if strings.TrimSpace(requestedProvider) != "" {
  1197. return []string{strings.TrimSpace(requestedProvider)}
  1198. }
  1199. return []string{BatchImageProviderGeminiAPI, BatchImageProviderVertex}
  1200. }
  1201. func batchImageModelsFromAccountMapping(account *Account) []string {
  1202. if account == nil {
  1203. return nil
  1204. }
  1205. mapping := account.GetModelMapping()
  1206. if len(mapping) == 0 {
  1207. return nil
  1208. }
  1209. models := make(map[string]struct{})
  1210. for model := range mapping {
  1211. model = strings.TrimSpace(model)
  1212. if model == "" {
  1213. continue
  1214. }
  1215. if strings.ContainsAny(model, "*?") {
  1216. for _, candidate := range defaultBatchImageModelCandidates() {
  1217. if matchWildcard(model, candidate) {
  1218. models[candidate] = struct{}{}
  1219. }
  1220. }
  1221. continue
  1222. }
  1223. models[model] = struct{}{}
  1224. }
  1225. out := make([]string, 0, len(models))
  1226. for model := range models {
  1227. out = append(out, model)
  1228. }
  1229. sort.Strings(out)
  1230. return out
  1231. }
  1232. func defaultBatchImageModelCandidates() []string {
  1233. return []string{
  1234. "gemini-2.0-flash-exp-image-generation",
  1235. "gemini-2.5-flash-image",
  1236. "gemini-3-pro-image",
  1237. "gemini-3-pro-image-preview",
  1238. "gemini-3.1-flash-image",
  1239. "gemini-3.1-flash-image-preview",
  1240. "gemini-3.1-flash-lite-image",
  1241. }
  1242. }
  1243. func batchImageGCSRef(provider, ref string) string {
  1244. if provider == BatchImageProviderVertex && strings.HasPrefix(strings.TrimSpace(ref), "gs://") {
  1245. return strings.TrimSpace(ref)
  1246. }
  1247. return ""
  1248. }
  1249. func batchImageProviderSubmitPublicError(err error) error {
  1250. reason := strings.TrimSpace(infraerrors.Reason(err))
  1251. switch reason {
  1252. case "VERTEX_MANAGED_GCS_BUCKET_MISSING":
  1253. return ErrBatchImageVertexGCSBucketMissing
  1254. case "BATCH_IMAGE_PROVIDER_MISSING_API_KEY":
  1255. return ErrBatchImageProviderMissingAPIKey
  1256. case "BATCH_IMAGE_PROVIDER_MISSING_SERVICE_ACCOUNT":
  1257. return ErrBatchImageProviderMissingServiceAccount
  1258. case "BATCH_IMAGE_PROVIDER_UNSUPPORTED_ACCOUNT":
  1259. return ErrBatchImageProviderUnsupportedAccount
  1260. default:
  1261. return ErrBatchImageProviderSubmitFailed
  1262. }
  1263. }
  1264. func batchImagePublicTaskName(job *BatchImageJob) string {
  1265. if job == nil {
  1266. return ""
  1267. }
  1268. if strings.TrimSpace(job.TaskName) != "" {
  1269. return strings.TrimSpace(job.TaskName)
  1270. }
  1271. return defaultBatchImageTaskName(job.CreatedAt)
  1272. }
  1273. func defaultBatchImageTaskName(now time.Time) string {
  1274. if now.IsZero() {
  1275. now = time.Now()
  1276. }
  1277. return now.Format("2006-01-02 15:04:05")
  1278. }
  1279. func batchImageProviderSubmitRecordCode(err error) string {
  1280. reason := strings.TrimSpace(infraerrors.Reason(err))
  1281. if reason == "" || reason == "BATCH_IMAGE_PROVIDER_SUBMIT_FAILED" {
  1282. return "PROVIDER_SUBMIT_FAILED"
  1283. }
  1284. return reason
  1285. }
  1286. func parseBatchImageListTime(raw string) *time.Time {
  1287. raw = strings.TrimSpace(raw)
  1288. if raw == "" {
  1289. return nil
  1290. }
  1291. if unix, err := strconv.ParseInt(raw, 10, 64); err == nil && unix > 0 {
  1292. t := time.Unix(unix, 0)
  1293. return &t
  1294. }
  1295. if t, err := time.Parse(time.RFC3339, raw); err == nil {
  1296. return &t
  1297. }
  1298. if t, err := time.Parse("2006-01-02", raw); err == nil {
  1299. return &t
  1300. }
  1301. return nil
  1302. }
  1303. func sanitizeBatchImageMetadata(in map[string]string) map[string]string {
  1304. if len(in) == 0 {
  1305. return nil
  1306. }
  1307. keys := make([]string, 0, len(in))
  1308. for k := range in {
  1309. keys = append(keys, k)
  1310. }
  1311. sort.Strings(keys)
  1312. out := make(map[string]string, len(keys))
  1313. for _, k := range keys {
  1314. key := strings.TrimSpace(k)
  1315. if key == "" || len(key) > 64 {
  1316. continue
  1317. }
  1318. value := strings.TrimSpace(in[k])
  1319. if len(value) > 256 {
  1320. value = value[:256]
  1321. }
  1322. out[key] = value
  1323. if len(out) >= 20 {
  1324. break
  1325. }
  1326. }
  1327. return out
  1328. }
  1329. func sanitizeBatchImagePublicMessage(message string) string {
  1330. message = strings.TrimSpace(message)
  1331. for _, marker := range []string{"gs://", "files/", "projects/"} {
  1332. if strings.Contains(message, marker) {
  1333. message = "upstream provider operation failed"
  1334. break
  1335. }
  1336. }
  1337. if len(message) > maxBatchImagePublicErrorChars {
  1338. message = message[:maxBatchImagePublicErrorChars]
  1339. }
  1340. return message
  1341. }
  1342. func batchImageUnixPtr(t *time.Time) *int64 {
  1343. if t == nil {
  1344. return nil
  1345. }
  1346. v := t.Unix()
  1347. return &v
  1348. }
  1349. func parseBatchImageCursor(cursor string) int {
  1350. offset, err := strconv.Atoi(strings.TrimSpace(cursor))
  1351. if err != nil || offset < 0 {
  1352. return 0
  1353. }
  1354. return offset
  1355. }