package warmupcontent import ( "context" "fmt" "strings" "time" "github.com/getsentry/sentry-go" "github.com/google/uuid" "github.com/rs/zerolog/log" "github.com/warmbly/warmbly/internal/models" "github.com/warmbly/warmbly/internal/pkg/generation" "github.com/warmbly/warmbly/internal/pkg/humanlint" "github.com/warmbly/warmbly/internal/pkg/warmlint" ) // themeForIndex mirrors runJob's theme selection: a pinned theme wins, otherwise // rotate through the default theme set so a batch spans varied topics. func themeForIndex(pinned string, i int) string { if pinned != "" { return pinned } return defaultThemes[i%len(defaultThemes)] } // GenerateBatch submits an async OpenAI Batch API run. It builds N // chat-completion requests (theme rotation identical to the sync runJob), // uploads them as a batch, and persists a job row in mode='batch' with the // OpenAI batch/file identifiers. It returns immediately — the batch is ingested // later by PollBatches when OpenAI finishes processing (up to the completion // window, typically 24h). func (s *service) GenerateBatch(ctx context.Context, req GenerateRequest) (uuid.UUID, error) { if s.gen == nil { return uuid.Nil, ErrNotConfigured } if req.PoolType == "" { req.PoolType = "premium" } if req.Trigger == "" { req.Trigger = "manual" } if req.Count <= 0 { req.Count = 100 } if req.Count > maxPerBatch { req.Count = maxPerBatch } settings, _ := s.repo.GetGenerationSettings(ctx) if req.Model == "" && settings != nil { req.Model = settings.Model } maxMessages := 4 if settings != nil { maxMessages = settings.MaxMessagesPerThread } if req.MaxMessages > 0 { maxMessages = req.MaxMessages } // Respect the daily generation cap (shared with the sync/scheduled paths) so // a huge batch can't blow past the admin's budget for the day. if settings != nil && settings.DailyGenerationCap > 0 { remaining := dailyRemaining(ctx, s.repo, settings.DailyGenerationCap) if remaining <= 0 { return uuid.Nil, fmt.Errorf("daily generation cap reached") } if req.Count > remaining { req.Count = remaining } } window := req.CompletionWindow if window == "" { window = "24h" } job := &models.WarmupGenerationJob{ ID: uuid.New(), RequestedBy: req.RequestedBy, Trigger: req.Trigger, Mode: models.WarmupGenerationModeBatch, PoolType: req.PoolType, Segment: req.Segment, Theme: req.Theme, Model: req.Model, RequestedCount: req.Count, Status: "pending", CompletionWindow: window, } // custom_id → theme so results map back after the (unordered) batch returns. requests := make([]generation.BatchRequest, 0, req.Count) for i := 0; i < req.Count; i++ { theme := themeForIndex(req.Theme, i) requests = append(requests, generation.BatchRequest{ CustomID: fmt.Sprintf("%s-%d", job.ID.String(), i), Theme: theme, Model: req.Model, MaxMessages: maxMessages, }) } batchID, inputFileID, err := s.gen.SubmitBatch(ctx, requests, window) if err != nil { // Persist a failed job row for visibility rather than dropping silently. now := time.Now() job.Status = "failed" job.Error = err.Error() job.StartedAt = &now job.FinishedAt = &now if cerr := s.repo.CreateGenerationJob(ctx, job); cerr != nil { return uuid.Nil, cerr } return uuid.Nil, err } now := time.Now() job.Status = "running" job.BatchStatus = "submitted" job.BatchID = batchID job.BatchInputFileID = inputFileID job.StartedAt = &now if err := s.repo.CreateGenerationJob(ctx, job); err != nil { return uuid.Nil, err } log.Info(). Str("job_id", job.ID.String()). Str("batch_id", batchID). Int("requested", req.Count). Str("completion_window", window). Msg("warmup batch generation submitted") return job.ID, nil } // PollBatches reconciles every in-flight batch job against OpenAI. Completed // batches are downloaded and ingested (clean + lint + cache, mirroring runJob); // failed/expired/cancelled batches mark the job failed; otherwise the latest // batch status is persisted so the admin UI reflects progress. func (s *service) PollBatches(ctx context.Context) error { if s.gen == nil { return nil } jobs, err := s.repo.ListActiveBatchJobs(ctx) if err != nil { return err } for i := range jobs { job := &jobs[i] if err := s.pollBatchJob(ctx, job); err != nil { sentry.CaptureException(err) log.Warn().Err(err).Str("job_id", job.ID.String()).Str("batch_id", job.BatchID). Msg("warmup batch generation: poll failed") } } return nil } // pollBatchJob reconciles a single batch job. func (s *service) pollBatchJob(ctx context.Context, job *models.WarmupGenerationJob) error { status, outputFileID, counts, err := s.gen.GetBatch(ctx, job.BatchID) if err != nil { return err } job.BatchStatus = status switch status { case "completed": job.BatchOutputFileID = outputFileID return s.ingestBatch(ctx, job, outputFileID, counts) case "failed", "expired", "cancelled": now := time.Now() job.Status = "failed" job.FinishedAt = &now if job.Error == "" { job.Error = fmt.Sprintf("batch %s", status) } return s.repo.UpdateGenerationJob(ctx, job) default: // validating | in_progress | finalizing | cancelling | submitted — // still running; persist the latest status for visibility. return s.repo.UpdateGenerationJob(ctx, job) } } // ingestBatch downloads a completed batch's output, cleans + lints + caches each // conversation (mirroring runJob), and finalises the job row. func (s *service) ingestBatch(ctx context.Context, job *models.WarmupGenerationJob, outputFileID string, counts generation.BatchCounts) error { results, err := s.gen.FetchBatchResults(ctx, outputFileID) if err != nil { now := time.Now() job.Status = "failed" job.FinishedAt = &now job.Error = err.Error() _ = s.repo.UpdateGenerationJob(ctx, job) return err } // Reset the per-ingest counters; the output file is the source of truth. job.GeneratedCount = 0 job.LintRejectedCount = 0 job.FailedCount = 0 for i := range results { r := &results[i] if r.Err != "" || r.Conversation == nil { job.FailedCount++ if r.Err != "" { log.Debug().Str("job_id", job.ID.String()).Str("custom_id", r.CustomID).Str("err", r.Err). Msg("warmup batch generation: result line failed") } continue } theme := themeForCustomID(r.CustomID, job.Theme) subject := strings.TrimSpace(r.Conversation.Subject) description := strings.TrimSpace(r.Conversation.Description) messages := cleanMessages(r.Conversation.Messages) if description == "" || subject == "" { job.FailedCount++ continue } // Humanize before the gates (identical to the sync runJob path). subject, description, messages = humanizeThread(subject, description, messages) lintBody := description if len(messages) > 0 { lintBody += "\n\n" + strings.Join(messages, "\n") } if humanlint.LooksRobotic(lintBody) { job.LintRejectedCount++ continue } if err := warmlint.Check(subject, lintBody, false); err != nil { job.LintRejectedCount++ continue } record := &models.WarmupConversation{ ID: uuid.New(), PoolType: job.PoolType, Segment: job.Segment, Source: models.WarmupContentSourceAI, Theme: theme, Subject: subject, Description: description, Messages: messages, Status: "active", LintPassed: true, GeneratedByJob: &job.ID, } if err := s.repo.InsertConversation(ctx, record); err != nil { job.FailedCount++ sentry.CaptureException(err) continue } job.GeneratedCount++ } // The output file already contains one line per request (success or error), // so iterating results above accounts for every line — including provider // failures, which surface as error lines. counts.Failed is therefore not // added on top (that would double-count); it's reconciled only if the output // reported fewer lines than the batch's total, which shouldn't normally // happen but guards against a truncated/partial download. if missing := counts.Total - len(results); missing > 0 { job.FailedCount += missing } now := time.Now() job.FinishedAt = &now job.Status = "completed" if job.GeneratedCount == 0 && (job.FailedCount > 0 || job.Error != "") { job.Status = "failed" if job.Error == "" { job.Error = fmt.Sprintf("all %d batch results failed", job.FailedCount) } } if err := s.repo.UpdateGenerationJob(ctx, job); err != nil { return err } log.Info(). Str("job_id", job.ID.String()). Str("batch_id", job.BatchID). Int("generated", job.GeneratedCount). Int("lint_rejected", job.LintRejectedCount). Int("failed", job.FailedCount). Msg("warmup batch generation ingested") return nil } // themeForCustomID recovers the theme for a result. The custom_id is // "-"; the index re-derives the rotated theme so cached rows carry // the right topic even though the batch output is unordered. func themeForCustomID(customID, pinnedTheme string) string { if pinnedTheme != "" { return pinnedTheme } idx := strings.LastIndexByte(customID, '-') if idx < 0 || idx+1 >= len(customID) { return defaultThemes[0] } n := 0 for _, ch := range customID[idx+1:] { if ch < '0' || ch > '9' { return defaultThemes[0] } n = n*10 + int(ch-'0') } return defaultThemes[n%len(defaultThemes)] } // CancelBatch cancels an in-flight batch job both on OpenAI and locally. func (s *service) CancelBatch(ctx context.Context, jobID uuid.UUID) error { if s.gen == nil { return ErrNotConfigured } job, err := s.repo.GetGenerationJob(ctx, jobID) if err != nil { return err } if job == nil { return fmt.Errorf("generation job not found") } if job.Mode != models.WarmupGenerationModeBatch || job.BatchID == "" { return fmt.Errorf("job is not a batch job") } if job.Status == "completed" || job.Status == "failed" { return fmt.Errorf("job already finished") } if err := s.gen.CancelBatch(ctx, job.BatchID); err != nil { return err } now := time.Now() job.Status = "failed" job.BatchStatus = "cancelling" job.FinishedAt = &now if job.Error == "" { job.Error = "cancelled by admin" } return s.repo.UpdateGenerationJob(ctx, job) }