package app import ( "context" "database/sql" "encoding/json" "errors" "fmt" "strings" "time" "github.com/jackc/pgx/v5" "github.com/geo-platform/tenant-api/internal/shared/auth" "github.com/geo-platform/tenant-api/internal/shared/response" "github.com/geo-platform/tenant-api/internal/tenant/repository" ) type RegenerateArticleResponse struct { ArticleID int64 `json:"article_id"` TaskID int64 `json:"task_id"` } type articleRegenerationRecord struct { ArticleID int64 SourceType string TemplateID *int64 KolPromptID *int64 PromptRuleID *int64 GenerateStatus string WizardStateJSON []byte SourceTaskID int64 TaskType string RequestHash *string InputParamsJSON []byte InputParams map[string]interface{} } func (s *ArticleService) Regenerate(ctx context.Context, articleID int64) (*RegenerateArticleResponse, error) { actor := auth.MustActor(ctx) if s == nil || s.llm == nil || s.rabbitMQ == nil { return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is not configured") } if err := s.llm.Validate(); err != nil { return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error()) } quotaRepo := repository.NewQuotaRepository(s.pool) balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation") if err != nil || balance < 1 { return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan") } tx, err := s.pool.Begin(ctx) if err != nil { return nil, fmt.Errorf("begin regeneration tx: %w", err) } defer rollbackGenerationTx(tx) record, err := loadArticleForRegeneration(ctx, tx, actor.TenantID, articleID) if err != nil { return nil, err } if err := validateArticleRegeneration(record); err != nil { return nil, err } quotaTx := repository.NewQuotaRepository(tx) auditTx := repository.NewAuditRepository(tx) newBalance := balance - 1 reason := "generation_reserve" referenceType := "generation_task" if _, err := quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ TenantID: actor.TenantID, OperatorID: actor.UserID, QuotaType: "article_generation", Delta: -1, BalanceAfter: newBalance, Reason: &reason, ReferenceType: &referenceType, }); err != nil { return nil, fmt.Errorf("reserve regeneration quota: %w", err) } reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{ TenantID: actor.TenantID, OperatorID: actor.UserID, QuotaType: "article_generation", ResourceType: "article", ResourceID: nil, ReservedAmount: 1, ExpireAt: time.Now().Add(time.Hour), }) if err != nil { return nil, fmt.Errorf("create regeneration reservation: %w", err) } if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil { return nil, fmt.Errorf("update regeneration reservation resource: %w", err) } title := regenerationInitialTitle(record) if err := updateGeneratedArticleState(ctx, tx, articleID, actor.TenantID, title, "generating", record.WizardStateJSON); err != nil { return nil, fmt.Errorf("mark article regenerating: %w", err) } requestHash := record.RequestHash if requestHash == nil { if hash := strings.TrimSpace(extractString(record.InputParams, "rendered_prompt_hash")); hash != "" { requestHash = &hash } } taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{ OperatorID: actor.UserID, TenantID: actor.TenantID, ArticleID: &articleID, QuotaReservationID: &reservationID, TaskType: record.TaskType, RequestHash: requestHash, InputParamsJSON: record.InputParamsJSON, }) if err != nil { return nil, fmt.Errorf("create regeneration task: %w", err) } if record.TaskType == kolGenerationTaskType { if _, err := auditTx.CreateTaskRecord(ctx, repository.TaskRecordInput{ TenantID: actor.TenantID, OperatorID: actor.UserID, TaskType: kolGenerationTaskRecordType, ResourceType: kolGenerationTaskRecordScope, ResourceID: taskID, Status: "queued", }); err != nil { return nil, fmt.Errorf("create kol regeneration task record: %w", err) } if err := createKolUsageForRegeneration(ctx, tx, actor.TenantID, taskID, record.InputParams); err != nil { return nil, fmt.Errorf("create kol regeneration usage: %w", err) } } if err := tx.Commit(ctx); err != nil { return nil, fmt.Errorf("commit regeneration task: %w", err) } invalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID) task := ClaimedGenerationTask{ ID: taskID, TenantID: actor.TenantID, OperatorID: actor.UserID, ArticleID: articleID, QuotaReservationID: reservationID, TaskType: record.TaskType, InputParams: cloneInputParams(record.InputParams), } if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{ TaskID: taskID, TaskType: record.TaskType, TenantID: actor.TenantID, ArticleID: articleID, }); err != nil { RecordGenerationTaskEvent(GenerationTaskEventQueueFailure, GenerationTaskLogContext{ TaskID: taskID, TenantID: actor.TenantID, OperatorID: actor.UserID, ArticleID: articleID, ReservationID: reservationID, TaskType: record.TaskType, Stage: "queue_enqueue", StatusBefore: "queued", StatusAfter: "failed", Result: GenerationTaskResultFailure, }) s.failRegenerationQueue(context.Background(), task, title, err) return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry") } runtimeCfg := s.runtimeConfig() if runtimeCfg.StreamEnabled && s.streams != nil { s.streams.Start(articleID, taskID, title) } return &RegenerateArticleResponse{ArticleID: articleID, TaskID: taskID}, nil } func loadArticleForRegeneration(ctx context.Context, tx pgx.Tx, tenantID, articleID int64) (*articleRegenerationRecord, error) { record := &articleRegenerationRecord{} var templateID, kolPromptID, promptRuleID sql.NullInt64 err := tx.QueryRow(ctx, ` SELECT id, source_type, template_id, kol_prompt_id, prompt_rule_id, generate_status, wizard_state_json FROM articles WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL FOR UPDATE `, articleID, tenantID).Scan( &record.ArticleID, &record.SourceType, &templateID, &kolPromptID, &promptRuleID, &record.GenerateStatus, &record.WizardStateJSON, ) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return nil, response.ErrNotFound(40411, "article_not_found", "article not found") } return nil, response.ErrInternal(50010, "query_failed", "failed to fetch article") } record.TemplateID = nullableSQLInt64(templateID) record.KolPromptID = nullableSQLInt64(kolPromptID) record.PromptRuleID = nullableSQLInt64(promptRuleID) var requestHash sql.NullString err = tx.QueryRow(ctx, ` SELECT id, task_type, request_hash, input_params_json FROM generation_tasks WHERE article_id = $1 AND tenant_id = $2 AND status = 'failed' ORDER BY (COALESCE(input_params_json, '{}'::jsonb) = '{}'::jsonb), created_at DESC, id DESC LIMIT 1 `, articleID, tenantID).Scan( &record.SourceTaskID, &record.TaskType, &requestHash, &record.InputParamsJSON, ) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return nil, response.ErrConflict(40914, "generation_task_missing", "failed generation task is missing") } return nil, response.ErrInternal(50010, "query_failed", "failed to fetch generation task") } if requestHash.Valid { record.RequestHash = &requestHash.String } if len(record.InputParamsJSON) == 0 { record.InputParamsJSON = []byte("{}") } params := map[string]interface{}{} if err := json.Unmarshal(record.InputParamsJSON, ¶ms); err != nil { return nil, response.ErrConflict(40915, "generation_params_invalid", "failed generation params are invalid") } record.InputParams = params return record, nil } func validateArticleRegeneration(record *articleRegenerationRecord) error { if record == nil { return response.ErrNotFound(40411, "article_not_found", "article not found") } if record.GenerateStatus != "failed" { return response.ErrConflict(40913, "article_not_regeneratable", "article cannot be regenerated in its current state") } switch record.SourceType { case "template": if record.TaskType == "template" && record.TemplateID != nil { return nil } case "custom_generation": if record.TaskType == "custom_generation" { return nil } case articleImitationSourceType: if record.TaskType == articleImitationTaskType { return nil } case kolGenerationTaskType: if record.TaskType == kolGenerationTaskType { return nil } } return response.ErrConflict(40913, "article_not_regeneratable", "article cannot be regenerated in its current state") } func regenerationInitialTitle(record *articleRegenerationRecord) string { if record == nil { return "Generated Article" } switch record.TaskType { case "template": return resolveArticleTitle(record.InputParams, "") case "custom_generation": if title := strings.TrimSpace(extractString(record.InputParams, "task_name")); title != "" { return title } return promptRuleGeneratingTitlePlaceholder case articleImitationTaskType: return buildImitationInitialTitle(extractString(record.InputParams, "source_title")) case kolGenerationTaskType: if title := strings.TrimSpace(extractString(record.InputParams, "prompt_name")); title != "" { return title } } if title := strings.TrimSpace(extractString(record.InputParams, "title")); title != "" { return title } return "Generated Article" } func (s *ArticleService) failRegenerationQueue(ctx context.Context, task ClaimedGenerationTask, title string, taskErr error) { if task.TaskType == kolGenerationTaskType { _ = HandleKolGenerationTaskFailure(ctx, s.pool, s.cache, task, "queue_enqueue", taskErr) } else { _ = HandleClaimedGenerationTaskFailure(ctx, s.pool, nil, nil, task, taskErr) } runtimeCfg := s.runtimeConfig() if runtimeCfg.StreamEnabled && s.streams != nil { s.streams.Fail(task.ArticleID, task.ID, title, formatGenerationFailure("queue_enqueue", taskErr)) } } func createKolUsageForRegeneration(ctx context.Context, tx pgx.Tx, tenantID, taskID int64, params map[string]interface{}) error { packageID := int64(extractInt(params, "package_id")) promptID := int64(extractInt(params, "prompt_id")) if packageID <= 0 || promptID <= 0 { return nil } promptRevisionNo := extractInt(params, "prompt_revision_no") if promptRevisionNo <= 0 { promptRevisionNo = 1 } subscriptionID := optionalInt64FromParams(params, "subscription_id") subscriptionPromptID := optionalInt64FromParams(params, "subscription_prompt_id") variablesSnapshot, err := marshalOptionalJSONField(params, "variables") if err != nil { return err } schemaSnapshot, err := marshalOptionalJSONField(params, "schema_snapshot") if err != nil { return err } renderedPromptHash := strings.TrimSpace(extractString(params, "rendered_prompt_hash")) var renderedPromptHashPtr *string if renderedPromptHash != "" { renderedPromptHashPtr = &renderedPromptHash } _, err = repository.NewKolUsageRepository(tx).Create(ctx, repository.CreateKolUsageLogInput{ TenantID: tenantID, SubscriptionID: subscriptionID, SubscriptionPromptID: subscriptionPromptID, PackageID: packageID, PromptID: promptID, PromptRevisionNo: promptRevisionNo, GenerationTaskID: &taskID, VariablesSnapshot: variablesSnapshot, SchemaSnapshot: schemaSnapshot, RenderedPromptHash: renderedPromptHashPtr, }) return err } func marshalOptionalJSONField(params map[string]interface{}, key string) ([]byte, error) { if len(params) == 0 { return nil, nil } value, ok := params[key] if !ok || value == nil { return nil, nil } return json.Marshal(value) } func optionalInt64FromParams(params map[string]interface{}, key string) *int64 { value := int64(extractInt(params, key)) if value <= 0 { return nil } return &value } func nullableSQLInt64(value sql.NullInt64) *int64 { if !value.Valid { return nil } result := value.Int64 return &result }