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))
}
server/internal/worker/generate/kol_generation_worker.go
+227
View File
@@ -0,0 +1,227 @@
package generate
import (
"context"
"errors"
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"go.uber.org/zap"
sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
"github.com/geo-platform/tenant-api/internal/shared/config"
"github.com/geo-platform/tenant-api/internal/shared/contentstats"
"github.com/geo-platform/tenant-api/internal/shared/llm"
tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
"github.com/geo-platform/tenant-api/internal/tenant/repository"
)
const kolGenerationTaskType = "kol"
type KolGenerationWorker struct {
pool *pgxpool.Pool
llm llm.Client
logger *zap.Logger
articleTimeout time.Duration
maxOutputTokens int64
cache sharedcache.Cache
}
func NewKolGenerationWorker(
pool *pgxpool.Pool,
llmClient llm.Client,
logger *zap.Logger,
cfg config.GenerationConfig,
llmMaxOutputTokens int64,
) *KolGenerationWorker {
articleTimeout := cfg.ArticleTimeout
if articleTimeout <= 0 {
articleTimeout = 8 * time.Minute
}
return &KolGenerationWorker{
pool: pool,
llm: llmClient,
logger: logger,
articleTimeout: articleTimeout,
maxOutputTokens: llmMaxOutputTokens,
}
}
func (w *KolGenerationWorker) WithCache(c sharedcache.Cache) *KolGenerationWorker {
w.cache = c
return w
}
func (w *KolGenerationWorker) Process(ctx context.Context, task tenantapp.ClaimedGenerationTask) {
renderedPrompt := strings.TrimSpace(kolStringInput(task.InputParams, "rendered_prompt"))
if renderedPrompt == "" {
w.HandleTaskFailure(ctx, task, "task_load", errMissingRenderedPrompt)
return
}
now := time.Now()
auditRepo := repository.NewAuditRepository(w.pool)
if err := auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
ID: task.ID,
TenantID: task.TenantID,
Status: "running",
StartedAt: &now,
}); err != nil {
w.HandleTaskFailure(ctx, task, "task_load", err)
return
}
if err := auditRepo.UpdateTaskRecordStatus(ctx, repository.TaskRecordInput{
TenantID: task.TenantID,
TaskType: "article_generation",
ResourceType: "generation_task",
ResourceID: task.ID,
Status: "running",
}); err != nil {
w.HandleTaskFailure(ctx, task, "task_load", err)
return
}
result, err := w.llm.Generate(ctx, llm.GenerateRequest{
Prompt: renderedPrompt,
Timeout: w.articleTimeout,
MaxOutputTokens: w.maxOutputTokens,
}, nil)
if err != nil {
w.HandleTaskFailure(ctx, task, "llm_generate", err)
return
}
content := strings.TrimSpace(result.Content)
if content == "" {
w.HandleTaskFailure(ctx, task, "llm_generate", errEmptyGeneratedContent)
return
}
title := kolResolveArticleTitle(content)
wordCount := contentstats.CountWords(content)
tx, err := w.pool.Begin(ctx)
if err != nil {
w.HandleTaskFailure(ctx, task, "persist_begin_tx", err)
return
}
defer func() {
_ = tx.Rollback(ctx)
}()
articleRepo := repository.NewArticleRepository(tx)
quotaRepo := repository.NewQuotaRepository(tx)
auditTx := repository.NewAuditRepository(tx)
usageRepo := repository.NewKolUsageRepository(tx)
versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{
ArticleID: task.ArticleID,
VersionNo: 1,
Title: title,
HTMLContent: "",
MarkdownContent: content,
WordCount: wordCount,
SourceLabel: result.Model,
})
if err != nil {
w.rollbackAndFail(ctx, tx, task, "persist_create_version", err)
return
}
if err := articleRepo.UpdateArticleCurrentVersion(ctx, task.ArticleID, task.TenantID, versionID); err != nil {
w.rollbackAndFail(ctx, tx, task, "persist_current_version", err)
return
}
if err := articleRepo.UpdateArticleGenerateStatus(ctx, task.ArticleID, task.TenantID, "completed"); err != nil {
w.rollbackAndFail(ctx, tx, task, "persist_article_status", err)
return
}
completedAt := time.Now()
if err := auditTx.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
ID: task.ID,
TenantID: task.TenantID,
Status: "completed",
StartedAt: &now,
CompletedAt: &completedAt,
}); err != nil {
w.rollbackAndFail(ctx, tx, task, "persist_task_status", err)
return
}
if err := auditTx.UpdateTaskRecordStatus(ctx, repository.TaskRecordInput{
TenantID: task.TenantID,
TaskType: "article_generation",
ResourceType: "generation_task",
ResourceID: task.ID,
Status: "completed",
}); err != nil {
w.rollbackAndFail(ctx, tx, task, "persist_task_record_status", err)
return
}
if err := usageRepo.MarkCompletedByTask(ctx, task.TenantID, task.ID, task.ArticleID); err != nil {
w.rollbackAndFail(ctx, tx, task, "persist_usage_status", err)
return
}
if err := quotaRepo.ConfirmReservation(ctx, task.QuotaReservationID, task.TenantID); err != nil {
w.rollbackAndFail(ctx, tx, task, "confirm_reservation", err)
return
}
if err := tx.Commit(ctx); err != nil {
w.rollbackAndFail(ctx, tx, task, "persist_commit", err)
return
}
cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
tenantapp.InvalidateArticleCaches(cleanupCtx, w.cache, task.TenantID, &task.ArticleID)
}
func (w *KolGenerationWorker) HandleTaskFailure(ctx context.Context, task tenantapp.ClaimedGenerationTask, stage string, taskErr error) {
if err := tenantapp.HandleKolGenerationTaskFailure(ctx, w.pool, w.cache, task, stage, taskErr); err != nil && w.logger != nil {
w.logger.Warn("kol generation task failure handling failed",
zap.Error(err),
zap.Int64("task_id", task.ID),
zap.String("stage", stage),
)
}
}
func (w *KolGenerationWorker) rollbackAndFail(ctx context.Context, tx pgx.Tx, task tenantapp.ClaimedGenerationTask, stage string, taskErr error) {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
_ = tx.Rollback(cleanupCtx)
w.HandleTaskFailure(ctx, task, stage, taskErr)
}
var (
errMissingRenderedPrompt = errors.New("kol generation task missing rendered prompt")
errEmptyGeneratedContent = errors.New("kol generation returned empty content")
)
func kolStringInput(input map[string]interface{}, key string) string {
if input == nil {
return ""
}
value, ok := input[key]
if !ok || value == nil {
return ""
}
text, ok := value.(string)
if !ok {
return ""
}
return text
}
func kolResolveArticleTitle(markdown string) string {
for _, line := range strings.Split(markdown, "\n") {
line = strings.TrimSpace(line)
if strings.HasPrefix(line, "#") {
return strings.TrimSpace(strings.TrimLeft(line, "#"))
}
}
return "Generated Article"
}