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

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