454 lines
14 KiB
Go
454 lines
14 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
|
|
sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
|
|
"github.com/geo-platform/tenant-api/internal/shared/llm"
|
|
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
|
)
|
|
|
|
type ClaimedGenerationTask struct {
|
|
ID int64
|
|
TenantID int64
|
|
OperatorID int64
|
|
ArticleID int64
|
|
QuotaReservationID int64
|
|
TaskType string
|
|
InputParams map[string]interface{}
|
|
}
|
|
|
|
const generationCleanupTimeout = 30 * time.Second
|
|
|
|
var errGenerationArticleUnavailable = errors.New("generation article is unavailable")
|
|
|
|
type generationDBTX interface {
|
|
QueryRow(context.Context, string, ...interface{}) pgx.Row
|
|
}
|
|
|
|
type promptRuleTaskContext struct {
|
|
PromptRuleID int64
|
|
BrandID int64
|
|
}
|
|
|
|
func newGenerationCleanupContext() (context.Context, context.CancelFunc) {
|
|
return context.WithTimeout(context.Background(), generationCleanupTimeout)
|
|
}
|
|
|
|
func rollbackGenerationTx(tx pgx.Tx) {
|
|
if tx == nil {
|
|
return
|
|
}
|
|
|
|
ctx, cancel := newGenerationCleanupContext()
|
|
defer cancel()
|
|
|
|
_ = tx.Rollback(ctx)
|
|
}
|
|
|
|
func generationStringPtr(value string) *string {
|
|
return &value
|
|
}
|
|
|
|
func EnsureGenerationArticleWritable(ctx context.Context, db generationDBTX, articleID, tenantID int64) error {
|
|
if articleID <= 0 || tenantID <= 0 {
|
|
return errGenerationArticleUnavailable
|
|
}
|
|
|
|
var lockedArticleID int64
|
|
if err := db.QueryRow(ctx, `
|
|
SELECT id
|
|
FROM articles
|
|
WHERE id = $1
|
|
AND tenant_id = $2
|
|
AND deleted_at IS NULL
|
|
FOR UPDATE
|
|
`, articleID, tenantID).Scan(&lockedArticleID); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func IsGenerationArticleUnavailable(err error) bool {
|
|
return errors.Is(err, errGenerationArticleUnavailable) || errors.Is(err, pgx.ErrNoRows)
|
|
}
|
|
|
|
func HandleGenerationTaskArticleUnavailable(
|
|
ctx context.Context,
|
|
pool *pgxpool.Pool,
|
|
cache sharedcache.Cache,
|
|
task ClaimedGenerationTask,
|
|
taskType string,
|
|
updateKolTaskRecord bool,
|
|
) error {
|
|
cleanupCtx, cancel := newGenerationCleanupContext()
|
|
defer cancel()
|
|
|
|
const errMsg = "文章已删除，生成结果已丢弃"
|
|
now := time.Now()
|
|
auditRepo := repository.NewAuditRepository(pool)
|
|
quotaRepo := repository.NewQuotaRepository(pool)
|
|
usageRepo := repository.NewKolUsageRepository(pool)
|
|
|
|
var cleanupErr error
|
|
if task.ID > 0 {
|
|
if err := auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{
|
|
ID: task.ID,
|
|
TenantID: task.TenantID,
|
|
Status: "cancelled",
|
|
ErrorMessage: generationStringPtr(errMsg),
|
|
CompletedAt: &now,
|
|
}); err != nil {
|
|
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("cancel generation task: %w", err))
|
|
}
|
|
if updateKolTaskRecord {
|
|
_ = auditRepo.UpdateTaskRecordStatus(cleanupCtx, repository.TaskRecordInput{
|
|
TenantID: task.TenantID,
|
|
TaskType: kolGenerationTaskRecordType,
|
|
ResourceType: kolGenerationTaskRecordScope,
|
|
ResourceID: task.ID,
|
|
Status: "cancelled",
|
|
ErrorMessage: generationStringPtr(errMsg),
|
|
})
|
|
_ = usageRepo.MarkFailedByTask(cleanupCtx, task.TenantID, task.ID, errMsg)
|
|
}
|
|
}
|
|
|
|
if task.QuotaReservationID > 0 {
|
|
var balance int
|
|
balanceErr := error(nil)
|
|
if task.OperatorID > 0 {
|
|
balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
|
}
|
|
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.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 if task.OperatorID > 0 {
|
|
reason := "generation_cancel_refund"
|
|
referenceType := "generation_task"
|
|
taskReferenceID := task.ID
|
|
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
|
|
TenantID: task.TenantID,
|
|
OperatorID: task.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 cancel refund ledger: %w", err))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if task.ArticleID > 0 {
|
|
invalidateArticleCaches(cleanupCtx, cache, task.TenantID, &task.ArticleID)
|
|
}
|
|
|
|
RecordGenerationTaskEvent(GenerationTaskEventTransition, GenerationTaskLogContext{
|
|
TaskID: task.ID,
|
|
TenantID: task.TenantID,
|
|
OperatorID: task.OperatorID,
|
|
ArticleID: task.ArticleID,
|
|
ReservationID: task.QuotaReservationID,
|
|
TaskType: taskType,
|
|
Stage: "article_deleted",
|
|
StatusAfter: "cancelled",
|
|
Result: GenerationTaskResultSkipped,
|
|
})
|
|
return cleanupErr
|
|
}
|
|
|
|
func HandleClaimedGenerationTaskFailure(
|
|
ctx context.Context,
|
|
pool *pgxpool.Pool,
|
|
templateService *TemplateService,
|
|
promptRuleService *PromptRuleGenerationService,
|
|
task ClaimedGenerationTask,
|
|
taskErr error,
|
|
) error {
|
|
cleanupCtx, cancel := newGenerationCleanupContext()
|
|
defer cancel()
|
|
|
|
errMsg := formatGenerationFailure("task_load", taskErr)
|
|
if taskErr != nil && strings.Contains(taskErr.Error(), "[worker_recovery]") {
|
|
errMsg = strings.TrimSpace(taskErr.Error())
|
|
}
|
|
now := time.Now()
|
|
|
|
auditRepo := repository.NewAuditRepository(pool)
|
|
articleRepo := repository.NewArticleRepository(pool)
|
|
quotaRepo := repository.NewQuotaRepository(pool)
|
|
|
|
var failureErr error
|
|
if err := auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{
|
|
ID: task.ID,
|
|
TenantID: task.TenantID,
|
|
Status: "failed",
|
|
ErrorMessage: &errMsg,
|
|
CompletedAt: &now,
|
|
}); err != nil {
|
|
failureErr = errors.Join(failureErr, fmt.Errorf("update generation task failed status: %w", err))
|
|
}
|
|
if task.ArticleID > 0 {
|
|
if err := articleRepo.UpdateArticleGenerateStatus(cleanupCtx, task.ArticleID, task.TenantID, "failed"); err != nil {
|
|
failureErr = errors.Join(failureErr, fmt.Errorf("update article failed status: %w", err))
|
|
}
|
|
}
|
|
|
|
if task.QuotaReservationID > 0 {
|
|
var balance int
|
|
balanceErr := error(nil)
|
|
if task.OperatorID > 0 {
|
|
balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
|
}
|
|
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
|
|
if err != nil {
|
|
failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err))
|
|
} else if refunded {
|
|
if balanceErr != nil {
|
|
failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
|
|
} else if task.OperatorID > 0 {
|
|
reason := "generation_refund"
|
|
referenceType := "generation_task"
|
|
taskReferenceID := task.ID
|
|
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
|
|
TenantID: task.TenantID,
|
|
OperatorID: task.OperatorID,
|
|
QuotaType: "article_generation",
|
|
Delta: 1,
|
|
BalanceAfter: balance + 1,
|
|
Reason: &reason,
|
|
ReferenceType: &referenceType,
|
|
ReferenceID: &taskReferenceID,
|
|
}); err != nil {
|
|
failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if task.ArticleID > 0 {
|
|
switch {
|
|
case templateService != nil && templateService.cache != nil:
|
|
invalidateArticleCaches(cleanupCtx, templateService.cache, task.TenantID, &task.ArticleID)
|
|
case promptRuleService != nil && promptRuleService.cache != nil:
|
|
invalidateArticleCaches(cleanupCtx, promptRuleService.cache, task.TenantID, &task.ArticleID)
|
|
}
|
|
}
|
|
|
|
title := resolveArticleTitle(task.InputParams, "")
|
|
if strings.TrimSpace(extractString(task.InputParams, "task_name")) != "" {
|
|
title = strings.TrimSpace(extractString(task.InputParams, "task_name"))
|
|
}
|
|
if title == "" {
|
|
title = promptRuleGeneratingTitlePlaceholder
|
|
}
|
|
|
|
if task.TaskType == "template" && templateService != nil &&
|
|
templateService.runtimeConfig().StreamEnabled && templateService.streams != nil {
|
|
templateService.streams.Fail(task.ArticleID, task.ID, title, errMsg)
|
|
}
|
|
if task.TaskType == "custom_generation" && promptRuleService != nil &&
|
|
promptRuleService.runtimeConfig().StreamEnabled && promptRuleService.streams != nil {
|
|
promptRuleService.streams.Fail(task.ArticleID, task.ID, title, errMsg)
|
|
}
|
|
|
|
return failureErr
|
|
}
|
|
|
|
func (s *TemplateService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) {
|
|
job := generationJob{
|
|
OperatorID: task.OperatorID,
|
|
TenantID: task.TenantID,
|
|
ArticleID: task.ArticleID,
|
|
TaskID: task.ID,
|
|
ReservationID: task.QuotaReservationID,
|
|
Params: cloneInputParams(task.InputParams),
|
|
InitialTitle: resolveArticleTitle(task.InputParams, ""),
|
|
WebSearch: llm.WebSearchOptions{
|
|
Enabled: extractBool(task.InputParams, "enable_web_search"),
|
|
},
|
|
}
|
|
if limit := extractInt(task.InputParams, "web_search_limit"); limit > 0 {
|
|
job.WebSearch.Limit = int32(limit)
|
|
}
|
|
|
|
templateKey, templateName, promptTemplate := extractTemplateSnapshot(task.InputParams)
|
|
if templateKey == "" || templateName == "" {
|
|
record, err := s.loadTemplateRecordForTask(ctx, task)
|
|
if err != nil {
|
|
s.failGeneration(context.Background(), job, job.InitialTitle, "task_load", err)
|
|
return
|
|
}
|
|
templateKey = record.TemplateKey
|
|
templateName = record.TemplateName
|
|
promptTemplate = record.PromptTemplate
|
|
}
|
|
job.TemplateKey = templateKey
|
|
job.TemplateName = templateName
|
|
job.PromptTemplate = promptTemplate
|
|
|
|
s.executeGeneration(ctx, job)
|
|
}
|
|
|
|
func (s *TemplateService) loadTemplateRecordForTask(ctx context.Context, task ClaimedGenerationTask) (*repository.TemplateRecord, error) {
|
|
var templateID int64
|
|
if err := s.pool.QueryRow(ctx, `
|
|
SELECT template_id
|
|
FROM articles
|
|
WHERE id = $1
|
|
AND tenant_id = $2
|
|
AND source_type = 'template'
|
|
AND deleted_at IS NULL
|
|
`, task.ArticleID, task.TenantID).Scan(&templateID); err != nil {
|
|
return nil, fmt.Errorf("load task template id: %w", err)
|
|
}
|
|
|
|
record, err := s.templates.GetTemplateByID(ctx, templateID, task.TenantID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("load task template: %w", err)
|
|
}
|
|
return record, nil
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) {
|
|
inputParams := cloneInputParams(task.InputParams)
|
|
initialTitle := strings.TrimSpace(extractString(inputParams, "task_name"))
|
|
if initialTitle == "" {
|
|
initialTitle = promptRuleGeneratingTitlePlaceholder
|
|
}
|
|
|
|
job := promptRuleGenerationJob{
|
|
OperatorID: task.OperatorID,
|
|
TenantID: task.TenantID,
|
|
ArticleID: task.ArticleID,
|
|
TaskID: task.ID,
|
|
ReservationID: task.QuotaReservationID,
|
|
InputParams: inputParams,
|
|
InitialTitle: initialTitle,
|
|
WebSearch: llm.WebSearchOptions{
|
|
Enabled: extractBool(inputParams, "enable_web_search"),
|
|
},
|
|
}
|
|
|
|
rule := extractPromptRuleSnapshot(inputParams)
|
|
if rule == nil || int64(extractInt(inputParams, "prompt_rule_id")) <= 0 || int64(extractInt(inputParams, "brand_id")) <= 0 {
|
|
taskContext, err := s.resolvePromptRuleTaskContext(ctx, task)
|
|
if err != nil {
|
|
s.failGeneration(context.Background(), job, initialTitle, "task_load", err)
|
|
return
|
|
}
|
|
inputParams["prompt_rule_id"] = taskContext.PromptRuleID
|
|
inputParams["brand_id"] = taskContext.BrandID
|
|
}
|
|
if rule == nil {
|
|
promptRuleID := int64(extractInt(inputParams, "prompt_rule_id"))
|
|
brandID := int64(extractInt(inputParams, "brand_id"))
|
|
var err error
|
|
rule, err = s.loadPromptRule(ctx, task.TenantID, brandID, promptRuleID)
|
|
if err != nil {
|
|
s.failGeneration(context.Background(), job, initialTitle, "task_load", err)
|
|
return
|
|
}
|
|
}
|
|
|
|
knowledgePrompt := strings.TrimSpace(extractString(inputParams, generationPayloadKnowledgePrompt))
|
|
if knowledgePrompt == "" {
|
|
if knowledgeGroupIDs := extractKnowledgeGroupIDs(inputParams["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 {
|
|
if s.knowledge == nil {
|
|
s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", fmt.Errorf("knowledge service is not configured"))
|
|
return
|
|
}
|
|
|
|
knowledgeContext, err := s.knowledge.ResolveContext(
|
|
ctx,
|
|
task.TenantID,
|
|
int64(extractInt(inputParams, "brand_id")),
|
|
knowledgeGroupIDs,
|
|
buildPromptRuleKnowledgeQuery(rule, inputParams),
|
|
)
|
|
if err != nil {
|
|
s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", err)
|
|
return
|
|
}
|
|
knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
|
|
}
|
|
}
|
|
|
|
job.Prompt = buildPromptRuleGenerationPrompt(rule, inputParams, knowledgePrompt)
|
|
s.executeGeneration(ctx, job)
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) resolvePromptRuleTaskContext(ctx context.Context, task ClaimedGenerationTask) (promptRuleTaskContext, error) {
|
|
resolved := promptRuleTaskContext{
|
|
PromptRuleID: int64(extractInt(task.InputParams, "prompt_rule_id")),
|
|
BrandID: int64(extractInt(task.InputParams, "brand_id")),
|
|
}
|
|
if resolved.PromptRuleID > 0 && resolved.BrandID > 0 {
|
|
return resolved, nil
|
|
}
|
|
|
|
if task.ArticleID <= 0 {
|
|
return resolved, fmt.Errorf("article context missing for prompt rule task")
|
|
}
|
|
|
|
var articlePromptRuleID int64
|
|
var articleBrandID int64
|
|
if err := s.pool.QueryRow(ctx, `
|
|
SELECT COALESCE(prompt_rule_id, 0), brand_id
|
|
FROM articles
|
|
WHERE id = $1
|
|
AND tenant_id = $2
|
|
AND source_type = 'custom_generation'
|
|
AND deleted_at IS NULL
|
|
`, task.ArticleID, task.TenantID).Scan(&articlePromptRuleID, &articleBrandID); err != nil {
|
|
return resolved, fmt.Errorf("load task prompt rule context: %w", err)
|
|
}
|
|
|
|
if resolved.PromptRuleID <= 0 {
|
|
resolved.PromptRuleID = articlePromptRuleID
|
|
}
|
|
if resolved.BrandID <= 0 {
|
|
resolved.BrandID = articleBrandID
|
|
}
|
|
if resolved.PromptRuleID <= 0 {
|
|
return resolved, fmt.Errorf("prompt rule context missing for prompt rule task")
|
|
}
|
|
if resolved.BrandID <= 0 {
|
|
return resolved, fmt.Errorf("brand context missing for prompt rule task")
|
|
}
|
|
return resolved, nil
|
|
}
|
|
|
|
func (s *PromptRuleGenerationService) resolvePromptRuleIDForTask(ctx context.Context, task ClaimedGenerationTask) (int64, error) {
|
|
if promptRuleID := int64(extractInt(task.InputParams, "prompt_rule_id")); promptRuleID > 0 {
|
|
return promptRuleID, nil
|
|
}
|
|
|
|
var promptRuleID int64
|
|
if err := s.pool.QueryRow(ctx, `
|
|
SELECT prompt_rule_id
|
|
FROM articles
|
|
WHERE id = $1
|
|
AND tenant_id = $2
|
|
AND source_type = 'custom_generation'
|
|
AND deleted_at IS NULL
|
|
`, task.ArticleID, task.TenantID).Scan(&promptRuleID); err != nil {
|
|
return 0, fmt.Errorf("load task prompt rule id: %w", err)
|
|
}
|
|
return promptRuleID, nil
|
|
}
|