feat(generation): migrate article generation and template-assist to RabbitMQ queues
Replace in-process generation with queue-based async execution via RabbitMQ. Add generation task runtime for lifecycle management, article generation payload/queue abstractions, and template-assist queue publisher. Refactor prompt generation service to support batch generation with per-item error tracking. Update stream hub to bridge queue events to SSE. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,222 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user