geo/server/internal/worker/assist/kol_assist_worker.go

package assist

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

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

	sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
	"github.com/geo-platform/tenant-api/internal/shared/llm"
	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

var errKolAssistConsumerClosed = errors.New("kol assist consumer channel closed")

type KolAssistWorker struct {
	pool              *pgxpool.Pool
	rabbitMQ          *rabbitmq.Client
	repo              repository.KolAssistRepository
	llm               llm.Client
	logger            *zap.Logger
	cache             sharedcache.Cache
	retryInterval     time.Duration
	processTimeout    time.Duration
	consumerPrefix    string
	workerConcurrency int
}

func NewKolAssistWorker(
	pool *pgxpool.Pool,
	rabbitMQClient *rabbitmq.Client,
	repo repository.KolAssistRepository,
	llmClient llm.Client,
	logger *zap.Logger,
	cache sharedcache.Cache,
	workerConcurrency int,
	processTimeout time.Duration,
) *KolAssistWorker {
	if workerConcurrency <= 0 {
		workerConcurrency = 1
	}
	if processTimeout <= 0 {
		processTimeout = 5 * time.Minute
	}

	return &KolAssistWorker{
		pool:              pool,
		rabbitMQ:          rabbitMQClient,
		repo:              repo,
		llm:               llmClient,
		logger:            logger,
		cache:             cache,
		retryInterval:     5 * time.Second,
		processTimeout:    processTimeout,
		consumerPrefix:    fmt.Sprintf("kol-assist-%d", os.Getpid()),
		workerConcurrency: workerConcurrency,
	}
}

func (w *KolAssistWorker) Start(ctx context.Context) {
	if w == nil || w.rabbitMQ == nil || w.repo == nil || w.llm == 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 *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJob) error {
	task, err := w.repo.Get(ctx, job.TenantID, job.TaskID)
	if err != nil {
		return fmt.Errorf("load kol assist task: %w", err)
	}
	if task == nil {
		return nil
	}
	switch task.Status {
	case "completed":
		w.completeAIPoints(ctx, job, "")
		return nil
	case "failed":
		errorMessage := "kol assist task failed"
		if task.ErrorMessage != nil && *task.ErrorMessage != "" {
			errorMessage = *task.ErrorMessage
		}
		w.refundAIPoints(ctx, job, errorMessage)
		return nil
	}

	if err := w.repo.MarkStarted(ctx, job.TenantID, job.TaskID); err != nil {
		return fmt.Errorf("mark kol assist started: %w", err)
	}

	if err := w.processTask(ctx, job, task); err != nil {
		_ = w.repo.MarkFailed(context.Background(), job.TenantID, job.TaskID, err.Error())
		w.refundAIPoints(context.Background(), job, err.Error())
		if w.logger != nil {
			w.logger.Warn("kol assist task failed",
				zap.Error(err),
				zap.String("task_id", job.TaskID),
				zap.Int64("tenant_id", job.TenantID),
			)
		}
	}

	return nil
}

func (w *KolAssistWorker) 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, errKolAssistConsumerClosed) {
				w.logger.Info("kol assist consumer channel closed, retrying",
					zap.String("consumer", consumerName),
				)
			} else {
				w.logger.Warn("kol assist worker stopped unexpectedly",
					zap.Error(err),
					zap.String("consumer", consumerName),
				)
			}
		}

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

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

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

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

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

	if err := w.Process(ctx, *job); err != nil {
		requeue := !delivery.Redelivered
		w.rejectDelivery(delivery, requeue, err, job.TaskID, job.TenantID)
		return
	}

	if err := delivery.Ack(false); err != nil && w.logger != nil {
		w.logger.Warn("kol assist ack failed",
			zap.Error(err),
			zap.String("task_id", job.TaskID),
			zap.Int64("tenant_id", job.TenantID),
		)
	}
}

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

func (w *KolAssistWorker) processTask(ctx context.Context, job tenantapp.KolAssistJob, task *repository.KolAssistTask) error {
	var req tenantapp.AssistRequest
	if err := json.Unmarshal(task.RequestJSON, &req); err != nil {
		return fmt.Errorf("decode kol assist request: %w", err)
	}

	generateReq, err := tenantapp.BuildKolAssistGenerateRequest(req)
	if err != nil {
		return err
	}

	result, err := w.llm.Generate(ctx, generateReq, nil)
	if err != nil {
		return fmt.Errorf("generate kol assist content: %w", err)
	}

	parsed, err := tenantapp.ParseKolAssistResponse(result.Content)
	if err != nil {
		return err
	}
	parsed, err = tenantapp.FinalizeKolAssistResult(req, parsed)
	if err != nil {
		return err
	}

	resultJSON, err := json.Marshal(parsed)
	if err != nil {
		return fmt.Errorf("marshal kol assist result: %w", err)
	}

	if err := w.repo.MarkCompleted(ctx, job.TenantID, job.TaskID, resultJSON); err != nil {
		return fmt.Errorf("mark kol assist completed: %w", err)
	}
	w.completeAIPoints(ctx, job, result.Model)
	return nil
}

func (w *KolAssistWorker) completeAIPoints(ctx context.Context, job tenantapp.KolAssistJob, model string) {
	resourceType := "kol_assist_task"
	resourceUID := job.TaskID
	if err := tenantapp.CompletePendingAIPointsByResource(ctx, w.pool, job.TenantID, resourceType, nil, &resourceUID, model); err != nil && w.logger != nil {
		w.logger.Warn("kol assist ai points confirmation failed",
			zap.Error(err),
			zap.String("task_id", job.TaskID),
			zap.Int64("tenant_id", job.TenantID),
		)
	}
}

func (w *KolAssistWorker) refundAIPoints(ctx context.Context, job tenantapp.KolAssistJob, errorMessage string) {
	resourceType := "kol_assist_task"
	resourceUID := job.TaskID
	if err := tenantapp.RefundPendingAIPointsByResource(ctx, w.pool, w.cache, job.TenantID, resourceType, nil, &resourceUID, errorMessage); err != nil && w.logger != nil {
		w.logger.Warn("kol assist ai points refund failed",
			zap.Error(err),
			zap.String("task_id", job.TaskID),
			zap.Int64("tenant_id", job.TenantID),
		)
	}
}