geo/server/internal/worker/generate/kol_generation_worker.go

package generate

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"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
}

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) 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,
			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)

	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)
}

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 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 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"
}