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:
@@ -13,6 +13,7 @@ import (
|
||||
"github.com/geo-platform/tenant-api/internal/shared/auth"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/llm"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/response"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/stream"
|
||||
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
||||
@@ -22,13 +23,12 @@ type TemplateService struct {
|
||||
pool *pgxpool.Pool
|
||||
templates repository.TemplateRepository
|
||||
llm llm.Client
|
||||
rabbitMQ *rabbitmq.Client
|
||||
knowledge *KnowledgeService
|
||||
streams *stream.GenerationHub
|
||||
streamEnabled bool
|
||||
articleTimeout time.Duration
|
||||
maxOutputTokens int64
|
||||
jobs chan generationJob
|
||||
assistJobs chan assistJob
|
||||
}
|
||||
|
||||
type generationJob struct {
|
||||
@@ -49,21 +49,12 @@ func NewTemplateService(
|
||||
pool *pgxpool.Pool,
|
||||
templates repository.TemplateRepository,
|
||||
llmClient llm.Client,
|
||||
rabbitMQClient *rabbitmq.Client,
|
||||
knowledge *KnowledgeService,
|
||||
streams *stream.GenerationHub,
|
||||
cfg config.GenerationConfig,
|
||||
llmMaxOutputTokens int64,
|
||||
) *TemplateService {
|
||||
queueSize := cfg.QueueSize
|
||||
if queueSize <= 0 {
|
||||
queueSize = 128
|
||||
}
|
||||
|
||||
workerCount := cfg.WorkerConcurrency
|
||||
if workerCount <= 0 {
|
||||
workerCount = 1
|
||||
}
|
||||
|
||||
articleTimeout := cfg.ArticleTimeout
|
||||
if articleTimeout <= 0 {
|
||||
articleTimeout = 8 * time.Minute
|
||||
@@ -73,18 +64,12 @@ func NewTemplateService(
|
||||
pool: pool,
|
||||
templates: templates,
|
||||
llm: llmClient,
|
||||
rabbitMQ: rabbitMQClient,
|
||||
knowledge: knowledge,
|
||||
streams: streams,
|
||||
streamEnabled: cfg.StreamEnabled,
|
||||
articleTimeout: articleTimeout,
|
||||
maxOutputTokens: llmMaxOutputTokens,
|
||||
jobs: make(chan generationJob, queueSize),
|
||||
assistJobs: make(chan assistJob, queueSize),
|
||||
}
|
||||
|
||||
for i := 0; i < workerCount; i++ {
|
||||
go svc.runGenerationWorker()
|
||||
go svc.runAssistWorker()
|
||||
}
|
||||
|
||||
return svc
|
||||
@@ -272,7 +257,13 @@ func (s *TemplateService) Generate(ctx context.Context, templateID int64, req Ge
|
||||
return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan")
|
||||
}
|
||||
|
||||
inputJSON, _ := json.Marshal(req.InputParams)
|
||||
taskInputParams := cloneInputParams(req.InputParams)
|
||||
taskInputParams["enable_web_search"] = req.EnableWebSearch
|
||||
if req.WebSearchLimit != nil {
|
||||
taskInputParams["web_search_limit"] = *req.WebSearchLimit
|
||||
}
|
||||
taskInputParams = attachTemplateSnapshot(taskInputParams, templateRecord)
|
||||
inputJSON, _ := json.Marshal(taskInputParams)
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
@@ -369,26 +360,22 @@ func (s *TemplateService) Generate(ctx context.Context, templateID int64, req Ge
|
||||
}
|
||||
|
||||
job := generationJob{
|
||||
OperatorID: actor.UserID,
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: articleID,
|
||||
TaskID: taskID,
|
||||
ReservationID: reservationID,
|
||||
TemplateKey: templateRecord.TemplateKey,
|
||||
TemplateName: templateRecord.TemplateName,
|
||||
PromptTemplate: templateRecord.PromptTemplate,
|
||||
Params: cloneInputParams(req.InputParams),
|
||||
InitialTitle: initialTitle,
|
||||
WebSearch: llm.WebSearchOptions{
|
||||
Enabled: req.EnableWebSearch,
|
||||
},
|
||||
}
|
||||
if req.WebSearchLimit != nil {
|
||||
job.WebSearch.Limit = *req.WebSearchLimit
|
||||
OperatorID: actor.UserID,
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: articleID,
|
||||
TaskID: taskID,
|
||||
ReservationID: reservationID,
|
||||
Params: taskInputParams,
|
||||
InitialTitle: initialTitle,
|
||||
}
|
||||
|
||||
if err := s.enqueueGeneration(job); err != nil {
|
||||
s.failGeneration(context.Background(), job, job.InitialTitle, "queue_enqueue", err)
|
||||
if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
|
||||
TaskID: taskID,
|
||||
TaskType: "template",
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: articleID,
|
||||
}); err != nil {
|
||||
s.failGeneration(context.Background(), job, initialTitle, "queue_enqueue", err)
|
||||
return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry")
|
||||
}
|
||||
|
||||
@@ -427,21 +414,6 @@ func updateGeneratedArticleState(
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TemplateService) enqueueGeneration(job generationJob) error {
|
||||
select {
|
||||
case s.jobs <- job:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("generation queue is full")
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TemplateService) runGenerationWorker() {
|
||||
for job := range s.jobs {
|
||||
s.executeGeneration(context.Background(), job)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TemplateService) executeGeneration(ctx context.Context, job generationJob) {
|
||||
now := time.Now()
|
||||
title := strings.TrimSpace(job.InitialTitle)
|
||||
@@ -622,6 +594,8 @@ func formatGenerationFailure(stage string, genErr error) string {
|
||||
|
||||
func generationFailureStageLabel(stage string) string {
|
||||
switch stage {
|
||||
case "task_load":
|
||||
return "加载生成任务失败"
|
||||
case "queue_enqueue":
|
||||
return "任务入队失败"
|
||||
case "llm_generate":
|
||||
|
||||
Reference in New Issue
Block a user