1005 lines
32 KiB
Go
1005 lines
32 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
"go.uber.org/zap"
|
|
|
|
"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"
|
|
)
|
|
|
|
type PromptRuleGenerationService struct {
|
|
pool *pgxpool.Pool
|
|
llm llm.Client
|
|
rabbitMQ *rabbitmq.Client
|
|
knowledge *KnowledgeService
|
|
streams *stream.GenerationHub
|
|
cfg generationRuntimeConfig
|
|
configProvider runtimeConfigProvider
|
|
cache sharedcache.Cache
|
|
publishJobs *PublishJobService
|
|
logger *zap.Logger
|
|
}
|
|
|
|
type promptRuleGenerationJob struct {
|
|
OperatorID int64
|
|
TenantID int64
|
|
ArticleID int64
|
|
TaskID int64
|
|
ReservationID int64
|
|
Prompt string
|
|
InputParams map[string]interface{}
|
|
InitialTitle string
|
|
WebSearch llm.WebSearchOptions
|
|
}
|
|
|
|
type GenerateFromRuleRequest struct {
|
|
PromptRuleID int64 `json:"prompt_rule_id" binding:"required"`
|
|
BrandID *int64 `json:"brand_id"`
|
|
ExtraParams map[string]interface{} `json:"extra_params"`
|
|
}
|
|
|
|
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
|
|
WorkspaceID int64
|
|
AutoPublish bool
|
|
PublishAccountIDs []string
|
|
CoverAssetURL *string
|
|
CoverImageAssetID *int64
|
|
EnableWebSearch bool
|
|
GenerateCount int
|
|
ScheduledFor time.Time
|
|
}
|
|
|
|
type promptRuleGenerationRecord struct {
|
|
ID int64
|
|
Name string
|
|
PromptContent string
|
|
Scene *string
|
|
DefaultTone *string
|
|
DefaultWordCount *int
|
|
KnowledgeGroupIDs []int64
|
|
KnowledgeGroups []PromptRuleKnowledgeGroup
|
|
Status string
|
|
}
|
|
|
|
const promptRuleGeneratingTitlePlaceholder = "标题生成中"
|
|
|
|
type enqueuePromptRuleGenerationInput struct {
|
|
OperatorID *int64
|
|
TenantID int64
|
|
PromptRuleID int64
|
|
BrandID *int64
|
|
ExtraParams map[string]interface{}
|
|
FallbackName string
|
|
TaskBatchID *string
|
|
}
|
|
|
|
func NewPromptRuleGenerationService(
|
|
pool *pgxpool.Pool,
|
|
llmClient llm.Client,
|
|
rabbitMQClient *rabbitmq.Client,
|
|
knowledge *KnowledgeService,
|
|
streams *stream.GenerationHub,
|
|
cfg config.GenerationConfig,
|
|
llmMaxOutputTokens int64,
|
|
) *PromptRuleGenerationService {
|
|
svc := &PromptRuleGenerationService{
|
|
pool: pool,
|
|
llm: llmClient,
|
|
rabbitMQ: rabbitMQClient,
|
|
knowledge: knowledge,
|
|
streams: streams,
|
|
cfg: generationRuntimeFromConfig(cfg, config.LLMConfig{MaxOutputTokens: llmMaxOutputTokens}),
|
|
}
|
|
|
|
return svc
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) WithCache(c sharedcache.Cache) *PromptRuleGenerationService {
|
|
s.cache = c
|
|
return s
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) WithPublishJobService(publishJobs *PublishJobService) *PromptRuleGenerationService {
|
|
s.publishJobs = publishJobs
|
|
return s
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) WithLogger(logger *zap.Logger) *PromptRuleGenerationService {
|
|
s.logger = logger
|
|
return s
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) WithConfigProvider(provider runtimeConfigProvider) *PromptRuleGenerationService {
|
|
s.configProvider = provider
|
|
return s
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) runtimeConfig() generationRuntimeConfig {
|
|
if s != nil && s.configProvider != nil {
|
|
if cfg := s.configProvider.Current(); cfg != nil {
|
|
return generationRuntimeFromConfig(cfg.Generation, cfg.LLM)
|
|
}
|
|
}
|
|
if s == nil {
|
|
return generationRuntimeFromConfig(config.GenerationConfig{}, config.LLMConfig{})
|
|
}
|
|
return s.cfg
|
|
}
|
|
|
|
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)
|
|
taskBatchID := "instant-" + uuid.NewString()
|
|
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,
|
|
ExtraParams: clonePromptRuleExtraParams(extra),
|
|
FallbackName: "即时任务",
|
|
TaskBatchID: &taskBatchID,
|
|
})
|
|
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.WorkspaceID > 0 {
|
|
extra["workspace_id"] = input.WorkspaceID
|
|
}
|
|
if input.AutoPublish {
|
|
extra["schedule_auto_publish"] = true
|
|
extra["schedule_publish_account_ids"] = input.PublishAccountIDs
|
|
if input.CoverAssetURL != nil {
|
|
extra["schedule_cover_asset_url"] = strings.TrimSpace(*input.CoverAssetURL)
|
|
}
|
|
if input.CoverImageAssetID != nil {
|
|
extra["schedule_cover_image_asset_id"] = *input.CoverImageAssetID
|
|
}
|
|
}
|
|
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,
|
|
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, input.TenantID, input.PromptRuleID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if rule.Status != "enabled" {
|
|
return nil, response.ErrBadRequest(40017, "prompt_rule_disabled", "prompt rule is disabled")
|
|
}
|
|
|
|
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"))
|
|
if taskName == "" {
|
|
taskName = strings.TrimSpace(rule.Name)
|
|
}
|
|
if taskName == "" {
|
|
taskName = strings.TrimSpace(input.FallbackName)
|
|
}
|
|
if taskName == "" {
|
|
taskName = "内容任务"
|
|
}
|
|
|
|
quotaRepo := repository.NewQuotaRepository(s.pool)
|
|
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": input.PromptRuleID,
|
|
"task_name": taskName,
|
|
"enable_web_search": extractBool(extra, "enable_web_search"),
|
|
"generate_count": generateCount,
|
|
}
|
|
if input.TaskBatchID != nil && strings.TrimSpace(*input.TaskBatchID) != "" {
|
|
inputParams["task_batch_id"] = strings.TrimSpace(*input.TaskBatchID)
|
|
}
|
|
if workspaceID := extractInt(extra, "workspace_id"); workspaceID > 0 {
|
|
inputParams["workspace_id"] = workspaceID
|
|
}
|
|
if generationMode != "" {
|
|
inputParams["generation_mode"] = generationMode
|
|
}
|
|
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
|
|
}
|
|
if extractBool(extra, "schedule_auto_publish") {
|
|
inputParams["schedule_auto_publish"] = true
|
|
inputParams["schedule_publish_account_ids"] = extractStringList(extra["schedule_publish_account_ids"], 64)
|
|
if coverURL := strings.TrimSpace(extractString(extra, "schedule_cover_asset_url")); coverURL != "" {
|
|
inputParams["schedule_cover_asset_url"] = coverURL
|
|
}
|
|
if coverImageAssetID := extractInt(extra, "schedule_cover_image_asset_id"); coverImageAssetID > 0 {
|
|
inputParams["schedule_cover_image_asset_id"] = coverImageAssetID
|
|
}
|
|
}
|
|
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 {
|
|
if s.knowledge == nil {
|
|
return nil, response.ErrServiceUnavailable(50352, "knowledge_unavailable", "knowledge service is not configured")
|
|
}
|
|
knowledgeContext, resolveErr := s.knowledge.ResolveContext(
|
|
ctx,
|
|
input.TenantID,
|
|
rule.KnowledgeGroupIDs,
|
|
buildPromptRuleKnowledgeQuery(rule, inputParams),
|
|
)
|
|
if resolveErr != nil {
|
|
return nil, resolveErr
|
|
}
|
|
knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
|
|
if knowledgePrompt != "" {
|
|
inputParams[generationPayloadKnowledgePrompt] = knowledgePrompt
|
|
}
|
|
}
|
|
|
|
inputJSON, _ := json.Marshal(inputParams)
|
|
var coverAssetURL *string
|
|
var coverImageAssetID NullableInt64Input
|
|
if coverURL := strings.TrimSpace(extractString(inputParams, "schedule_cover_asset_url")); coverURL != "" {
|
|
coverAssetURL = &coverURL
|
|
}
|
|
if rawCoverImageAssetID := int64(extractInt(inputParams, "schedule_cover_image_asset_id")); rawCoverImageAssetID > 0 {
|
|
coverImageAssetID.Set = true
|
|
coverImageAssetID.Value = &rawCoverImageAssetID
|
|
}
|
|
wizardStateJSON, _ := mergeArticleWizardState(nil, promptRuleGeneratingTitlePlaceholder, coverAssetURL, coverImageAssetID)
|
|
|
|
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)
|
|
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: input.TenantID,
|
|
OperatorID: operatorID,
|
|
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: input.TenantID,
|
|
OperatorID: operatorID,
|
|
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, prompt_rule_id, generate_status, publish_status, wizard_state_json)
|
|
VALUES ($1, 'custom_generation', $2, 'generating', 'unpublished', $3)
|
|
RETURNING id
|
|
`, 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, 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: operatorID,
|
|
TenantID: input.TenantID,
|
|
ArticleID: &articleID,
|
|
TaskBatchID: input.TaskBatchID,
|
|
QuotaReservationID: &reservationID,
|
|
TaskType: "custom_generation",
|
|
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, input.TenantID, &articleID)
|
|
s.invalidateScheduleTaskCacheForGeneration(ctx, input.TenantID, inputParams)
|
|
|
|
job := promptRuleGenerationJob{
|
|
OperatorID: operatorID,
|
|
TenantID: input.TenantID,
|
|
ArticleID: articleID,
|
|
TaskID: taskID,
|
|
ReservationID: reservationID,
|
|
InputParams: inputParams,
|
|
InitialTitle: promptRuleGeneratingTitlePlaceholder,
|
|
}
|
|
|
|
if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
|
|
TaskID: taskID,
|
|
TaskType: "custom_generation",
|
|
TenantID: input.TenantID,
|
|
ArticleID: articleID,
|
|
}); err != nil {
|
|
RecordGenerationTaskEvent(GenerationTaskEventQueueFailure, promptRuleGenerationJobLogContext(job, "queue_enqueue", "queued", "failed", GenerationTaskResultFailure))
|
|
s.failGeneration(context.Background(), job, promptRuleGeneratingTitlePlaceholder, "queue_enqueue", err)
|
|
return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry")
|
|
}
|
|
|
|
runtimeCfg := s.runtimeConfig()
|
|
if runtimeCfg.StreamEnabled && s.streams != nil {
|
|
s.streams.Start(articleID, taskID, promptRuleGeneratingTitlePlaceholder)
|
|
}
|
|
|
|
return &GenerateFromRuleResponse{ArticleID: articleID, TaskID: taskID}, nil
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) loadPromptRule(ctx context.Context, tenantID, promptRuleID int64) (*promptRuleGenerationRecord, error) {
|
|
record := &promptRuleGenerationRecord{}
|
|
err := s.pool.QueryRow(ctx, `
|
|
SELECT id, name, prompt_content, scene, default_tone, default_word_count, status
|
|
FROM prompt_rules
|
|
WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL
|
|
`, promptRuleID, tenantID).Scan(
|
|
&record.ID,
|
|
&record.Name,
|
|
&record.PromptContent,
|
|
&record.Scene,
|
|
&record.DefaultTone,
|
|
&record.DefaultWordCount,
|
|
&record.Status,
|
|
)
|
|
if err != nil {
|
|
return nil, response.ErrNotFound(40431, "rule_not_found", "prompt rule not found")
|
|
}
|
|
knowledgeRows, err := s.pool.Query(ctx, `
|
|
SELECT kg.id, kg.name
|
|
FROM prompt_rule_knowledge_groups prkg
|
|
JOIN knowledge_groups kg ON kg.id = prkg.knowledge_group_id
|
|
WHERE prkg.prompt_rule_id = $1
|
|
AND kg.tenant_id = $2
|
|
AND kg.deleted_at IS NULL
|
|
ORDER BY kg.sort_order, kg.created_at, kg.id
|
|
`, promptRuleID, tenantID)
|
|
if err != nil {
|
|
return nil, response.ErrInternal(50010, "knowledge_group_query_failed", "failed to query knowledge groups")
|
|
}
|
|
defer knowledgeRows.Close()
|
|
|
|
for knowledgeRows.Next() {
|
|
var group PromptRuleKnowledgeGroup
|
|
if err := knowledgeRows.Scan(&group.ID, &group.Name); err != nil {
|
|
return nil, response.ErrInternal(50010, "knowledge_group_query_failed", err.Error())
|
|
}
|
|
record.KnowledgeGroups = append(record.KnowledgeGroups, group)
|
|
record.KnowledgeGroupIDs = append(record.KnowledgeGroupIDs, group.ID)
|
|
}
|
|
return record, nil
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) executeGeneration(ctx context.Context, job promptRuleGenerationJob) {
|
|
now := time.Now()
|
|
title := strings.TrimSpace(job.InitialTitle)
|
|
if title == "" {
|
|
title = "Generated Article"
|
|
}
|
|
|
|
auditRepo := repository.NewAuditRepository(s.pool)
|
|
if err := auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
|
|
ID: job.TaskID,
|
|
TenantID: job.TenantID,
|
|
Status: "running",
|
|
StartedAt: &now,
|
|
}); err != nil {
|
|
s.failGeneration(ctx, job, title, "task_status_running", err)
|
|
return
|
|
}
|
|
RecordGenerationTaskEvent(GenerationTaskEventTransition, promptRuleGenerationJobLogContext(job, "task_status_running", "", GenerationTaskStatusRunning, GenerationTaskResultSuccess))
|
|
LogGenerationTaskInfo(s.logger, "article generation task marked running", promptRuleGenerationJobLogContext(job, "task_status_running", "", GenerationTaskStatusRunning, GenerationTaskResultSuccess))
|
|
|
|
runtimeCfg := s.runtimeConfig()
|
|
req := llm.GenerateRequest{
|
|
Prompt: job.Prompt,
|
|
Timeout: runtimeCfg.ArticleTimeout,
|
|
MaxOutputTokens: runtimeCfg.MaxOutputTokens,
|
|
}
|
|
if job.WebSearch.Enabled {
|
|
req.WebSearch = &job.WebSearch
|
|
}
|
|
|
|
result, err := s.llm.Generate(ctx, req, func(delta string) {
|
|
if runtimeCfg.StreamEnabled && s.streams != nil && delta != "" {
|
|
s.streams.AppendDelta(job.ArticleID, job.TaskID, title, delta)
|
|
}
|
|
})
|
|
if err != nil {
|
|
s.failGeneration(ctx, job, title, "llm_generate", err)
|
|
return
|
|
}
|
|
|
|
content := strings.TrimSpace(result.Content)
|
|
if content == "" {
|
|
s.failGeneration(ctx, job, title, "llm_generate", fmt.Errorf("模型返回空内容"))
|
|
return
|
|
}
|
|
|
|
title = resolveArticleTitle(job.InputParams, content)
|
|
wordCount := estimateWordCount(content)
|
|
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
s.failGeneration(ctx, job, title, "persist_begin_tx", err)
|
|
return
|
|
}
|
|
defer rollbackGenerationTx(tx)
|
|
|
|
failWithRollback := func(stage string, err error) {
|
|
rollbackGenerationTx(tx)
|
|
s.failGeneration(ctx, job, title, stage, err)
|
|
}
|
|
|
|
articleRepo := repository.NewArticleRepository(tx)
|
|
quotaRepo := repository.NewQuotaRepository(tx)
|
|
auditTx := repository.NewAuditRepository(tx)
|
|
|
|
versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{
|
|
ArticleID: job.ArticleID,
|
|
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)
|
|
s.invalidateScheduleTaskCacheForGeneration(cleanupCtx, job.TenantID, job.InputParams)
|
|
|
|
if runtimeCfg.StreamEnabled && s.streams != nil {
|
|
s.streams.Complete(job.ArticleID, job.TaskID, title, content)
|
|
}
|
|
|
|
s.enqueueScheduleAutoPublish(cleanupCtx, job, title)
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) enqueueScheduleAutoPublish(ctx context.Context, job promptRuleGenerationJob, title string) {
|
|
if s == nil || s.publishJobs == nil || !extractBool(job.InputParams, "schedule_auto_publish") {
|
|
return
|
|
}
|
|
|
|
workspaceID := int64(extractInt(job.InputParams, "workspace_id"))
|
|
accountIDs := extractStringList(job.InputParams["schedule_publish_account_ids"], 64)
|
|
coverURL := strings.TrimSpace(extractString(job.InputParams, "schedule_cover_asset_url"))
|
|
if workspaceID <= 0 || job.OperatorID <= 0 || len(accountIDs) == 0 {
|
|
s.logAutoPublishWarn("schedule auto publish skipped because configuration is incomplete", nil,
|
|
zap.Int64("generation_task_id", job.TaskID),
|
|
zap.Int64("article_id", job.ArticleID),
|
|
zap.Int64("workspace_id", workspaceID),
|
|
zap.Int("account_count", len(accountIDs)),
|
|
zap.Bool("has_cover", coverURL != ""),
|
|
)
|
|
return
|
|
}
|
|
|
|
accounts := make([]CreatePublishJobAccountRequest, 0, len(accountIDs))
|
|
for _, accountID := range accountIDs {
|
|
accountID = strings.TrimSpace(accountID)
|
|
if accountID == "" {
|
|
continue
|
|
}
|
|
accounts = append(accounts, CreatePublishJobAccountRequest{AccountID: accountID})
|
|
}
|
|
if len(accounts) == 0 {
|
|
return
|
|
}
|
|
|
|
publishTitle := strings.TrimSpace(title)
|
|
if publishTitle == "" {
|
|
publishTitle = "定时生成文章"
|
|
}
|
|
|
|
if _, err := s.publishJobs.Create(ctx, auth.Actor{
|
|
TenantID: job.TenantID,
|
|
UserID: job.OperatorID,
|
|
PrimaryWorkspaceID: workspaceID,
|
|
}, CreatePublishJobRequest{
|
|
Title: publishTitle,
|
|
ContentRef: map[string]any{
|
|
"article_id": job.ArticleID,
|
|
},
|
|
Accounts: accounts,
|
|
}); err != nil {
|
|
s.logAutoPublishWarn("schedule auto publish enqueue failed", err,
|
|
zap.Int64("generation_task_id", job.TaskID),
|
|
zap.Int64("article_id", job.ArticleID),
|
|
zap.Int64("workspace_id", workspaceID),
|
|
zap.Int("account_count", len(accounts)),
|
|
)
|
|
}
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) logAutoPublishWarn(message string, err error, fields ...zap.Field) {
|
|
if s == nil || s.logger == nil {
|
|
return
|
|
}
|
|
if err != nil {
|
|
fields = append(fields, zap.Error(err))
|
|
}
|
|
s.logger.Warn(message, fields...)
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) failGeneration(ctx context.Context, job promptRuleGenerationJob, title, stage string, genErr error) {
|
|
cleanupCtx, cancel := newGenerationCleanupContext()
|
|
defer cancel()
|
|
|
|
errMsg := formatGenerationFailure(stage, genErr)
|
|
now := time.Now()
|
|
logCtx := promptRuleGenerationJobLogContext(job, stage, "", "failed", GenerationTaskResultFailure)
|
|
RecordGenerationTaskEvent(GenerationTaskEventFailure, logCtx)
|
|
LogGenerationTaskWarn(s.logger, "prompt rule generation failed", logCtx, zap.Error(genErr))
|
|
failureTitle := strings.TrimSpace(extractString(job.InputParams, "task_name"))
|
|
if failureTitle == "" {
|
|
failureTitle = title
|
|
}
|
|
|
|
auditRepo := repository.NewAuditRepository(s.pool)
|
|
articleRepo := repository.NewArticleRepository(s.pool)
|
|
quotaRepo := repository.NewQuotaRepository(s.pool)
|
|
|
|
var cleanupErr error
|
|
if err := auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{
|
|
ID: job.TaskID,
|
|
TenantID: job.TenantID,
|
|
Status: "failed",
|
|
ErrorMessage: &errMsg,
|
|
CompletedAt: &now,
|
|
}); err != nil && s.logger != nil {
|
|
s.logger.Warn("prompt rule generation task failed status update failed",
|
|
zap.Error(err),
|
|
zap.Int64("task_id", job.TaskID),
|
|
zap.Int64("article_id", job.ArticleID),
|
|
zap.String("stage", stage),
|
|
)
|
|
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("update generation task failed status: %w", err))
|
|
}
|
|
if err := articleRepo.UpdateArticleGenerateStatus(cleanupCtx, job.ArticleID, job.TenantID, "failed"); err != nil && s.logger != nil {
|
|
s.logger.Warn("prompt rule generation article failed status update failed",
|
|
zap.Error(err),
|
|
zap.Int64("task_id", job.TaskID),
|
|
zap.Int64("article_id", job.ArticleID),
|
|
zap.String("stage", stage),
|
|
)
|
|
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("update article failed status: %w", err))
|
|
}
|
|
if failureTitle != "" {
|
|
_, _ = 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
|
|
`, failureTitle, job.ArticleID, job.TenantID)
|
|
}
|
|
|
|
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
|
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
|
if err != nil {
|
|
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
|
|
} else if refunded {
|
|
if balanceErr != nil {
|
|
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
|
|
} else {
|
|
reason := "generation_refund"
|
|
referenceType := "generation_task"
|
|
taskReferenceID := job.TaskID
|
|
if _, err := 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,
|
|
}); err != nil {
|
|
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
|
|
}
|
|
}
|
|
}
|
|
invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
|
|
s.invalidateScheduleTaskCacheForGeneration(cleanupCtx, job.TenantID, job.InputParams)
|
|
if cleanupErr != nil {
|
|
RecordGenerationTaskEvent(GenerationTaskEventCleanupFailure, logCtx)
|
|
LogGenerationTaskError(s.logger, "prompt rule generation failure cleanup failed", logCtx, zap.Error(cleanupErr))
|
|
}
|
|
|
|
runtimeCfg := s.runtimeConfig()
|
|
if runtimeCfg.StreamEnabled && s.streams != nil {
|
|
s.streams.Fail(job.ArticleID, job.TaskID, failureTitle, errMsg)
|
|
}
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) invalidateScheduleTaskCacheForGeneration(ctx context.Context, tenantID int64, inputParams map[string]interface{}) {
|
|
if s == nil || s.cache == nil || strings.TrimSpace(extractString(inputParams, "generation_mode")) != "schedule" {
|
|
return
|
|
}
|
|
scheduleTaskID := int64(extractInt(inputParams, "schedule_task_id"))
|
|
if scheduleTaskID <= 0 {
|
|
invalidateScheduleTaskCaches(ctx, s.cache, tenantID, nil)
|
|
return
|
|
}
|
|
invalidateScheduleTaskCaches(ctx, s.cache, tenantID, &scheduleTaskID)
|
|
}
|
|
|
|
func promptRuleGenerationJobLogContext(job promptRuleGenerationJob, stage, statusBefore, statusAfter, result string) GenerationTaskLogContext {
|
|
return GenerationTaskLogContext{
|
|
TaskID: job.TaskID,
|
|
TenantID: job.TenantID,
|
|
OperatorID: job.OperatorID,
|
|
ArticleID: job.ArticleID,
|
|
ReservationID: job.ReservationID,
|
|
TaskType: "custom_generation",
|
|
Stage: stage,
|
|
StatusBefore: statusBefore,
|
|
StatusAfter: statusAfter,
|
|
Result: result,
|
|
}
|
|
}
|
|
|
|
func buildPromptRuleGenerationPrompt(rule *promptRuleGenerationRecord, params map[string]interface{}, knowledgePrompt string) string {
|
|
sections := []string{strings.TrimSpace(rule.PromptContent)}
|
|
|
|
var supplements []string
|
|
if taskName := strings.TrimSpace(extractString(params, "task_name")); taskName != "" {
|
|
supplements = append(supplements, prompts.PromptRuleTaskNameSupplement(taskName))
|
|
}
|
|
if rule.Scene != nil && strings.TrimSpace(*rule.Scene) != "" {
|
|
supplements = append(supplements, prompts.PromptRuleSceneSupplement(strings.TrimSpace(*rule.Scene)))
|
|
}
|
|
if rule.DefaultTone != nil && strings.TrimSpace(*rule.DefaultTone) != "" {
|
|
supplements = append(supplements, prompts.PromptRuleToneSupplement(strings.TrimSpace(*rule.DefaultTone)))
|
|
}
|
|
if rule.DefaultWordCount != nil && *rule.DefaultWordCount > 0 {
|
|
supplements = append(supplements, prompts.PromptRuleWordCountSupplement(*rule.DefaultWordCount))
|
|
}
|
|
|
|
if len(supplements) > 0 {
|
|
sections = append(sections, prompts.PromptRuleSupplementSection(supplements))
|
|
}
|
|
|
|
if strings.TrimSpace(knowledgePrompt) != "" {
|
|
sections = append(sections, knowledgePrompt)
|
|
}
|
|
|
|
sections = append(sections, prompts.PromptRuleOutputRequirementsSection())
|
|
return strings.Join(sections, "\n\n")
|
|
}
|
|
|
|
func buildPromptRuleKnowledgeQuery(rule *promptRuleGenerationRecord, params map[string]interface{}) string {
|
|
parts := make([]string, 0, 6)
|
|
if rule != nil {
|
|
if text := strings.TrimSpace(rule.Name); text != "" {
|
|
parts = append(parts, text)
|
|
}
|
|
if rule.Scene != nil && strings.TrimSpace(*rule.Scene) != "" {
|
|
parts = append(parts, strings.TrimSpace(*rule.Scene))
|
|
}
|
|
}
|
|
if taskName := strings.TrimSpace(extractString(params, "task_name")); taskName != "" {
|
|
parts = append(parts, taskName)
|
|
}
|
|
for _, key := range []string{
|
|
"title",
|
|
"topic",
|
|
"subject",
|
|
"brand_name",
|
|
"product_name",
|
|
"primary_keyword",
|
|
"key_points",
|
|
"extra_requirements",
|
|
} {
|
|
if text := strings.TrimSpace(extractString(params, key)); text != "" {
|
|
parts = append(parts, text)
|
|
}
|
|
}
|
|
if keywords := extractStringList(params["keywords"], 16); len(keywords) > 0 {
|
|
parts = append(parts, strings.Join(keywords, "\n"))
|
|
}
|
|
if rule != nil {
|
|
if text := strings.TrimSpace(rule.PromptContent); text != "" {
|
|
parts = append(parts, text)
|
|
}
|
|
}
|
|
return strings.Join(parts, "\n")
|
|
}
|
|
|
|
func extractGenerateCount(extra map[string]interface{}) int {
|
|
value := extractInt(extra, "generate_count")
|
|
if value <= 0 {
|
|
return 1
|
|
}
|
|
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
|
|
}
|
|
raw, ok := extra[key]
|
|
if !ok || raw == nil {
|
|
return false
|
|
}
|
|
switch typed := raw.(type) {
|
|
case bool:
|
|
return typed
|
|
case string:
|
|
normalized := strings.TrimSpace(strings.ToLower(typed))
|
|
return normalized == "true" || normalized == "1" || normalized == "yes"
|
|
case float64:
|
|
return typed != 0
|
|
case int:
|
|
return typed != 0
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func extractInt(extra map[string]interface{}, key string) int {
|
|
if len(extra) == 0 {
|
|
return 0
|
|
}
|
|
raw, ok := extra[key]
|
|
if !ok || raw == nil {
|
|
return 0
|
|
}
|
|
switch typed := raw.(type) {
|
|
case int:
|
|
return typed
|
|
case int32:
|
|
return int(typed)
|
|
case int64:
|
|
return int(typed)
|
|
case float64:
|
|
return int(typed)
|
|
case string:
|
|
value, _ := strconv.Atoi(strings.TrimSpace(typed))
|
|
return value
|
|
default:
|
|
return 0
|
|
}
|
|
}
|
|
|
|
func extractString(payload map[string]interface{}, key string) string {
|
|
if len(payload) == 0 {
|
|
return ""
|
|
}
|
|
raw, ok := payload[key]
|
|
if !ok || raw == nil {
|
|
return ""
|
|
}
|
|
switch typed := raw.(type) {
|
|
case string:
|
|
return typed
|
|
default:
|
|
return fmt.Sprintf("%v", typed)
|
|
}
|
|
}
|
|
|
|
func derefString(value *string) string {
|
|
if value == nil {
|
|
return ""
|
|
}
|
|
return *value
|
|
}
|