geo/server/internal/tenant/app/generation_task_runtime.go

package app

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/geo-platform/tenant-api/internal/shared/llm"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

type ClaimedGenerationTask struct {
	ID                 int64
	TenantID           int64
	OperatorID         int64
	ArticleID          int64
	QuotaReservationID int64
	TaskType           string
	InputParams        map[string]interface{}
}

func HandleClaimedGenerationTaskFailure(
	ctx context.Context,
	pool *pgxpool.Pool,
	templateService *TemplateService,
	promptRuleService *PromptRuleGenerationService,
	task ClaimedGenerationTask,
	taskErr error,
) error {
	errMsg := formatGenerationFailure("task_load", taskErr)
	now := time.Now()

	auditRepo := repository.NewAuditRepository(pool)
	articleRepo := repository.NewArticleRepository(pool)
	quotaRepo := repository.NewQuotaRepository(pool)

	_ = auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
		ID:           task.ID,
		TenantID:     task.TenantID,
		Status:       "failed",
		ErrorMessage: &errMsg,
		CompletedAt:  &now,
	})
	if task.ArticleID > 0 {
		_ = articleRepo.UpdateArticleGenerateStatus(ctx, task.ArticleID, task.TenantID, "failed")
	}

	if task.QuotaReservationID > 0 {
		balance, balanceErr := quotaRepo.GetCurrentBalance(ctx, task.TenantID, "article_generation")
		if balanceErr == nil && task.OperatorID > 0 {
			reason := "generation_refund"
			referenceType := "generation_task"
			taskReferenceID := task.ID
			_, _ = quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
				TenantID:      task.TenantID,
				OperatorID:    task.OperatorID,
				QuotaType:     "article_generation",
				Delta:         1,
				BalanceAfter:  balance + 1,
				Reason:        &reason,
				ReferenceType: &referenceType,
				ReferenceID:   &taskReferenceID,
			})
		}

		_ = quotaRepo.RefundReservation(ctx, task.QuotaReservationID, task.TenantID)
	}

	title := resolveArticleTitle(task.InputParams, "")
	if strings.TrimSpace(extractString(task.InputParams, "task_name")) != "" {
		title = strings.TrimSpace(extractString(task.InputParams, "task_name"))
	}
	if title == "" {
		title = promptRuleGeneratingTitlePlaceholder
	}

	if task.TaskType == "template" && templateService != nil && templateService.streamEnabled {
		templateService.streams.Fail(task.ArticleID, task.ID, title, errMsg)
	}
	if task.TaskType == "custom_generation" && promptRuleService != nil && promptRuleService.streamEnabled {
		promptRuleService.streams.Fail(task.ArticleID, task.ID, title, errMsg)
	}

	return nil
}

func (s *TemplateService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) {
	job := generationJob{
		OperatorID:    task.OperatorID,
		TenantID:      task.TenantID,
		ArticleID:     task.ArticleID,
		TaskID:        task.ID,
		ReservationID: task.QuotaReservationID,
		Params:        cloneInputParams(task.InputParams),
		InitialTitle:  resolveArticleTitle(task.InputParams, ""),
		WebSearch: llm.WebSearchOptions{
			Enabled: extractBool(task.InputParams, "enable_web_search"),
		},
	}
	if limit := extractInt(task.InputParams, "web_search_limit"); limit > 0 {
		job.WebSearch.Limit = int32(limit)
	}

	templateKey, templateName, promptTemplate := extractTemplateSnapshot(task.InputParams)
	if templateKey == "" || templateName == "" {
		record, err := s.loadTemplateRecordForTask(ctx, task)
		if err != nil {
			s.failGeneration(context.Background(), job, job.InitialTitle, "task_load", err)
			return
		}
		templateKey = record.TemplateKey
		templateName = record.TemplateName
		promptTemplate = record.PromptTemplate
	}
	job.TemplateKey = templateKey
	job.TemplateName = templateName
	job.PromptTemplate = promptTemplate

	s.executeGeneration(ctx, job)
}

func (s *TemplateService) loadTemplateRecordForTask(ctx context.Context, task ClaimedGenerationTask) (*repository.TemplateRecord, error) {
	var templateID int64
	if err := s.pool.QueryRow(ctx, `
		SELECT template_id
		FROM articles
		WHERE id = $1
		  AND tenant_id = $2
		  AND source_type = 'template'
		  AND deleted_at IS NULL
	`, task.ArticleID, task.TenantID).Scan(&templateID); err != nil {
		return nil, fmt.Errorf("load task template id: %w", err)
	}

	record, err := s.templates.GetTemplateByID(ctx, templateID, task.TenantID)
	if err != nil {
		return nil, fmt.Errorf("load task template: %w", err)
	}
	return record, nil
}

func (s *PromptRuleGenerationService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) {
	inputParams := cloneInputParams(task.InputParams)
	initialTitle := strings.TrimSpace(extractString(inputParams, "task_name"))
	if initialTitle == "" {
		initialTitle = promptRuleGeneratingTitlePlaceholder
	}

	job := promptRuleGenerationJob{
		OperatorID:    task.OperatorID,
		TenantID:      task.TenantID,
		ArticleID:     task.ArticleID,
		TaskID:        task.ID,
		ReservationID: task.QuotaReservationID,
		InputParams:   inputParams,
		InitialTitle:  initialTitle,
		WebSearch: llm.WebSearchOptions{
			Enabled: extractBool(inputParams, "enable_web_search"),
		},
	}

	rule := extractPromptRuleSnapshot(inputParams)
	if rule == nil {
		promptRuleID, err := s.resolvePromptRuleIDForTask(ctx, task)
		if err != nil {
			s.failGeneration(context.Background(), job, initialTitle, "task_load", err)
			return
		}
		rule, err = s.loadPromptRule(ctx, task.TenantID, promptRuleID)
		if err != nil {
			s.failGeneration(context.Background(), job, initialTitle, "task_load", err)
			return
		}
	}

	knowledgePrompt := strings.TrimSpace(extractString(inputParams, generationPayloadKnowledgePrompt))
	if knowledgePrompt == "" {
		if knowledgeGroupIDs := extractKnowledgeGroupIDs(inputParams["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 {
			if s.knowledge == nil {
				s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", fmt.Errorf("knowledge service is not configured"))
				return
			}

			knowledgeContext, err := s.knowledge.ResolveContext(
				ctx,
				task.TenantID,
				knowledgeGroupIDs,
				buildPromptRuleKnowledgeQuery(rule, inputParams),
			)
			if err != nil {
				s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", err)
				return
			}
			knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
		}
	}

	job.Prompt = buildPromptRuleGenerationPrompt(rule, inputParams, knowledgePrompt)
	s.executeGeneration(ctx, job)
}

func (s *PromptRuleGenerationService) resolvePromptRuleIDForTask(ctx context.Context, task ClaimedGenerationTask) (int64, error) {
	if promptRuleID := int64(extractInt(task.InputParams, "prompt_rule_id")); promptRuleID > 0 {
		return promptRuleID, nil
	}

	var promptRuleID int64
	if err := s.pool.QueryRow(ctx, `
		SELECT prompt_rule_id
		FROM articles
		WHERE id = $1
		  AND tenant_id = $2
		  AND source_type = 'custom_generation'
		  AND deleted_at IS NULL
	`, task.ArticleID, task.TenantID).Scan(&promptRuleID); err != nil {
		return 0, fmt.Errorf("load task prompt rule id: %w", err)
	}
	return promptRuleID, nil
}