| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440 |
- package service
- import (
- "context"
- "crypto/sha256"
- "encoding/hex"
- "encoding/json"
- "errors"
- "fmt"
- "sort"
- "strconv"
- "strings"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- "go.uber.org/zap"
- )
- const (
- defaultBatchImageMaxItems = 200
- defaultBatchImageMaxOutputImages = 200
- defaultBatchImageMaxOutputCount = 4
- defaultBatchImageMaxPromptChars = 8000
- defaultBatchImageResponseMime = "image/png"
- defaultBatchImageImageSize = "1K"
- defaultBatchImageDiscountMultiplier = 0.5
- defaultBatchImageHoldMultiplier = 0.6
- maxBatchImagePublicErrorChars = 500
- maxBatchImageReferenceImageBytes = 10 * 1024 * 1024
- defaultBatchImageMaxReferenceImages = 1000
- defaultBatchImageMaxReferenceBytes = 128 * 1024 * 1024
- )
- type BatchImageAccountSelectionRepository interface {
- GetByID(ctx context.Context, id int64) (*Account, error)
- ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error)
- ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error)
- }
- type BatchImageGroupPricingRepository interface {
- GetByIDLite(ctx context.Context, id int64) (*Group, error)
- }
- type BatchImageUserGroupRateRepository interface {
- GetByUserAndGroup(ctx context.Context, userID, groupID int64) (*float64, error)
- }
- type BatchImageSubmitRequest struct {
- Model string `json:"model"`
- TaskName string `json:"task_name"`
- ParentBatchID string `json:"parent_batch_id"`
- Provider string `json:"provider"`
- Items []BatchImageSubmitItem `json:"items"`
- ResponseMimeType string `json:"response_mime_type"`
- AspectRatio string `json:"aspect_ratio"`
- ImageSize string `json:"image_size"`
- Metadata map[string]string `json:"metadata"`
- SessionID *string `json:"-"`
- }
- type BatchImageSubmitItem struct {
- CustomID string `json:"custom_id"`
- Prompt string `json:"prompt"`
- OutputCount int `json:"output_count,omitempty"`
- ReferenceImages []BatchImageReferenceInput `json:"reference_images,omitempty"`
- }
- type BatchImageReferenceInput struct {
- ID string `json:"id,omitempty"`
- Type string `json:"type,omitempty"`
- MimeType string `json:"mime_type"`
- Data []byte `json:"data,omitempty"`
- FileURI string `json:"file_uri,omitempty"`
- }
- type BatchImageOwner struct {
- UserID int64
- APIKeyID int64
- GroupID *int64
- }
- type BatchImagePublicService struct {
- Repo BatchImageRepository
- AccountRepo BatchImageAccountSelectionRepository
- GroupRepo BatchImageGroupPricingRepository
- UserGroupRateRepo BatchImageUserGroupRateRepository
- Queue BatchImageQueue
- ProviderRegistry *BatchImageProviderRegistry
- Pricing BatchImagePricingResolver
- BillingRepo UsageBillingRepository
- AuthCache APIKeyAuthCacheInvalidator
- Config *config.Config
- }
- type BatchImagePricingSnapshot struct {
- BaseUnitPrice float64
- GroupRateMultiplier float64
- AccountRateMultiplier float64
- BatchDiscountMultiplier float64
- HoldMultiplier float64
- BillableUnitPrice float64
- HoldUnitPrice float64
- EstimatedCost float64
- HoldAmount float64
- }
- type BatchImagePublicBatch struct {
- ID string `json:"id"`
- Object string `json:"object"`
- TaskName string `json:"task_name"`
- ParentBatchID *string `json:"parent_batch_id,omitempty"`
- Status string `json:"status"`
- Model string `json:"model"`
- Provider string `json:"provider"`
- ItemCount int `json:"item_count"`
- SuccessCount int `json:"success_count"`
- FailCount int `json:"fail_count"`
- EstimatedCost float64 `json:"estimated_cost"`
- HoldAmount float64 `json:"hold_amount"`
- ActualCost *float64 `json:"actual_cost"`
- CreatedAt int64 `json:"created_at"`
- SubmittedAt *int64 `json:"submitted_at"`
- SettledAt *int64 `json:"settled_at"`
- DownloadedAt *int64 `json:"downloaded_at,omitempty"`
- OutputDeletedAt *int64 `json:"output_deleted_at,omitempty"`
- }
- type BatchImagePublicItem struct {
- CustomID string `json:"custom_id"`
- Status string `json:"status"`
- PromptPreview *string `json:"prompt_preview,omitempty"`
- MimeType *string `json:"mime_type"`
- FileExtension *string `json:"file_extension"`
- ImageCount int `json:"image_count"`
- Error *BatchImagePublicError `json:"error"`
- }
- type BatchImagePublicError struct {
- Code string `json:"code"`
- Message string `json:"message"`
- Source string `json:"source,omitempty"`
- }
- type BatchImagePublicItemsResponse struct {
- Object string `json:"object"`
- Data []BatchImagePublicItem `json:"data"`
- HasMore bool `json:"has_more"`
- }
- type BatchImagePublicListResponse struct {
- Object string `json:"object"`
- Data []*BatchImagePublicBatch `json:"data"`
- HasMore bool `json:"has_more"`
- }
- type BatchImagePublicModel struct {
- ID string `json:"id"`
- Object string `json:"object"`
- Provider string `json:"provider"`
- }
- type BatchImagePublicModelsResponse struct {
- Object string `json:"object"`
- Data []BatchImagePublicModel `json:"data"`
- }
- type BatchImageJobsQuery struct {
- Status string
- TaskName string
- Downloaded string
- From string
- To string
- Limit int
- Cursor string
- }
- type BatchImageItemsQuery struct {
- Status string
- Limit int
- Cursor string
- }
- func NewBatchImagePublicService(repo BatchImageRepository, accountRepo AccountRepository, groupRepo GroupRepository, userGroupRateRepo UserGroupRateRepository, queue BatchImageQueue, pricing *BatchImageModelPricingResolver, billingRepo UsageBillingRepository, authCache APIKeyAuthCacheInvalidator, cfg *config.Config) *BatchImagePublicService {
- return &BatchImagePublicService{
- Repo: repo,
- AccountRepo: accountRepo,
- GroupRepo: groupRepo,
- UserGroupRateRepo: userGroupRateRepo,
- Queue: queue,
- ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg),
- Pricing: pricing,
- BillingRepo: billingRepo,
- AuthCache: authCache,
- Config: cfg,
- }
- }
- func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOwner, req BatchImageSubmitRequest, idempotencyKey string) (*BatchImagePublicBatch, error) {
- if !s.enabled() {
- return nil, ErrBatchImageDisabled
- }
- normalized, err := s.validateSubmitRequest(req)
- if err != nil {
- return nil, err
- }
- // 与 ListModels 使用同一鉴权谓词(AllowBatchImageGeneration + Platform==Gemini),
- // 避免两个入口校验口径不一致留下防御纵深缺口。
- if err := s.ensureGroupAllowsBatchImage(ctx, owner.GroupID); err != nil {
- return nil, err
- }
- requestHash := HashBatchImageSubmitRequest(normalized)
- idempotencyKey = strings.TrimSpace(idempotencyKey)
- if idempotencyKey != "" {
- existing, err := s.Repo.GetBatchImageJobByIdempotencyKey(ctx, owner.UserID, owner.APIKeyID, idempotencyKey)
- if err == nil {
- if batchImageDerefString(existing.RequestHash) != requestHash {
- return nil, ErrBatchImageIdempotencyConflict
- }
- if existing.Status == BatchImageJobStatusSubmitted && s.Queue != nil {
- if enqueueErr := s.Queue.Enqueue(ctx, existing.BatchID); enqueueErr != nil && !errors.Is(enqueueErr, ErrBatchImageAlreadyQueued) {
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, existing.BatchID, "QUEUE_FAILED", sanitizeBatchImagePublicMessage(enqueueErr.Error()), false)
- return nil, ErrBatchImageQueueFailed
- }
- }
- return BatchImageJobToPublic(existing), nil
- }
- if !errors.Is(err, ErrBatchImageJobNotFound) {
- return nil, err
- }
- }
- provider, account, err := s.selectProviderAndAccount(ctx, owner, normalized.Provider, normalized.Model)
- if err != nil {
- return nil, err
- }
- pricingSnapshot, err := s.resolvePricingSnapshot(ctx, owner, normalized, provider.Name(), account)
- if err != nil {
- return nil, err
- }
- parentBatchID := batchImageOptionalStringPtr(normalized.ParentBatchID)
- if parentBatchID != nil {
- parent, parentErr := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, *parentBatchID)
- if parentErr != nil {
- return nil, parentErr
- }
- if parent.ParentBatchID != nil && strings.TrimSpace(*parent.ParentBatchID) != "" {
- parentBatchID = batchImageOptionalStringPtr(*parent.ParentBatchID)
- }
- }
- batchID, err := NewBatchImageID()
- if err != nil {
- return nil, err
- }
- apiKeyID := owner.APIKeyID
- accountID := account.ID
- holdID := BatchImageHoldRequestID(batchID)
- holdAmount := pricingSnapshot.HoldAmount
- job, err := s.Repo.CreateBatchImageJob(ctx, CreateBatchImageJobParams{
- BatchID: batchID,
- UserID: owner.UserID,
- APIKeyID: &apiKeyID,
- AccountID: &accountID,
- Provider: provider.Name(),
- Model: normalized.Model,
- TaskName: normalized.TaskName,
- ParentBatchID: parentBatchID,
- Status: BatchImageJobStatusCreated,
- ItemCount: len(normalized.Items),
- EstimatedCost: pricingSnapshot.EstimatedCost,
- HoldAmount: &holdAmount,
- BaseUnitPrice: pricingSnapshot.BaseUnitPrice,
- GroupRateMultiplier: pricingSnapshot.GroupRateMultiplier,
- AccountRateMultiplier: pricingSnapshot.AccountRateMultiplier,
- BatchDiscountMultiplier: pricingSnapshot.BatchDiscountMultiplier,
- HoldMultiplier: pricingSnapshot.HoldMultiplier,
- BillableUnitPrice: pricingSnapshot.BillableUnitPrice,
- HoldUnitPrice: pricingSnapshot.HoldUnitPrice,
- PricingSnapshotVersion: 1,
- Currency: "USD",
- HoldID: &holdID,
- IdempotencyKey: batchImageOptionalStringPtr(idempotencyKey),
- RequestHash: batchImageStringPtr(requestHash),
- SessionID: normalized.SessionID,
- })
- if err != nil {
- return nil, err
- }
- if err := reserveBatchImageBalanceHold(ctx, s.BillingRepo, job, requestHash); err != nil {
- code := "BILLING_HOLD_FAILED"
- if errors.Is(err, ErrBatchImageInsufficientBalance) {
- code = "INSUFFICIENT_BALANCE"
- }
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, code, sanitizeBatchImagePublicMessage(err.Error()), true)
- s.hidePreUpstreamSubmitFailure(ctx, owner, job)
- return nil, err
- }
- s.invalidateAuthCache(ctx, owner.UserID)
- if err := s.createPendingItems(ctx, job.BatchID, requestHash, normalized.Items); err != nil {
- if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
- return nil, releaseErr
- }
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "ITEM_CREATE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true)
- s.hidePreUpstreamSubmitFailure(ctx, owner, job)
- return nil, ErrBatchImageQueueFailed
- }
- input := BatchImageInput{
- BatchID: job.BatchID,
- Model: normalized.Model,
- DisplayName: job.BatchID,
- ResponseMimeType: normalized.ResponseMimeType,
- AspectRatio: normalized.AspectRatio,
- ImageSize: normalized.ImageSize,
- Metadata: normalized.Metadata,
- Items: make([]BatchImageInputItem, 0, len(normalized.Items)),
- }
- for _, item := range normalized.Items {
- refs := make([]BatchImageReference, 0, len(item.ReferenceImages))
- for _, ref := range item.ReferenceImages {
- refs = append(refs, BatchImageReference(ref))
- }
- input.Items = append(input.Items, BatchImageInputItem{
- CustomID: item.CustomID,
- Prompt: item.Prompt,
- ReferenceImages: refs,
- })
- }
- // 上游提交(上传参考图 + 创建批任务)可能长达数分钟且不刷新 updated_at,
- // 会被 stale 恢复扫描误判为滞留并退款。提交前转入 uploading 刷新时间戳,
- // 提交期间用心跳持续续期。
- if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusUploading, BatchImageTransitionOptions{
- EventType: "upload_started",
- EventPayload: map[string]any{"batch_id": job.BatchID},
- }); err != nil {
- if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
- return nil, releaseErr
- }
- // 并发 Cancel 等导致的非法转换:job 已处于终态,不再覆盖其状态。
- if !errors.Is(err, ErrBatchImageInvalidTransition) {
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "UPLOAD_TRANSITION_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true)
- s.hidePreUpstreamSubmitFailure(ctx, owner, job)
- }
- return nil, err
- }
- job.Status = BatchImageJobStatusUploading
- hbCtx, hbCancel := context.WithCancel(ctx)
- hbDone := make(chan struct{})
- go s.runSubmitHeartbeat(hbCtx, job.BatchID, hbDone)
- providerJob, err := provider.Submit(ctx, job, account, input)
- hbCancel()
- <-hbDone
- if err != nil {
- if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
- return nil, releaseErr
- }
- publicErr := batchImageProviderSubmitPublicError(err)
- reason := batchImageProviderSubmitRecordCode(publicErr)
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, reason, sanitizeBatchImagePublicMessage(err.Error()), true)
- s.hidePreUpstreamSubmitFailure(ctx, owner, job)
- return nil, publicErr
- }
- if providerJob == nil || strings.TrimSpace(providerJob.ProviderJobName) == "" {
- if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
- return nil, releaseErr
- }
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "PROVIDER_SUBMIT_FAILED", "provider job name missing", true)
- s.hidePreUpstreamSubmitFailure(ctx, owner, job)
- return nil, ErrBatchImageProviderSubmitFailed
- }
- if err := s.Repo.UpdateBatchImageJobProviderSubmit(ctx, UpdateBatchImageJobProviderSubmitParams{
- BatchID: job.BatchID,
- ProviderJobName: providerJob.ProviderJobName,
- ProviderInputRef: providerJob.ProviderInputRef,
- ProviderOutputRef: providerJob.ProviderOutputRef,
- GCSInputURI: batchImageGCSRef(provider.Name(), providerJob.ProviderInputRef),
- GCSOutputURI: batchImageGCSRef(provider.Name(), providerJob.ProviderOutputRef),
- EventPayload: map[string]any{"provider": provider.Name()},
- }); err != nil {
- // job 可能已被恢复扫描转 failed 并退款:上游批任务已创建成功,
- // 必须尽力取消并清理输入,否则上游照常产生成本(孤儿任务)。
- s.abortOrphanProviderJob(ctx, provider, job, account, providerJob)
- return nil, err
- }
- if s.Queue != nil {
- if err := s.Queue.Enqueue(ctx, job.BatchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "QUEUE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), false)
- return nil, ErrBatchImageQueueFailed
- }
- }
- created, err := s.Repo.GetBatchImageJobByBatchID(ctx, job.BatchID)
- if err != nil {
- return nil, err
- }
- return BatchImageJobToPublic(created), nil
- }
- func (s *BatchImagePublicService) releaseFailedSubmitHold(ctx context.Context, job *BatchImageJob, requestHash string) error {
- if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, requestHash); err != nil {
- _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "BILLING_RELEASE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true)
- s.enqueueBillingRetry(ctx, job.BatchID)
- return ErrBatchImageBillingHoldFailed
- }
- s.invalidateAuthCache(ctx, job.UserID)
- return nil
- }
- // runSubmitHeartbeat 在 provider.Submit 期间周期性刷新 job 的 updated_at,
- // 使 stale 恢复扫描能区分"仍在慢提交"与"进程死亡后的滞留"。
- func (s *BatchImagePublicService) runSubmitHeartbeat(ctx context.Context, batchID string, done chan<- struct{}) {
- defer close(done)
- interval := s.submitHeartbeatInterval()
- ticker := time.NewTicker(interval)
- defer ticker.Stop()
- for {
- select {
- case <-ctx.Done():
- return
- case <-ticker.C:
- if err := s.Repo.TouchBatchImageJobSubmitting(ctx, batchID); err != nil && ctx.Err() == nil {
- logger.L().Warn("batch_image.submit_heartbeat_failed",
- zap.String("batch_id", batchID),
- zap.Error(err),
- )
- }
- }
- }
- }
- func (s *BatchImagePublicService) submitHeartbeatInterval() time.Duration {
- staleAfter := 10 * time.Minute
- if s != nil && s.Config != nil && s.Config.BatchImage.StaleActiveAfterSeconds > 0 {
- staleAfter = time.Duration(s.Config.BatchImage.StaleActiveAfterSeconds) * time.Second
- }
- interval := staleAfter / 3
- if interval < 15*time.Second {
- interval = 15 * time.Second
- }
- return interval
- }
- // abortOrphanProviderJob 在上游任务创建成功但本地状态推进失败时,
- // 尽力取消上游批任务并清理已上传的输入文件,避免孤儿任务持续产生成本。
- func (s *BatchImagePublicService) abortOrphanProviderJob(ctx context.Context, provider BatchImageProvider, job *BatchImageJob, account *Account, providerJob *BatchProviderJob) {
- if s == nil || provider == nil || job == nil || providerJob == nil {
- return
- }
- orphan := *job
- orphan.ProviderJobName = batchImageOptionalStringPtr(providerJob.ProviderJobName)
- orphan.ProviderInputRef = batchImageOptionalStringPtr(providerJob.ProviderInputRef)
- orphan.GCSInputURI = batchImageOptionalStringPtr(batchImageGCSRef(provider.Name(), providerJob.ProviderInputRef))
- if err := provider.Cancel(ctx, &orphan, account); err != nil {
- logger.L().Warn("batch_image.orphan_provider_job_cancel_failed",
- zap.String("batch_id", job.BatchID),
- zap.String("provider", provider.Name()),
- zap.Error(err),
- )
- }
- if err := provider.Cleanup(ctx, &orphan, account, CleanupTargetInput); err != nil {
- logger.L().Warn("batch_image.orphan_provider_job_cleanup_failed",
- zap.String("batch_id", job.BatchID),
- zap.String("provider", provider.Name()),
- zap.Error(err),
- )
- }
- if err := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "provider_job_aborted_after_submit", map[string]any{
- "batch_id": job.BatchID,
- "provider": provider.Name(),
- }); err != nil {
- logger.L().Warn("batch_image.orphan_provider_job_event_failed",
- zap.String("batch_id", job.BatchID),
- zap.Error(err),
- )
- }
- }
- func (s *BatchImagePublicService) createPendingItems(ctx context.Context, batchID, requestHash string, items []BatchImageSubmitItem) error {
- if s == nil || s.Repo == nil || len(items) == 0 {
- return nil
- }
- params := make([]CreateBatchImageItemParams, 0, len(items))
- for _, item := range items {
- preview := truncateBatchImageMessage(item.Prompt, s.maxPromptChars())
- params = append(params, CreateBatchImageItemParams{
- JobID: batchID,
- CustomID: item.CustomID,
- Status: BatchImageItemStatusPending,
- RequestHash: batchImageStringPtr(requestHash),
- PromptPreview: batchImageStringPtr(preview),
- ImageCount: 0,
- })
- }
- return s.Repo.BulkCreateBatchImageItems(ctx, params)
- }
- func (s *BatchImagePublicService) enqueueBillingRetry(ctx context.Context, batchID string) {
- if s == nil || s.Queue == nil {
- return
- }
- if err := s.Queue.Enqueue(ctx, batchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
- logger.L().Warn("batch_image.billing_retry_enqueue_failed",
- zap.String("batch_id", batchID),
- zap.Error(err),
- )
- if eventErr := s.Repo.AppendBatchImageEvent(ctx, batchID, "billing_retry_enqueue_failed", map[string]any{
- "batch_id": batchID,
- "error": sanitizeBatchImagePublicMessage(err.Error()),
- }); eventErr != nil {
- logger.L().Warn("batch_image.billing_retry_event_failed",
- zap.String("batch_id", batchID),
- zap.Error(eventErr),
- )
- }
- }
- }
- func (s *BatchImagePublicService) hidePreUpstreamSubmitFailure(ctx context.Context, owner BatchImageOwner, job *BatchImageJob) {
- if s == nil || s.Repo == nil || job == nil || job.ProviderJobName != nil {
- return
- }
- if err := s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now()); err != nil {
- logger.L().Warn("batch_image.hide_pre_upstream_failure_failed",
- zap.String("batch_id", job.BatchID),
- zap.Error(err),
- )
- }
- }
- func (s *BatchImagePublicService) Get(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) {
- job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
- if err != nil {
- return nil, err
- }
- return BatchImageJobToPublic(job), nil
- }
- func (s *BatchImagePublicService) List(ctx context.Context, owner BatchImageOwner, query BatchImageJobsQuery) (*BatchImagePublicListResponse, error) {
- filter := BatchImageJobFilter{Limit: query.Limit, Offset: parseBatchImageCursor(query.Cursor), ExcludeDeleted: true}
- filter.TaskNameLike = strings.TrimSpace(query.TaskName)
- switch strings.TrimSpace(query.Status) {
- case "", "all":
- case "queued":
- filter.Status = BatchImageJobStatusSubmitted
- case "processing_results":
- filter.Status = BatchImageJobStatusIndexing
- case "completed":
- filter.Status = BatchImageJobStatusCompleted
- case "failed":
- filter.Status = BatchImageJobStatusFailed
- case "cancelled":
- filter.Status = BatchImageJobStatusCancelled
- case "output_deleted":
- filter.Status = BatchImageJobStatusOutputDeleted
- default:
- filter.Status = strings.TrimSpace(query.Status)
- }
- switch strings.TrimSpace(strings.ToLower(query.Downloaded)) {
- case "", "all":
- case "true", "1", "yes", "downloaded":
- downloaded := true
- filter.Downloaded = &downloaded
- case "false", "0", "no", "not_downloaded":
- downloaded := false
- filter.Downloaded = &downloaded
- default:
- return nil, ErrBatchImageInvalidItems
- }
- if from := parseBatchImageListTime(query.From); from != nil {
- filter.CreatedAfter = from
- }
- if to := parseBatchImageListTime(query.To); to != nil {
- filter.CreatedBefore = to
- }
- if filter.Limit <= 0 || filter.Limit > 100 {
- filter.Limit = 20
- }
- jobs, err := s.Repo.ListBatchImageJobsForOwner(ctx, owner.UserID, owner.APIKeyID, filter)
- if err != nil {
- return nil, err
- }
- data := make([]*BatchImagePublicBatch, 0, len(jobs))
- for _, job := range jobs {
- data = append(data, BatchImageJobToPublic(job))
- }
- return &BatchImagePublicListResponse{
- Object: "list",
- Data: data,
- HasMore: len(data) == filter.Limit,
- }, nil
- }
- func (s *BatchImagePublicService) MarkDownloaded(ctx context.Context, owner BatchImageOwner, batchID string) error {
- job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
- if err != nil {
- return err
- }
- return s.Repo.MarkBatchImageDownloaded(ctx, job.BatchID, time.Now())
- }
- func (s *BatchImagePublicService) DeleteRecord(ctx context.Context, owner BatchImageOwner, batchID string) error {
- job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
- if err != nil {
- return err
- }
- if !isBatchImageProcessorDoneStatus(job.Status) {
- return ErrBatchImageRecordDeleteNotReady
- }
- return s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now())
- }
- func (s *BatchImagePublicService) ListModels(ctx context.Context, owner BatchImageOwner) (*BatchImagePublicModelsResponse, error) {
- if !s.enabled() {
- return nil, ErrBatchImageDisabled
- }
- if s.Pricing == nil {
- return nil, ErrBatchImageSettlementPricingMissing
- }
- if err := s.ensureGroupAllowsBatchImage(ctx, owner.GroupID); err != nil {
- return nil, err
- }
- modelsByProvider := make(map[string]map[string]struct{})
- for _, providerName := range batchImageProviderSelectionOrder("") {
- provider, ok := s.ProviderRegistry.Get(providerName)
- if !ok || provider == nil {
- continue
- }
- accounts, err := s.listCandidateAccounts(ctx, owner.GroupID, batchImageProviderPlatform(providerName))
- if err != nil {
- return nil, err
- }
- for i := range accounts {
- account := accounts[i]
- if !account.IsSchedulable() || !provider.SupportsAccount(&account) {
- continue
- }
- for _, model := range batchImageModelsFromAccountMapping(&account) {
- if _, err := s.Pricing.BatchImageUnitPrice(ctx, &BatchImageJob{Provider: providerName, Model: model}); err != nil {
- continue
- }
- if !account.IsModelSupported(model) {
- continue
- }
- if modelsByProvider[providerName] == nil {
- modelsByProvider[providerName] = make(map[string]struct{})
- }
- modelsByProvider[providerName][model] = struct{}{}
- }
- }
- }
- out := make([]BatchImagePublicModel, 0)
- for _, providerName := range batchImageProviderSelectionOrder("") {
- models := make([]string, 0, len(modelsByProvider[providerName]))
- for model := range modelsByProvider[providerName] {
- models = append(models, model)
- }
- sort.Strings(models)
- for _, model := range models {
- out = append(out, BatchImagePublicModel{
- ID: model,
- Object: "image.batch.model",
- Provider: providerName,
- })
- }
- }
- return &BatchImagePublicModelsResponse{Object: "list", Data: out}, nil
- }
- func (s *BatchImagePublicService) ListItems(ctx context.Context, owner BatchImageOwner, batchID string, query BatchImageItemsQuery) (*BatchImagePublicItemsResponse, error) {
- filter := BatchImageItemFilter{Limit: query.Limit, Offset: parseBatchImageCursor(query.Cursor)}
- switch strings.TrimSpace(query.Status) {
- case "", "all":
- case "succeeded", "success":
- filter.Status = BatchImageItemStatusSuccess
- case "pending":
- filter.Status = BatchImageItemStatusPending
- case "failed":
- filter.Status = BatchImageItemStatusFailed
- default:
- return nil, ErrBatchImageInvalidItems
- }
- if filter.Limit <= 0 || filter.Limit > 500 {
- filter.Limit = 100
- }
- items, err := s.Repo.ListBatchImageItemsForOwner(ctx, owner.UserID, owner.APIKeyID, batchID, filter)
- if err != nil {
- return nil, err
- }
- data := make([]BatchImagePublicItem, 0, len(items))
- for _, item := range items {
- data = append(data, BatchImageItemToPublic(item))
- }
- return &BatchImagePublicItemsResponse{
- Object: "list",
- Data: data,
- HasMore: len(data) == filter.Limit,
- }, nil
- }
- func (s *BatchImagePublicService) Cancel(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) {
- job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
- if err != nil {
- return nil, err
- }
- if isBatchImageProcessorDoneStatus(job.Status) {
- if job.Status == BatchImageJobStatusFailed || job.Status == BatchImageJobStatusCancelled {
- if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil {
- s.enqueueBillingRetry(ctx, job.BatchID)
- return nil, ErrBatchImageCancelFailed
- }
- s.invalidateAuthCache(ctx, owner.UserID)
- }
- return BatchImageJobToPublic(job), nil
- }
- if job.ProviderJobName != nil && strings.TrimSpace(*job.ProviderJobName) != "" {
- provider, ok := s.ProviderRegistry.Get(job.Provider)
- if !ok || provider == nil {
- return nil, ErrBatchImageUnsupportedProvider
- }
- if job.AccountID == nil {
- return nil, ErrBatchImageCancelFailed
- }
- account, err := s.AccountRepo.GetByID(ctx, *job.AccountID)
- if err != nil {
- return nil, ErrBatchImageCancelFailed
- }
- if err := provider.Cancel(ctx, job, account); err != nil {
- return nil, ErrBatchImageCancelFailed
- }
- if eventErr := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "job_cancel_requested", map[string]any{"batch_id": job.BatchID}); eventErr != nil {
- logger.L().Warn("batch_image.cancel_event_failed",
- zap.String("batch_id", job.BatchID),
- zap.Error(eventErr),
- )
- }
- if s.Queue != nil {
- if err := s.Queue.Enqueue(ctx, job.BatchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
- return nil, ErrBatchImageCancelFailed
- }
- }
- updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
- if err != nil {
- return nil, err
- }
- return BatchImageJobToPublic(updated), nil
- }
- if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusCancelled, BatchImageTransitionOptions{
- EventType: "job_cancelled",
- EventPayload: map[string]any{"batch_id": job.BatchID},
- }); err != nil {
- return nil, err
- }
- if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil {
- s.enqueueBillingRetry(ctx, job.BatchID)
- return nil, ErrBatchImageCancelFailed
- }
- s.invalidateAuthCache(ctx, owner.UserID)
- updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
- if err != nil {
- return nil, err
- }
- return BatchImageJobToPublic(updated), nil
- }
- func (s *BatchImagePublicService) validateSubmitRequest(req BatchImageSubmitRequest) (BatchImageSubmitRequest, error) {
- req.Model = strings.TrimSpace(req.Model)
- req.TaskName = strings.TrimSpace(req.TaskName)
- req.ParentBatchID = strings.TrimSpace(req.ParentBatchID)
- req.Provider = strings.TrimSpace(req.Provider)
- req.ResponseMimeType = strings.TrimSpace(req.ResponseMimeType)
- req.AspectRatio = strings.TrimSpace(req.AspectRatio)
- req.ImageSize = strings.TrimSpace(req.ImageSize)
- if req.Model == "" {
- return req, ErrBatchImageInvalidModel
- }
- if req.TaskName == "" {
- req.TaskName = defaultBatchImageTaskName(time.Now())
- }
- if len(req.TaskName) > 255 {
- req.TaskName = truncateBatchImageMessage(req.TaskName, 255)
- }
- if req.Provider != "" && !IsSupportedBatchImageProvider(req.Provider) {
- return req, ErrBatchImageUnsupportedProvider
- }
- if len(req.Items) == 0 {
- return req, ErrBatchImageInvalidItems
- }
- maxItems := s.maxItems()
- if len(req.Items) > maxItems {
- return req, ErrBatchImageInvalidItems
- }
- if req.ResponseMimeType == "" {
- req.ResponseMimeType = s.defaultResponseMimeType()
- }
- if req.ImageSize == "" {
- req.ImageSize = s.defaultImageSize()
- }
- if !strings.EqualFold(req.ImageSize, defaultBatchImageImageSize) {
- return req, ErrBatchImageInvalidItems
- }
- req.ImageSize = defaultBatchImageImageSize
- req.Metadata = sanitizeBatchImageMetadata(req.Metadata)
- seen := make(map[string]struct{}, len(req.Items))
- totalReferenceImages := 0
- totalInlineReferenceBytes := 0
- totalOutputImages := 0
- expandedItems := make([]BatchImageSubmitItem, 0, len(req.Items))
- for i := range req.Items {
- req.Items[i].CustomID = strings.TrimSpace(req.Items[i].CustomID)
- if req.Items[i].CustomID == "" {
- req.Items[i].CustomID = fmt.Sprintf("item_%06d", i+1)
- }
- outputCount := req.Items[i].OutputCount
- if outputCount == 0 {
- outputCount = 1
- }
- if outputCount < 1 || outputCount > s.maxOutputImagesPerItem() {
- return req, ErrBatchImageInvalidItems
- }
- totalOutputImages += outputCount
- if totalOutputImages > s.maxOutputImagesPerJob() {
- return req, ErrBatchImageTooManyOutputImages
- }
- req.Items[i].Prompt = strings.TrimSpace(req.Items[i].Prompt)
- if req.Items[i].Prompt == "" {
- return req, ErrBatchImageInvalidItems
- }
- if len(req.Items[i].Prompt) > s.maxPromptChars() {
- return req, ErrBatchImagePromptTooLong
- }
- referenceCount, inlineReferenceBytes, err := normalizeBatchImageReferenceInputs(req.Model, &req.Items[i])
- if err != nil {
- return req, err
- }
- totalReferenceImages += referenceCount * outputCount
- if totalReferenceImages > s.maxReferenceImagesPerJob() {
- return req, ErrBatchImageTooManyReferenceImages
- }
- totalInlineReferenceBytes += inlineReferenceBytes * outputCount
- if totalInlineReferenceBytes > s.maxReferenceInlineBytesPerJob() {
- return req, ErrBatchImageReferenceImagesTooLarge
- }
- for repeatIndex := 1; repeatIndex <= outputCount; repeatIndex++ {
- expanded := req.Items[i]
- expanded.OutputCount = 0
- if outputCount > 1 {
- expanded.CustomID = fmt.Sprintf("%s_%0*d", req.Items[i].CustomID, batchImageRepeatSuffixWidth(outputCount), repeatIndex)
- }
- if _, ok := seen[expanded.CustomID]; ok {
- return req, ErrBatchImageDuplicateCustomIDInRequest
- }
- seen[expanded.CustomID] = struct{}{}
- expandedItems = append(expandedItems, expanded)
- }
- }
- req.Items = expandedItems
- return req, nil
- }
- func normalizeBatchImageReferenceInputs(model string, item *BatchImageSubmitItem) (int, int, error) {
- if item == nil || len(item.ReferenceImages) == 0 {
- return 0, 0, nil
- }
- maxRefs := maxBatchImageReferenceImagesForModel(model)
- if maxRefs <= 0 || len(item.ReferenceImages) > maxRefs {
- return 0, 0, ErrBatchImageTooManyReferenceImages
- }
- out := make([]BatchImageReferenceInput, 0, len(item.ReferenceImages))
- inlineBytes := 0
- for _, ref := range item.ReferenceImages {
- ref.ID = truncateBatchImageMessage(strings.TrimSpace(ref.ID), 80)
- ref.Type = truncateBatchImageMessage(strings.TrimSpace(ref.Type), 40)
- ref.MimeType = normalizeBatchImageReferenceMimeType(ref.MimeType)
- ref.FileURI = strings.TrimSpace(ref.FileURI)
- if ref.MimeType == "" {
- return 0, 0, ErrBatchImageInvalidReferenceImage
- }
- if len(ref.Data) == 0 && ref.FileURI == "" {
- return 0, 0, ErrBatchImageInvalidReferenceImage
- }
- if len(ref.Data) > 0 && ref.FileURI != "" {
- return 0, 0, ErrBatchImageInvalidReferenceImage
- }
- if len(ref.Data) > maxBatchImageReferenceImageBytes {
- return 0, 0, ErrBatchImageInvalidReferenceImage
- }
- if ref.FileURI != "" && !strings.HasPrefix(ref.FileURI, "gs://") {
- return 0, 0, ErrBatchImageInvalidReferenceImage
- }
- inlineBytes += len(ref.Data)
- out = append(out, ref)
- }
- item.ReferenceImages = out
- return len(out), inlineBytes, nil
- }
- func normalizeBatchImageReferenceMimeType(v string) string {
- switch strings.ToLower(strings.TrimSpace(v)) {
- case "image/jpeg", "image/jpg":
- return "image/jpeg"
- case "image/png":
- return "image/png"
- case "image/webp":
- return "image/webp"
- default:
- return ""
- }
- }
- func batchImageRepeatSuffixWidth(count int) int {
- if count < 10 {
- return 2
- }
- return len(strconv.Itoa(count))
- }
- func maxBatchImageReferenceImagesForModel(model string) int {
- model = strings.ToLower(strings.TrimSpace(model))
- if strings.Contains(model, "pro-image") {
- return 14
- }
- if strings.Contains(model, "flash-image") {
- return 3
- }
- return 0
- }
- func (s *BatchImagePublicService) selectProviderAndAccount(ctx context.Context, owner BatchImageOwner, requestedProvider, model string) (BatchImageProvider, *Account, error) {
- providers := batchImageProviderSelectionOrder(requestedProvider)
- for _, providerName := range providers {
- provider, ok := s.ProviderRegistry.Get(providerName)
- if !ok || provider == nil {
- continue
- }
- accounts, err := s.listCandidateAccounts(ctx, owner.GroupID, batchImageProviderPlatform(providerName))
- if err != nil {
- return nil, nil, err
- }
- sort.SliceStable(accounts, func(i, j int) bool {
- if accounts[i].Priority != accounts[j].Priority {
- return accounts[i].Priority > accounts[j].Priority
- }
- return accounts[i].ID < accounts[j].ID
- })
- for i := range accounts {
- account := accounts[i]
- if !account.IsSchedulable() || !account.IsModelSupported(model) {
- continue
- }
- if provider.SupportsAccount(&account) {
- return provider, &account, nil
- }
- }
- }
- if requestedProvider != "" {
- return nil, nil, ErrBatchImageNoAccountAvailable
- }
- return nil, nil, ErrBatchImageNoAccountAvailable
- }
- func (s *BatchImagePublicService) listCandidateAccounts(ctx context.Context, groupID *int64, platform string) ([]Account, error) {
- if s.AccountRepo == nil {
- return nil, ErrBatchImageNoAccountAvailable
- }
- if groupID != nil && *groupID > 0 {
- return s.AccountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, platform)
- }
- return s.AccountRepo.ListSchedulableByPlatform(ctx, platform)
- }
- func (s *BatchImagePublicService) ensureGroupAllowsBatchImage(ctx context.Context, groupID *int64) error {
- if groupID == nil || *groupID <= 0 {
- return nil
- }
- if s.GroupRepo == nil {
- return ErrBatchImageSettlementPricingMissing
- }
- group, err := s.GroupRepo.GetByIDLite(ctx, *groupID)
- if err != nil || group == nil {
- return ErrBatchImageSettlementPricingMissing
- }
- if !group.AllowBatchImageGeneration {
- return ErrBatchImageGroupDisabled
- }
- if group.Platform != PlatformGemini {
- return ErrBatchImageGroupDisabled
- }
- return nil
- }
- func (s *BatchImagePublicService) resolvePricingSnapshot(ctx context.Context, owner BatchImageOwner, req BatchImageSubmitRequest, provider string, account *Account) (*BatchImagePricingSnapshot, error) {
- unit := -1.0
- groupMultiplier := 1.0
- discountMultiplier := defaultBatchImageDiscountMultiplier
- holdMultiplier := defaultBatchImageHoldMultiplier
- if owner.GroupID != nil && *owner.GroupID > 0 {
- if s.GroupRepo == nil {
- return nil, ErrBatchImageSettlementPricingMissing
- }
- group, err := s.GroupRepo.GetByIDLite(ctx, *owner.GroupID)
- if err != nil || group == nil {
- return nil, ErrBatchImageSettlementPricingMissing
- }
- if !group.AllowBatchImageGeneration {
- return nil, ErrBatchImageGroupDisabled
- }
- groupDefaultMultiplier := group.RateMultiplier
- if groupDefaultMultiplier < 0 {
- groupDefaultMultiplier = 0
- }
- effectiveGroupMultiplier := groupDefaultMultiplier
- if s.UserGroupRateRepo != nil {
- userRate, rateErr := s.UserGroupRateRepo.GetByUserAndGroup(ctx, owner.UserID, group.ID)
- if rateErr != nil {
- return nil, ErrBatchImageSettlementPricingMissing
- }
- if userRate != nil {
- effectiveGroupMultiplier = *userRate
- }
- }
- groupMultiplier = effectiveGroupMultiplier
- if group.ImageRateIndependent {
- groupMultiplier = group.ImageRateMultiplier
- }
- if groupMultiplier < 0 {
- groupMultiplier = 0
- }
- discountMultiplier = group.BatchImageDiscountMultiplier
- if discountMultiplier < 0 {
- discountMultiplier = 0
- }
- if group.BatchImageHoldMultiplier >= 0 {
- holdMultiplier = group.BatchImageHoldMultiplier
- }
- if configuredUnit := group.GetImagePrice(req.ImageSize); configuredUnit != nil && *configuredUnit >= 0 {
- unit = *configuredUnit
- }
- }
- if unit < 0 {
- if s.Pricing == nil {
- return nil, ErrBatchImageSettlementPricingMissing
- }
- resolvedUnit, err := s.Pricing.BatchImageUnitPrice(ctx, &BatchImageJob{Provider: provider, Model: req.Model})
- if err != nil || resolvedUnit < 0 {
- return nil, ErrBatchImageSettlementPricingMissing
- }
- unit = resolvedUnit
- }
- // 定价不变式:hold 比例不得低于 discount 比例,否则成功率足够高时
- // actualCost > holdAmount,结算永远失败、冻结余额无法解冻。
- // 管理端已校验新配置,此处兜底钳制存量脏数据。
- if holdMultiplier < discountMultiplier {
- logger.L().Warn("batch_image.hold_multiplier_below_discount_clamped",
- zap.Float64("hold_multiplier", holdMultiplier),
- zap.Float64("discount_multiplier", discountMultiplier),
- )
- holdMultiplier = discountMultiplier
- }
- accountMultiplier := 1.0
- if account != nil {
- accountMultiplier = account.BillingRateMultiplier()
- }
- if accountMultiplier < 0 {
- accountMultiplier = 0
- }
- standardUnitPrice := unit * groupMultiplier * accountMultiplier
- billableUnitPrice := standardUnitPrice * discountMultiplier
- holdUnitPrice := standardUnitPrice * holdMultiplier
- return &BatchImagePricingSnapshot{
- BaseUnitPrice: unit,
- GroupRateMultiplier: groupMultiplier,
- AccountRateMultiplier: accountMultiplier,
- BatchDiscountMultiplier: discountMultiplier,
- HoldMultiplier: holdMultiplier,
- BillableUnitPrice: billableUnitPrice,
- HoldUnitPrice: holdUnitPrice,
- EstimatedCost: billableUnitPrice * float64(len(req.Items)),
- HoldAmount: holdUnitPrice * float64(len(req.Items)),
- }, nil
- }
- func (s *BatchImagePublicService) enabled() bool {
- return s != nil && s.Repo != nil && s.AccountRepo != nil && s.Config != nil && s.Config.BatchImage.Enabled
- }
- func (s *BatchImagePublicService) invalidateAuthCache(ctx context.Context, userID int64) {
- if s != nil && s.AuthCache != nil && userID > 0 {
- s.AuthCache.InvalidateAuthCacheByUserID(ctx, userID)
- }
- }
- func (s *BatchImagePublicService) maxItems() int {
- if s != nil && s.Config != nil && s.Config.BatchImage.MaxItemsPerJobDefault > 0 {
- return s.Config.BatchImage.MaxItemsPerJobDefault
- }
- return defaultBatchImageMaxItems
- }
- func (s *BatchImagePublicService) maxOutputImagesPerJob() int {
- if s != nil && s.Config != nil && s.Config.BatchImage.MaxOutputImagesPerJob > 0 {
- return s.Config.BatchImage.MaxOutputImagesPerJob
- }
- return defaultBatchImageMaxOutputImages
- }
- func (s *BatchImagePublicService) maxOutputImagesPerItem() int {
- if s != nil && s.Config != nil && s.Config.BatchImage.MaxOutputImagesPerItem > 0 {
- return s.Config.BatchImage.MaxOutputImagesPerItem
- }
- return defaultBatchImageMaxOutputCount
- }
- func (s *BatchImagePublicService) maxPromptChars() int {
- if s != nil && s.Config != nil && s.Config.BatchImage.MaxPromptCharsPerItem > 0 {
- return s.Config.BatchImage.MaxPromptCharsPerItem
- }
- return defaultBatchImageMaxPromptChars
- }
- func (s *BatchImagePublicService) maxReferenceImagesPerJob() int {
- if s != nil && s.Config != nil && s.Config.BatchImage.MaxReferenceImagesPerJob > 0 {
- return s.Config.BatchImage.MaxReferenceImagesPerJob
- }
- return defaultBatchImageMaxReferenceImages
- }
- func (s *BatchImagePublicService) maxReferenceInlineBytesPerJob() int {
- if s != nil && s.Config != nil && s.Config.BatchImage.MaxReferenceInlineBytesPerJob > 0 {
- return s.Config.BatchImage.MaxReferenceInlineBytesPerJob
- }
- return defaultBatchImageMaxReferenceBytes
- }
- func (s *BatchImagePublicService) defaultResponseMimeType() string {
- if s != nil && s.Config != nil && strings.TrimSpace(s.Config.BatchImage.DefaultResponseMimeType) != "" {
- return strings.TrimSpace(s.Config.BatchImage.DefaultResponseMimeType)
- }
- return defaultBatchImageResponseMime
- }
- func (s *BatchImagePublicService) defaultImageSize() string {
- if s != nil && s.Config != nil && strings.TrimSpace(s.Config.BatchImage.DefaultImageSize) != "" {
- return strings.TrimSpace(s.Config.BatchImage.DefaultImageSize)
- }
- return defaultBatchImageImageSize
- }
- func BatchImageJobToPublic(job *BatchImageJob) *BatchImagePublicBatch {
- if job == nil {
- return nil
- }
- holdAmount := job.EstimatedCost
- if job.HoldAmount != nil {
- holdAmount = *job.HoldAmount
- }
- return &BatchImagePublicBatch{
- ID: job.BatchID,
- Object: "image.batch",
- TaskName: batchImagePublicTaskName(job),
- ParentBatchID: job.ParentBatchID,
- Status: PublicBatchImageStatus(job.Status),
- Model: job.Model,
- Provider: job.Provider,
- ItemCount: job.ItemCount,
- SuccessCount: job.SuccessCount,
- FailCount: job.FailCount,
- EstimatedCost: job.EstimatedCost,
- HoldAmount: holdAmount,
- ActualCost: job.ActualCost,
- CreatedAt: job.CreatedAt.Unix(),
- SubmittedAt: batchImageUnixPtr(job.SubmittedAt),
- SettledAt: batchImageUnixPtr(job.SettledAt),
- DownloadedAt: batchImageUnixPtr(job.DownloadedAt),
- OutputDeletedAt: batchImageUnixPtr(job.OutputDeletedAt),
- }
- }
- func BatchImageItemToPublic(item *BatchImageItem) BatchImagePublicItem {
- out := BatchImagePublicItem{
- CustomID: item.CustomID,
- Status: "failed",
- PromptPreview: item.PromptPreview,
- MimeType: item.MimeType,
- FileExtension: item.FileExtension,
- ImageCount: item.ImageCount,
- }
- if item.Status == BatchImageItemStatusPending {
- out.Status = "pending"
- return out
- }
- if item.Status == BatchImageItemStatusSuccess {
- out.Status = "succeeded"
- return out
- }
- out.Error = &BatchImagePublicError{
- Code: batchImageDerefString(item.ErrorCode),
- Message: sanitizeBatchImagePublicMessage(batchImageDerefString(item.ErrorMessage)),
- Source: batchImageItemErrorSource(item),
- }
- return out
- }
- func batchImageItemErrorSource(item *BatchImageItem) string {
- if item == nil || item.ErrorCode == nil {
- return ""
- }
- code := strings.TrimSpace(*item.ErrorCode)
- if batchImageDerefString(item.ProviderSourceObject) != "" {
- return "provider"
- }
- switch code {
- case "EMPTY_IMAGE_OUTPUT", "PROVIDER_ITEM_FAILED":
- return "provider"
- case "INDEX_OUTPUT_MISSING", "INDEX_PARSE_FAILED", "DUPLICATE_CUSTOM_ID_IN_OUTPUT":
- return "system"
- default:
- return ""
- }
- }
- func PublicBatchImageStatus(status string) string {
- switch status {
- case BatchImageJobStatusCreated, BatchImageJobStatusUploading, BatchImageJobStatusSubmitted:
- return "queued"
- case BatchImageJobStatusRunning:
- return "running"
- case BatchImageJobStatusIndexing:
- return "processing_results"
- case BatchImageJobStatusSettling:
- return "settling"
- case BatchImageJobStatusCompleted:
- return "completed"
- case BatchImageJobStatusFailed:
- return "failed"
- case BatchImageJobStatusCancelled:
- return "cancelled"
- case BatchImageJobStatusOutputDeleted:
- return "output_deleted"
- default:
- return status
- }
- }
- func HashBatchImageSubmitRequest(req BatchImageSubmitRequest) string {
- req.Metadata = sanitizeBatchImageMetadata(req.Metadata)
- b, _ := json.Marshal(req)
- sum := sha256.Sum256(b)
- return hex.EncodeToString(sum[:])
- }
- func batchImageProviderPlatform(provider string) string {
- switch provider {
- case BatchImageProviderGeminiAPI, BatchImageProviderVertex:
- return PlatformGemini
- default:
- return PlatformGemini
- }
- }
- func batchImageProviderSelectionOrder(requestedProvider string) []string {
- if strings.TrimSpace(requestedProvider) != "" {
- return []string{strings.TrimSpace(requestedProvider)}
- }
- return []string{BatchImageProviderGeminiAPI, BatchImageProviderVertex}
- }
- func batchImageModelsFromAccountMapping(account *Account) []string {
- if account == nil {
- return nil
- }
- mapping := account.GetModelMapping()
- if len(mapping) == 0 {
- return nil
- }
- models := make(map[string]struct{})
- for model := range mapping {
- model = strings.TrimSpace(model)
- if model == "" {
- continue
- }
- if strings.ContainsAny(model, "*?") {
- for _, candidate := range defaultBatchImageModelCandidates() {
- if matchWildcard(model, candidate) {
- models[candidate] = struct{}{}
- }
- }
- continue
- }
- models[model] = struct{}{}
- }
- out := make([]string, 0, len(models))
- for model := range models {
- out = append(out, model)
- }
- sort.Strings(out)
- return out
- }
- func defaultBatchImageModelCandidates() []string {
- return []string{
- "gemini-2.0-flash-exp-image-generation",
- "gemini-2.5-flash-image",
- "gemini-3-pro-image",
- "gemini-3-pro-image-preview",
- "gemini-3.1-flash-image",
- "gemini-3.1-flash-image-preview",
- "gemini-3.1-flash-lite-image",
- }
- }
- func batchImageGCSRef(provider, ref string) string {
- if provider == BatchImageProviderVertex && strings.HasPrefix(strings.TrimSpace(ref), "gs://") {
- return strings.TrimSpace(ref)
- }
- return ""
- }
- func batchImageProviderSubmitPublicError(err error) error {
- reason := strings.TrimSpace(infraerrors.Reason(err))
- switch reason {
- case "VERTEX_MANAGED_GCS_BUCKET_MISSING":
- return ErrBatchImageVertexGCSBucketMissing
- case "BATCH_IMAGE_PROVIDER_MISSING_API_KEY":
- return ErrBatchImageProviderMissingAPIKey
- case "BATCH_IMAGE_PROVIDER_MISSING_SERVICE_ACCOUNT":
- return ErrBatchImageProviderMissingServiceAccount
- case "BATCH_IMAGE_PROVIDER_UNSUPPORTED_ACCOUNT":
- return ErrBatchImageProviderUnsupportedAccount
- default:
- return ErrBatchImageProviderSubmitFailed
- }
- }
- func batchImagePublicTaskName(job *BatchImageJob) string {
- if job == nil {
- return ""
- }
- if strings.TrimSpace(job.TaskName) != "" {
- return strings.TrimSpace(job.TaskName)
- }
- return defaultBatchImageTaskName(job.CreatedAt)
- }
- func defaultBatchImageTaskName(now time.Time) string {
- if now.IsZero() {
- now = time.Now()
- }
- return now.Format("2006-01-02 15:04:05")
- }
- func batchImageProviderSubmitRecordCode(err error) string {
- reason := strings.TrimSpace(infraerrors.Reason(err))
- if reason == "" || reason == "BATCH_IMAGE_PROVIDER_SUBMIT_FAILED" {
- return "PROVIDER_SUBMIT_FAILED"
- }
- return reason
- }
- func parseBatchImageListTime(raw string) *time.Time {
- raw = strings.TrimSpace(raw)
- if raw == "" {
- return nil
- }
- if unix, err := strconv.ParseInt(raw, 10, 64); err == nil && unix > 0 {
- t := time.Unix(unix, 0)
- return &t
- }
- if t, err := time.Parse(time.RFC3339, raw); err == nil {
- return &t
- }
- if t, err := time.Parse("2006-01-02", raw); err == nil {
- return &t
- }
- return nil
- }
- func sanitizeBatchImageMetadata(in map[string]string) map[string]string {
- if len(in) == 0 {
- return nil
- }
- keys := make([]string, 0, len(in))
- for k := range in {
- keys = append(keys, k)
- }
- sort.Strings(keys)
- out := make(map[string]string, len(keys))
- for _, k := range keys {
- key := strings.TrimSpace(k)
- if key == "" || len(key) > 64 {
- continue
- }
- value := strings.TrimSpace(in[k])
- if len(value) > 256 {
- value = value[:256]
- }
- out[key] = value
- if len(out) >= 20 {
- break
- }
- }
- return out
- }
- func sanitizeBatchImagePublicMessage(message string) string {
- message = strings.TrimSpace(message)
- for _, marker := range []string{"gs://", "files/", "projects/"} {
- if strings.Contains(message, marker) {
- message = "upstream provider operation failed"
- break
- }
- }
- if len(message) > maxBatchImagePublicErrorChars {
- message = message[:maxBatchImagePublicErrorChars]
- }
- return message
- }
- func batchImageUnixPtr(t *time.Time) *int64 {
- if t == nil {
- return nil
- }
- v := t.Unix()
- return &v
- }
- func parseBatchImageCursor(cursor string) int {
- offset, err := strconv.Atoi(strings.TrimSpace(cursor))
- if err != nil || offset < 0 {
- return 0
- }
- return offset
- }
|