root/geo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(generation): migrate article generation and template-assist to RabbitMQ queues

Browse Source 
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:
liangxu
2026-04-15 14:20:26 +08:00
parent e8f48c6d37
commit ce13331e26
14 changed files with 951 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/tenant/app/prompt_generate_service.go
+208 -61
View File
@@ -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/prompts"
@@ -22,12 +23,12 @@ import (
type PromptRuleGenerationService struct {
pool *pgxpool.Pool
llm llm.Client
rabbitMQ *rabbitmq.Client
knowledge *KnowledgeService
streams *stream.GenerationHub
streamEnabled bool
articleTimeout time.Duration
maxOutputTokens int64
jobs chan promptRuleGenerationJob
}
type promptRuleGenerationJob struct {
@@ -50,10 +51,34 @@ type GenerateFromRuleRequest struct {
}
type GenerateFromRuleResponse struct {
ArticleID int64 `json:"article_id"`
TaskID int64 `json:"task_id"`
ArticleIDs []int64 `json:"article_ids"`
TaskIDs []int64 `json:"task_ids"`
Items []GenerateFromRuleItem `json:"items"`
RequestedGenerateCount int `json:"requested_generate_count"`
SuccessCount int `json:"success_count"`
FailedCount int `json:"failed_count"`
}
type GenerateFromRuleItem struct {
ArticleID int64 `json:"article_id"`
TaskID int64 `json:"task_id"`
}
type SchedulePromptRuleGenerationInput struct {
ScheduleTaskID int64
OperatorID *int64
TenantID int64
PromptRuleID int64
BrandID *int64
Name string
TargetPlatform *string
EnableWebSearch bool
GenerateCount int
ScheduledFor time.Time
}
type promptRuleGenerationRecord struct {
ID int64
Name string
@@ -68,24 +93,25 @@ type promptRuleGenerationRecord struct {
const promptRuleGeneratingTitlePlaceholder = "标题生成中"
type enqueuePromptRuleGenerationInput struct {
OperatorID *int64
TenantID int64
PromptRuleID int64
BrandID *int64
TargetPlatform *string
ExtraParams map[string]interface{}
FallbackName string
}
func NewPromptRuleGenerationService(
pool *pgxpool.Pool,
llmClient llm.Client,
rabbitMQClient *rabbitmq.Client,
knowledge *KnowledgeService,
streams *stream.GenerationHub,
cfg config.GenerationConfig,
llmMaxOutputTokens int64,
) *PromptRuleGenerationService {
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
@@ -94,16 +120,12 @@ func NewPromptRuleGenerationService(
svc := &PromptRuleGenerationService{
pool: pool,
llm: llmClient,
rabbitMQ: rabbitMQClient,
knowledge: knowledge,
streams: streams,
streamEnabled: cfg.StreamEnabled,
articleTimeout: articleTimeout,
maxOutputTokens: llmMaxOutputTokens,
jobs: make(chan promptRuleGenerationJob, queueSize),
}
for i := 0; i < workerCount; i++ {
go svc.runGenerationWorker()
}
return svc
@@ -111,12 +133,86 @@ func NewPromptRuleGenerationService(
func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req GenerateFromRuleRequest) (*GenerateFromRuleResponse, error) {
actor := auth.MustActor(ctx)
extra := clonePromptRuleExtraParams(req.ExtraParams)
if strings.TrimSpace(extractString(extra, "generation_mode")) == "" {
extra["generation_mode"] = "instant"
}
generateCount, err := normalizePromptRuleGenerateCount(extractGenerateCount(extra))
if err != nil {
return nil, err
}
quotaRepo := repository.NewQuotaRepository(s.pool)
balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation")
if err != nil || balance < generateCount {
return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan")
}
results := make([]GenerateFromRuleItem, 0, generateCount)
var lastErr error
for idx := 0; idx < generateCount; idx++ {
result, enqueueErr := s.enqueuePromptRuleGeneration(ctx, enqueuePromptRuleGenerationInput{
OperatorID: &actor.UserID,
TenantID: actor.TenantID,
PromptRuleID: req.PromptRuleID,
BrandID: req.BrandID,
TargetPlatform: req.TargetPlatform,
ExtraParams: clonePromptRuleExtraParams(extra),
FallbackName: "即时任务",
})
if enqueueErr != nil {
lastErr = enqueueErr
continue
}
results = append(results, GenerateFromRuleItem{
ArticleID: result.ArticleID,
TaskID: result.TaskID,
})
}
if len(results) == 0 {
if lastErr != nil {
return nil, lastErr
}
return nil, response.ErrInternal(50012, "create_failed", "failed to create generation task")
}
return buildGenerateFromRuleResponse(results, generateCount), nil
}
func (s *PromptRuleGenerationService) EnqueueScheduledGeneration(ctx context.Context, input SchedulePromptRuleGenerationInput) (*GenerateFromRuleResponse, error) {
extra := map[string]interface{}{
"generation_mode": "schedule",
"schedule_task_id": input.ScheduleTaskID,
"enable_web_search": input.EnableWebSearch,
}
if input.GenerateCount > 0 {
extra["generate_count"] = input.GenerateCount
}
if name := strings.TrimSpace(input.Name); name != "" {
extra["task_name"] = name
}
if !input.ScheduledFor.IsZero() {
extra["scheduled_for"] = input.ScheduledFor.UTC().Format(time.RFC3339)
}
return s.enqueuePromptRuleGeneration(ctx, enqueuePromptRuleGenerationInput{
OperatorID: input.OperatorID,
TenantID: input.TenantID,
PromptRuleID: input.PromptRuleID,
BrandID: input.BrandID,
TargetPlatform: input.TargetPlatform,
ExtraParams: extra,
FallbackName: "定时任务",
})
}
func (s *PromptRuleGenerationService) enqueuePromptRuleGeneration(ctx context.Context, input enqueuePromptRuleGenerationInput) (*GenerateFromRuleResponse, error) {
if err := s.llm.Validate(); err != nil {
return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
}
rule, err := s.loadPromptRule(ctx, actor.TenantID, req.PromptRuleID)
rule, err := s.loadPromptRule(ctx, input.TenantID, input.PromptRuleID)
if err != nil {
return nil, err
}
@@ -124,9 +220,11 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
return nil, response.ErrBadRequest(40017, "prompt_rule_disabled", "prompt rule is disabled")
}
extra := req.ExtraParams
if extractGenerateCount(extra) > 1 {
return nil, response.ErrBadRequest(40018, "generate_count_not_supported", "prompt rule instant generation currently supports generate_count = 1 only")
extra := input.ExtraParams
generationMode := strings.TrimSpace(extractString(extra, "generation_mode"))
generateCount, err := normalizePromptRuleGenerateCount(extractGenerateCount(extra))
if err != nil {
return nil, err
}
taskName := strings.TrimSpace(extractString(extra, "task_name"))
@@ -134,34 +232,46 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
taskName = strings.TrimSpace(rule.Name)
}
if taskName == "" {
taskName = "即时任务"
taskName = strings.TrimSpace(input.FallbackName)
}
if taskName == "" {
taskName = "内容任务"
}
quotaRepo := repository.NewQuotaRepository(s.pool)
balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation")
balance, err := quotaRepo.GetCurrentBalance(ctx, input.TenantID, "article_generation")
if err != nil || balance < 1 {
return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan")
}
inputParams := map[string]interface{}{
"prompt_rule_id": req.PromptRuleID,
"prompt_rule_id": input.PromptRuleID,
"task_name": taskName,
"target_platform": strings.TrimSpace(derefString(req.TargetPlatform)),
"target_platform": strings.TrimSpace(derefString(input.TargetPlatform)),
"enable_web_search": extractBool(extra, "enable_web_search"),
"generate_count": generateCount,
}
if generationMode := strings.TrimSpace(extractString(extra, "generation_mode")); generationMode != "" {
if generationMode != "" {
inputParams["generation_mode"] = generationMode
}
platformIDs := parsePlatformIDs(derefString(req.TargetPlatform))
if scheduleTaskID := extractInt(extra, "schedule_task_id"); scheduleTaskID > 0 {
inputParams["schedule_task_id"] = scheduleTaskID
}
if scheduledFor := strings.TrimSpace(extractString(extra, "scheduled_for")); scheduledFor != "" {
inputParams["scheduled_for"] = scheduledFor
}
platformIDs := parsePlatformIDs(derefString(input.TargetPlatform))
if len(platformIDs) > 0 {
inputParams["target_platforms"] = platformIDs
}
if req.BrandID != nil {
inputParams["brand_id"] = *req.BrandID
if input.BrandID != nil {
inputParams["brand_id"] = *input.BrandID
}
if len(rule.KnowledgeGroupIDs) > 0 {
inputParams["knowledge_group_ids"] = rule.KnowledgeGroupIDs
}
inputParams = attachPromptRuleSnapshot(inputParams, rule)
knowledgePrompt := ""
if len(rule.KnowledgeGroupIDs) > 0 {
@@ -170,7 +280,7 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
}
knowledgeContext, resolveErr := s.knowledge.ResolveContext(
ctx,
actor.TenantID,
input.TenantID,
rule.KnowledgeGroupIDs,
buildPromptRuleKnowledgeQuery(rule, inputParams),
)
@@ -178,6 +288,9 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
return nil, resolveErr
}
knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
if knowledgePrompt != "" {
inputParams[generationPayloadKnowledgePrompt] = knowledgePrompt
}
}
inputJSON, _ := json.Marshal(inputParams)
@@ -193,13 +306,17 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
quotaTx := repository.NewQuotaRepository(tx)
auditTx := repository.NewAuditRepository(tx)
operatorID := int64(0)
if input.OperatorID != nil {
operatorID = *input.OperatorID
}
newBalance := balance - 1
reason := "generation_reserve"
referenceType := "generation_task"
_, err = quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
TenantID: actor.TenantID,
OperatorID: actor.UserID,
TenantID: input.TenantID,
OperatorID: operatorID,
QuotaType: "article_generation",
Delta: -1,
BalanceAfter: newBalance,
@@ -211,8 +328,8 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
}
reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{
TenantID: actor.TenantID,
OperatorID: actor.UserID,
TenantID: input.TenantID,
OperatorID: operatorID,
QuotaType: "article_generation",
ResourceType: "article",
ResourceID: nil,
@@ -228,18 +345,18 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
INSERT INTO articles (tenant_id, source_type, prompt_rule_id, generate_status, publish_status, wizard_state_json)
VALUES ($1, 'custom_generation', $2, 'generating', 'unpublished', $3)
RETURNING id
`, actor.TenantID, req.PromptRuleID, wizardStateJSON).Scan(&articleID)
`, input.TenantID, input.PromptRuleID, wizardStateJSON).Scan(&articleID)
if err != nil {
return nil, response.ErrInternal(50012, "create_failed", "failed to create prompt rule article")
}
if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil {
if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, input.TenantID, articleID); err != nil {
return nil, response.ErrInternal(50012, "create_failed", "failed to update quota reservation resource")
}
taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{
OperatorID: actor.UserID,
TenantID: actor.TenantID,
OperatorID: operatorID,
TenantID: input.TenantID,
ArticleID: &articleID,
QuotaReservationID: &reservationID,
TaskType: "custom_generation",
@@ -254,25 +371,26 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
}
job := promptRuleGenerationJob{
OperatorID: actor.UserID,
TenantID: actor.TenantID,
OperatorID: operatorID,
TenantID: input.TenantID,
ArticleID: articleID,
TaskID: taskID,
ReservationID: reservationID,
Prompt: buildPromptRuleGenerationPrompt(rule, inputParams, knowledgePrompt),
InputParams: inputParams,
InitialTitle: promptRuleGeneratingTitlePlaceholder,
WebSearch: llm.WebSearchOptions{
Enabled: extractBool(extra, "enable_web_search"),
},
}
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: "custom_generation",
TenantID: input.TenantID,
ArticleID: articleID,
}); err != nil {
s.failGeneration(context.Background(), job, promptRuleGeneratingTitlePlaceholder, "queue_enqueue", err)
return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry")
}
if s.streamEnabled {
if s.streamEnabled && s.streams != nil {
s.streams.Start(articleID, taskID, promptRuleGeneratingTitlePlaceholder)
}
@@ -322,21 +440,6 @@ func (s *PromptRuleGenerationService) loadPromptRule(ctx context.Context, tenant
return record, nil
}
func (s *PromptRuleGenerationService) enqueueGeneration(job promptRuleGenerationJob) error {
select {
case s.jobs <- job:
return nil
default:
return fmt.Errorf("generation queue is full")
}
}
func (s *PromptRuleGenerationService) runGenerationWorker() {
for job := range s.jobs {
s.executeGeneration(context.Background(), job)
}
}
func (s *PromptRuleGenerationService) executeGeneration(ctx context.Context, job promptRuleGenerationJob) {
now := time.Now()
title := strings.TrimSpace(job.InitialTitle)
@@ -557,6 +660,50 @@ func extractGenerateCount(extra map[string]interface{}) int {
return value
}
func normalizePromptRuleGenerateCount(value int) (int, error) {
if value < 1 || value > 20 {
return 0, response.ErrBadRequest(40022, "invalid_generate_count", "generate_count must be between 1 and 20")
}
return value, nil
}
func clonePromptRuleExtraParams(source map[string]interface{}) map[string]interface{} {
if len(source) == 0 {
return map[string]interface{}{}
}
cloned := make(map[string]interface{}, len(source))
for key, value := range source {
cloned[key] = value
}
return cloned
}
func buildGenerateFromRuleResponse(items []GenerateFromRuleItem, requestedCount int) *GenerateFromRuleResponse {
responseItems := make([]GenerateFromRuleItem, 0, len(items))
articleIDs := make([]int64, 0, len(items))
taskIDs := make([]int64, 0, len(items))
for _, item := range items {
responseItems = append(responseItems, item)
articleIDs = append(articleIDs, item.ArticleID)
taskIDs = append(taskIDs, item.TaskID)
}
result := &GenerateFromRuleResponse{
ArticleIDs: articleIDs,
TaskIDs: taskIDs,
Items: responseItems,
RequestedGenerateCount: requestedCount,
SuccessCount: len(items),
FailedCount: requestedCount - len(items),
}
if len(items) > 0 {
result.ArticleID = items[0].ArticleID
result.TaskID = items[0].TaskID
}
return result
}
func extractBool(extra map[string]interface{}, key string) bool {
if len(extra) == 0 {
return false