feat(kol): generation service + worker branch
This commit is contained in:
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user