feat(migrations): Harden task audit tracking and optimize article templates
- Added migration to harden task audit tracking by modifying audit_logs and related tables. - Introduced operator_id to several tables for better tracking of actions. - Updated article_templates with new prompt templates for various article types, enhancing content generation. - Created prompt_rules and schedule_tasks tables to manage content generation rules and scheduling. - Added foreign key constraints to articles for better data integrity.
This commit is contained in:
@@ -0,0 +1,538 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/auth"
|
||||
"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/response"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/stream"
|
||||
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
||||
)
|
||||
|
||||
type PromptRuleGenerationService struct {
|
||||
pool *pgxpool.Pool
|
||||
llm llm.Client
|
||||
streams *stream.GenerationHub
|
||||
streamEnabled bool
|
||||
articleTimeout time.Duration
|
||||
maxOutputTokens int64
|
||||
jobs chan promptRuleGenerationJob
|
||||
}
|
||||
|
||||
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"`
|
||||
TargetPlatform *string `json:"target_platform"`
|
||||
ExtraParams map[string]interface{} `json:"extra_params"`
|
||||
}
|
||||
|
||||
type GenerateFromRuleResponse struct {
|
||||
ArticleID int64 `json:"article_id"`
|
||||
TaskID int64 `json:"task_id"`
|
||||
}
|
||||
|
||||
type promptRuleGenerationRecord struct {
|
||||
ID int64
|
||||
Name string
|
||||
PromptContent string
|
||||
Scene *string
|
||||
DefaultTone *string
|
||||
DefaultWordCount *int
|
||||
Status string
|
||||
}
|
||||
|
||||
func NewPromptRuleGenerationService(
|
||||
pool *pgxpool.Pool,
|
||||
llmClient llm.Client,
|
||||
streams *stream.GenerationHub,
|
||||
cfg config.GenerationConfig,
|
||||
llmMaxOutputTokens int64,
|
||||
) *PromptRuleGenerationService {
|
||||
queueSize := cfg.QueueSize
|
||||
if queueSize <= 0 {
|
||||
queueSize = 128
|
||||
}
|
||||
|
||||
workerCount := cfg.WorkerConcurrency
|
||||
if workerCount <= 0 {
|
||||
workerCount = 1
|
||||
}
|
||||
|
||||
articleTimeout := cfg.ArticleTimeout
|
||||
if articleTimeout <= 0 {
|
||||
articleTimeout = 8 * time.Minute
|
||||
}
|
||||
|
||||
svc := &PromptRuleGenerationService{
|
||||
pool: pool,
|
||||
llm: llmClient,
|
||||
streams: streams,
|
||||
streamEnabled: cfg.StreamEnabled,
|
||||
articleTimeout: articleTimeout,
|
||||
maxOutputTokens: llmMaxOutputTokens,
|
||||
jobs: make(chan promptRuleGenerationJob, queueSize),
|
||||
}
|
||||
|
||||
for i := 0; i < workerCount; i++ {
|
||||
go svc.runGenerationWorker()
|
||||
}
|
||||
|
||||
return svc
|
||||
}
|
||||
|
||||
func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req GenerateFromRuleRequest) (*GenerateFromRuleResponse, error) {
|
||||
actor := auth.MustActor(ctx)
|
||||
|
||||
if err := s.llm.Validate(); err != nil {
|
||||
return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
|
||||
}
|
||||
|
||||
rule, err := s.loadPromptRule(ctx, actor.TenantID, req.PromptRuleID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if rule.Status != "enabled" {
|
||||
return nil, response.ErrBadRequest(40017, "prompt_rule_disabled", "prompt rule is disabled")
|
||||
}
|
||||
|
||||
extra := req.ExtraParams
|
||||
if extractGenerateCount(extra) > 1 {
|
||||
return nil, response.ErrBadRequest(40018, "generate_count_not_supported", "prompt rule instant generation currently supports generate_count = 1 only")
|
||||
}
|
||||
|
||||
initialTitle := strings.TrimSpace(extractString(extra, "task_name"))
|
||||
if initialTitle == "" {
|
||||
initialTitle = strings.TrimSpace(rule.Name)
|
||||
}
|
||||
if initialTitle == "" {
|
||||
initialTitle = "Generated Article"
|
||||
}
|
||||
|
||||
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 := map[string]interface{}{
|
||||
"title": initialTitle,
|
||||
"prompt_rule_id": req.PromptRuleID,
|
||||
"target_platform": strings.TrimSpace(derefString(req.TargetPlatform)),
|
||||
"enable_web_search": extractBool(extra, "enable_web_search"),
|
||||
}
|
||||
platformIDs := parsePlatformIDs(derefString(req.TargetPlatform))
|
||||
if len(platformIDs) > 0 {
|
||||
inputParams["target_platforms"] = platformIDs
|
||||
}
|
||||
if req.BrandID != nil {
|
||||
inputParams["brand_id"] = *req.BrandID
|
||||
}
|
||||
|
||||
inputJSON, _ := json.Marshal(inputParams)
|
||||
wizardStateJSON, _ := mergeArticleWizardState(nil, initialTitle, platformIDs)
|
||||
|
||||
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, prompt_rule_id, generate_status, publish_status, wizard_state_json)
|
||||
VALUES ($1, 'prompt_rule', $2, 'generating', 'unpublished', $3)
|
||||
RETURNING id
|
||||
`, actor.TenantID, req.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, 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: "prompt_rule",
|
||||
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")
|
||||
}
|
||||
|
||||
job := promptRuleGenerationJob{
|
||||
OperatorID: actor.UserID,
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: articleID,
|
||||
TaskID: taskID,
|
||||
ReservationID: reservationID,
|
||||
Prompt: buildPromptRuleGenerationPrompt(rule, inputParams),
|
||||
InputParams: inputParams,
|
||||
InitialTitle: initialTitle,
|
||||
WebSearch: llm.WebSearchOptions{
|
||||
Enabled: extractBool(extra, "enable_web_search"),
|
||||
},
|
||||
}
|
||||
|
||||
if err := s.enqueueGeneration(job); err != nil {
|
||||
s.failGeneration(context.Background(), job, job.InitialTitle, "queue_enqueue", err)
|
||||
return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry")
|
||||
}
|
||||
|
||||
if s.streamEnabled {
|
||||
s.streams.Start(articleID, taskID, initialTitle)
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (s *PromptRuleGenerationService) enqueueGeneration(job promptRuleGenerationJob) error {
|
||||
select {
|
||||
case s.jobs <- job:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("generation queue is full")
|
||||
}
|
||||
}
|
||||
|
||||
func (s *PromptRuleGenerationService) runGenerationWorker() {
|
||||
for job := range s.jobs {
|
||||
s.executeGeneration(context.Background(), job)
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
_ = auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
|
||||
ID: job.TaskID,
|
||||
TenantID: job.TenantID,
|
||||
Status: "running",
|
||||
StartedAt: &now,
|
||||
})
|
||||
|
||||
req := llm.GenerateRequest{
|
||||
Prompt: job.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 && 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 func() {
|
||||
_ = tx.Rollback(ctx)
|
||||
}()
|
||||
|
||||
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 {
|
||||
s.failGeneration(ctx, job, title, "persist_create_version", err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := articleRepo.UpdateArticleCurrentVersion(ctx, job.ArticleID, job.TenantID, versionID); err != nil {
|
||||
s.failGeneration(ctx, job, title, "persist_current_version", err)
|
||||
return
|
||||
}
|
||||
if err := articleRepo.UpdateArticleGenerateStatus(ctx, job.ArticleID, job.TenantID, "completed"); err != nil {
|
||||
s.failGeneration(ctx, job, title, "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 {
|
||||
s.failGeneration(ctx, job, title, "persist_task_status", err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := quotaRepo.ConfirmReservation(ctx, job.ReservationID, job.TenantID); err != nil {
|
||||
s.failGeneration(ctx, job, title, "confirm_reservation", err)
|
||||
return
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
s.failGeneration(ctx, job, title, "persist_commit", err)
|
||||
return
|
||||
}
|
||||
|
||||
if s.streamEnabled {
|
||||
s.streams.Complete(job.ArticleID, job.TaskID, title, content)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *PromptRuleGenerationService) failGeneration(ctx context.Context, job promptRuleGenerationJob, title, stage string, genErr error) {
|
||||
errMsg := formatGenerationFailure(stage, genErr)
|
||||
now := time.Now()
|
||||
|
||||
auditRepo := repository.NewAuditRepository(s.pool)
|
||||
articleRepo := repository.NewArticleRepository(s.pool)
|
||||
quotaRepo := repository.NewQuotaRepository(s.pool)
|
||||
|
||||
_ = auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
|
||||
ID: job.TaskID,
|
||||
TenantID: job.TenantID,
|
||||
Status: "failed",
|
||||
ErrorMessage: &errMsg,
|
||||
CompletedAt: &now,
|
||||
})
|
||||
_ = articleRepo.UpdateArticleGenerateStatus(ctx, job.ArticleID, job.TenantID, "failed")
|
||||
|
||||
balance, balanceErr := quotaRepo.GetCurrentBalance(ctx, job.TenantID, "article_generation")
|
||||
if balanceErr == nil {
|
||||
reason := "generation_refund"
|
||||
referenceType := "generation_task"
|
||||
taskReferenceID := job.TaskID
|
||||
_, _ = quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||
TenantID: job.TenantID,
|
||||
OperatorID: job.OperatorID,
|
||||
QuotaType: "article_generation",
|
||||
Delta: 1,
|
||||
BalanceAfter: balance + 1,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
ReferenceID: &taskReferenceID,
|
||||
})
|
||||
}
|
||||
|
||||
_ = quotaRepo.RefundReservation(ctx, job.ReservationID, job.TenantID)
|
||||
|
||||
if s.streamEnabled {
|
||||
s.streams.Fail(job.ArticleID, job.TaskID, title, errMsg)
|
||||
}
|
||||
}
|
||||
|
||||
func buildPromptRuleGenerationPrompt(rule *promptRuleGenerationRecord, params map[string]interface{}) string {
|
||||
sections := []string{strings.TrimSpace(rule.PromptContent)}
|
||||
|
||||
var supplements []string
|
||||
if title := strings.TrimSpace(extractString(params, "title")); title != "" {
|
||||
supplements = append(supplements, "任务名称："+title)
|
||||
}
|
||||
if rule.Scene != nil && strings.TrimSpace(*rule.Scene) != "" {
|
||||
supplements = append(supplements, "适用场景："+strings.TrimSpace(*rule.Scene))
|
||||
}
|
||||
if rule.DefaultTone != nil && strings.TrimSpace(*rule.DefaultTone) != "" {
|
||||
supplements = append(supplements, "默认语气："+strings.TrimSpace(*rule.DefaultTone))
|
||||
}
|
||||
if rule.DefaultWordCount != nil && *rule.DefaultWordCount > 0 {
|
||||
supplements = append(supplements, "建议字数："+strconv.Itoa(*rule.DefaultWordCount)+" 字左右")
|
||||
}
|
||||
if target := strings.TrimSpace(extractString(params, "target_platform")); target != "" {
|
||||
supplements = append(supplements, "目标发布平台："+strings.ReplaceAll(target, ",", "、"))
|
||||
}
|
||||
|
||||
if len(supplements) > 0 {
|
||||
sections = append(sections, "补充要求：\n- "+strings.Join(supplements, "\n- "))
|
||||
}
|
||||
|
||||
sections = append(sections, "输出要求：\n1. 直接输出完整可发布的中文 Markdown 文章。\n2. 标题使用一级标题。\n3. 不要输出你的思考过程、解释或致歉。")
|
||||
return strings.Join(sections, "\n\n")
|
||||
}
|
||||
|
||||
func extractGenerateCount(extra map[string]interface{}) int {
|
||||
value := extractInt(extra, "generate_count")
|
||||
if value <= 0 {
|
||||
return 1
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user