From 019410fffe7fb59b13cd63d6cc993d5950717f98 Mon Sep 17 00:00:00 2001 From: liangxu Date: Fri, 17 Apr 2026 14:52:08 +0800 Subject: [PATCH] feat(kol): generation service + worker branch --- server/cmd/worker-generate/main.go | 9 + server/internal/tenant/app/cache_support.go | 4 + .../tenant/app/kol_generation_service.go | 401 ++++++++++++++++++ .../generate/article_generation_worker.go | 21 +- .../worker/generate/kol_generation_worker.go | 227 ++++++++++ 5 files changed, 660 insertions(+), 2 deletions(-) create mode 100644 server/internal/tenant/app/kol_generation_service.go create mode 100644 server/internal/worker/generate/kol_generation_worker.go diff --git a/server/cmd/worker-generate/main.go b/server/cmd/worker-generate/main.go index 7e58449..aa321e6 100644 --- a/server/cmd/worker-generate/main.go +++ b/server/cmd/worker-generate/main.go @@ -65,6 +65,14 @@ func main() { app.Config.LLM.MaxOutputTokens, ).WithCache(app.Cache) + kolWorker := generateworker.NewKolGenerationWorker( + app.DB, + app.LLM, + app.Logger, + generationCfg, + app.Config.LLM.MaxOutputTokens, + ).WithCache(app.Cache) + ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -73,6 +81,7 @@ func main() { app.RabbitMQ, templateSvc, promptRuleSvc, + kolWorker, app.Logger, app.Config.Generation, ).Start(ctx) diff --git a/server/internal/tenant/app/cache_support.go b/server/internal/tenant/app/cache_support.go index 70d8018..9e8dd0b 100644 --- a/server/internal/tenant/app/cache_support.go +++ b/server/internal/tenant/app/cache_support.go @@ -215,3 +215,7 @@ func invalidateArticleCaches(ctx context.Context, c sharedcache.Cache, tenantID } invalidateWorkspaceCaches(ctx, c, tenantID) } + +func InvalidateArticleCaches(ctx context.Context, c sharedcache.Cache, tenantID int64, articleID *int64) { + invalidateArticleCaches(ctx, c, tenantID, articleID) +} diff --git a/server/internal/tenant/app/kol_generation_service.go b/server/internal/tenant/app/kol_generation_service.go new file mode 100644 index 0000000..2705b41 --- /dev/null +++ b/server/internal/tenant/app/kol_generation_service.go @@ -0,0 +1,401 @@ +package app + +import ( + "context" + "encoding/json" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/geo-platform/tenant-api/internal/shared/auth" + sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache" + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" + "github.com/geo-platform/tenant-api/internal/shared/response" + "github.com/geo-platform/tenant-api/internal/tenant/repository" +) + +const ( + kolGenerationTaskType = "kol" + kolGenerationTaskRecordType = "article_generation" + kolGenerationTaskRecordScope = "generation_task" +) + +var ErrKolGenerationForbidden = response.ErrForbidden(40373, "kol_generation_forbidden", "subscription prompt is not available") + +type KolGenerationSubmitRequest struct { + Variables map[string]any `json:"variables"` +} + +type KolGenerationSubmitResponse struct { + ArticleID int64 `json:"article_id"` + GenerationTaskID int64 `json:"generation_task_id"` + UsageLogID int64 `json:"usage_log_id"` +} + +type KolSubscriptionPromptSchemaResponse struct { + SubscriptionPromptID int64 `json:"subscription_prompt_id"` + PackageID int64 `json:"package_id"` + PackageName string `json:"package_name"` + PromptID int64 `json:"prompt_id"` + PromptName string `json:"prompt_name"` + PlatformHint *string `json:"platform_hint"` + SchemaJSON *KolSchemaJSON `json:"schema_json"` + CardConfigJSON map[string]interface{} `json:"card_config_json"` +} + +type kolGenerationTaskInput struct { + SubscriptionPromptID int64 `json:"subscription_prompt_id"` + SubscriptionID int64 `json:"subscription_id"` + PackageID int64 `json:"package_id"` + PromptID int64 `json:"prompt_id"` + PromptRevisionNo int `json:"prompt_revision_no"` + Variables map[string]any `json:"variables"` + SchemaSnapshot json.RawMessage `json:"schema_snapshot"` + CardConfigSnapshot json.RawMessage `json:"card_config_snapshot"` + RenderedPromptHash string `json:"rendered_prompt_hash"` + RenderedPrompt string `json:"rendered_prompt"` +} + +type KolGenerationService struct { + pool *pgxpool.Pool + subRepo repository.KolSubscriptionRepository + promptRepo repository.KolPromptRepository + usageRepo repository.KolUsageRepository + articleRepo repository.ArticleRepository + taskRepo repository.AuditRepository + quotaRepo repository.QuotaRepository + asset *KolPromptAsset + rabbitMQ *rabbitmq.Client + cache sharedcache.Cache +} + +func NewKolGenerationService( + pool *pgxpool.Pool, + subRepo repository.KolSubscriptionRepository, + promptRepo repository.KolPromptRepository, + usageRepo repository.KolUsageRepository, + articleRepo repository.ArticleRepository, + taskRepo repository.AuditRepository, + quotaRepo repository.QuotaRepository, + asset *KolPromptAsset, + rabbitMQ *rabbitmq.Client, +) *KolGenerationService { + return &KolGenerationService{ + pool: pool, + subRepo: subRepo, + promptRepo: promptRepo, + usageRepo: usageRepo, + articleRepo: articleRepo, + taskRepo: taskRepo, + quotaRepo: quotaRepo, + asset: asset, + rabbitMQ: rabbitMQ, + } +} + +func (s *KolGenerationService) WithCache(c sharedcache.Cache) *KolGenerationService { + s.cache = c + return s +} + +func (s *KolGenerationService) GetSchema(ctx context.Context, actor auth.Actor, subPromptID int64) (*KolSubscriptionPromptSchemaResponse, error) { + row, revision, err := s.loadAuthorizedRevision(ctx, actor, subPromptID) + if err != nil { + return nil, err + } + + schema, err := decodeKolSchema(revision.SchemaJSON) + if err != nil { + return nil, response.ErrInternal(50096, "kol_generation_schema_decode_failed", err.Error()) + } + cardConfig, err := decodeKolCardConfig(revision.CardConfigJSON) + if err != nil { + return nil, response.ErrInternal(50097, "kol_generation_card_config_decode_failed", err.Error()) + } + + return &KolSubscriptionPromptSchemaResponse{ + SubscriptionPromptID: row.ID, + PackageID: row.PackageID, + PackageName: row.PackageName, + PromptID: row.PromptID, + PromptName: row.PromptName, + PlatformHint: row.PlatformHint, + SchemaJSON: schema, + CardConfigJSON: cardConfig, + }, nil +} + +func (s *KolGenerationService) Submit(ctx context.Context, actor auth.Actor, subPromptID int64, variables map[string]any) (*KolGenerationSubmitResponse, error) { + row, revision, err := s.loadAuthorizedRevision(ctx, actor, subPromptID) + if err != nil { + return nil, err + } + + if variables == nil { + variables = map[string]any{} + } + + schema, err := decodeKolSchema(revision.SchemaJSON) + if err != nil { + return nil, response.ErrInternal(50096, "kol_generation_schema_decode_failed", err.Error()) + } + + content, err := s.asset.Get(ctx, revision.PromptAssetKey) + if err != nil { + return nil, response.ErrInternal(50098, "kol_generation_prompt_asset_load_failed", err.Error()) + } + + renderedPrompt, err := RenderPrompt(content, *schema, variables) + if err != nil { + return nil, response.ErrBadRequest(40071, "kol_generation_variables_invalid", err.Error()) + } + renderedPromptHash := HashPrompt(renderedPrompt) + + balance, err := s.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") + } + + inputJSON, err := json.Marshal(kolGenerationTaskInput{ + SubscriptionPromptID: subPromptID, + SubscriptionID: row.SubscriptionID, + PackageID: row.PackageID, + PromptID: row.PromptID, + PromptRevisionNo: *row.PublishedRevisionNo, + Variables: variables, + SchemaSnapshot: json.RawMessage(revision.SchemaJSON), + CardConfigSnapshot: json.RawMessage(revision.CardConfigJSON), + RenderedPromptHash: renderedPromptHash, + RenderedPrompt: renderedPrompt, + }) + if err != nil { + return nil, response.ErrInternal(50099, "kol_generation_input_encode_failed", err.Error()) + } + + variablesJSON, err := json.Marshal(variables) + if err != nil { + return nil, response.ErrBadRequest(40072, "kol_generation_variables_invalid", "variables must be valid json") + } + + tx, err := s.pool.Begin(ctx) + if err != nil { + return nil, response.ErrInternal(50100, "kol_generation_tx_failed", "failed to begin generation transaction") + } + defer rollbackGenerationTx(tx) + + quotaTx := repository.NewQuotaRepository(tx) + articleTx := repository.NewArticleRepository(tx) + auditTx := repository.NewAuditRepository(tx) + usageTx := repository.NewKolUsageRepository(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, response.ErrInternal(50101, "kol_generation_quota_reserve_failed", err.Error()) + } + + 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, response.ErrInternal(50102, "kol_generation_reservation_create_failed", err.Error()) + } + + articleID, err := articleTx.CreateArticle(ctx, repository.CreateArticleInput{ + TenantID: actor.TenantID, + SourceType: kolGenerationTaskType, + KolPromptID: &row.PromptID, + }) + if err != nil { + return nil, response.ErrInternal(50103, "kol_generation_article_create_failed", err.Error()) + } + + if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil { + return nil, response.ErrInternal(50104, "kol_generation_reservation_update_failed", err.Error()) + } + + taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{ + OperatorID: actor.UserID, + TenantID: actor.TenantID, + ArticleID: &articleID, + QuotaReservationID: &reservationID, + TaskType: kolGenerationTaskType, + RequestHash: &renderedPromptHash, + InputParamsJSON: inputJSON, + }) + if err != nil { + return nil, response.ErrInternal(50105, "kol_generation_task_create_failed", err.Error()) + } + + 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, response.ErrInternal(50106, "kol_generation_task_record_create_failed", err.Error()) + } + + usage, err := usageTx.Create(ctx, repository.CreateKolUsageLogInput{ + TenantID: actor.TenantID, + SubscriptionID: &row.SubscriptionID, + SubscriptionPromptID: &row.ID, + PackageID: row.PackageID, + PromptID: row.PromptID, + PromptRevisionNo: *row.PublishedRevisionNo, + GenerationTaskID: &taskID, + VariablesSnapshot: variablesJSON, + SchemaSnapshot: revision.SchemaJSON, + RenderedPromptHash: &renderedPromptHash, + }) + if err != nil { + return nil, response.ErrInternal(50107, "kol_generation_usage_create_failed", err.Error()) + } + + if err := tx.Commit(ctx); err != nil { + return nil, response.ErrInternal(50108, "kol_generation_commit_failed", err.Error()) + } + + InvalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID) + + task := ClaimedGenerationTask{ + ID: taskID, + TenantID: actor.TenantID, + OperatorID: actor.UserID, + ArticleID: articleID, + QuotaReservationID: reservationID, + TaskType: kolGenerationTaskType, + } + if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{ + TaskID: taskID, + TaskType: kolGenerationTaskType, + TenantID: actor.TenantID, + ArticleID: articleID, + }); err != nil { + _ = HandleKolGenerationTaskFailure(context.Background(), s.pool, s.cache, task, "queue_enqueue", err) + return nil, response.ErrServiceUnavailable(50343, "generation_queue_unavailable", "generation queue is busy, please retry") + } + + return &KolGenerationSubmitResponse{ + ArticleID: articleID, + GenerationTaskID: taskID, + UsageLogID: usage.ID, + }, nil +} + +func (s *KolGenerationService) loadAuthorizedRevision(ctx context.Context, actor auth.Actor, subPromptID int64) (*repository.KolSubscriptionPrompt, *repository.KolPromptRevision, error) { + row, err := s.subRepo.GetSubscriptionPromptByID(ctx, actor.TenantID, subPromptID) + if err != nil { + return nil, nil, response.ErrInternal(50109, "kol_subscription_prompt_query_failed", err.Error()) + } + if row == nil { + return nil, nil, ErrKolGenerationForbidden + } + + now := time.Now() + switch { + case row.Status != "active": + return nil, nil, ErrKolGenerationForbidden + case row.SubscriptionStatus != "active": + return nil, nil, ErrKolGenerationForbidden + case row.SubscriptionEndAt != nil && row.SubscriptionEndAt.Before(now): + return nil, nil, ErrKolGenerationForbidden + case row.PackageStatus != "published": + return nil, nil, ErrKolGenerationForbidden + case row.PromptStatus != "active" || row.PublishedRevisionNo == nil: + return nil, nil, ErrKolGenerationForbidden + } + + revision, err := s.promptRepo.GetRevision(ctx, row.CreatorTenantID, row.PromptID, *row.PublishedRevisionNo) + if err != nil { + return nil, nil, response.ErrInternal(50110, "kol_prompt_revision_query_failed", err.Error()) + } + if revision == nil { + return nil, nil, response.ErrInternal(50111, "kol_prompt_revision_missing", "published prompt revision not found") + } + + return row, revision, nil +} + +func HandleKolGenerationTaskFailure( + ctx context.Context, + pool *pgxpool.Pool, + cache sharedcache.Cache, + task ClaimedGenerationTask, + stage string, + taskErr error, +) error { + cleanupCtx, cancel := newGenerationCleanupContext() + defer cancel() + + errMsg := formatGenerationFailure(stage, taskErr) + now := time.Now() + + auditRepo := repository.NewAuditRepository(pool) + articleRepo := repository.NewArticleRepository(pool) + quotaRepo := repository.NewQuotaRepository(pool) + usageRepo := repository.NewKolUsageRepository(pool) + + if task.ID > 0 { + _ = auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{ + ID: task.ID, + TenantID: task.TenantID, + Status: "failed", + ErrorMessage: &errMsg, + CompletedAt: &now, + }) + _ = auditRepo.UpdateTaskRecordStatus(cleanupCtx, repository.TaskRecordInput{ + TenantID: task.TenantID, + TaskType: kolGenerationTaskRecordType, + ResourceType: kolGenerationTaskRecordScope, + ResourceID: task.ID, + Status: "failed", + ErrorMessage: &errMsg, + }) + _ = usageRepo.MarkFailedByTask(cleanupCtx, task.TenantID, task.ID, errMsg) + } + + if task.ArticleID > 0 { + _ = articleRepo.UpdateArticleGenerateStatus(cleanupCtx, task.ArticleID, task.TenantID, "failed") + } + + if task.QuotaReservationID > 0 { + balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation") + if balanceErr == nil && task.OperatorID > 0 { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := task.ID + _, _ = quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ + TenantID: task.TenantID, + OperatorID: task.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }) + } + _ = quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID) + } + + InvalidateArticleCaches(cleanupCtx, cache, task.TenantID, &task.ArticleID) + return nil +} diff --git a/server/internal/worker/generate/article_generation_worker.go b/server/internal/worker/generate/article_generation_worker.go index fc6a558..81f0177 100644 --- a/server/internal/worker/generate/article_generation_worker.go +++ b/server/internal/worker/generate/article_generation_worker.go @@ -28,6 +28,7 @@ type ArticleGenerationWorker struct { rabbitMQ *rabbitmq.Client templateService *tenantapp.TemplateService promptRuleService *tenantapp.PromptRuleGenerationService + kolWorker *KolGenerationWorker logger *zap.Logger retryInterval time.Duration processTimeout time.Duration @@ -40,6 +41,7 @@ func NewArticleGenerationWorker( rabbitMQClient *rabbitmq.Client, templateService *tenantapp.TemplateService, promptRuleService *tenantapp.PromptRuleGenerationService, + kolWorker *KolGenerationWorker, logger *zap.Logger, cfg config.GenerationConfig, ) *ArticleGenerationWorker { @@ -58,6 +60,7 @@ func NewArticleGenerationWorker( rabbitMQ: rabbitMQClient, templateService: templateService, promptRuleService: promptRuleService, + kolWorker: kolWorker, logger: logger, retryInterval: defaultArticleGenerationWorkerRetryInterval, processTimeout: processTimeout, @@ -142,7 +145,11 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver return } if task.ID > 0 { - _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, err) + if task.TaskType == kolGenerationTaskType && w.kolWorker != nil { + w.kolWorker.HandleTaskFailure(ctx, task, "task_load", err) + } else { + _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, err) + } if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil { w.logger.Warn("article generation ack failed after task load failure", zap.Error(ackErr), @@ -167,11 +174,21 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver switch { case task.ArticleID <= 0 || task.QuotaReservationID <= 0: - _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("generation task references are incomplete")) + if task.TaskType == kolGenerationTaskType && w.kolWorker != nil { + w.kolWorker.HandleTaskFailure(ctx, task, "task_load", fmt.Errorf("generation task references are incomplete")) + } else { + _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("generation task references are incomplete")) + } case task.TaskType == "template": w.templateService.ProcessQueuedTask(ctx, task) case task.TaskType == "custom_generation": w.promptRuleService.ProcessQueuedTask(ctx, task) + case task.TaskType == kolGenerationTaskType: + if w.kolWorker != nil { + w.kolWorker.Process(ctx, task) + } else { + _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("kol generation worker is not configured")) + } default: _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("unsupported generation task type: %s", task.TaskType)) } diff --git a/server/internal/worker/generate/kol_generation_worker.go b/server/internal/worker/generate/kol_generation_worker.go new file mode 100644 index 0000000..9000f84 --- /dev/null +++ b/server/internal/worker/generate/kol_generation_worker.go @@ -0,0 +1,227 @@ +package generate + +import ( + "context" + "errors" + "strings" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "go.uber.org/zap" + + sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache" + "github.com/geo-platform/tenant-api/internal/shared/config" + "github.com/geo-platform/tenant-api/internal/shared/contentstats" + "github.com/geo-platform/tenant-api/internal/shared/llm" + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" + "github.com/geo-platform/tenant-api/internal/tenant/repository" +) + +const kolGenerationTaskType = "kol" + +type KolGenerationWorker struct { + pool *pgxpool.Pool + llm llm.Client + logger *zap.Logger + articleTimeout time.Duration + maxOutputTokens int64 + cache sharedcache.Cache +} + +func NewKolGenerationWorker( + pool *pgxpool.Pool, + llmClient llm.Client, + logger *zap.Logger, + cfg config.GenerationConfig, + llmMaxOutputTokens int64, +) *KolGenerationWorker { + articleTimeout := cfg.ArticleTimeout + if articleTimeout <= 0 { + articleTimeout = 8 * time.Minute + } + + return &KolGenerationWorker{ + pool: pool, + llm: llmClient, + logger: logger, + articleTimeout: articleTimeout, + maxOutputTokens: llmMaxOutputTokens, + } +} + +func (w *KolGenerationWorker) WithCache(c sharedcache.Cache) *KolGenerationWorker { + w.cache = c + return w +} + +func (w *KolGenerationWorker) Process(ctx context.Context, task tenantapp.ClaimedGenerationTask) { + renderedPrompt := strings.TrimSpace(kolStringInput(task.InputParams, "rendered_prompt")) + if renderedPrompt == "" { + w.HandleTaskFailure(ctx, task, "task_load", errMissingRenderedPrompt) + return + } + + now := time.Now() + auditRepo := repository.NewAuditRepository(w.pool) + if err := auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{ + ID: task.ID, + TenantID: task.TenantID, + Status: "running", + StartedAt: &now, + }); err != nil { + w.HandleTaskFailure(ctx, task, "task_load", err) + return + } + if err := auditRepo.UpdateTaskRecordStatus(ctx, repository.TaskRecordInput{ + TenantID: task.TenantID, + TaskType: "article_generation", + ResourceType: "generation_task", + ResourceID: task.ID, + Status: "running", + }); err != nil { + w.HandleTaskFailure(ctx, task, "task_load", err) + return + } + + result, err := w.llm.Generate(ctx, llm.GenerateRequest{ + Prompt: renderedPrompt, + Timeout: w.articleTimeout, + MaxOutputTokens: w.maxOutputTokens, + }, nil) + if err != nil { + w.HandleTaskFailure(ctx, task, "llm_generate", err) + return + } + + content := strings.TrimSpace(result.Content) + if content == "" { + w.HandleTaskFailure(ctx, task, "llm_generate", errEmptyGeneratedContent) + return + } + + title := kolResolveArticleTitle(content) + wordCount := contentstats.CountWords(content) + + tx, err := w.pool.Begin(ctx) + if err != nil { + w.HandleTaskFailure(ctx, task, "persist_begin_tx", err) + return + } + defer func() { + _ = tx.Rollback(ctx) + }() + + articleRepo := repository.NewArticleRepository(tx) + quotaRepo := repository.NewQuotaRepository(tx) + auditTx := repository.NewAuditRepository(tx) + usageRepo := repository.NewKolUsageRepository(tx) + + versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{ + ArticleID: task.ArticleID, + VersionNo: 1, + Title: title, + HTMLContent: "", + MarkdownContent: content, + WordCount: wordCount, + SourceLabel: result.Model, + }) + if err != nil { + w.rollbackAndFail(ctx, tx, task, "persist_create_version", err) + return + } + + if err := articleRepo.UpdateArticleCurrentVersion(ctx, task.ArticleID, task.TenantID, versionID); err != nil { + w.rollbackAndFail(ctx, tx, task, "persist_current_version", err) + return + } + if err := articleRepo.UpdateArticleGenerateStatus(ctx, task.ArticleID, task.TenantID, "completed"); err != nil { + w.rollbackAndFail(ctx, tx, task, "persist_article_status", err) + return + } + + completedAt := time.Now() + if err := auditTx.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{ + ID: task.ID, + TenantID: task.TenantID, + Status: "completed", + StartedAt: &now, + CompletedAt: &completedAt, + }); err != nil { + w.rollbackAndFail(ctx, tx, task, "persist_task_status", err) + return + } + if err := auditTx.UpdateTaskRecordStatus(ctx, repository.TaskRecordInput{ + TenantID: task.TenantID, + TaskType: "article_generation", + ResourceType: "generation_task", + ResourceID: task.ID, + Status: "completed", + }); err != nil { + w.rollbackAndFail(ctx, tx, task, "persist_task_record_status", err) + return + } + if err := usageRepo.MarkCompletedByTask(ctx, task.TenantID, task.ID, task.ArticleID); err != nil { + w.rollbackAndFail(ctx, tx, task, "persist_usage_status", err) + return + } + if err := quotaRepo.ConfirmReservation(ctx, task.QuotaReservationID, task.TenantID); err != nil { + w.rollbackAndFail(ctx, tx, task, "confirm_reservation", err) + return + } + if err := tx.Commit(ctx); err != nil { + w.rollbackAndFail(ctx, tx, task, "persist_commit", err) + return + } + + cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + tenantapp.InvalidateArticleCaches(cleanupCtx, w.cache, task.TenantID, &task.ArticleID) +} + +func (w *KolGenerationWorker) HandleTaskFailure(ctx context.Context, task tenantapp.ClaimedGenerationTask, stage string, taskErr error) { + if err := tenantapp.HandleKolGenerationTaskFailure(ctx, w.pool, w.cache, task, stage, taskErr); err != nil && w.logger != nil { + w.logger.Warn("kol generation task failure handling failed", + zap.Error(err), + zap.Int64("task_id", task.ID), + zap.String("stage", stage), + ) + } +} + +func (w *KolGenerationWorker) rollbackAndFail(ctx context.Context, tx pgx.Tx, task tenantapp.ClaimedGenerationTask, stage string, taskErr error) { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _ = tx.Rollback(cleanupCtx) + w.HandleTaskFailure(ctx, task, stage, taskErr) +} + +var ( + errMissingRenderedPrompt = errors.New("kol generation task missing rendered prompt") + errEmptyGeneratedContent = errors.New("kol generation returned empty content") +) + +func kolStringInput(input map[string]interface{}, key string) string { + if input == nil { + return "" + } + value, ok := input[key] + if !ok || value == nil { + return "" + } + text, ok := value.(string) + if !ok { + return "" + } + return text +} + +func kolResolveArticleTitle(markdown string) string { + for _, line := range strings.Split(markdown, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "#") { + return strings.TrimSpace(strings.TrimLeft(line, "#")) + } + } + return "Generated Article" +}