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

package generate

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"

	"github.com/geo-platform/tenant-api/internal/shared/config"
	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
)

const defaultArticleGenerationWorkerRetryInterval = 5 * time.Second

var errArticleGenerationConsumerClosed = errors.New("article generation consumer channel closed")

type ArticleGenerationWorker struct {
	pool              *pgxpool.Pool
	rabbitMQ          *rabbitmq.Client
	templateService   *tenantapp.TemplateService
	promptRuleService *tenantapp.PromptRuleGenerationService
	kolWorker         *KolGenerationWorker
	logger            *zap.Logger
	retryInterval     time.Duration
	processTimeout    time.Duration
	consumerPrefix    string
	workerConcurrency int
}

func NewArticleGenerationWorker(
	pool *pgxpool.Pool,
	rabbitMQClient *rabbitmq.Client,
	templateService *tenantapp.TemplateService,
	promptRuleService *tenantapp.PromptRuleGenerationService,
	kolWorker *KolGenerationWorker,
	logger *zap.Logger,
	cfg config.GenerationConfig,
) *ArticleGenerationWorker {
	workerConcurrency := cfg.WorkerConcurrency
	if workerConcurrency <= 0 {
		workerConcurrency = 1
	}

	processTimeout := cfg.ArticleTimeout + 2*time.Minute
	if processTimeout <= 2*time.Minute {
		processTimeout = 10 * time.Minute
	}

	return &ArticleGenerationWorker{
		pool:              pool,
		rabbitMQ:          rabbitMQClient,
		templateService:   templateService,
		promptRuleService: promptRuleService,
		kolWorker:         kolWorker,
		logger:            logger,
		retryInterval:     defaultArticleGenerationWorkerRetryInterval,
		processTimeout:    processTimeout,
		consumerPrefix:    fmt.Sprintf("worker-generate-%d", os.Getpid()),
		workerConcurrency: workerConcurrency,
	}
}

func (w *ArticleGenerationWorker) Start(ctx context.Context) {
	if w == nil || w.pool == nil || w.rabbitMQ == nil {
		return
	}

	for i := 0; i < w.workerConcurrency; i++ {
		consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1)
		go w.run(ctx, consumerName)
	}
}

func (w *ArticleGenerationWorker) run(ctx context.Context, consumerName string) {
	for {
		if err := w.consumeOnce(ctx, consumerName); err != nil && ctx.Err() == nil && w.logger != nil {
			if errors.Is(err, errArticleGenerationConsumerClosed) {
				w.logger.Info("article generation consumer channel closed, retrying",
					zap.String("consumer", consumerName),
				)
			} else {
				w.logger.Warn("article generation worker stopped unexpectedly",
					zap.Error(err),
					zap.String("consumer", consumerName),
				)
			}
		}

		select {
		case <-ctx.Done():
			return
		case <-time.After(w.retryInterval):
		}
	}
}

func (w *ArticleGenerationWorker) consumeOnce(ctx context.Context, consumerName string) error {
	deliveries, ch, err := w.rabbitMQ.ConsumeGenerationTask(consumerName)
	if err != nil {
		return err
	}
	defer ch.Close()

	for {
		select {
		case <-ctx.Done():
			return nil
		case delivery, ok := <-deliveries:
			if !ok {
				return errArticleGenerationConsumerClosed
			}
			w.handleDelivery(ctx, delivery)
		}
	}
}

func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, delivery amqp.Delivery) {
	envelope, err := tenantapp.DecodeArticleGenerationTaskEnvelope(delivery.Body)
	if err != nil {
		w.rejectDelivery(delivery, false, err, 0)
		return
	}

	ctx, cancel := context.WithTimeout(parent, w.processTimeout)
	defer cancel()

	task, terminal, err := w.loadTask(ctx, envelope.TaskID)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil {
				w.logger.Warn("article generation ack failed for missing task",
					zap.Error(ackErr),
					zap.Int64("task_id", envelope.TaskID),
				)
			}
			return
		}
		if task.ID > 0 {
			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),
					zap.Int64("task_id", envelope.TaskID),
				)
			}
			return
		}
		requeue := !delivery.Redelivered
		w.rejectDelivery(delivery, requeue, err, envelope.TaskID)
		return
	}
	if terminal {
		if err := delivery.Ack(false); err != nil && w.logger != nil {
			w.logger.Warn("article generation ack failed for terminal task",
				zap.Error(err),
				zap.Int64("task_id", envelope.TaskID),
			)
		}
		return
	}

	switch {
	case task.ArticleID <= 0 || task.QuotaReservationID <= 0:
		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))
	}

	if err := delivery.Ack(false); err != nil && w.logger != nil {
		w.logger.Warn("article generation ack failed",
			zap.Error(err),
			zap.Int64("task_id", task.ID),
			zap.Int64("article_id", task.ArticleID),
		)
	}
}

func (w *ArticleGenerationWorker) loadTask(ctx context.Context, taskID int64) (tenantapp.ClaimedGenerationTask, bool, error) {
	var (
		task            tenantapp.ClaimedGenerationTask
		status          string
		operatorID      sql.NullInt64
		articleID       sql.NullInt64
		reservationID   sql.NullInt64
		inputParamsJSON []byte
	)

	err := w.pool.QueryRow(ctx, `
		SELECT id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, status
		FROM generation_tasks
		WHERE id = $1
	`, taskID).Scan(
		&task.ID,
		&task.TenantID,
		&operatorID,
		&articleID,
		&reservationID,
		&task.TaskType,
		&inputParamsJSON,
		&status,
	)
	if err != nil {
		return tenantapp.ClaimedGenerationTask{}, false, err
	}

	if status == "completed" || status == "failed" {
		return tenantapp.ClaimedGenerationTask{}, true, nil
	}

	task.OperatorID = operatorID.Int64
	task.ArticleID = articleID.Int64
	task.QuotaReservationID = reservationID.Int64

	params, err := parseJSONMap(inputParamsJSON)
	if err != nil {
		task.InputParams = map[string]interface{}{}
		return task, false, fmt.Errorf("decode generation input params: %w", err)
	}
	task.InputParams = params

	return task, false, nil
}

func (w *ArticleGenerationWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID int64) {
	if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil {
		w.logger.Warn("article generation nack failed",
			zap.Error(nackErr),
			zap.Int64("task_id", taskID),
		)
	}
	if w.logger != nil {
		w.logger.Warn("article generation processing failed",
			zap.Error(err),
			zap.Bool("requeue", requeue),
			zap.Bool("redelivered", delivery.Redelivered),
			zap.Int64("task_id", taskID),
		)
	}
}

func parseJSONMap(raw []byte) (map[string]interface{}, error) {
	if len(raw) == 0 {
		return map[string]interface{}{}, nil
	}

	payload := make(map[string]interface{})
	if err := json.Unmarshal(raw, &payload); err != nil {
		return nil, err
	}
	return payload, nil
}