508 lines
14 KiB
Go
508 lines
14 KiB
Go
package generate
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sort"
|
|
"strconv"
|
|
"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
|
|
knowledge *tenantapp.KnowledgeService
|
|
logger *zap.Logger
|
|
cfg config.GenerationConfig
|
|
llmCfg config.LLMConfig
|
|
provider config.Provider
|
|
cache sharedcache.Cache
|
|
publishJobs *tenantapp.PublishJobService
|
|
enterpriseSites *tenantapp.EnterpriseSiteService
|
|
}
|
|
|
|
func NewKolGenerationWorker(
|
|
pool *pgxpool.Pool,
|
|
llmClient llm.Client,
|
|
knowledgeSvc *tenantapp.KnowledgeService,
|
|
logger *zap.Logger,
|
|
cfg config.GenerationConfig,
|
|
llmMaxOutputTokens int64,
|
|
) *KolGenerationWorker {
|
|
return &KolGenerationWorker{
|
|
pool: pool,
|
|
llm: llmClient,
|
|
knowledge: knowledgeSvc,
|
|
logger: logger,
|
|
cfg: cfg,
|
|
llmCfg: config.LLMConfig{MaxOutputTokens: llmMaxOutputTokens},
|
|
}
|
|
}
|
|
|
|
func (w *KolGenerationWorker) WithCache(c sharedcache.Cache) *KolGenerationWorker {
|
|
w.cache = c
|
|
return w
|
|
}
|
|
|
|
func (w *KolGenerationWorker) WithPublishJobService(svc *tenantapp.PublishJobService) *KolGenerationWorker {
|
|
w.publishJobs = svc
|
|
return w
|
|
}
|
|
|
|
func (w *KolGenerationWorker) WithEnterpriseSiteService(svc *tenantapp.EnterpriseSiteService) *KolGenerationWorker {
|
|
w.enterpriseSites = svc
|
|
return w
|
|
}
|
|
|
|
func (w *KolGenerationWorker) WithConfigProvider(provider config.Provider) *KolGenerationWorker {
|
|
w.provider = provider
|
|
return w
|
|
}
|
|
|
|
func (w *KolGenerationWorker) runtimeConfig() (time.Duration, int64) {
|
|
if w != nil && w.provider != nil {
|
|
if cfg := w.provider.Current(); cfg != nil {
|
|
return kolArticleTimeout(cfg.Generation), cfg.LLM.MaxOutputTokens
|
|
}
|
|
}
|
|
if w == nil {
|
|
return 8 * time.Minute, 0
|
|
}
|
|
return kolArticleTimeout(w.cfg), w.llmCfg.MaxOutputTokens
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
knowledgePrompt := ""
|
|
if knowledgeGroupIDs := kolKnowledgeGroupIDs(task.InputParams["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 {
|
|
if w.knowledge == nil {
|
|
w.HandleTaskFailure(ctx, task, "knowledge_resolve", fmt.Errorf("knowledge service is not configured"))
|
|
return
|
|
}
|
|
|
|
knowledgeContext, err := w.knowledge.ResolveContext(
|
|
ctx,
|
|
task.TenantID,
|
|
int64(kolIntInput(task.InputParams, "brand_id")),
|
|
knowledgeGroupIDs,
|
|
buildKolKnowledgeQuery(task.InputParams),
|
|
)
|
|
if err != nil {
|
|
w.HandleTaskFailure(ctx, task, "knowledge_resolve", err)
|
|
return
|
|
}
|
|
knowledgePrompt = w.knowledge.RenderPromptSection(knowledgeContext)
|
|
}
|
|
|
|
prompt := renderedPrompt
|
|
if strings.TrimSpace(knowledgePrompt) != "" {
|
|
prompt = strings.TrimSpace(renderedPrompt + "\n\n" + knowledgePrompt)
|
|
}
|
|
|
|
articleTimeout, maxOutputTokens := w.runtimeConfig()
|
|
generateReq := llm.GenerateRequest{
|
|
Prompt: prompt,
|
|
Timeout: articleTimeout,
|
|
MaxOutputTokens: maxOutputTokens,
|
|
}
|
|
if kolBoolInput(task.InputParams, "enable_web_search") {
|
|
generateReq.WebSearch = &llm.WebSearchOptions{Enabled: true}
|
|
}
|
|
|
|
result, err := w.llm.Generate(ctx, generateReq, 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)
|
|
|
|
if err := tenantapp.EnsureGenerationArticleWritable(ctx, tx, task.ArticleID, task.TenantID); err != nil {
|
|
if tenantapp.IsGenerationArticleUnavailable(err) {
|
|
w.rollbackAndCancelUnavailableArticle(ctx, tx, task)
|
|
return
|
|
}
|
|
w.rollbackAndFail(ctx, tx, task, "persist_article_status", err)
|
|
return
|
|
}
|
|
|
|
versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{
|
|
ArticleID: task.ArticleID,
|
|
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)
|
|
tenantapp.StartScheduleAutoPublish(tenantapp.ScheduleAutoPublishInput{
|
|
TenantID: task.TenantID,
|
|
WorkspaceID: int64(kolIntInput(task.InputParams, "workspace_id")),
|
|
OperatorID: task.OperatorID,
|
|
BrandID: int64(kolIntInput(task.InputParams, "brand_id")),
|
|
ArticleID: task.ArticleID,
|
|
GenerationTaskID: task.ID,
|
|
Title: title,
|
|
InputParams: task.InputParams,
|
|
PublishJobs: w.publishJobs,
|
|
EnterpriseSites: w.enterpriseSites,
|
|
Logger: w.logger,
|
|
})
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
func (w *KolGenerationWorker) rollbackAndCancelUnavailableArticle(ctx context.Context, tx pgx.Tx, task tenantapp.ClaimedGenerationTask) {
|
|
cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
_ = tx.Rollback(cleanupCtx)
|
|
if err := tenantapp.HandleGenerationTaskArticleUnavailable(ctx, w.pool, w.cache, task, kolGenerationTaskType, true); err != nil && w.logger != nil {
|
|
w.logger.Warn("kol generation article unavailable cleanup failed",
|
|
zap.Error(err),
|
|
zap.Int64("task_id", task.ID),
|
|
zap.Int64("article_id", task.ArticleID),
|
|
)
|
|
}
|
|
}
|
|
|
|
var (
|
|
errMissingRenderedPrompt = errors.New("kol generation task missing rendered prompt")
|
|
errEmptyGeneratedContent = errors.New("kol generation returned empty content")
|
|
)
|
|
|
|
func kolArticleTimeout(cfg config.GenerationConfig) time.Duration {
|
|
if cfg.ArticleTimeout > 0 {
|
|
return cfg.ArticleTimeout
|
|
}
|
|
return 8 * time.Minute
|
|
}
|
|
|
|
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 kolBoolInput(input map[string]interface{}, key string) bool {
|
|
if input == nil {
|
|
return false
|
|
}
|
|
value, ok := input[key]
|
|
if !ok || value == nil {
|
|
return false
|
|
}
|
|
enabled, ok := value.(bool)
|
|
return ok && enabled
|
|
}
|
|
|
|
func kolIntInput(input map[string]interface{}, key string) int {
|
|
if input == nil {
|
|
return 0
|
|
}
|
|
value, ok := input[key]
|
|
if !ok || value == nil {
|
|
return 0
|
|
}
|
|
switch typed := value.(type) {
|
|
case int:
|
|
return typed
|
|
case int32:
|
|
return int(typed)
|
|
case int64:
|
|
return int(typed)
|
|
case float64:
|
|
return int(typed)
|
|
case string:
|
|
parsed, _ := strconv.Atoi(strings.TrimSpace(typed))
|
|
return parsed
|
|
default:
|
|
return 0
|
|
}
|
|
}
|
|
|
|
func kolKnowledgeGroupIDs(value interface{}) []int64 {
|
|
switch typed := value.(type) {
|
|
case []int64:
|
|
return dedupeKolIDs(typed)
|
|
case []interface{}:
|
|
ids := make([]int64, 0, len(typed))
|
|
for _, item := range typed {
|
|
switch raw := item.(type) {
|
|
case int64:
|
|
ids = append(ids, raw)
|
|
case int32:
|
|
ids = append(ids, int64(raw))
|
|
case int:
|
|
ids = append(ids, int64(raw))
|
|
case float64:
|
|
ids = append(ids, int64(raw))
|
|
case float32:
|
|
ids = append(ids, int64(raw))
|
|
}
|
|
}
|
|
return dedupeKolIDs(ids)
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func dedupeKolIDs(ids []int64) []int64 {
|
|
if len(ids) == 0 {
|
|
return nil
|
|
}
|
|
|
|
seen := make(map[int64]struct{}, len(ids))
|
|
result := make([]int64, 0, len(ids))
|
|
for _, id := range ids {
|
|
if id <= 0 {
|
|
continue
|
|
}
|
|
if _, exists := seen[id]; exists {
|
|
continue
|
|
}
|
|
seen[id] = struct{}{}
|
|
result = append(result, id)
|
|
}
|
|
sort.Slice(result, func(i, j int) bool {
|
|
return result[i] < result[j]
|
|
})
|
|
return result
|
|
}
|
|
|
|
func buildKolKnowledgeQuery(input map[string]interface{}) string {
|
|
promptName := strings.TrimSpace(kolStringInput(input, "prompt_name"))
|
|
platformHint := strings.TrimSpace(kolStringInput(input, "platform_hint"))
|
|
variables, _ := input["variables"].(map[string]interface{})
|
|
keywords := make([]string, 0)
|
|
keyPoints := make([]string, 0)
|
|
if len(variables) > 0 {
|
|
keys := make([]string, 0, len(variables))
|
|
for key := range variables {
|
|
keys = append(keys, key)
|
|
}
|
|
sort.Strings(keys)
|
|
for _, key := range keys {
|
|
valueText := strings.TrimSpace(kolValueString(variables[key]))
|
|
if valueText == "" {
|
|
continue
|
|
}
|
|
switch {
|
|
case strings.Contains(strings.ToLower(key), "keyword") || strings.Contains(key, "关键词"):
|
|
keywords = append(keywords, valueText)
|
|
case strings.Contains(strings.ToLower(key), "brand") || strings.Contains(key, "品牌"):
|
|
keyPoints = append(keyPoints, fmt.Sprintf("%s: %s", key, valueText))
|
|
default:
|
|
keyPoints = append(keyPoints, fmt.Sprintf("%s: %s", key, valueText))
|
|
}
|
|
}
|
|
}
|
|
|
|
renderedPrompt := strings.TrimSpace(kolStringInput(input, "rendered_prompt"))
|
|
query := tenantapp.BuildKnowledgeQuery(tenantapp.KnowledgeQueryInput{
|
|
TaskType: firstNonEmptyKolText(platformHint, "KOL 内容生成"),
|
|
TemplateName: promptName,
|
|
Title: firstNonEmptyKolText(promptName, renderedPrompt),
|
|
Topic: firstNonEmptyKolText(promptName, renderedPrompt),
|
|
PrimaryKeyword: firstNonEmptyKolText(strings.Join(keywords, "、"), promptName),
|
|
Keywords: keywords,
|
|
KeyPoints: strings.Join(keyPoints, "\n"),
|
|
Extra: renderedPrompt,
|
|
})
|
|
if strings.TrimSpace(query) != "" {
|
|
return kolLimitText(query, 2000)
|
|
}
|
|
return kolLimitText(renderedPrompt, 2000)
|
|
}
|
|
|
|
func firstNonEmptyKolText(values ...string) string {
|
|
for _, value := range values {
|
|
if text := strings.TrimSpace(value); text != "" {
|
|
return text
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func kolValueString(value interface{}) string {
|
|
switch typed := value.(type) {
|
|
case string:
|
|
return typed
|
|
case bool:
|
|
if typed {
|
|
return "true"
|
|
}
|
|
return "false"
|
|
case []interface{}:
|
|
items := make([]string, 0, len(typed))
|
|
for _, item := range typed {
|
|
text := strings.TrimSpace(kolValueString(item))
|
|
if text != "" {
|
|
items = append(items, text)
|
|
}
|
|
}
|
|
return strings.Join(items, ", ")
|
|
default:
|
|
return fmt.Sprintf("%v", typed)
|
|
}
|
|
}
|
|
|
|
func kolLimitText(value string, maxRunes int) string {
|
|
text := strings.TrimSpace(value)
|
|
if text == "" || maxRunes <= 0 {
|
|
return text
|
|
}
|
|
|
|
runes := []rune(text)
|
|
if len(runes) <= maxRunes {
|
|
return text
|
|
}
|
|
return strings.TrimSpace(string(runes[:maxRunes]))
|
|
}
|
|
|
|
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"
|
|
}
|