feat(article): support regenerating failed articles
This commit is contained in:
@@ -0,0 +1,386 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/auth"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/response"
|
||||
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
||||
)
|
||||
|
||||
type RegenerateArticleResponse struct {
|
||||
ArticleID int64 `json:"article_id"`
|
||||
TaskID int64 `json:"task_id"`
|
||||
}
|
||||
|
||||
type articleRegenerationRecord struct {
|
||||
ArticleID int64
|
||||
SourceType string
|
||||
TemplateID *int64
|
||||
KolPromptID *int64
|
||||
PromptRuleID *int64
|
||||
GenerateStatus string
|
||||
WizardStateJSON []byte
|
||||
SourceTaskID int64
|
||||
TaskType string
|
||||
RequestHash *string
|
||||
InputParamsJSON []byte
|
||||
InputParams map[string]interface{}
|
||||
}
|
||||
|
||||
func (s *ArticleService) Regenerate(ctx context.Context, articleID int64) (*RegenerateArticleResponse, error) {
|
||||
actor := auth.MustActor(ctx)
|
||||
if s == nil || s.llm == nil || s.rabbitMQ == nil {
|
||||
return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is not configured")
|
||||
}
|
||||
if err := s.llm.Validate(); err != nil {
|
||||
return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("begin regeneration tx: %w", err)
|
||||
}
|
||||
defer rollbackGenerationTx(tx)
|
||||
|
||||
record, err := loadArticleForRegeneration(ctx, tx, actor.TenantID, articleID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := validateArticleRegeneration(record); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
quotaTx := repository.NewQuotaRepository(tx)
|
||||
auditTx := repository.NewAuditRepository(tx)
|
||||
|
||||
newBalance := balance - 1
|
||||
reason := "generation_reserve"
|
||||
referenceType := "generation_task"
|
||||
if _, err := quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
QuotaType: "article_generation",
|
||||
Delta: -1,
|
||||
BalanceAfter: newBalance,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
}); err != nil {
|
||||
return nil, fmt.Errorf("reserve regeneration quota: %w", err)
|
||||
}
|
||||
|
||||
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, fmt.Errorf("create regeneration reservation: %w", err)
|
||||
}
|
||||
if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil {
|
||||
return nil, fmt.Errorf("update regeneration reservation resource: %w", err)
|
||||
}
|
||||
|
||||
title := regenerationInitialTitle(record)
|
||||
if err := updateGeneratedArticleState(ctx, tx, articleID, actor.TenantID, title, "generating", record.WizardStateJSON); err != nil {
|
||||
return nil, fmt.Errorf("mark article regenerating: %w", err)
|
||||
}
|
||||
|
||||
requestHash := record.RequestHash
|
||||
if requestHash == nil {
|
||||
if hash := strings.TrimSpace(extractString(record.InputParams, "rendered_prompt_hash")); hash != "" {
|
||||
requestHash = &hash
|
||||
}
|
||||
}
|
||||
|
||||
taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{
|
||||
OperatorID: actor.UserID,
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: &articleID,
|
||||
QuotaReservationID: &reservationID,
|
||||
TaskType: record.TaskType,
|
||||
RequestHash: requestHash,
|
||||
InputParamsJSON: record.InputParamsJSON,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create regeneration task: %w", err)
|
||||
}
|
||||
|
||||
if record.TaskType == kolGenerationTaskType {
|
||||
if _, err := auditTx.CreateTaskRecord(ctx, repository.TaskRecordInput{
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
TaskType: kolGenerationTaskRecordType,
|
||||
ResourceType: kolGenerationTaskRecordScope,
|
||||
ResourceID: taskID,
|
||||
Status: "queued",
|
||||
}); err != nil {
|
||||
return nil, fmt.Errorf("create kol regeneration task record: %w", err)
|
||||
}
|
||||
if err := createKolUsageForRegeneration(ctx, tx, actor.TenantID, taskID, record.InputParams); err != nil {
|
||||
return nil, fmt.Errorf("create kol regeneration usage: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return nil, fmt.Errorf("commit regeneration task: %w", err)
|
||||
}
|
||||
invalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID)
|
||||
|
||||
task := ClaimedGenerationTask{
|
||||
ID: taskID,
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
ArticleID: articleID,
|
||||
QuotaReservationID: reservationID,
|
||||
TaskType: record.TaskType,
|
||||
InputParams: cloneInputParams(record.InputParams),
|
||||
}
|
||||
if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
|
||||
TaskID: taskID,
|
||||
TaskType: record.TaskType,
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: articleID,
|
||||
}); err != nil {
|
||||
RecordGenerationTaskEvent(GenerationTaskEventQueueFailure, GenerationTaskLogContext{
|
||||
TaskID: taskID,
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
ArticleID: articleID,
|
||||
ReservationID: reservationID,
|
||||
TaskType: record.TaskType,
|
||||
Stage: "queue_enqueue",
|
||||
StatusBefore: "queued",
|
||||
StatusAfter: "failed",
|
||||
Result: GenerationTaskResultFailure,
|
||||
})
|
||||
s.failRegenerationQueue(context.Background(), task, title, 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, title)
|
||||
}
|
||||
|
||||
return &RegenerateArticleResponse{ArticleID: articleID, TaskID: taskID}, nil
|
||||
}
|
||||
|
||||
func loadArticleForRegeneration(ctx context.Context, tx pgx.Tx, tenantID, articleID int64) (*articleRegenerationRecord, error) {
|
||||
record := &articleRegenerationRecord{}
|
||||
var templateID, kolPromptID, promptRuleID sql.NullInt64
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT id, source_type, template_id, kol_prompt_id, prompt_rule_id, generate_status, wizard_state_json
|
||||
FROM articles
|
||||
WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL
|
||||
FOR UPDATE
|
||||
`, articleID, tenantID).Scan(
|
||||
&record.ArticleID,
|
||||
&record.SourceType,
|
||||
&templateID,
|
||||
&kolPromptID,
|
||||
&promptRuleID,
|
||||
&record.GenerateStatus,
|
||||
&record.WizardStateJSON,
|
||||
)
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, response.ErrNotFound(40411, "article_not_found", "article not found")
|
||||
}
|
||||
return nil, response.ErrInternal(50010, "query_failed", "failed to fetch article")
|
||||
}
|
||||
record.TemplateID = nullableSQLInt64(templateID)
|
||||
record.KolPromptID = nullableSQLInt64(kolPromptID)
|
||||
record.PromptRuleID = nullableSQLInt64(promptRuleID)
|
||||
|
||||
var requestHash sql.NullString
|
||||
err = tx.QueryRow(ctx, `
|
||||
SELECT id, task_type, request_hash, input_params_json
|
||||
FROM generation_tasks
|
||||
WHERE article_id = $1
|
||||
AND tenant_id = $2
|
||||
AND status = 'failed'
|
||||
ORDER BY (COALESCE(input_params_json, '{}'::jsonb) = '{}'::jsonb), created_at DESC, id DESC
|
||||
LIMIT 1
|
||||
`, articleID, tenantID).Scan(
|
||||
&record.SourceTaskID,
|
||||
&record.TaskType,
|
||||
&requestHash,
|
||||
&record.InputParamsJSON,
|
||||
)
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, response.ErrConflict(40914, "generation_task_missing", "failed generation task is missing")
|
||||
}
|
||||
return nil, response.ErrInternal(50010, "query_failed", "failed to fetch generation task")
|
||||
}
|
||||
if requestHash.Valid {
|
||||
record.RequestHash = &requestHash.String
|
||||
}
|
||||
if len(record.InputParamsJSON) == 0 {
|
||||
record.InputParamsJSON = []byte("{}")
|
||||
}
|
||||
params := map[string]interface{}{}
|
||||
if err := json.Unmarshal(record.InputParamsJSON, ¶ms); err != nil {
|
||||
return nil, response.ErrConflict(40915, "generation_params_invalid", "failed generation params are invalid")
|
||||
}
|
||||
record.InputParams = params
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func validateArticleRegeneration(record *articleRegenerationRecord) error {
|
||||
if record == nil {
|
||||
return response.ErrNotFound(40411, "article_not_found", "article not found")
|
||||
}
|
||||
if record.GenerateStatus != "failed" {
|
||||
return response.ErrConflict(40913, "article_not_regeneratable", "article cannot be regenerated in its current state")
|
||||
}
|
||||
|
||||
switch record.SourceType {
|
||||
case "template":
|
||||
if record.TaskType == "template" && record.TemplateID != nil {
|
||||
return nil
|
||||
}
|
||||
case "custom_generation":
|
||||
if record.TaskType == "custom_generation" {
|
||||
return nil
|
||||
}
|
||||
case articleImitationSourceType:
|
||||
if record.TaskType == articleImitationTaskType {
|
||||
return nil
|
||||
}
|
||||
case kolGenerationTaskType:
|
||||
if record.TaskType == kolGenerationTaskType {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return response.ErrConflict(40913, "article_not_regeneratable", "article cannot be regenerated in its current state")
|
||||
}
|
||||
|
||||
func regenerationInitialTitle(record *articleRegenerationRecord) string {
|
||||
if record == nil {
|
||||
return "Generated Article"
|
||||
}
|
||||
|
||||
switch record.TaskType {
|
||||
case "template":
|
||||
return resolveArticleTitle(record.InputParams, "")
|
||||
case "custom_generation":
|
||||
if title := strings.TrimSpace(extractString(record.InputParams, "task_name")); title != "" {
|
||||
return title
|
||||
}
|
||||
return promptRuleGeneratingTitlePlaceholder
|
||||
case articleImitationTaskType:
|
||||
return buildImitationInitialTitle(extractString(record.InputParams, "source_title"))
|
||||
case kolGenerationTaskType:
|
||||
if title := strings.TrimSpace(extractString(record.InputParams, "prompt_name")); title != "" {
|
||||
return title
|
||||
}
|
||||
}
|
||||
|
||||
if title := strings.TrimSpace(extractString(record.InputParams, "title")); title != "" {
|
||||
return title
|
||||
}
|
||||
return "Generated Article"
|
||||
}
|
||||
|
||||
func (s *ArticleService) failRegenerationQueue(ctx context.Context, task ClaimedGenerationTask, title string, taskErr error) {
|
||||
if task.TaskType == kolGenerationTaskType {
|
||||
_ = HandleKolGenerationTaskFailure(ctx, s.pool, s.cache, task, "queue_enqueue", taskErr)
|
||||
} else {
|
||||
_ = HandleClaimedGenerationTaskFailure(ctx, s.pool, nil, nil, task, taskErr)
|
||||
}
|
||||
|
||||
runtimeCfg := s.runtimeConfig()
|
||||
if runtimeCfg.StreamEnabled && s.streams != nil {
|
||||
s.streams.Fail(task.ArticleID, task.ID, title, formatGenerationFailure("queue_enqueue", taskErr))
|
||||
}
|
||||
}
|
||||
|
||||
func createKolUsageForRegeneration(ctx context.Context, tx pgx.Tx, tenantID, taskID int64, params map[string]interface{}) error {
|
||||
packageID := int64(extractInt(params, "package_id"))
|
||||
promptID := int64(extractInt(params, "prompt_id"))
|
||||
if packageID <= 0 || promptID <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
promptRevisionNo := extractInt(params, "prompt_revision_no")
|
||||
if promptRevisionNo <= 0 {
|
||||
promptRevisionNo = 1
|
||||
}
|
||||
subscriptionID := optionalInt64FromParams(params, "subscription_id")
|
||||
subscriptionPromptID := optionalInt64FromParams(params, "subscription_prompt_id")
|
||||
variablesSnapshot, err := marshalOptionalJSONField(params, "variables")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
schemaSnapshot, err := marshalOptionalJSONField(params, "schema_snapshot")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
renderedPromptHash := strings.TrimSpace(extractString(params, "rendered_prompt_hash"))
|
||||
var renderedPromptHashPtr *string
|
||||
if renderedPromptHash != "" {
|
||||
renderedPromptHashPtr = &renderedPromptHash
|
||||
}
|
||||
|
||||
_, err = repository.NewKolUsageRepository(tx).Create(ctx, repository.CreateKolUsageLogInput{
|
||||
TenantID: tenantID,
|
||||
SubscriptionID: subscriptionID,
|
||||
SubscriptionPromptID: subscriptionPromptID,
|
||||
PackageID: packageID,
|
||||
PromptID: promptID,
|
||||
PromptRevisionNo: promptRevisionNo,
|
||||
GenerationTaskID: &taskID,
|
||||
VariablesSnapshot: variablesSnapshot,
|
||||
SchemaSnapshot: schemaSnapshot,
|
||||
RenderedPromptHash: renderedPromptHashPtr,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func marshalOptionalJSONField(params map[string]interface{}, key string) ([]byte, error) {
|
||||
if len(params) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
value, ok := params[key]
|
||||
if !ok || value == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return json.Marshal(value)
|
||||
}
|
||||
|
||||
func optionalInt64FromParams(params map[string]interface{}, key string) *int64 {
|
||||
value := int64(extractInt(params, key))
|
||||
if value <= 0 {
|
||||
return nil
|
||||
}
|
||||
return &value
|
||||
}
|
||||
|
||||
func nullableSQLInt64(value sql.NullInt64) *int64 {
|
||||
if !value.Valid {
|
||||
return nil
|
||||
}
|
||||
result := value.Int64
|
||||
return &result
|
||||
}
|
||||
Reference in New Issue
Block a user