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

feat(kol): generation service + worker branch

Browse Source
This commit is contained in:
liangxu
2026-04-17 14:52:08 +08:00
parent 26bc8a3749
commit 019410fffe
5 changed files with 660 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/worker/generate/article_generation_worker.go
+19 -2
View File
@@ -28,6 +28,7 @@ type ArticleGenerationWorker struct {
rabbitMQ *rabbitmq.Client
templateService *tenantapp.TemplateService
promptRuleService *tenantapp.PromptRuleGenerationService
kolWorker *KolGenerationWorker
logger *zap.Logger
retryInterval time.Duration
processTimeout time.Duration
@@ -40,6 +41,7 @@ func NewArticleGenerationWorker(
rabbitMQClient *rabbitmq.Client,
templateService *tenantapp.TemplateService,
promptRuleService *tenantapp.PromptRuleGenerationService,
kolWorker *KolGenerationWorker,
logger *zap.Logger,
cfg config.GenerationConfig,
) *ArticleGenerationWorker {
@@ -58,6 +60,7 @@ func NewArticleGenerationWorker(
rabbitMQ: rabbitMQClient,
templateService: templateService,
promptRuleService: promptRuleService,
kolWorker: kolWorker,
logger: logger,
retryInterval: defaultArticleGenerationWorkerRetryInterval,
processTimeout: processTimeout,
@@ -142,7 +145,11 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver
return
}
if task.ID > 0 {
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, err)
if task.TaskType == kolGenerationTaskType && w.kolWorker != nil {
w.kolWorker.HandleTaskFailure(ctx, task, "task_load", err)
} else {
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, err)
}
if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil {
w.logger.Warn("article generation ack failed after task load failure",
zap.Error(ackErr),
@@ -167,11 +174,21 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver
switch {
case task.ArticleID <= 0 || task.QuotaReservationID <= 0:
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("generation task references are incomplete"))
if task.TaskType == kolGenerationTaskType && w.kolWorker != nil {
w.kolWorker.HandleTaskFailure(ctx, task, "task_load", fmt.Errorf("generation task references are incomplete"))
} else {
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("generation task references are incomplete"))
}
case task.TaskType == "template":
w.templateService.ProcessQueuedTask(ctx, task)
case task.TaskType == "custom_generation":
w.promptRuleService.ProcessQueuedTask(ctx, task)
case task.TaskType == kolGenerationTaskType:
if w.kolWorker != nil {
w.kolWorker.Process(ctx, task)
} else {
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("kol generation worker is not configured"))
}
default:
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("unsupported generation task type: %s", task.TaskType))
}