root
01fa2b8309
Desktop Client Build / Resolve Build Metadata (push) Successful in 28s
Frontend CI / Frontend (push) Successful in 3m57s
Backend CI / Backend (push) Successful in 17m3s
Desktop Client Build / Build Desktop Client (push) Successful in 26m8s
Desktop Client Build / Publish Client Artifacts to NAS (push) Successful in 51s
485 lines
17 KiB
Go
485 lines
17 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"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"`
|
|
EnableWebSearch bool `json:"enable_web_search"`
|
|
KnowledgeGroupIDs []int64 `json:"knowledge_group_ids"`
|
|
}
|
|
|
|
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"`
|
|
PromptName string `json:"prompt_name"`
|
|
PlatformHint string `json:"platform_hint"`
|
|
PromptRevisionNo int `json:"prompt_revision_no"`
|
|
Variables map[string]any `json:"variables"`
|
|
EnableWebSearch bool `json:"enable_web_search"`
|
|
KnowledgeGroupIDs []int64 `json:"knowledge_group_ids"`
|
|
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, prompt, err := s.loadAuthorizedPrompt(ctx, actor, subPromptID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
schema, err := decodeKolSchema(prompt.SchemaJSON)
|
|
if err != nil {
|
|
return nil, response.ErrInternal(50096, "kol_generation_schema_decode_failed", err.Error())
|
|
}
|
|
cardConfig, err := decodeKolCardConfig(prompt.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, req KolGenerationSubmitRequest) (*KolGenerationSubmitResponse, error) {
|
|
row, prompt, err := s.loadAuthorizedPrompt(ctx, actor, subPromptID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
variables := req.Variables
|
|
if variables == nil {
|
|
variables = map[string]any{}
|
|
}
|
|
|
|
schema, err := decodeKolSchema(prompt.SchemaJSON)
|
|
if err != nil {
|
|
return nil, response.ErrInternal(50096, "kol_generation_schema_decode_failed", err.Error())
|
|
}
|
|
cardConfig, err := decodeKolCardConfig(prompt.CardConfigJSON)
|
|
if err != nil {
|
|
return nil, response.ErrInternal(50097, "kol_generation_card_config_decode_failed", err.Error())
|
|
}
|
|
|
|
if prompt.PromptAssetKey == nil || *prompt.PromptAssetKey == "" {
|
|
return nil, response.ErrInternal(50098, "kol_generation_prompt_asset_missing", "published prompt content not found")
|
|
}
|
|
|
|
content, err := s.asset.Get(ctx, *prompt.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)
|
|
if req.EnableWebSearch && !kolCardConfigAllowsWebSearch(cardConfig) {
|
|
return nil, response.ErrBadRequest(40073, "kol_generation_web_search_disabled", "web search is disabled for this prompt")
|
|
}
|
|
|
|
knowledgeGroupIDs := normalizeInt64IDs(req.KnowledgeGroupIDs)
|
|
if len(knowledgeGroupIDs) > 0 && !kolCardConfigAllowsUserKnowledge(cardConfig) {
|
|
return nil, response.ErrBadRequest(40074, "kol_generation_knowledge_disabled", "knowledge selection is disabled for this prompt")
|
|
}
|
|
|
|
enableWebSearch := kolCardConfigAllowsWebSearch(cardConfig) && req.EnableWebSearch
|
|
promptRevisionNo := 1
|
|
if row.PublishedRevisionNo != nil && *row.PublishedRevisionNo > 0 {
|
|
promptRevisionNo = *row.PublishedRevisionNo
|
|
}
|
|
|
|
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,
|
|
PromptName: row.PromptName,
|
|
PlatformHint: derefString(row.PlatformHint),
|
|
PromptRevisionNo: promptRevisionNo,
|
|
Variables: variables,
|
|
EnableWebSearch: enableWebSearch,
|
|
KnowledgeGroupIDs: knowledgeGroupIDs,
|
|
SchemaSnapshot: json.RawMessage(prompt.SchemaJSON),
|
|
CardConfigSnapshot: json.RawMessage(prompt.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 := articleTx.UpdateArticleGenerateStatus(ctx, articleID, actor.TenantID, "generating"); err != nil {
|
|
return nil, response.ErrInternal(50112, "kol_generation_article_state_update_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: promptRevisionNo,
|
|
GenerationTaskID: &taskID,
|
|
VariablesSnapshot: variablesJSON,
|
|
SchemaSnapshot: prompt.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 {
|
|
logCtx := GenerationTaskLogContext{
|
|
TaskID: task.ID,
|
|
TenantID: task.TenantID,
|
|
OperatorID: task.OperatorID,
|
|
ArticleID: task.ArticleID,
|
|
ReservationID: task.QuotaReservationID,
|
|
TaskType: kolGenerationTaskType,
|
|
Stage: "queue_enqueue",
|
|
StatusBefore: "queued",
|
|
StatusAfter: "failed",
|
|
Result: GenerationTaskResultFailure,
|
|
}
|
|
RecordGenerationTaskEvent(GenerationTaskEventQueueFailure, logCtx)
|
|
_ = 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) loadAuthorizedPrompt(ctx context.Context, actor auth.Actor, subPromptID int64) (*repository.KolSubscriptionPrompt, *repository.KolPrompt, 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":
|
|
return nil, nil, ErrKolGenerationForbidden
|
|
}
|
|
|
|
prompt, err := s.promptRepo.GetByID(ctx, row.CreatorTenantID, row.PromptID)
|
|
if err != nil {
|
|
return nil, nil, response.ErrInternal(50110, "kol_prompt_query_failed", err.Error())
|
|
}
|
|
if prompt == nil || prompt.PromptAssetKey == nil || *prompt.PromptAssetKey == "" {
|
|
return nil, nil, response.ErrInternal(50111, "kol_prompt_content_missing", "published prompt content not found")
|
|
}
|
|
|
|
return row, prompt, 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()
|
|
logCtx := GenerationTaskLogContext{
|
|
TaskID: task.ID,
|
|
TenantID: task.TenantID,
|
|
OperatorID: task.OperatorID,
|
|
ArticleID: task.ArticleID,
|
|
ReservationID: task.QuotaReservationID,
|
|
TaskType: kolGenerationTaskType,
|
|
Stage: stage,
|
|
StatusAfter: "failed",
|
|
Result: GenerationTaskResultFailure,
|
|
}
|
|
RecordGenerationTaskEvent(GenerationTaskEventFailure, logCtx)
|
|
|
|
auditRepo := repository.NewAuditRepository(pool)
|
|
articleRepo := repository.NewArticleRepository(pool)
|
|
quotaRepo := repository.NewQuotaRepository(pool)
|
|
usageRepo := repository.NewKolUsageRepository(pool)
|
|
|
|
var failureErr error
|
|
if task.ID > 0 {
|
|
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))
|
|
}
|
|
_ = 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 {
|
|
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))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
InvalidateArticleCaches(cleanupCtx, cache, task.TenantID, &task.ArticleID)
|
|
if failureErr != nil {
|
|
RecordGenerationTaskEvent(GenerationTaskEventCleanupFailure, logCtx)
|
|
}
|
|
return failureErr
|
|
}
|