feat(kol): generation service + worker branch
This commit is contained in:
@@ -65,6 +65,14 @@ func main() {
|
||||
app.Config.LLM.MaxOutputTokens,
|
||||
).WithCache(app.Cache)
|
||||
|
||||
kolWorker := generateworker.NewKolGenerationWorker(
|
||||
app.DB,
|
||||
app.LLM,
|
||||
app.Logger,
|
||||
generationCfg,
|
||||
app.Config.LLM.MaxOutputTokens,
|
||||
).WithCache(app.Cache)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
@@ -73,6 +81,7 @@ func main() {
|
||||
app.RabbitMQ,
|
||||
templateSvc,
|
||||
promptRuleSvc,
|
||||
kolWorker,
|
||||
app.Logger,
|
||||
app.Config.Generation,
|
||||
).Start(ctx)
|
||||
|
||||
@@ -215,3 +215,7 @@ func invalidateArticleCaches(ctx context.Context, c sharedcache.Cache, tenantID
|
||||
}
|
||||
invalidateWorkspaceCaches(ctx, c, tenantID)
|
||||
}
|
||||
|
||||
func InvalidateArticleCaches(ctx context.Context, c sharedcache.Cache, tenantID int64, articleID *int64) {
|
||||
invalidateArticleCaches(ctx, c, tenantID, articleID)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"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/messaging/rabbitmq"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/response"
|
||||
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
||||
)
|
||||
|
||||
const (
|
||||
kolGenerationTaskType = "kol"
|
||||
kolGenerationTaskRecordType = "article_generation"
|
||||
kolGenerationTaskRecordScope = "generation_task"
|
||||
)
|
||||
|
||||
var ErrKolGenerationForbidden = response.ErrForbidden(40373, "kol_generation_forbidden", "subscription prompt is not available")
|
||||
|
||||
type KolGenerationSubmitRequest struct {
|
||||
Variables map[string]any `json:"variables"`
|
||||
}
|
||||
|
||||
type KolGenerationSubmitResponse struct {
|
||||
ArticleID int64 `json:"article_id"`
|
||||
GenerationTaskID int64 `json:"generation_task_id"`
|
||||
UsageLogID int64 `json:"usage_log_id"`
|
||||
}
|
||||
|
||||
type KolSubscriptionPromptSchemaResponse struct {
|
||||
SubscriptionPromptID int64 `json:"subscription_prompt_id"`
|
||||
PackageID int64 `json:"package_id"`
|
||||
PackageName string `json:"package_name"`
|
||||
PromptID int64 `json:"prompt_id"`
|
||||
PromptName string `json:"prompt_name"`
|
||||
PlatformHint *string `json:"platform_hint"`
|
||||
SchemaJSON *KolSchemaJSON `json:"schema_json"`
|
||||
CardConfigJSON map[string]interface{} `json:"card_config_json"`
|
||||
}
|
||||
|
||||
type kolGenerationTaskInput struct {
|
||||
SubscriptionPromptID int64 `json:"subscription_prompt_id"`
|
||||
SubscriptionID int64 `json:"subscription_id"`
|
||||
PackageID int64 `json:"package_id"`
|
||||
PromptID int64 `json:"prompt_id"`
|
||||
PromptRevisionNo int `json:"prompt_revision_no"`
|
||||
Variables map[string]any `json:"variables"`
|
||||
SchemaSnapshot json.RawMessage `json:"schema_snapshot"`
|
||||
CardConfigSnapshot json.RawMessage `json:"card_config_snapshot"`
|
||||
RenderedPromptHash string `json:"rendered_prompt_hash"`
|
||||
RenderedPrompt string `json:"rendered_prompt"`
|
||||
}
|
||||
|
||||
type KolGenerationService struct {
|
||||
pool *pgxpool.Pool
|
||||
subRepo repository.KolSubscriptionRepository
|
||||
promptRepo repository.KolPromptRepository
|
||||
usageRepo repository.KolUsageRepository
|
||||
articleRepo repository.ArticleRepository
|
||||
taskRepo repository.AuditRepository
|
||||
quotaRepo repository.QuotaRepository
|
||||
asset *KolPromptAsset
|
||||
rabbitMQ *rabbitmq.Client
|
||||
cache sharedcache.Cache
|
||||
}
|
||||
|
||||
func NewKolGenerationService(
|
||||
pool *pgxpool.Pool,
|
||||
subRepo repository.KolSubscriptionRepository,
|
||||
promptRepo repository.KolPromptRepository,
|
||||
usageRepo repository.KolUsageRepository,
|
||||
articleRepo repository.ArticleRepository,
|
||||
taskRepo repository.AuditRepository,
|
||||
quotaRepo repository.QuotaRepository,
|
||||
asset *KolPromptAsset,
|
||||
rabbitMQ *rabbitmq.Client,
|
||||
) *KolGenerationService {
|
||||
return &KolGenerationService{
|
||||
pool: pool,
|
||||
subRepo: subRepo,
|
||||
promptRepo: promptRepo,
|
||||
usageRepo: usageRepo,
|
||||
articleRepo: articleRepo,
|
||||
taskRepo: taskRepo,
|
||||
quotaRepo: quotaRepo,
|
||||
asset: asset,
|
||||
rabbitMQ: rabbitMQ,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *KolGenerationService) WithCache(c sharedcache.Cache) *KolGenerationService {
|
||||
s.cache = c
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *KolGenerationService) GetSchema(ctx context.Context, actor auth.Actor, subPromptID int64) (*KolSubscriptionPromptSchemaResponse, error) {
|
||||
row, revision, err := s.loadAuthorizedRevision(ctx, actor, subPromptID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
schema, err := decodeKolSchema(revision.SchemaJSON)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50096, "kol_generation_schema_decode_failed", err.Error())
|
||||
}
|
||||
cardConfig, err := decodeKolCardConfig(revision.CardConfigJSON)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50097, "kol_generation_card_config_decode_failed", err.Error())
|
||||
}
|
||||
|
||||
return &KolSubscriptionPromptSchemaResponse{
|
||||
SubscriptionPromptID: row.ID,
|
||||
PackageID: row.PackageID,
|
||||
PackageName: row.PackageName,
|
||||
PromptID: row.PromptID,
|
||||
PromptName: row.PromptName,
|
||||
PlatformHint: row.PlatformHint,
|
||||
SchemaJSON: schema,
|
||||
CardConfigJSON: cardConfig,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *KolGenerationService) Submit(ctx context.Context, actor auth.Actor, subPromptID int64, variables map[string]any) (*KolGenerationSubmitResponse, error) {
|
||||
row, revision, err := s.loadAuthorizedRevision(ctx, actor, subPromptID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if variables == nil {
|
||||
variables = map[string]any{}
|
||||
}
|
||||
|
||||
schema, err := decodeKolSchema(revision.SchemaJSON)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50096, "kol_generation_schema_decode_failed", err.Error())
|
||||
}
|
||||
|
||||
content, err := s.asset.Get(ctx, revision.PromptAssetKey)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50098, "kol_generation_prompt_asset_load_failed", err.Error())
|
||||
}
|
||||
|
||||
renderedPrompt, err := RenderPrompt(content, *schema, variables)
|
||||
if err != nil {
|
||||
return nil, response.ErrBadRequest(40071, "kol_generation_variables_invalid", err.Error())
|
||||
}
|
||||
renderedPromptHash := HashPrompt(renderedPrompt)
|
||||
|
||||
balance, err := s.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")
|
||||
}
|
||||
|
||||
inputJSON, err := json.Marshal(kolGenerationTaskInput{
|
||||
SubscriptionPromptID: subPromptID,
|
||||
SubscriptionID: row.SubscriptionID,
|
||||
PackageID: row.PackageID,
|
||||
PromptID: row.PromptID,
|
||||
PromptRevisionNo: *row.PublishedRevisionNo,
|
||||
Variables: variables,
|
||||
SchemaSnapshot: json.RawMessage(revision.SchemaJSON),
|
||||
CardConfigSnapshot: json.RawMessage(revision.CardConfigJSON),
|
||||
RenderedPromptHash: renderedPromptHash,
|
||||
RenderedPrompt: renderedPrompt,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50099, "kol_generation_input_encode_failed", err.Error())
|
||||
}
|
||||
|
||||
variablesJSON, err := json.Marshal(variables)
|
||||
if err != nil {
|
||||
return nil, response.ErrBadRequest(40072, "kol_generation_variables_invalid", "variables must be valid json")
|
||||
}
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50100, "kol_generation_tx_failed", "failed to begin generation transaction")
|
||||
}
|
||||
defer rollbackGenerationTx(tx)
|
||||
|
||||
quotaTx := repository.NewQuotaRepository(tx)
|
||||
articleTx := repository.NewArticleRepository(tx)
|
||||
auditTx := repository.NewAuditRepository(tx)
|
||||
usageTx := repository.NewKolUsageRepository(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, response.ErrInternal(50101, "kol_generation_quota_reserve_failed", err.Error())
|
||||
}
|
||||
|
||||
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(50102, "kol_generation_reservation_create_failed", err.Error())
|
||||
}
|
||||
|
||||
articleID, err := articleTx.CreateArticle(ctx, repository.CreateArticleInput{
|
||||
TenantID: actor.TenantID,
|
||||
SourceType: kolGenerationTaskType,
|
||||
KolPromptID: &row.PromptID,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50103, "kol_generation_article_create_failed", err.Error())
|
||||
}
|
||||
|
||||
if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil {
|
||||
return nil, response.ErrInternal(50104, "kol_generation_reservation_update_failed", err.Error())
|
||||
}
|
||||
|
||||
taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{
|
||||
OperatorID: actor.UserID,
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: &articleID,
|
||||
QuotaReservationID: &reservationID,
|
||||
TaskType: kolGenerationTaskType,
|
||||
RequestHash: &renderedPromptHash,
|
||||
InputParamsJSON: inputJSON,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50105, "kol_generation_task_create_failed", err.Error())
|
||||
}
|
||||
|
||||
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, response.ErrInternal(50106, "kol_generation_task_record_create_failed", err.Error())
|
||||
}
|
||||
|
||||
usage, err := usageTx.Create(ctx, repository.CreateKolUsageLogInput{
|
||||
TenantID: actor.TenantID,
|
||||
SubscriptionID: &row.SubscriptionID,
|
||||
SubscriptionPromptID: &row.ID,
|
||||
PackageID: row.PackageID,
|
||||
PromptID: row.PromptID,
|
||||
PromptRevisionNo: *row.PublishedRevisionNo,
|
||||
GenerationTaskID: &taskID,
|
||||
VariablesSnapshot: variablesJSON,
|
||||
SchemaSnapshot: revision.SchemaJSON,
|
||||
RenderedPromptHash: &renderedPromptHash,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50107, "kol_generation_usage_create_failed", err.Error())
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return nil, response.ErrInternal(50108, "kol_generation_commit_failed", err.Error())
|
||||
}
|
||||
|
||||
InvalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID)
|
||||
|
||||
task := ClaimedGenerationTask{
|
||||
ID: taskID,
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
ArticleID: articleID,
|
||||
QuotaReservationID: reservationID,
|
||||
TaskType: kolGenerationTaskType,
|
||||
}
|
||||
if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
|
||||
TaskID: taskID,
|
||||
TaskType: kolGenerationTaskType,
|
||||
TenantID: actor.TenantID,
|
||||
ArticleID: articleID,
|
||||
}); err != nil {
|
||||
_ = HandleKolGenerationTaskFailure(context.Background(), s.pool, s.cache, task, "queue_enqueue", err)
|
||||
return nil, response.ErrServiceUnavailable(50343, "generation_queue_unavailable", "generation queue is busy, please retry")
|
||||
}
|
||||
|
||||
return &KolGenerationSubmitResponse{
|
||||
ArticleID: articleID,
|
||||
GenerationTaskID: taskID,
|
||||
UsageLogID: usage.ID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *KolGenerationService) loadAuthorizedRevision(ctx context.Context, actor auth.Actor, subPromptID int64) (*repository.KolSubscriptionPrompt, *repository.KolPromptRevision, error) {
|
||||
row, err := s.subRepo.GetSubscriptionPromptByID(ctx, actor.TenantID, subPromptID)
|
||||
if err != nil {
|
||||
return nil, nil, response.ErrInternal(50109, "kol_subscription_prompt_query_failed", err.Error())
|
||||
}
|
||||
if row == nil {
|
||||
return nil, nil, ErrKolGenerationForbidden
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
switch {
|
||||
case row.Status != "active":
|
||||
return nil, nil, ErrKolGenerationForbidden
|
||||
case row.SubscriptionStatus != "active":
|
||||
return nil, nil, ErrKolGenerationForbidden
|
||||
case row.SubscriptionEndAt != nil && row.SubscriptionEndAt.Before(now):
|
||||
return nil, nil, ErrKolGenerationForbidden
|
||||
case row.PackageStatus != "published":
|
||||
return nil, nil, ErrKolGenerationForbidden
|
||||
case row.PromptStatus != "active" || row.PublishedRevisionNo == nil:
|
||||
return nil, nil, ErrKolGenerationForbidden
|
||||
}
|
||||
|
||||
revision, err := s.promptRepo.GetRevision(ctx, row.CreatorTenantID, row.PromptID, *row.PublishedRevisionNo)
|
||||
if err != nil {
|
||||
return nil, nil, response.ErrInternal(50110, "kol_prompt_revision_query_failed", err.Error())
|
||||
}
|
||||
if revision == nil {
|
||||
return nil, nil, response.ErrInternal(50111, "kol_prompt_revision_missing", "published prompt revision not found")
|
||||
}
|
||||
|
||||
return row, revision, nil
|
||||
}
|
||||
|
||||
func HandleKolGenerationTaskFailure(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
cache sharedcache.Cache,
|
||||
task ClaimedGenerationTask,
|
||||
stage string,
|
||||
taskErr error,
|
||||
) error {
|
||||
cleanupCtx, cancel := newGenerationCleanupContext()
|
||||
defer cancel()
|
||||
|
||||
errMsg := formatGenerationFailure(stage, taskErr)
|
||||
now := time.Now()
|
||||
|
||||
auditRepo := repository.NewAuditRepository(pool)
|
||||
articleRepo := repository.NewArticleRepository(pool)
|
||||
quotaRepo := repository.NewQuotaRepository(pool)
|
||||
usageRepo := repository.NewKolUsageRepository(pool)
|
||||
|
||||
if task.ID > 0 {
|
||||
_ = auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{
|
||||
ID: task.ID,
|
||||
TenantID: task.TenantID,
|
||||
Status: "failed",
|
||||
ErrorMessage: &errMsg,
|
||||
CompletedAt: &now,
|
||||
})
|
||||
_ = auditRepo.UpdateTaskRecordStatus(cleanupCtx, repository.TaskRecordInput{
|
||||
TenantID: task.TenantID,
|
||||
TaskType: kolGenerationTaskRecordType,
|
||||
ResourceType: kolGenerationTaskRecordScope,
|
||||
ResourceID: task.ID,
|
||||
Status: "failed",
|
||||
ErrorMessage: &errMsg,
|
||||
})
|
||||
_ = usageRepo.MarkFailedByTask(cleanupCtx, task.TenantID, task.ID, errMsg)
|
||||
}
|
||||
|
||||
if task.ArticleID > 0 {
|
||||
_ = articleRepo.UpdateArticleGenerateStatus(cleanupCtx, task.ArticleID, task.TenantID, "failed")
|
||||
}
|
||||
|
||||
if task.QuotaReservationID > 0 {
|
||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
||||
if balanceErr == nil && task.OperatorID > 0 {
|
||||
reason := "generation_refund"
|
||||
referenceType := "generation_task"
|
||||
taskReferenceID := task.ID
|
||||
_, _ = 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,
|
||||
})
|
||||
}
|
||||
_ = quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
|
||||
}
|
||||
|
||||
InvalidateArticleCaches(cleanupCtx, cache, task.TenantID, &task.ArticleID)
|
||||
return nil
|
||||
}
|
||||
@@ -28,6 +28,7 @@ type ArticleGenerationWorker struct {
|
||||
rabbitMQ *rabbitmq.Client
|
||||
templateService *tenantapp.TemplateService
|
||||
promptRuleService *tenantapp.PromptRuleGenerationService
|
||||
kolWorker *KolGenerationWorker
|
||||
logger *zap.Logger
|
||||
retryInterval time.Duration
|
||||
processTimeout time.Duration
|
||||
@@ -40,6 +41,7 @@ func NewArticleGenerationWorker(
|
||||
rabbitMQClient *rabbitmq.Client,
|
||||
templateService *tenantapp.TemplateService,
|
||||
promptRuleService *tenantapp.PromptRuleGenerationService,
|
||||
kolWorker *KolGenerationWorker,
|
||||
logger *zap.Logger,
|
||||
cfg config.GenerationConfig,
|
||||
) *ArticleGenerationWorker {
|
||||
@@ -58,6 +60,7 @@ func NewArticleGenerationWorker(
|
||||
rabbitMQ: rabbitMQClient,
|
||||
templateService: templateService,
|
||||
promptRuleService: promptRuleService,
|
||||
kolWorker: kolWorker,
|
||||
logger: logger,
|
||||
retryInterval: defaultArticleGenerationWorkerRetryInterval,
|
||||
processTimeout: processTimeout,
|
||||
@@ -142,7 +145,11 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver
|
||||
return
|
||||
}
|
||||
if task.ID > 0 {
|
||||
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, err)
|
||||
if task.TaskType == kolGenerationTaskType && w.kolWorker != nil {
|
||||
w.kolWorker.HandleTaskFailure(ctx, task, "task_load", err)
|
||||
} else {
|
||||
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, err)
|
||||
}
|
||||
if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation ack failed after task load failure",
|
||||
zap.Error(ackErr),
|
||||
@@ -167,11 +174,21 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver
|
||||
|
||||
switch {
|
||||
case task.ArticleID <= 0 || task.QuotaReservationID <= 0:
|
||||
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("generation task references are incomplete"))
|
||||
if task.TaskType == kolGenerationTaskType && w.kolWorker != nil {
|
||||
w.kolWorker.HandleTaskFailure(ctx, task, "task_load", fmt.Errorf("generation task references are incomplete"))
|
||||
} else {
|
||||
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("generation task references are incomplete"))
|
||||
}
|
||||
case task.TaskType == "template":
|
||||
w.templateService.ProcessQueuedTask(ctx, task)
|
||||
case task.TaskType == "custom_generation":
|
||||
w.promptRuleService.ProcessQueuedTask(ctx, task)
|
||||
case task.TaskType == kolGenerationTaskType:
|
||||
if w.kolWorker != nil {
|
||||
w.kolWorker.Process(ctx, task)
|
||||
} else {
|
||||
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("kol generation worker is not configured"))
|
||||
}
|
||||
default:
|
||||
_ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("unsupported generation task type: %s", task.TaskType))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
package generate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"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
|
||||
logger *zap.Logger
|
||||
articleTimeout time.Duration
|
||||
maxOutputTokens int64
|
||||
cache sharedcache.Cache
|
||||
}
|
||||
|
||||
func NewKolGenerationWorker(
|
||||
pool *pgxpool.Pool,
|
||||
llmClient llm.Client,
|
||||
logger *zap.Logger,
|
||||
cfg config.GenerationConfig,
|
||||
llmMaxOutputTokens int64,
|
||||
) *KolGenerationWorker {
|
||||
articleTimeout := cfg.ArticleTimeout
|
||||
if articleTimeout <= 0 {
|
||||
articleTimeout = 8 * time.Minute
|
||||
}
|
||||
|
||||
return &KolGenerationWorker{
|
||||
pool: pool,
|
||||
llm: llmClient,
|
||||
logger: logger,
|
||||
articleTimeout: articleTimeout,
|
||||
maxOutputTokens: llmMaxOutputTokens,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *KolGenerationWorker) WithCache(c sharedcache.Cache) *KolGenerationWorker {
|
||||
w.cache = c
|
||||
return w
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
result, err := w.llm.Generate(ctx, llm.GenerateRequest{
|
||||
Prompt: renderedPrompt,
|
||||
Timeout: w.articleTimeout,
|
||||
MaxOutputTokens: w.maxOutputTokens,
|
||||
}, 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)
|
||||
|
||||
versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{
|
||||
ArticleID: task.ArticleID,
|
||||
VersionNo: 1,
|
||||
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)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
var (
|
||||
errMissingRenderedPrompt = errors.New("kol generation task missing rendered prompt")
|
||||
errEmptyGeneratedContent = errors.New("kol generation returned empty content")
|
||||
)
|
||||
|
||||
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 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"
|
||||
}
|
||||
Reference in New Issue
Block a user