From 40d9a6cc634064095cd58381652596ef1dee7c8b Mon Sep 17 00:00:00 2001 From: liangxu Date: Sat, 25 Apr 2026 21:36:44 +0800 Subject: [PATCH] feat(imitation): add article imitation create flow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce 仿写创作: new list and create views, backend imitation service and prompt templates, worker task routing for imitation jobs, and one-click rewrite triggers from question citation sources. Surface source article URL/title on article list/detail, and restrict failed articles to delete-only. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/components/ArticleDetailDrawer.vue | 56 +- .../src/components/ArticleSourceMeta.vue | 14 +- .../src/components/PublishArticleModal.vue | 7 - apps/admin-web/src/i18n/messages/en-US.ts | 78 ++ apps/admin-web/src/i18n/messages/zh-CN.ts | 78 ++ apps/admin-web/src/layouts/AppShell.vue | 133 +++- apps/admin-web/src/lib/api.ts | 8 + .../admin-web/src/lib/article-list-actions.ts | 12 +- apps/admin-web/src/lib/display.ts | 3 + apps/admin-web/src/lib/errors.ts | 2 +- apps/admin-web/src/router/index.ts | 20 + .../src/views/ImitationGenerateView.vue | 531 ++++++++++++++ apps/admin-web/src/views/ImitationView.vue | 632 ++++++++++++++++ .../src/views/KolMarketplaceView.vue | 165 +++-- .../src/views/TemplateWizardView.vue | 2 +- apps/admin-web/src/views/TemplatesView.vue | 2 +- .../src/views/TrackingQuestionDetailView.vue | 126 +++- apps/admin-web/src/views/WorkspaceView.vue | 2 +- packages/shared-types/src/index.ts | 29 + server/cmd/worker-generate/main.go | 11 + server/configs/prompts.yml | 41 ++ .../tenant/app/article_imitation_service.go | 686 ++++++++++++++++++ server/internal/tenant/app/article_service.go | 36 +- .../internal/tenant/app/template_service.go | 6 +- server/internal/tenant/prompts/loader.go | 3 + server/internal/tenant/prompts/runtime.go | 49 ++ .../tenant/transport/article_handler.go | 26 + server/internal/tenant/transport/router.go | 1 + .../generate/article_generation_worker.go | 39 +- 29 files changed, 2678 insertions(+), 120 deletions(-) create mode 100644 apps/admin-web/src/views/ImitationGenerateView.vue create mode 100644 apps/admin-web/src/views/ImitationView.vue create mode 100644 server/internal/tenant/app/article_imitation_service.go diff --git a/apps/admin-web/src/components/ArticleDetailDrawer.vue b/apps/admin-web/src/components/ArticleDetailDrawer.vue index ddcd6eb..89b75f7 100644 --- a/apps/admin-web/src/components/ArticleDetailDrawer.vue +++ b/apps/admin-web/src/components/ArticleDetailDrawer.vue @@ -1,5 +1,5 @@ + + + + diff --git a/apps/admin-web/src/views/ImitationView.vue b/apps/admin-web/src/views/ImitationView.vue new file mode 100644 index 0000000..5db7b72 --- /dev/null +++ b/apps/admin-web/src/views/ImitationView.vue @@ -0,0 +1,632 @@ + + + + + diff --git a/apps/admin-web/src/views/KolMarketplaceView.vue b/apps/admin-web/src/views/KolMarketplaceView.vue index 41620e0..b6b0b34 100644 --- a/apps/admin-web/src/views/KolMarketplaceView.vue +++ b/apps/admin-web/src/views/KolMarketplaceView.vue @@ -52,40 +52,44 @@ function handleCardClick(id: number) { @@ -1546,6 +1592,7 @@ const contentCitationColumns = computed(() => [ .tracking-question-source-card, .tracking-question-content-card { + position: relative; display: flex; flex-direction: column; gap: 6px; @@ -1553,6 +1600,7 @@ const contentCitationColumns = computed(() => [ border: 1px solid #edf2f7; border-radius: 10px; background: #f8fbff; + cursor: pointer; transition: all 0.2s ease; text-decoration: none; } @@ -1563,6 +1611,10 @@ const contentCitationColumns = computed(() => [ box-shadow: 0 4px 12px rgba(37, 99, 235, 0.05); } +.tracking-question-source-card { + padding-right: 58px; +} + .tracking-question-source-card.is-linked { border-color: #93c5fd; background: #eef6ff; @@ -1649,6 +1701,62 @@ const contentCitationColumns = computed(() => [ white-space: nowrap; } +.tracking-question-imitation-trigger { + display: inline-flex; + width: 28px; + height: 28px; + flex-shrink: 0; + align-items: center; + justify-content: center; + border: 1px solid #dbeafe; + border-radius: 8px; + background: #ffffff; + color: #2563eb; + cursor: pointer; + opacity: 0; + transform: translateY(1px); + transition: + opacity 0.18s ease, + transform 0.18s ease, + border-color 0.18s ease, + background-color 0.18s ease; +} + +.tracking-question-imitation-trigger:hover { + border-color: #93c5fd; + background: #eff6ff; +} + +.tracking-question-imitation-trigger--source { + position: absolute; + top: 50%; + right: 16px; + transform: translateY(calc(-50% + 1px)); +} + +.tracking-question-source-card:hover .tracking-question-imitation-trigger, +.tracking-question-source-card:focus-within .tracking-question-imitation-trigger { + opacity: 1; + transform: translateY(0); +} + +.tracking-question-source-card:hover .tracking-question-imitation-trigger--source, +.tracking-question-source-card:focus-within .tracking-question-imitation-trigger--source { + opacity: 1; + transform: translateY(-50%); +} + +.tracking-question-imitation-trigger--table-action { + margin: 0 auto; + opacity: 1; + transform: none; +} + +.tracking-question-table__empty-action { + color: #cbd5e1; + font-size: 12px; +} + .tracking-question-source-card__link { color: #94a3b8; font-size: 12px; diff --git a/apps/admin-web/src/views/WorkspaceView.vue b/apps/admin-web/src/views/WorkspaceView.vue index 0e5624a..0370320 100644 --- a/apps/admin-web/src/views/WorkspaceView.vue +++ b/apps/admin-web/src/views/WorkspaceView.vue @@ -190,7 +190,7 @@ async function openEditor(article: RecentArticle): Promise { try { const detail = await articlesApi.detail(article.id); - if (detail.generate_status !== "completed" && detail.source_type === "template" && detail.template_id) { + if (detail.generate_status === "draft" && detail.source_type === "template" && detail.template_id) { await router.push({ name: "article-wizard", query: { diff --git a/packages/shared-types/src/index.ts b/packages/shared-types/src/index.ts index 4b920b7..033a4a9 100644 --- a/packages/shared-types/src/index.ts +++ b/packages/shared-types/src/index.ts @@ -502,6 +502,8 @@ export interface ArticleListItem { generation_error_message: string | null; word_count: number; source_label: string | null; + source_url?: string | null; + source_title?: string | null; created_at: string; } @@ -528,6 +530,8 @@ export interface ArticleDetail { markdown_content: string | null; word_count: number; source_label: string | null; + source_url?: string | null; + source_title?: string | null; version_no: number | null; generation_error_message: string | null; wizard_state?: Record | null; @@ -538,6 +542,31 @@ export interface CreateArticleRequest { title?: string; } +export interface GenerateImitationRequest { + source_url: string; + source_title?: string; + locale?: string; + industry?: string; + brand_name?: string; + region?: string; + target_audience?: string; + content_goal?: string; + tone?: string; + length_goal?: string; + keywords?: string[]; + preserve_points?: string; + avoid_points?: string; + extra_requirements?: string; + enable_web_search?: boolean; + web_search_limit?: number; + knowledge_group_ids?: number[]; +} + +export interface GenerateImitationResponse { + article_id: number; + task_id: number; +} + export interface UpdateArticleRequest { title: string; markdown_content: string; diff --git a/server/cmd/worker-generate/main.go b/server/cmd/worker-generate/main.go index 2a441f0..b7f313a 100644 --- a/server/cmd/worker-generate/main.go +++ b/server/cmd/worker-generate/main.go @@ -65,6 +65,16 @@ func main() { app.Config.LLM.MaxOutputTokens, ).WithCache(app.Cache) + imitationSvc := tenantapp.NewArticleImitationService( + app.DB, + app.LLM, + nil, + knowledgeSvc, + app.GenerationStreams, + generationCfg, + app.Config.LLM.MaxOutputTokens, + ).WithCache(app.Cache) + kolWorker := generateworker.NewKolGenerationWorker( app.DB, app.LLM, @@ -82,6 +92,7 @@ func main() { app.RabbitMQ, templateSvc, promptRuleSvc, + imitationSvc, kolWorker, app.Logger, app.Config.Generation, diff --git a/server/configs/prompts.yml b/server/configs/prompts.yml index debd97d..0da95a8 100644 --- a/server/configs/prompts.yml +++ b/server/configs/prompts.yml @@ -208,6 +208,47 @@ runtime: 2. 标题使用一级标题，并由 AI 根据正文内容自行生成，不要直接复用任务名称。 3. 不要输出你的思考过程、解释或致歉。 4. 保持客观中立，避免排名、绝对化承诺、权威背书、夸大宣传和刻意贬损同行。 + article_imitation_prompt_template: | + 你是一名资深内容策略编辑，请基于给定源文章做高质量仿写创作。 + + 核心要求： + - 保留源文章中的事实信息、论证逻辑和可复用观点，但必须重写标题、结构、表达和段落组织。 + - 不要逐句同义替换，不要复制源文的独特措辞，不要编造源文、知识库和补充要求之外的事实。 + - 标题也必须仿写改写：第一行必须是新的一级标题，需仿照源标题的信息焦点、受众意图和表达节奏，但不能直接复用源标题，也不能只做简单词语替换。 + - 如果给出了地域、品牌、目标读者或内容目标，标题和正文都要自然贴合这些上下文。 + %s + + 输出格式：Markdown。第一行必须是重写后的一级标题，正文结构清晰，可直接进入文章编辑器。 + + ## 源文章 + %s + + ## 仿写设置 + %s + + 如果某项设置为空，请根据源文章和行业常识自行选择最贴合的写法。 + + %s + + ## 源文章内容 + %s + + 请开始输出仿写后的完整文章。 + article_imitation_locale_instructions: + zh-CN: "输出语言：中文简体。" + en-US: "Output language: English." + article_imitation_setting_labels: + industry: "行业/场景" + brand_name: "品牌/主体" + region: "地域" + target_audience: "目标读者" + content_goal: "内容目标" + tone: "语气风格" + length_goal: "篇幅要求" + keywords: "关键词" + preserve_points: "必须保留/强调" + avoid_points: "需要规避" + extra_requirements: "其他要求" knowledge_prompt_intro_lines: - "以下为可参考的知识库内容，只用于补充事实背景、术语定义、产品信息或表达素材：" - "- 优先吸收并改写，但地址、电话、官网、邮箱、营业时间、价格、日期，以及主营业务、经营范围、服务范围、核心产品等业务事实如需引用，必须严格依据原文，不得写成知识库里没有的业务内容。" diff --git a/server/internal/tenant/app/article_imitation_service.go b/server/internal/tenant/app/article_imitation_service.go new file mode 100644 index 0000000..fd5c36d --- /dev/null +++ b/server/internal/tenant/app/article_imitation_service.go @@ -0,0 +1,686 @@ +package app + +import ( + "context" + "encoding/json" + "fmt" + "net/url" + "strings" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/geo-platform/tenant-api/internal/shared/auth" + 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/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" + "github.com/geo-platform/tenant-api/internal/tenant/repository" +) + +const ( + articleImitationTaskType = "imitation" + articleImitationSourceType = "imitation" + articleImitationFallbackTitle = "仿写文章生成中" + maxImitationSourceContentRunes = 24000 + defaultImitationWebSearchLimit = int32(5) +) + +type ArticleImitationService struct { + pool *pgxpool.Pool + llm llm.Client + rabbitMQ *rabbitmq.Client + knowledge *KnowledgeService + streams *stream.GenerationHub + streamEnabled bool + articleTimeout time.Duration + maxOutputTokens int64 + cache sharedcache.Cache +} + +type GenerateImitationRequest struct { + SourceURL string `json:"source_url" binding:"required"` + SourceTitle string `json:"source_title"` + Locale string `json:"locale"` + Industry string `json:"industry"` + BrandName string `json:"brand_name"` + Region string `json:"region"` + TargetAudience string `json:"target_audience"` + ContentGoal string `json:"content_goal"` + Tone string `json:"tone"` + LengthGoal string `json:"length_goal"` + Keywords []string `json:"keywords"` + PreservePoints string `json:"preserve_points"` + AvoidPoints string `json:"avoid_points"` + ExtraRequirements string `json:"extra_requirements"` + EnableWebSearch bool `json:"enable_web_search"` + WebSearchLimit *int32 `json:"web_search_limit"` + KnowledgeGroupIDs []int64 `json:"knowledge_group_ids"` +} + +type GenerateImitationResponse struct { + ArticleID int64 `json:"article_id"` + TaskID int64 `json:"task_id"` +} + +type articleImitationJob struct { + OperatorID int64 + TenantID int64 + ArticleID int64 + TaskID int64 + ReservationID int64 + InputParams map[string]interface{} + InitialTitle string + WebSearch llm.WebSearchOptions +} + +func NewArticleImitationService( + pool *pgxpool.Pool, + llmClient llm.Client, + rabbitMQClient *rabbitmq.Client, + knowledge *KnowledgeService, + streams *stream.GenerationHub, + cfg config.GenerationConfig, + llmMaxOutputTokens int64, +) *ArticleImitationService { + articleTimeout := cfg.ArticleTimeout + if articleTimeout <= 0 { + articleTimeout = 8 * time.Minute + } + + return &ArticleImitationService{ + pool: pool, + llm: llmClient, + rabbitMQ: rabbitMQClient, + knowledge: knowledge, + streams: streams, + streamEnabled: cfg.StreamEnabled, + articleTimeout: articleTimeout, + maxOutputTokens: llmMaxOutputTokens, + } +} + +func (s *ArticleImitationService) WithCache(c sharedcache.Cache) *ArticleImitationService { + s.cache = c + return s +} + +func (s *ArticleImitationService) Generate(ctx context.Context, req GenerateImitationRequest) (*GenerateImitationResponse, error) { + actor := auth.MustActor(ctx) + if err := s.llm.Validate(); err != nil { + return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error()) + } + + sourceURL, err := normalizeImitationSourceURL(req.SourceURL) + if err != nil { + return nil, err + } + + quotaRepo := repository.NewQuotaRepository(s.pool) + balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation") + if err != nil || balance < 1 { + return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan") + } + + inputParams := buildImitationInputParams(sourceURL, req) + inputJSON, _ := json.Marshal(inputParams) + initialTitle := buildImitationInitialTitle(req.SourceTitle) + wizardStateJSON, _ := json.Marshal(map[string]interface{}{ + "title": initialTitle, + "source_url": sourceURL, + "source_title": strings.TrimSpace(req.SourceTitle), + "current_step": 0, + }) + + tx, err := s.pool.Begin(ctx) + if err != nil { + return nil, response.ErrInternal(50012, "create_failed", "failed to begin generation transaction") + } + defer func() { + _ = tx.Rollback(ctx) + }() + + quotaTx := repository.NewQuotaRepository(tx) + auditTx := repository.NewAuditRepository(tx) + + newBalance := balance - 1 + reason := "generation_reserve" + referenceType := "generation_task" + _, err = quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ + TenantID: actor.TenantID, + OperatorID: actor.UserID, + QuotaType: "article_generation", + Delta: -1, + BalanceAfter: newBalance, + Reason: &reason, + ReferenceType: &referenceType, + }) + if err != nil { + return nil, response.ErrInternal(50012, "create_failed", "failed to reserve generation quota") + } + + reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{ + TenantID: actor.TenantID, + OperatorID: actor.UserID, + QuotaType: "article_generation", + ResourceType: "article", + ResourceID: nil, + ReservedAmount: 1, + ExpireAt: time.Now().Add(time.Hour), + }) + if err != nil { + return nil, response.ErrInternal(50012, "create_failed", "failed to create quota reservation") + } + + var articleID int64 + err = tx.QueryRow(ctx, ` + INSERT INTO articles (tenant_id, source_type, generate_status, publish_status, wizard_state_json) + VALUES ($1, $2, 'generating', 'unpublished', $3) + RETURNING id + `, actor.TenantID, articleImitationSourceType, wizardStateJSON).Scan(&articleID) + if err != nil { + return nil, response.ErrInternal(50012, "create_failed", "failed to create imitation article") + } + + if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.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, + ArticleID: &articleID, + QuotaReservationID: &reservationID, + TaskType: articleImitationTaskType, + InputParamsJSON: inputJSON, + }) + if err != nil { + return nil, response.ErrInternal(50012, "create_failed", "failed to create generation task") + } + + if err := tx.Commit(ctx); err != nil { + return nil, response.ErrInternal(50012, "create_failed", "failed to commit generation task") + } + invalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID) + + job := articleImitationJob{ + OperatorID: actor.UserID, + TenantID: actor.TenantID, + ArticleID: articleID, + TaskID: taskID, + ReservationID: reservationID, + InputParams: inputParams, + InitialTitle: initialTitle, + WebSearch: llm.WebSearchOptions{ + Enabled: req.EnableWebSearch, + }, + } + if req.WebSearchLimit != nil && *req.WebSearchLimit > 0 { + job.WebSearch.Limit = *req.WebSearchLimit + } else if req.EnableWebSearch { + job.WebSearch.Limit = defaultImitationWebSearchLimit + } + + if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{ + TaskID: taskID, + TaskType: articleImitationTaskType, + TenantID: actor.TenantID, + ArticleID: articleID, + }); err != nil { + s.failGeneration(context.Background(), job, "queue_enqueue", err) + return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry") + } + + if s.streamEnabled && s.streams != nil { + s.streams.Start(articleID, taskID, initialTitle) + } + + return &GenerateImitationResponse{ArticleID: articleID, TaskID: taskID}, nil +} + +func (s *ArticleImitationService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) { + inputParams := cloneInputParams(task.InputParams) + job := articleImitationJob{ + OperatorID: task.OperatorID, + TenantID: task.TenantID, + ArticleID: task.ArticleID, + TaskID: task.ID, + ReservationID: task.QuotaReservationID, + InputParams: inputParams, + InitialTitle: buildImitationInitialTitle(extractString(inputParams, "source_title")), + WebSearch: llm.WebSearchOptions{ + Enabled: extractBool(inputParams, "enable_web_search"), + }, + } + if limit := extractInt(inputParams, "web_search_limit"); limit > 0 { + job.WebSearch.Limit = int32(limit) + } else if job.WebSearch.Enabled { + job.WebSearch.Limit = defaultImitationWebSearchLimit + } + + s.executeGeneration(ctx, job) +} + +func (s *ArticleImitationService) HandleTaskFailure(ctx context.Context, task ClaimedGenerationTask, stage string, taskErr error) { + job := articleImitationJob{ + OperatorID: task.OperatorID, + TenantID: task.TenantID, + ArticleID: task.ArticleID, + TaskID: task.ID, + ReservationID: task.QuotaReservationID, + InputParams: cloneInputParams(task.InputParams), + InitialTitle: buildImitationInitialTitle(extractString(task.InputParams, "source_title")), + } + s.failGeneration(ctx, job, stage, taskErr) +} + +func (s *ArticleImitationService) executeGeneration(ctx context.Context, job articleImitationJob) { + now := time.Now() + title := strings.TrimSpace(job.InitialTitle) + if title == "" { + title = articleImitationFallbackTitle + } + + auditRepo := repository.NewAuditRepository(s.pool) + _ = auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{ + ID: job.TaskID, + TenantID: job.TenantID, + Status: "running", + StartedAt: &now, + }) + + sourceURL := strings.TrimSpace(extractString(job.InputParams, "source_url")) + if sourceURL == "" { + s.failGeneration(ctx, job, "task_load", fmt.Errorf("source_url is required")) + return + } + if s.knowledge == nil { + s.failGeneration(ctx, job, "source_fetch", fmt.Errorf("knowledge service is not configured")) + return + } + + parsed, err := s.knowledge.parseWebsiteKnowledge(ctx, sourceURL) + if err != nil { + s.failGeneration(ctx, job, "source_fetch", err) + return + } + + sourceContent := truncateRunes(strings.TrimSpace(parsed.Markdown), maxImitationSourceContentRunes) + if sourceContent == "" { + sourceContent = truncateRunes(strings.TrimSpace(parsed.Text), maxImitationSourceContentRunes) + } + if sourceContent == "" { + s.failGeneration(ctx, job, "source_fetch", fmt.Errorf("source article content is empty")) + return + } + + knowledgePrompt := "" + if knowledgeGroupIDs := extractKnowledgeGroupIDs(job.InputParams["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 { + if s.knowledge == nil { + s.failGeneration(ctx, job, "knowledge_resolve", fmt.Errorf("knowledge service is not configured")) + return + } + knowledgeContext, resolveErr := s.knowledge.ResolveContext( + ctx, + job.TenantID, + knowledgeGroupIDs, + buildImitationKnowledgeQuery(job.InputParams), + ) + if resolveErr != nil { + s.failGeneration(ctx, job, "knowledge_resolve", resolveErr) + return + } + knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext) + } + + prompt := buildImitationGenerationPrompt(job.InputParams, sourceContent, knowledgePrompt) + req := llm.GenerateRequest{ + Prompt: prompt, + Timeout: s.articleTimeout, + MaxOutputTokens: s.maxOutputTokens, + } + if job.WebSearch.Enabled { + req.WebSearch = &job.WebSearch + } + + result, err := s.llm.Generate(ctx, req, func(delta string) { + if s.streamEnabled && s.streams != nil && delta != "" { + s.streams.AppendDelta(job.ArticleID, job.TaskID, title, delta) + } + }) + if err != nil { + s.failGeneration(ctx, job, "llm_generate", err) + return + } + + content := strings.TrimSpace(result.Content) + if content == "" { + s.failGeneration(ctx, job, "llm_generate", fmt.Errorf("模型返回空内容")) + return + } + + title = resolveGeneratedImitationTitle(content, job.InputParams) + wordCount := estimateWordCount(content) + + tx, err := s.pool.Begin(ctx) + if err != nil { + s.failGeneration(ctx, job, "persist_begin_tx", err) + return + } + defer rollbackGenerationTx(tx) + + failWithRollback := func(stage string, err error) { + rollbackGenerationTx(tx) + s.failGeneration(ctx, job, stage, err) + } + + articleRepo := repository.NewArticleRepository(tx) + quotaRepo := repository.NewQuotaRepository(tx) + auditTx := repository.NewAuditRepository(tx) + + versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{ + ArticleID: job.ArticleID, + VersionNo: 1, + Title: title, + HTMLContent: "", + MarkdownContent: content, + WordCount: wordCount, + SourceLabel: result.Model, + }) + if err != nil { + failWithRollback("persist_create_version", err) + return + } + + if err := articleRepo.UpdateArticleCurrentVersion(ctx, job.ArticleID, job.TenantID, versionID); err != nil { + failWithRollback("persist_current_version", err) + return + } + if err := articleRepo.UpdateArticleGenerateStatus(ctx, job.ArticleID, job.TenantID, "completed"); err != nil { + failWithRollback("persist_article_status", err) + return + } + + completed := time.Now() + if err := auditTx.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{ + ID: job.TaskID, + TenantID: job.TenantID, + Status: "completed", + StartedAt: &now, + CompletedAt: &completed, + }); err != nil { + failWithRollback("persist_task_status", err) + return + } + + if err := quotaRepo.ConfirmReservation(ctx, job.ReservationID, job.TenantID); err != nil { + failWithRollback("confirm_reservation", err) + return + } + if err := tx.Commit(ctx); err != nil { + failWithRollback("persist_commit", err) + return + } + + cleanupCtx, cancel := newGenerationCleanupContext() + defer cancel() + invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID) + + if s.streamEnabled && s.streams != nil { + s.streams.Complete(job.ArticleID, job.TaskID, title, content) + } +} + +func (s *ArticleImitationService) failGeneration(ctx context.Context, job articleImitationJob, stage string, genErr error) { + cleanupCtx, cancel := newGenerationCleanupContext() + defer cancel() + + errMsg := formatGenerationFailure(stage, genErr) + now := time.Now() + title := strings.TrimSpace(job.InitialTitle) + if title == "" { + title = articleImitationFallbackTitle + } + + auditRepo := repository.NewAuditRepository(s.pool) + articleRepo := repository.NewArticleRepository(s.pool) + quotaRepo := repository.NewQuotaRepository(s.pool) + + _ = auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{ + ID: job.TaskID, + TenantID: job.TenantID, + Status: "failed", + ErrorMessage: &errMsg, + CompletedAt: &now, + }) + if job.ArticleID > 0 { + _ = articleRepo.UpdateArticleGenerateStatus(cleanupCtx, job.ArticleID, job.TenantID, "failed") + _, _ = s.pool.Exec(cleanupCtx, ` + UPDATE articles + SET wizard_state_json = jsonb_set(COALESCE(wizard_state_json, '{}'::jsonb), '{title}', to_jsonb($1::text), true), + updated_at = NOW() + WHERE id = $2 AND tenant_id = $3 + `, title, job.ArticleID, job.TenantID) + } + + if job.ReservationID > 0 { + balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation") + if balanceErr == nil { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := job.TaskID + _, _ = quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ + TenantID: job.TenantID, + OperatorID: job.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }) + } + _ = quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID) + } + + if job.ArticleID > 0 { + invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID) + } + + if s.streamEnabled && s.streams != nil && job.ArticleID > 0 { + s.streams.Fail(job.ArticleID, job.TaskID, title, errMsg) + } +} + +func normalizeImitationSourceURL(raw string) (string, error) { + trimmed := strings.TrimSpace(raw) + if trimmed == "" { + return "", response.ErrBadRequest(40001, "source_url_required", "source_url is required") + } + parsed, err := url.Parse(trimmed) + if err != nil || parsed.Scheme == "" || parsed.Host == "" { + return "", response.ErrBadRequest(40001, "invalid_source_url", "source_url must be a valid URL") + } + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return "", response.ErrBadRequest(40001, "invalid_source_url", "source_url must use http or https") + } + return parsed.String(), nil +} + +func buildImitationInputParams(sourceURL string, req GenerateImitationRequest) map[string]interface{} { + params := map[string]interface{}{ + "source_url": sourceURL, + "source_title": strings.TrimSpace(req.SourceTitle), + "locale": normalizeImitationLocale(req.Locale), + "industry": strings.TrimSpace(req.Industry), + "brand_name": strings.TrimSpace(req.BrandName), + "region": strings.TrimSpace(req.Region), + "target_audience": strings.TrimSpace(req.TargetAudience), + "content_goal": strings.TrimSpace(req.ContentGoal), + "tone": strings.TrimSpace(req.Tone), + "length_goal": strings.TrimSpace(req.LengthGoal), + "keywords": normalizeImitationKeywords(req.Keywords), + "preserve_points": strings.TrimSpace(req.PreservePoints), + "avoid_points": strings.TrimSpace(req.AvoidPoints), + "extra_requirements": strings.TrimSpace(req.ExtraRequirements), + "enable_web_search": req.EnableWebSearch, + "knowledge_group_ids": normalizeInt64IDs(req.KnowledgeGroupIDs), + } + if req.WebSearchLimit != nil && *req.WebSearchLimit > 0 { + params["web_search_limit"] = *req.WebSearchLimit + } + return params +} + +func buildImitationInitialTitle(sourceTitle string) string { + title := strings.TrimSpace(sourceTitle) + if title == "" { + return articleImitationFallbackTitle + } + return "仿写：" + title +} + +func buildImitationGenerationPrompt(params map[string]interface{}, sourceContent string, knowledgePrompt string) string { + locale := normalizeImitationLocale(extractString(params, "locale")) + sourceTitle := strings.TrimSpace(extractString(params, "source_title")) + sourceURL := strings.TrimSpace(extractString(params, "source_url")) + keywords := extractStringList(params["keywords"], 16) + knowledgePrompt = strings.TrimSpace(knowledgePrompt) + + var sourceMeta strings.Builder + if sourceTitle != "" { + appendRawPromptLine(&sourceMeta, "标题", sourceTitle) + } + if sourceURL != "" { + appendRawPromptLine(&sourceMeta, "URL", sourceURL) + } + + var settings strings.Builder + appendPromptLine(&settings, "industry", extractString(params, "industry")) + appendPromptLine(&settings, "brand_name", extractString(params, "brand_name")) + appendPromptLine(&settings, "region", extractString(params, "region")) + appendPromptLine(&settings, "target_audience", extractString(params, "target_audience")) + appendPromptLine(&settings, "content_goal", extractString(params, "content_goal")) + appendPromptLine(&settings, "tone", extractString(params, "tone")) + appendPromptLine(&settings, "length_goal", extractString(params, "length_goal")) + if len(keywords) > 0 { + appendPromptLine(&settings, "keywords", strings.Join(keywords, "、")) + } + appendPromptLine(&settings, "preserve_points", extractString(params, "preserve_points")) + appendPromptLine(&settings, "avoid_points", extractString(params, "avoid_points")) + appendPromptLine(&settings, "extra_requirements", extractString(params, "extra_requirements")) + + knowledgeSection := "" + if knowledgePrompt != "" { + knowledgeSection = "## 知识库参考\n" + knowledgePrompt + } + + return prompts.ArticleImitationPrompt( + locale, + sourceMeta.String(), + settings.String(), + knowledgeSection, + sourceContent, + ) +} + +func buildImitationKnowledgeQuery(params map[string]interface{}) string { + if len(params) == 0 { + return "" + } + + parts := make([]string, 0, 8) + appendValue := func(value string) { + value = strings.TrimSpace(value) + if value == "" { + return + } + parts = append(parts, value) + } + + appendValue(extractString(params, "source_title")) + appendValue(extractString(params, "industry")) + appendValue(extractString(params, "brand_name")) + appendValue(extractString(params, "region")) + appendValue(extractString(params, "target_audience")) + appendValue(extractString(params, "content_goal")) + appendValue(extractString(params, "tone")) + appendValue(strings.Join(extractStringList(params["keywords"], 16), "\n")) + appendValue(extractString(params, "extra_requirements")) + + return strings.Join(parts, "\n") +} + +func appendPromptLine(b *strings.Builder, label, value string) { + value = strings.TrimSpace(value) + if value == "" { + return + } + b.WriteString("- " + prompts.ArticleImitationSettingLabel(label) + "：" + value + "\n") +} + +func appendRawPromptLine(b *strings.Builder, label, value string) { + value = strings.TrimSpace(value) + if value == "" { + return + } + b.WriteString("- " + label + "：" + value + "\n") +} + +func resolveGeneratedImitationTitle(markdown string, params map[string]interface{}) string { + for _, line := range strings.Split(markdown, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "#") { + title := strings.TrimSpace(strings.TrimLeft(line, "#")) + if title != "" { + return title + } + } + } + return buildImitationInitialTitle(extractString(params, "source_title")) +} + +func normalizeImitationLocale(value string) string { + if strings.EqualFold(strings.TrimSpace(value), "en-US") { + return "en-US" + } + return "zh-CN" +} + +func normalizeImitationKeywords(values []string) []string { + if len(values) == 0 { + return []string{} + } + seen := make(map[string]struct{}, len(values)) + keywords := make([]string, 0, len(values)) + for _, item := range values { + item = strings.TrimSpace(item) + if item == "" { + continue + } + key := strings.ToLower(item) + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + keywords = append(keywords, item) + if len(keywords) >= 16 { + break + } + } + return keywords +} + +func truncateRunes(value string, maxRunes int) string { + if maxRunes <= 0 { + return value + } + runes := []rune(value) + if len(runes) <= maxRunes { + return value + } + return string(runes[:maxRunes]) +} diff --git a/server/internal/tenant/app/article_service.go b/server/internal/tenant/app/article_service.go index 9f94f2a..662644d 100644 --- a/server/internal/tenant/app/article_service.go +++ b/server/internal/tenant/app/article_service.go @@ -94,6 +94,8 @@ type ArticleListItem struct { GenerationErrorMessage *string `json:"generation_error_message"` WordCount int `json:"word_count"` SourceLabel *string `json:"source_label"` + SourceURL *string `json:"source_url,omitempty"` + SourceTitle *string `json:"source_title,omitempty"` CreatedAt time.Time `json:"created_at"` } @@ -127,6 +129,8 @@ type ArticleDetailResponse struct { MarkdownContent *string `json:"markdown_content"` WordCount int `json:"word_count"` SourceLabel *string `json:"source_label"` + SourceURL *string `json:"source_url,omitempty"` + SourceTitle *string `json:"source_title,omitempty"` VersionNo *int `json:"version_no"` GenerationErrorMessage *string `json:"generation_error_message"` WizardState map[string]interface{} `json:"wizard_state,omitempty"` @@ -205,7 +209,11 @@ func (s *ArticleService) loadArticles(ctx context.Context, tenantID int64, param argIdx++ } if params.Keyword != nil { - countQuery += ` AND COALESCE(av.title, NULLIF(a.wizard_state_json ->> 'title', ''), NULLIF(gt.input_params_json ->> 'title', '')) ILIKE '%' || $` + strconv.Itoa(argIdx) + ` || '%'` + countQuery += ` AND ( + COALESCE(av.title, NULLIF(a.wizard_state_json ->> 'title', ''), NULLIF(gt.input_params_json ->> 'title', '')) ILIKE '%' || $` + strconv.Itoa(argIdx) + ` || '%' + OR COALESCE(NULLIF(a.wizard_state_json ->> 'source_title', ''), NULLIF(gt.input_params_json ->> 'source_title', '')) ILIKE '%' || $` + strconv.Itoa(argIdx) + ` || '%' + OR COALESCE(NULLIF(a.wizard_state_json ->> 'source_url', ''), NULLIF(gt.input_params_json ->> 'source_url', '')) ILIKE '%' || $` + strconv.Itoa(argIdx) + ` || '%' + )` countArgs = append(countArgs, *params.Keyword) argIdx++ } @@ -274,7 +282,11 @@ func (s *ArticleService) loadArticles(ctx context.Context, tenantID int64, param listIdx++ } if params.Keyword != nil { - listQuery += ` AND COALESCE(av.title, NULLIF(a.wizard_state_json ->> 'title', ''), NULLIF(gt.input_params_json ->> 'title', '')) ILIKE '%' || $` + strconv.Itoa(listIdx) + ` || '%'` + listQuery += ` AND ( + COALESCE(av.title, NULLIF(a.wizard_state_json ->> 'title', ''), NULLIF(gt.input_params_json ->> 'title', '')) ILIKE '%' || $` + strconv.Itoa(listIdx) + ` || '%' + OR COALESCE(NULLIF(a.wizard_state_json ->> 'source_title', ''), NULLIF(gt.input_params_json ->> 'source_title', '')) ILIKE '%' || $` + strconv.Itoa(listIdx) + ` || '%' + OR COALESCE(NULLIF(a.wizard_state_json ->> 'source_url', ''), NULLIF(gt.input_params_json ->> 'source_url', '')) ILIKE '%' || $` + strconv.Itoa(listIdx) + ` || '%' + )` listArgs = append(listArgs, *params.Keyword) listIdx++ } @@ -311,6 +323,7 @@ func (s *ArticleService) loadArticles(ctx context.Context, tenantID int64, param item.SourceType = dbSourceType item.GenerationMode = resolveArticleGenerationMode(dbSourceType, inputParamsJSON) item.Platforms = resolveArticlePlatforms(wizardStateJSON, inputParamsJSON) + item.SourceURL, item.SourceTitle = resolveArticleImitationSourceInfo(dbSourceType, wizardStateJSON, inputParamsJSON) items = append(items, item) } return &ArticleListResponse{Items: items, Total: total, Page: params.Page, PageSize: params.PageSize}, nil @@ -361,6 +374,7 @@ func (s *ArticleService) loadArticleDetail(ctx context.Context, tenantID, articl } item.GenerationMode = resolveArticleGenerationMode(dbSourceType, inputParamsJSON) item.Platforms = resolveArticlePlatforms(wizardStateJSON, inputParamsJSON) + item.SourceURL, item.SourceTitle = resolveArticleImitationSourceInfo(dbSourceType, wizardStateJSON, inputParamsJSON) item.CoverAssetURL = resolveArticleCoverAssetURL(wizardStateJSON) item.CoverImageAssetID = resolveArticleCoverImageAssetID(wizardStateJSON) if item.MarkdownContent != nil || item.CoverImageAssetID != nil { @@ -1080,6 +1094,24 @@ func resolveArticleGenerationMode(sourceType string, inputParamsJSON []byte) *st return &mode } +func resolveArticleImitationSourceInfo(sourceType string, wizardStateJSON, inputParamsJSON []byte) (*string, *string) { + if sourceType != articleImitationSourceType { + return nil, nil + } + + sourceURL := strings.TrimSpace(extractStringFromJSONPayload(wizardStateJSON, "source_url")) + if sourceURL == "" { + sourceURL = strings.TrimSpace(extractStringFromJSONPayload(inputParamsJSON, "source_url")) + } + + sourceTitle := strings.TrimSpace(extractStringFromJSONPayload(wizardStateJSON, "source_title")) + if sourceTitle == "" { + sourceTitle = strings.TrimSpace(extractStringFromJSONPayload(inputParamsJSON, "source_title")) + } + + return nilIfEmptyString(sourceURL), nilIfEmptyString(sourceTitle) +} + func nilIfEmptyString(value string) *string { value = strings.TrimSpace(value) if value == "" { diff --git a/server/internal/tenant/app/template_service.go b/server/internal/tenant/app/template_service.go index ae9d91a..a5539c9 100644 --- a/server/internal/tenant/app/template_service.go +++ b/server/internal/tenant/app/template_service.go @@ -241,7 +241,7 @@ func (s *TemplateService) SaveDraft(ctx context.Context, templateID int64, req S AND template_id = $4 AND source_type = 'template' AND deleted_at IS NULL - AND generate_status IN ('draft', 'failed') + AND generate_status = 'draft' RETURNING id `, stateJSON, *req.ArticleID, actor.TenantID, templateID).Scan(&articleID) if err != nil { @@ -333,7 +333,7 @@ func (s *TemplateService) Generate(ctx context.Context, templateID int64, req Ge AND template_id = $3 AND source_type = 'template' AND deleted_at IS NULL - AND generate_status IN ('draft', 'failed') + AND generate_status = 'draft' FOR UPDATE `, *req.ArticleID, actor.TenantID, templateID).Scan(&articleID) if err != nil { @@ -630,6 +630,8 @@ func generationFailureStageLabel(stage string) string { return "AI 正文生成失败" case "knowledge_resolve": return "知识库检索失败" + case "source_fetch": + return "源文章抓取失败" case "persist_begin_tx": return "写入生成结果事务启动失败" case "persist_create_version": diff --git a/server/internal/tenant/prompts/loader.go b/server/internal/tenant/prompts/loader.go index c31bb0e..c7b4128 100644 --- a/server/internal/tenant/prompts/loader.go +++ b/server/internal/tenant/prompts/loader.go @@ -48,6 +48,9 @@ type runtimePromptsConfig struct { PromptRuleTargetPlatformSupplementFormat string `yaml:"prompt_rule_target_platform_supplement_format"` PromptRuleSupplementHeading string `yaml:"prompt_rule_supplement_heading"` PromptRuleOutputRequirementsSection string `yaml:"prompt_rule_output_requirements_section"` + ArticleImitationPromptTemplate string `yaml:"article_imitation_prompt_template"` + ArticleImitationLocaleInstructions map[string]string `yaml:"article_imitation_locale_instructions"` + ArticleImitationSettingLabels map[string]string `yaml:"article_imitation_setting_labels"` KnowledgePromptIntroLines []string `yaml:"knowledge_prompt_intro_lines"` KnowledgeSnippetLabelFormat string `yaml:"knowledge_snippet_label_format"` KnowledgeWebsiteMarkdown knowledgeWebsiteMarkdownConfig `yaml:"knowledge_website_markdown"` diff --git a/server/internal/tenant/prompts/runtime.go b/server/internal/tenant/prompts/runtime.go index 3766c39..94e9482 100644 --- a/server/internal/tenant/prompts/runtime.go +++ b/server/internal/tenant/prompts/runtime.go @@ -136,6 +136,55 @@ func PromptRuleOutputRequirementsSection() string { return trimPrompt(loadRuntimePrompts().PromptRuleOutputRequirementsSection) } +func ArticleImitationPrompt(locale, sourceMeta, settingsBlock, knowledgeSection, sourceContent string) string { + runtime := loadRuntimePrompts() + template := trimPrompt(runtime.ArticleImitationPromptTemplate) + if template == "" { + template = `你是一名资深内容策略编辑，请基于给定源文章做高质量仿写创作。 +%s + +## 源文章 +%s + +## 仿写设置 +%s + +%s + +## 源文章内容 +%s + +请开始输出仿写后的完整文章。` + } + return strings.TrimSpace(fmt.Sprintf( + template, + ArticleImitationLocaleInstruction(locale), + strings.TrimSpace(sourceMeta), + strings.TrimSpace(settingsBlock), + strings.TrimSpace(knowledgeSection), + strings.TrimSpace(sourceContent), + )) +} + +func ArticleImitationLocaleInstruction(locale string) string { + instructions := loadRuntimePrompts().ArticleImitationLocaleInstructions + if value, ok := instructions[strings.TrimSpace(locale)]; ok && strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + if strings.EqualFold(strings.TrimSpace(locale), "en-US") { + return "Output language: English." + } + return "输出语言：中文简体。" +} + +func ArticleImitationSettingLabel(key string) string { + labels := loadRuntimePrompts().ArticleImitationSettingLabels + if value, ok := labels[strings.TrimSpace(key)]; ok && strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + return strings.TrimSpace(key) +} + func KnowledgePromptIntroLines() []string { lines := loadRuntimePrompts().KnowledgePromptIntroLines return append([]string(nil), lines...) diff --git a/server/internal/tenant/transport/article_handler.go b/server/internal/tenant/transport/article_handler.go index 8ee62c0..96b6214 100644 --- a/server/internal/tenant/transport/article_handler.go +++ b/server/internal/tenant/transport/article_handler.go @@ -22,6 +22,7 @@ type ArticleHandler struct { svc *app.ArticleService selectionOptimize *app.ArticleSelectionOptimizeService promptGenerateSvc *app.PromptRuleGenerationService + imitationSvc *app.ArticleImitationService streams *stream.GenerationHub streamEnabled bool assets *AssetHandler @@ -57,6 +58,15 @@ func NewArticleHandler(a *bootstrap.App) *ArticleHandler { a.Config.Generation, a.Config.LLM.MaxOutputTokens, ).WithCache(a.Cache), + imitationSvc: app.NewArticleImitationService( + a.DB, + a.LLM, + a.RabbitMQ, + knowledgeSvc, + a.GenerationStreams, + a.Config.Generation, + a.Config.LLM.MaxOutputTokens, + ).WithCache(a.Cache), streams: a.GenerationStreams, streamEnabled: a.Config.Generation.StreamEnabled, assets: NewAssetHandler(a), @@ -162,6 +172,22 @@ func (h *ArticleHandler) GenerateFromRule(c *gin.Context) { response.SuccessWithStatus(c, http.StatusCreated, data) } +func (h *ArticleHandler) GenerateImitation(c *gin.Context) { + var req app.GenerateImitationRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.Error(c, response.ErrBadRequest(40001, "invalid_params", err.Error())) + return + } + + data, err := h.imitationSvc.Generate(c.Request.Context(), req) + if err != nil { + response.Error(c, err) + return + } + + response.SuccessWithStatus(c, http.StatusCreated, data) +} + func (h *ArticleHandler) Create(c *gin.Context) { var req app.CreateArticleRequest if err := c.ShouldBindJSON(&req); err != nil { diff --git a/server/internal/tenant/transport/router.go b/server/internal/tenant/transport/router.go index c20bdbc..43d7532 100644 --- a/server/internal/tenant/transport/router.go +++ b/server/internal/tenant/transport/router.go @@ -149,6 +149,7 @@ func RegisterRoutes(app *bootstrap.App) { articles.GET("", artHandler.List) articles.POST("", artHandler.Create) articles.POST("/generate-from-rule", artHandler.GenerateFromRule) + articles.POST("/imitations/generate", artHandler.GenerateImitation) articles.POST("/:id/images", artHandler.UploadImage) articles.POST("/:id/selection-optimize/stream", artHandler.OptimizeSelectionStream) articles.GET("/:id", artHandler.Detail) diff --git a/server/internal/worker/generate/article_generation_worker.go b/server/internal/worker/generate/article_generation_worker.go index 81f0177..8804964 100644 --- a/server/internal/worker/generate/article_generation_worker.go +++ b/server/internal/worker/generate/article_generation_worker.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "os" + "strings" "time" "github.com/jackc/pgx/v5" @@ -28,6 +29,7 @@ type ArticleGenerationWorker struct { rabbitMQ *rabbitmq.Client templateService *tenantapp.TemplateService promptRuleService *tenantapp.PromptRuleGenerationService + imitationService *tenantapp.ArticleImitationService kolWorker *KolGenerationWorker logger *zap.Logger retryInterval time.Duration @@ -41,6 +43,7 @@ func NewArticleGenerationWorker( rabbitMQClient *rabbitmq.Client, templateService *tenantapp.TemplateService, promptRuleService *tenantapp.PromptRuleGenerationService, + imitationService *tenantapp.ArticleImitationService, kolWorker *KolGenerationWorker, logger *zap.Logger, cfg config.GenerationConfig, @@ -60,6 +63,7 @@ func NewArticleGenerationWorker( rabbitMQ: rabbitMQClient, templateService: templateService, promptRuleService: promptRuleService, + imitationService: imitationService, kolWorker: kolWorker, logger: logger, retryInterval: defaultArticleGenerationWorkerRetryInterval, @@ -145,11 +149,7 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver return } if task.ID > 0 { - 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) - } + w.handleTaskFailure(ctx, task, "task_load", 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), @@ -172,25 +172,28 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver return } + task.TaskType = strings.TrimSpace(task.TaskType) switch { case task.ArticleID <= 0 || task.QuotaReservationID <= 0: - 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")) - } + w.handleTaskFailure(ctx, task, "task_load", 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 == "imitation": + if w.imitationService != nil { + w.imitationService.ProcessQueuedTask(ctx, task) + } else { + w.handleTaskFailure(ctx, task, "task_load", fmt.Errorf("imitation generation worker is not configured")) + } 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")) + w.handleTaskFailure(ctx, task, "task_load", 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)) + w.handleTaskFailure(ctx, task, "task_load", fmt.Errorf("unsupported generation task type: %s", task.TaskType)) } if err := delivery.Ack(false); err != nil && w.logger != nil { @@ -202,6 +205,18 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver } } +func (w *ArticleGenerationWorker) handleTaskFailure(ctx context.Context, task tenantapp.ClaimedGenerationTask, stage string, taskErr error) { + task.TaskType = strings.TrimSpace(task.TaskType) + switch { + case task.TaskType == kolGenerationTaskType && w.kolWorker != nil: + w.kolWorker.HandleTaskFailure(ctx, task, stage, taskErr) + case task.TaskType == "imitation" && w.imitationService != nil: + w.imitationService.HandleTaskFailure(ctx, task, stage, taskErr) + default: + _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, taskErr) + } +} + func (w *ArticleGenerationWorker) loadTask(ctx context.Context, taskID int64) (tenantapp.ClaimedGenerationTask, bool, error) { var ( task tenantapp.ClaimedGenerationTask