package generate import ( "context" "errors" "fmt" "sort" "strconv" "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 knowledge *tenantapp.KnowledgeService logger *zap.Logger cfg config.GenerationConfig llmCfg config.LLMConfig provider config.Provider cache sharedcache.Cache publishJobs *tenantapp.PublishJobService enterpriseSites *tenantapp.EnterpriseSiteService } func NewKolGenerationWorker( pool *pgxpool.Pool, llmClient llm.Client, knowledgeSvc *tenantapp.KnowledgeService, logger *zap.Logger, cfg config.GenerationConfig, llmMaxOutputTokens int64, ) *KolGenerationWorker { return &KolGenerationWorker{ pool: pool, llm: llmClient, knowledge: knowledgeSvc, logger: logger, cfg: cfg, llmCfg: config.LLMConfig{MaxOutputTokens: llmMaxOutputTokens}, } } func (w *KolGenerationWorker) WithCache(c sharedcache.Cache) *KolGenerationWorker { w.cache = c return w } func (w *KolGenerationWorker) WithPublishJobService(svc *tenantapp.PublishJobService) *KolGenerationWorker { w.publishJobs = svc return w } func (w *KolGenerationWorker) WithEnterpriseSiteService(svc *tenantapp.EnterpriseSiteService) *KolGenerationWorker { w.enterpriseSites = svc return w } func (w *KolGenerationWorker) WithConfigProvider(provider config.Provider) *KolGenerationWorker { w.provider = provider return w } func (w *KolGenerationWorker) runtimeConfig() (time.Duration, int64) { if w != nil && w.provider != nil { if cfg := w.provider.Current(); cfg != nil { return kolArticleTimeout(cfg.Generation), cfg.LLM.MaxOutputTokens } } if w == nil { return 8 * time.Minute, 0 } return kolArticleTimeout(w.cfg), w.llmCfg.MaxOutputTokens } 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 } knowledgePrompt := "" if knowledgeGroupIDs := kolKnowledgeGroupIDs(task.InputParams["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 { if w.knowledge == nil { w.HandleTaskFailure(ctx, task, "knowledge_resolve", fmt.Errorf("knowledge service is not configured")) return } knowledgeContext, err := w.knowledge.ResolveContext( ctx, task.TenantID, int64(kolIntInput(task.InputParams, "brand_id")), knowledgeGroupIDs, buildKolKnowledgeQuery(task.InputParams), ) if err != nil { w.HandleTaskFailure(ctx, task, "knowledge_resolve", err) return } knowledgePrompt = w.knowledge.RenderPromptSection(knowledgeContext) } prompt := renderedPrompt if strings.TrimSpace(knowledgePrompt) != "" { prompt = strings.TrimSpace(renderedPrompt + "\n\n" + knowledgePrompt) } articleTimeout, maxOutputTokens := w.runtimeConfig() generateReq := llm.GenerateRequest{ Prompt: prompt, Timeout: articleTimeout, MaxOutputTokens: maxOutputTokens, } if kolBoolInput(task.InputParams, "enable_web_search") { generateReq.WebSearch = &llm.WebSearchOptions{Enabled: true} } result, err := w.llm.Generate(ctx, generateReq, 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) if err := tenantapp.EnsureGenerationArticleWritable(ctx, tx, task.ArticleID, task.TenantID); err != nil { if tenantapp.IsGenerationArticleUnavailable(err) { w.rollbackAndCancelUnavailableArticle(ctx, tx, task) return } w.rollbackAndFail(ctx, tx, task, "persist_article_status", err) return } versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{ ArticleID: task.ArticleID, 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) tenantapp.StartScheduleAutoPublish(tenantapp.ScheduleAutoPublishInput{ TenantID: task.TenantID, WorkspaceID: int64(kolIntInput(task.InputParams, "workspace_id")), OperatorID: task.OperatorID, BrandID: int64(kolIntInput(task.InputParams, "brand_id")), ArticleID: task.ArticleID, GenerationTaskID: task.ID, Title: title, InputParams: task.InputParams, PublishJobs: w.publishJobs, EnterpriseSites: w.enterpriseSites, Logger: w.logger, }) } 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) } func (w *KolGenerationWorker) rollbackAndCancelUnavailableArticle(ctx context.Context, tx pgx.Tx, task tenantapp.ClaimedGenerationTask) { cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() _ = tx.Rollback(cleanupCtx) if err := tenantapp.HandleGenerationTaskArticleUnavailable(ctx, w.pool, w.cache, task, kolGenerationTaskType, true); err != nil && w.logger != nil { w.logger.Warn("kol generation article unavailable cleanup failed", zap.Error(err), zap.Int64("task_id", task.ID), zap.Int64("article_id", task.ArticleID), ) } } var ( errMissingRenderedPrompt = errors.New("kol generation task missing rendered prompt") errEmptyGeneratedContent = errors.New("kol generation returned empty content") ) func kolArticleTimeout(cfg config.GenerationConfig) time.Duration { if cfg.ArticleTimeout > 0 { return cfg.ArticleTimeout } return 8 * time.Minute } 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 kolBoolInput(input map[string]interface{}, key string) bool { if input == nil { return false } value, ok := input[key] if !ok || value == nil { return false } enabled, ok := value.(bool) return ok && enabled } func kolIntInput(input map[string]interface{}, key string) int { if input == nil { return 0 } value, ok := input[key] if !ok || value == nil { return 0 } switch typed := value.(type) { case int: return typed case int32: return int(typed) case int64: return int(typed) case float64: return int(typed) case string: parsed, _ := strconv.Atoi(strings.TrimSpace(typed)) return parsed default: return 0 } } func kolKnowledgeGroupIDs(value interface{}) []int64 { switch typed := value.(type) { case []int64: return dedupeKolIDs(typed) case []interface{}: ids := make([]int64, 0, len(typed)) for _, item := range typed { switch raw := item.(type) { case int64: ids = append(ids, raw) case int32: ids = append(ids, int64(raw)) case int: ids = append(ids, int64(raw)) case float64: ids = append(ids, int64(raw)) case float32: ids = append(ids, int64(raw)) } } return dedupeKolIDs(ids) default: return nil } } func dedupeKolIDs(ids []int64) []int64 { if len(ids) == 0 { return nil } seen := make(map[int64]struct{}, len(ids)) result := make([]int64, 0, len(ids)) for _, id := range ids { if id <= 0 { continue } if _, exists := seen[id]; exists { continue } seen[id] = struct{}{} result = append(result, id) } sort.Slice(result, func(i, j int) bool { return result[i] < result[j] }) return result } func buildKolKnowledgeQuery(input map[string]interface{}) string { promptName := strings.TrimSpace(kolStringInput(input, "prompt_name")) platformHint := strings.TrimSpace(kolStringInput(input, "platform_hint")) variables, _ := input["variables"].(map[string]interface{}) keywords := make([]string, 0) keyPoints := make([]string, 0) if len(variables) > 0 { keys := make([]string, 0, len(variables)) for key := range variables { keys = append(keys, key) } sort.Strings(keys) for _, key := range keys { valueText := strings.TrimSpace(kolValueString(variables[key])) if valueText == "" { continue } switch { case strings.Contains(strings.ToLower(key), "keyword") || strings.Contains(key, "关键词"): keywords = append(keywords, valueText) case strings.Contains(strings.ToLower(key), "brand") || strings.Contains(key, "品牌"): keyPoints = append(keyPoints, fmt.Sprintf("%s: %s", key, valueText)) default: keyPoints = append(keyPoints, fmt.Sprintf("%s: %s", key, valueText)) } } } renderedPrompt := strings.TrimSpace(kolStringInput(input, "rendered_prompt")) query := tenantapp.BuildKnowledgeQuery(tenantapp.KnowledgeQueryInput{ TaskType: firstNonEmptyKolText(platformHint, "KOL 内容生成"), TemplateName: promptName, Title: firstNonEmptyKolText(promptName, renderedPrompt), Topic: firstNonEmptyKolText(promptName, renderedPrompt), PrimaryKeyword: firstNonEmptyKolText(strings.Join(keywords, "、"), promptName), Keywords: keywords, KeyPoints: strings.Join(keyPoints, "\n"), Extra: renderedPrompt, }) if strings.TrimSpace(query) != "" { return kolLimitText(query, 2000) } return kolLimitText(renderedPrompt, 2000) } func firstNonEmptyKolText(values ...string) string { for _, value := range values { if text := strings.TrimSpace(value); text != "" { return text } } return "" } func kolValueString(value interface{}) string { switch typed := value.(type) { case string: return typed case bool: if typed { return "true" } return "false" case []interface{}: items := make([]string, 0, len(typed)) for _, item := range typed { text := strings.TrimSpace(kolValueString(item)) if text != "" { items = append(items, text) } } return strings.Join(items, ", ") default: return fmt.Sprintf("%v", typed) } } func kolLimitText(value string, maxRunes int) string { text := strings.TrimSpace(value) if text == "" || maxRunes <= 0 { return text } runes := []rune(text) if len(runes) <= maxRunes { return text } return strings.TrimSpace(string(runes[:maxRunes])) } 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" }