diff --git a/server/cmd/kol-assist-worker/main.go b/server/cmd/kol-assist-worker/main.go index 630a91d..969e0b5 100644 --- a/server/cmd/kol-assist-worker/main.go +++ b/server/cmd/kol-assist-worker/main.go @@ -27,10 +27,12 @@ func main() { defer cancel() assistworker.NewKolAssistWorker( + app.DB, app.RabbitMQ, app.KolAssists, app.LLM, app.Logger, + app.Cache, app.Config.Generation.WorkerConcurrency, app.Config.Generation.ArticleTimeout, ).Start(ctx) diff --git a/server/internal/tenant/app/ai_point_usage.go b/server/internal/tenant/app/ai_point_usage.go new file mode 100644 index 0000000..8823700 --- /dev/null +++ b/server/internal/tenant/app/ai_point_usage.go @@ -0,0 +1,416 @@ +package app + +import ( + "context" + "encoding/json" + "errors" + "strings" + "time" + "unicode/utf8" + + "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/response" + "github.com/geo-platform/tenant-api/internal/tenant/repository" +) + +const ( + AIUsageTypeArticleSelectionOptimize = "article_selection_optimize" + AIUsageTypeTemplateAnalyze = "template_analyze" + AIUsageTypeTemplateTitleGenerate = "template_title_generate" + AIUsageTypeTemplateOutlineGenerate = "template_outline_generate" + AIUsageTypeKolPromptGenerate = "kol_prompt_generate" + AIUsageTypeKolPromptOptimize = "kol_prompt_optimize" + + aiPointsQuotaType = "ai_points" +) + +type AIPointReserveInput struct { + TenantID int64 + OperatorID int64 + UsageType string + ResourceType *string + ResourceID *int64 + ResourceUID *string + MeteredText string + FixedPoints int + Metadata map[string]any +} + +type AIPointReservation struct { + ReservationID int64 + UsageLogID int64 + Points int + RequestChars int + BaseChars int +} + +func ReserveAIPoints( + ctx context.Context, + pool *pgxpool.Pool, + cache sharedcache.Cache, + input AIPointReserveInput, +) (*AIPointReservation, error) { + if pool == nil { + return nil, response.ErrServiceUnavailable(50344, "ai_points_unavailable", "ai points store is unavailable") + } + + input.UsageType = strings.TrimSpace(input.UsageType) + if input.UsageType == "" { + return nil, response.ErrBadRequest(40075, "ai_points_usage_type_required", "usage type is required") + } + + tx, err := pool.Begin(ctx) + if err != nil { + return nil, response.ErrInternal(50120, "ai_points_tx_failed", "failed to begin ai points transaction") + } + defer rollbackGenerationTx(tx) + + if err := lockTenantAIPoints(ctx, tx, input.TenantID); err != nil { + return nil, response.ErrInternal(50121, "ai_points_lock_failed", "failed to lock ai points quota") + } + + now := time.Now().UTC() + quotaRepo := repository.NewQuotaRepository(tx) + status, err := quotaRepo.GetAIQuotaStatus(ctx, input.TenantID, now) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are not available for this plan") + } + return nil, response.ErrInternal(50122, "ai_points_query_failed", "failed to query ai points") + } + + requestChars := countAIPointChars(input.MeteredText) + points := calculateAIPointCost(requestChars, status.BaseChars) + if input.FixedPoints > 0 { + points = input.FixedPoints + } + if status.Total <= 0 || status.Balance < points { + return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are insufficient, please upgrade your plan") + } + + resourceType := "ai_usage" + if input.ResourceType != nil && strings.TrimSpace(*input.ResourceType) != "" { + resourceType = strings.TrimSpace(*input.ResourceType) + } + + reservationID, err := quotaRepo.CreateQuotaReservation(ctx, repository.QuotaReservationInput{ + TenantID: input.TenantID, + OperatorID: input.OperatorID, + QuotaType: aiPointsQuotaType, + ResourceType: resourceType, + ResourceID: input.ResourceID, + ReservedAmount: points, + ExpireAt: now.Add(time.Hour), + }) + if err != nil { + return nil, response.ErrInternal(50123, "ai_points_reserve_failed", "failed to reserve ai points") + } + + metadata := input.Metadata + if metadata == nil { + metadata = map[string]any{} + } + metadataJSON, err := json.Marshal(metadata) + if err != nil { + return nil, response.ErrInternal(50124, "ai_points_metadata_failed", "failed to encode ai points metadata") + } + usageID, err := repository.NewAIPointUsageRepository(tx).Create(ctx, repository.CreateAIPointUsageLogInput{ + TenantID: input.TenantID, + OperatorID: input.OperatorID, + QuotaReservationID: reservationID, + UsageType: input.UsageType, + ResourceType: &resourceType, + ResourceID: input.ResourceID, + ResourceUID: input.ResourceUID, + RequestChars: requestChars, + BaseChars: status.BaseChars, + Points: points, + MetadataJSON: metadataJSON, + }) + if err != nil { + return nil, response.ErrInternal(50125, "ai_points_audit_failed", "failed to create ai points audit log") + } + + reason := "ai_points_reserve" + referenceType := "ai_point_usage" + balanceAfter := status.Balance - points + if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ + TenantID: input.TenantID, + OperatorID: input.OperatorID, + QuotaType: aiPointsQuotaType, + Delta: -points, + BalanceAfter: balanceAfter, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &usageID, + }); err != nil { + return nil, response.ErrInternal(50126, "ai_points_ledger_failed", "failed to write ai points ledger") + } + + if err := tx.Commit(ctx); err != nil { + return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to commit ai points reservation") + } + deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(input.TenantID)) + + return &AIPointReservation{ + ReservationID: reservationID, + UsageLogID: usageID, + Points: points, + RequestChars: requestChars, + BaseChars: status.BaseChars, + }, nil +} + +func CompleteAIPoints( + ctx context.Context, + pool *pgxpool.Pool, + tenantID int64, + reservation AIPointReservation, + model string, +) error { + if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 { + return nil + } + + tx, err := pool.Begin(ctx) + if err != nil { + return err + } + defer rollbackGenerationTx(tx) + + pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID) + if err != nil || !pending { + return err + } + if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { + return err + } + var modelPtr *string + if strings.TrimSpace(model) != "" { + normalized := strings.TrimSpace(model) + modelPtr = &normalized + } + if err := repository.NewAIPointUsageRepository(tx).MarkCompleted(ctx, tenantID, reservation.UsageLogID, modelPtr); err != nil { + return err + } + return tx.Commit(ctx) +} + +func RefundAIPoints( + ctx context.Context, + pool *pgxpool.Pool, + cache sharedcache.Cache, + tenantID int64, + operatorID int64, + reservation AIPointReservation, + errorMessage string, +) error { + if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 || reservation.Points <= 0 { + return nil + } + + tx, err := pool.Begin(ctx) + if err != nil { + return err + } + defer rollbackGenerationTx(tx) + + if err := lockTenantAIPoints(ctx, tx, tenantID); err != nil { + return err + } + pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID) + if err != nil || !pending { + return err + } + + status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC()) + if err != nil { + return err + } + + quotaRepo := repository.NewQuotaRepository(tx) + reason := "ai_points_refund" + referenceType := "ai_point_usage" + usageID := reservation.UsageLogID + if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ + TenantID: tenantID, + OperatorID: operatorID, + QuotaType: aiPointsQuotaType, + Delta: reservation.Points, + BalanceAfter: status.Balance + reservation.Points, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &usageID, + }); err != nil { + return err + } + if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { + return err + } + if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil { + return err + } + if err := tx.Commit(ctx); err != nil { + return err + } + deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(tenantID)) + return nil +} + +func CompletePendingAIPointsByResource( + ctx context.Context, + pool *pgxpool.Pool, + tenantID int64, + resourceType string, + resourceID *int64, + resourceUID *string, + model string, +) error { + reservation, err := findPendingAIPointReservation(ctx, pool, tenantID, resourceType, resourceID, resourceUID) + if err != nil || reservation == nil { + return err + } + return CompleteAIPoints(ctx, pool, tenantID, *reservation, model) +} + +func RefundPendingAIPointsByResource( + ctx context.Context, + pool *pgxpool.Pool, + cache sharedcache.Cache, + tenantID int64, + resourceType string, + resourceID *int64, + resourceUID *string, + errorMessage string, +) error { + reservation, operatorID, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID) + if err != nil || reservation == nil { + return err + } + return RefundAIPoints(ctx, pool, cache, tenantID, operatorID, *reservation, errorMessage) +} + +func findPendingAIPointReservation( + ctx context.Context, + pool *pgxpool.Pool, + tenantID int64, + resourceType string, + resourceID *int64, + resourceUID *string, +) (*AIPointReservation, error) { + reservation, _, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID) + return reservation, err +} + +func findPendingAIPointReservationWithOperator( + ctx context.Context, + pool *pgxpool.Pool, + tenantID int64, + resourceType string, + resourceID *int64, + resourceUID *string, +) (*AIPointReservation, int64, error) { + if pool == nil { + return nil, 0, nil + } + log, err := repository.NewAIPointUsageRepository(pool).FindPendingByResource(ctx, tenantID, resourceType, resourceID, resourceUID) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, 0, nil + } + return nil, 0, err + } + if log == nil || log.QuotaReservationID == nil { + return nil, 0, nil + } + operatorID := int64(0) + if log.OperatorID != nil { + operatorID = *log.OperatorID + } + return &AIPointReservation{ + ReservationID: *log.QuotaReservationID, + UsageLogID: log.ID, + Points: log.Points, + RequestChars: log.RequestChars, + BaseChars: log.BaseChars, + }, operatorID, nil +} + +func countAIPointChars(text string) int { + return utf8.RuneCountInString(strings.TrimSpace(text)) +} + +func calculateAIPointCost(chars, baseChars int) int { + if baseChars <= 0 { + baseChars = 1000 + } + if chars <= 0 { + return 1 + } + return chars/baseChars + 1 +} + +func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error { + _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(20260428, hashint8($1)::int)`, tenantID) + return err +} + +func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogID int64) (bool, error) { + var status string + err := tx.QueryRow(ctx, ` + SELECT status + FROM ai_point_usage_logs + WHERE id = $1 AND tenant_id = $2 + FOR UPDATE + `, usageLogID, tenantID).Scan(&status) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return false, nil + } + return false, err + } + return status == "pending", nil +} + +func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { + tag, err := tx.Exec(ctx, ` + UPDATE quota_reservations + SET status = 'confirmed', + consumed_amount = reserved_amount, + updated_at = NOW() + WHERE id = $1 + AND tenant_id = $2 + AND quota_type = $3 + AND status = 'pending' + `, reservationID, tenantID, aiPointsQuotaType) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return errors.New("ai points reservation is not pending") + } + return nil +} + +func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { + tag, err := tx.Exec(ctx, ` + UPDATE quota_reservations + SET status = 'refunded', + refunded_amount = reserved_amount, + updated_at = NOW() + WHERE id = $1 + AND tenant_id = $2 + AND quota_type = $3 + AND status = 'pending' + `, reservationID, tenantID, aiPointsQuotaType) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return errors.New("ai points reservation is not pending") + } + return nil +} diff --git a/server/internal/tenant/app/ai_point_usage_test.go b/server/internal/tenant/app/ai_point_usage_test.go new file mode 100644 index 0000000..c44221b --- /dev/null +++ b/server/internal/tenant/app/ai_point_usage_test.go @@ -0,0 +1,33 @@ +package app + +import "testing" + +func TestCalculateAIPointCost(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + chars int + baseChars int + want int + }{ + {name: "empty still charges base usage", chars: 0, baseChars: 1000, want: 1}, + {name: "single char", chars: 1, baseChars: 1000, want: 1}, + {name: "first base boundary", chars: 1000, baseChars: 1000, want: 2}, + {name: "over first base", chars: 1001, baseChars: 1000, want: 2}, + {name: "second base boundary", chars: 2000, baseChars: 1000, want: 3}, + {name: "over second base", chars: 2001, baseChars: 1000, want: 3}, + {name: "default base chars", chars: 1000, baseChars: 0, want: 2}, + {name: "custom base chars", chars: 1500, baseChars: 500, want: 4}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := calculateAIPointCost(tt.chars, tt.baseChars); got != tt.want { + t.Fatalf("calculateAIPointCost(%d, %d) = %d, want %d", tt.chars, tt.baseChars, got, tt.want) + } + }) + } +} diff --git a/server/internal/tenant/app/article_selection_optimize_service.go b/server/internal/tenant/app/article_selection_optimize_service.go index f68853f..4d2a4d0 100644 --- a/server/internal/tenant/app/article_selection_optimize_service.go +++ b/server/internal/tenant/app/article_selection_optimize_service.go @@ -11,6 +11,7 @@ import ( "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/llm" "github.com/geo-platform/tenant-api/internal/shared/response" ) @@ -38,6 +39,7 @@ type ArticleSelectionOptimizeService struct { pool *pgxpool.Pool llm llm.Client defaultMaxOutTokens int64 + cache sharedcache.Cache } func NewArticleSelectionOptimizeService( @@ -52,6 +54,11 @@ func NewArticleSelectionOptimizeService( } } +func (s *ArticleSelectionOptimizeService) WithCache(c sharedcache.Cache) *ArticleSelectionOptimizeService { + s.cache = c + return s +} + func (s *ArticleSelectionOptimizeService) ValidateRequest( ctx context.Context, articleID int64, @@ -70,9 +77,29 @@ func (s *ArticleSelectionOptimizeService) ValidateRequest( func (s *ArticleSelectionOptimizeService) OptimizeSelection( ctx context.Context, + articleID int64, req ArticleSelectionOptimizeRequest, onDelta func(string), ) (*ArticleSelectionOptimizeResult, error) { + actor := auth.MustActor(ctx) + resourceType := "article" + reservation, err := ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{ + TenantID: actor.TenantID, + OperatorID: actor.UserID, + UsageType: AIUsageTypeArticleSelectionOptimize, + ResourceType: &resourceType, + ResourceID: &articleID, + MeteredText: req.SelectedText, + Metadata: map[string]any{ + "article_id": articleID, + "instruction": strings.TrimSpace(req.Instruction), + "title_length": countAIPointChars(req.Title), + }, + }) + if err != nil { + return nil, err + } + result, err := s.llm.Generate( ctx, llm.GenerateRequest{ @@ -83,14 +110,27 @@ func (s *ArticleSelectionOptimizeService) OptimizeSelection( onDelta, ) if err != nil { + refundCtx, cancel := newGenerationCleanupContext() + _ = RefundAIPoints(refundCtx, s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, err.Error()) + cancel() return nil, response.ErrInternal(50019, "article_selection_optimize_failed", err.Error()) } content := strings.TrimSpace(result.Content) if content == "" { + refundCtx, cancel := newGenerationCleanupContext() + _ = RefundAIPoints(refundCtx, s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, "optimized content is empty") + cancel() return nil, response.ErrInternal(50019, "article_selection_optimize_failed", "optimized content is empty") } + completeCtx, cancel := newGenerationCleanupContext() + if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, result.Model); err != nil { + cancel() + return nil, response.ErrInternal(50019, "article_selection_optimize_failed", "failed to confirm ai points usage") + } + cancel() + return &ArticleSelectionOptimizeResult{ Content: content, Model: result.Model, diff --git a/server/internal/tenant/app/kol_assist_runtime.go b/server/internal/tenant/app/kol_assist_runtime.go index d85af78..6fc7710 100644 --- a/server/internal/tenant/app/kol_assist_runtime.go +++ b/server/internal/tenant/app/kol_assist_runtime.go @@ -47,26 +47,42 @@ func (s *KolAssistService) Run( req.Description = s.enrichGenerateDescription(ctx, req.Description) } + reservation, err := s.reserveStreamingAIPoints(ctx, actor, req) + if err != nil { + return nil, "", err + } + generateReq, err := BuildKolAssistGenerateRequest(req) if err != nil { + s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error()) return nil, "", response.ErrInternal(50087, "kol_assist_prompt_build_failed", err.Error()) } streamAccumulator := newKolAssistContentStream(onDelta) generated, err := s.llm.Generate(ctx, generateReq, streamAccumulator.ConsumeRawDelta) if err != nil { + s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error()) return nil, "", err } result, err := ParseKolAssistResponse(generated.Content) if err != nil { + s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error()) return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", err.Error()) } result, err = FinalizeKolAssistResult(req, result) if err != nil { + s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error()) return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", err.Error()) } + completeCtx, cancel := newGenerationCleanupContext() + if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, generated.Model); err != nil { + cancel() + return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", "failed to confirm ai points usage") + } + cancel() + return result, generated.Model, nil } diff --git a/server/internal/tenant/app/kol_assist_service.go b/server/internal/tenant/app/kol_assist_service.go index a33f2d0..dc54c48 100644 --- a/server/internal/tenant/app/kol_assist_service.go +++ b/server/internal/tenant/app/kol_assist_service.go @@ -6,8 +6,10 @@ import ( "strings" "github.com/google/uuid" + "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/config" "github.com/geo-platform/tenant-api/internal/shared/llm" "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" @@ -48,15 +50,18 @@ type KolAssistTaskResponse struct { } type KolAssistService struct { + pool *pgxpool.Pool profileSvc *KolProfileService repo repository.KolAssistRepository llm llm.Client messaging *rabbitmq.Client logger *zap.Logger urlResolver *kolAssistURLResolver + cache sharedcache.Cache } func NewKolAssistService( + pool *pgxpool.Pool, profileSvc *KolProfileService, repo repository.KolAssistRepository, llmClient llm.Client, @@ -65,6 +70,7 @@ func NewKolAssistService( logger *zap.Logger, ) *KolAssistService { return &KolAssistService{ + pool: pool, profileSvc: profileSvc, repo: repo, llm: llmClient, @@ -74,6 +80,11 @@ func NewKolAssistService( } } +func (s *KolAssistService) WithCache(c sharedcache.Cache) *KolAssistService { + s.cache = c + return s +} + func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req AssistRequest) (string, error) { profile, err := s.profileSvc.RequireActiveKol(ctx, actor) if err != nil { @@ -94,6 +105,11 @@ func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req Ass } taskID := uuid.NewString() + reservation, err := s.reserveAIPoints(ctx, actor, taskID, req) + if err != nil { + return "", err + } + if err := s.repo.Create(ctx, repository.CreateKolAssistInput{ ID: taskID, TenantID: profile.TenantID, @@ -103,6 +119,7 @@ func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req Ass TaskType: taskType, RequestJSON: requestJSON, }); err != nil { + s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error()) return "", response.ErrInternal(50084, "kol_assist_create_failed", err.Error()) } @@ -111,6 +128,7 @@ func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req Ass TenantID: profile.TenantID, }); err != nil { _ = s.repo.MarkFailed(context.Background(), profile.TenantID, taskID, err.Error()) + s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error()) return "", response.ErrServiceUnavailable(50342, "kol_assist_queue_unavailable", "assist queue is busy, please retry") } @@ -184,3 +202,75 @@ func normalizeAssistRequest(req AssistRequest) (AssistRequest, string, error) { return req, "ai_" + req.Mode, nil } + +func (s *KolAssistService) reserveAIPoints( + ctx context.Context, + actor auth.Actor, + taskID string, + req AssistRequest, +) (*AIPointReservation, error) { + usageType := AIUsageTypeKolPromptGenerate + if req.Mode == kolAssistModeOptimize { + usageType = AIUsageTypeKolPromptOptimize + } + resourceType := "kol_assist_task" + resourceUID := taskID + var resourceID *int64 + if req.PromptID != nil && *req.PromptID > 0 { + resourceID = req.PromptID + } + return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{ + TenantID: actor.TenantID, + OperatorID: actor.UserID, + UsageType: usageType, + ResourceType: &resourceType, + ResourceID: resourceID, + ResourceUID: &resourceUID, + MeteredText: buildKolAssistAIPointMeteredText(req), + Metadata: map[string]any{ + "task_id": taskID, + "mode": req.Mode, + "prompt_id": req.PromptID, + }, + }) +} + +func (s *KolAssistService) reserveStreamingAIPoints(ctx context.Context, actor auth.Actor, req AssistRequest) (*AIPointReservation, error) { + usageType := AIUsageTypeKolPromptGenerate + if req.Mode == kolAssistModeOptimize { + usageType = AIUsageTypeKolPromptOptimize + } + resourceType := "kol_prompt" + var resourceID *int64 + if req.PromptID != nil && *req.PromptID > 0 { + resourceID = req.PromptID + } + return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{ + TenantID: actor.TenantID, + OperatorID: actor.UserID, + UsageType: usageType, + ResourceType: &resourceType, + ResourceID: resourceID, + MeteredText: buildKolAssistAIPointMeteredText(req), + Metadata: map[string]any{ + "mode": req.Mode, + "prompt_id": req.PromptID, + "stream": true, + }, + }) +} + +func (s *KolAssistService) refundReservedAIPoints(ctx context.Context, actor auth.Actor, reservation *AIPointReservation, errorMessage string) { + if reservation == nil { + return + } + _ = RefundAIPoints(ctx, s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, errorMessage) +} + +func buildKolAssistAIPointMeteredText(req AssistRequest) string { + parts := []string{req.Description} + if req.Mode == kolAssistModeOptimize { + parts = append(parts, req.CurrentContent) + } + return strings.TrimSpace(strings.Join(parts, "\n")) +} diff --git a/server/internal/tenant/app/template_assist.go b/server/internal/tenant/app/template_assist.go index c122b4e..22a2440 100644 --- a/server/internal/tenant/app/template_assist.go +++ b/server/internal/tenant/app/template_assist.go @@ -144,6 +144,11 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6 } taskID := uuid.NewString() + reservation, err := s.reserveTemplateAssistAIPoints(ctx, actor.TenantID, actor.UserID, templateID, taskID, templateAssistTaskTypeAnalyze, buildAnalyzeAIPointMeteredText(req)) + if err != nil { + return nil, err + } + repo := repository.NewTemplateAssistTaskRepository(s.pool) if err := repo.CreateTask(ctx, repository.CreateTemplateAssistTaskInput{ ID: taskID, @@ -153,6 +158,7 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6 TemplateID: templateID, RequestJSON: payload, }); err != nil { + s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error()) return nil, response.ErrInternal(50021, "task_create_failed", err.Error()) } @@ -168,6 +174,7 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6 } if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil { _ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeAnalyze, actor.TenantID, err.Error(), time.Now()) + s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error()) return nil, response.ErrServiceUnavailable(50312, "analyze_queue_unavailable", "analyze queue is busy, please retry") } @@ -225,6 +232,11 @@ func (s *TemplateService) CreateTitleTask(ctx context.Context, templateID int64, } taskID := uuid.NewString() + reservation, err := s.reserveTemplateAssistAIPoints(ctx, actor.TenantID, actor.UserID, templateID, taskID, templateAssistTaskTypeTitle, buildTitleAIPointMeteredText(req)) + if err != nil { + return nil, err + } + repo := repository.NewTemplateAssistTaskRepository(s.pool) if err := repo.CreateTask(ctx, repository.CreateTemplateAssistTaskInput{ ID: taskID, @@ -234,6 +246,7 @@ func (s *TemplateService) CreateTitleTask(ctx context.Context, templateID int64, TemplateID: templateID, RequestJSON: payload, }); err != nil { + s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error()) return nil, response.ErrInternal(50025, "task_create_failed", err.Error()) } @@ -249,6 +262,7 @@ func (s *TemplateService) CreateTitleTask(ctx context.Context, templateID int64, } if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil { _ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeTitle, actor.TenantID, err.Error(), time.Now()) + s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error()) return nil, response.ErrServiceUnavailable(50313, "title_queue_unavailable", "title queue is busy, please retry") } @@ -306,6 +320,11 @@ func (s *TemplateService) CreateOutlineTask(ctx context.Context, templateID int6 } taskID := uuid.NewString() + reservation, err := s.reserveTemplateAssistAIPoints(ctx, actor.TenantID, actor.UserID, templateID, taskID, templateAssistTaskTypeOutline, buildOutlineAIPointMeteredText(req)) + if err != nil { + return nil, err + } + repo := repository.NewTemplateAssistTaskRepository(s.pool) if err := repo.CreateTask(ctx, repository.CreateTemplateAssistTaskInput{ ID: taskID, @@ -315,6 +334,7 @@ func (s *TemplateService) CreateOutlineTask(ctx context.Context, templateID int6 TemplateID: templateID, RequestJSON: payload, }); err != nil { + s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error()) return nil, response.ErrInternal(50029, "task_create_failed", err.Error()) } @@ -330,6 +350,7 @@ func (s *TemplateService) CreateOutlineTask(ctx context.Context, templateID int6 } if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil { _ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeOutline, actor.TenantID, err.Error(), time.Now()) + s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error()) return nil, response.ErrServiceUnavailable(50314, "outline_queue_unavailable", "outline queue is busy, please retry") } @@ -401,7 +422,24 @@ func (s *TemplateService) IsAssistTaskTerminal(ctx context.Context, job AssistJo if err != nil { return false, err } - return record.Status == "completed" || record.Status == "failed", nil + switch record.Status { + case "completed": + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, "") + return true, nil + case "failed": + if isTemplateAssistAIPointTask(job.TaskType) { + resourceType := "template_assist_task" + resourceUID := job.TaskID + errorMessage := "template assist task failed" + if record.ErrorMessage != nil && strings.TrimSpace(*record.ErrorMessage) != "" { + errorMessage = strings.TrimSpace(*record.ErrorMessage) + } + _ = RefundPendingAIPointsByResource(ctx, s.pool, s.cache, job.TenantID, resourceType, nil, &resourceUID, errorMessage) + } + return true, nil + default: + return false, nil + } } func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) { @@ -433,7 +471,9 @@ func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repositor if err := repo.CompleteTask(ctx, job.TaskID, job.TaskType, job.TenantID, resultJSON, time.Now()); err != nil { s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err) + return } + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model) } func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) { @@ -470,7 +510,9 @@ func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository. if err := repo.CompleteTask(ctx, job.TaskID, job.TaskType, job.TenantID, resultJSON, time.Now()); err != nil { s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err) + return } + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model) } func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) { @@ -512,12 +554,75 @@ func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repositor if err := repo.CompleteTask(ctx, job.TaskID, job.TaskType, job.TenantID, resultJSON, time.Now()); err != nil { s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err) + return } + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model) } func (s *TemplateService) failAssist(ctx context.Context, taskID, taskType string, tenantID int64, err error) { repo := repository.NewTemplateAssistTaskRepository(s.pool) _ = repo.FailTask(ctx, taskID, taskType, tenantID, err.Error(), time.Now()) + if isTemplateAssistAIPointTask(taskType) { + resourceType := "template_assist_task" + resourceUID := taskID + _ = RefundPendingAIPointsByResource(ctx, s.pool, s.cache, tenantID, resourceType, nil, &resourceUID, err.Error()) + } +} + +func (s *TemplateService) reserveTemplateAssistAIPoints( + ctx context.Context, + tenantID int64, + operatorID int64, + templateID int64, + taskID string, + taskType string, + meteredText string, +) (*AIPointReservation, error) { + usageType := AIUsageTypeTemplateTitleGenerate + fixedPoints := 0 + switch taskType { + case templateAssistTaskTypeAnalyze: + usageType = AIUsageTypeTemplateAnalyze + fixedPoints = 1 + case templateAssistTaskTypeOutline: + usageType = AIUsageTypeTemplateOutlineGenerate + } + resourceType := "template_assist_task" + resourceUID := taskID + return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{ + TenantID: tenantID, + OperatorID: operatorID, + UsageType: usageType, + ResourceType: &resourceType, + ResourceID: &templateID, + ResourceUID: &resourceUID, + MeteredText: meteredText, + FixedPoints: fixedPoints, + Metadata: map[string]any{ + "template_id": templateID, + "task_id": taskID, + "task_type": taskType, + }, + }) +} + +func (s *TemplateService) refundReservedAIPoints(ctx context.Context, tenantID, operatorID int64, reservation *AIPointReservation, errorMessage string) { + if reservation == nil { + return + } + _ = RefundAIPoints(ctx, s.pool, s.cache, tenantID, operatorID, *reservation, errorMessage) +} + +func (s *TemplateService) completeTemplateAssistAIPoints(ctx context.Context, tenantID int64, taskID string, model string) { + resourceType := "template_assist_task" + resourceUID := taskID + _ = CompletePendingAIPointsByResource(ctx, s.pool, tenantID, resourceType, nil, &resourceUID, model) +} + +func isTemplateAssistAIPointTask(taskType string) bool { + return taskType == templateAssistTaskTypeAnalyze || + taskType == templateAssistTaskTypeTitle || + taskType == templateAssistTaskTypeOutline } func normalizeAnalyzeTaskRequest(req AnalyzeTaskRequest) AnalyzeTaskRequest { @@ -661,16 +766,7 @@ func hasOutlineContext(req OutlineTaskRequest) bool { } func buildAnalyzePrompt(templateKey, templateName string, req AnalyzeTaskRequest, analyzePromptTemplate *string) string { - contextPayload := map[string]interface{}{ - "template_key": templateKey, - "template_name": templateName, - "locale": req.Locale, - "brand_name": req.BrandName, - "official_website": req.Website, - "input_params": req.InputParams, - "existing_keywords": req.ExistingKeywords, - "existing_competitors": req.ExistingCompetitors, - } + contextPayload := analyzePromptParams(templateKey, templateName, req) if analyzePromptTemplate != nil && strings.TrimSpace(*analyzePromptTemplate) != "" { basePrompt := strings.TrimSpace(renderPromptTemplate(analyzePromptTemplate, contextPayload)) @@ -684,6 +780,11 @@ func buildAnalyzePrompt(templateKey, templateName string, req AnalyzeTaskRequest return prompts.AnalyzeFallbackPrompt(string(contextJSON)) } +func buildAnalyzeAIPointMeteredText(req AnalyzeTaskRequest) string { + payload, _ := json.Marshal(analyzePromptParams("", "", req)) + return string(payload) +} + func buildTitlePrompt(templateKey, templateName string, req TitleTaskRequest, titlePromptTemplate *string) string { contextPayload := titlePromptParams(templateKey, templateName, req) if titlePromptTemplate != nil && strings.TrimSpace(*titlePromptTemplate) != "" { @@ -723,6 +824,29 @@ func buildOutlinePrompt(templateKey, templateName string, req OutlineTaskRequest return prompts.OutlineFallbackPrompt(string(contextJSON)) } +func buildTitleAIPointMeteredText(req TitleTaskRequest) string { + payload, _ := json.Marshal(titlePromptParams("", "", req)) + return string(payload) +} + +func buildOutlineAIPointMeteredText(req OutlineTaskRequest) string { + payload, _ := json.Marshal(outlinePromptParams("", "", req)) + return string(payload) +} + +func analyzePromptParams(templateKey, templateName string, req AnalyzeTaskRequest) map[string]interface{} { + return map[string]interface{}{ + "template_key": templateKey, + "template_name": templateName, + "locale": req.Locale, + "brand_name": req.BrandName, + "official_website": req.Website, + "input_params": req.InputParams, + "existing_keywords": req.ExistingKeywords, + "existing_competitors": req.ExistingCompetitors, + } +} + func titlePromptParams(templateKey, templateName string, req TitleTaskRequest) map[string]interface{} { params := map[string]interface{}{ "template_key": templateKey, diff --git a/server/internal/tenant/app/workspace_service.go b/server/internal/tenant/app/workspace_service.go index 29e65cc..18022aa 100644 --- a/server/internal/tenant/app/workspace_service.go +++ b/server/internal/tenant/app/workspace_service.go @@ -20,6 +20,7 @@ type WorkspaceService struct { repo repository.WorkspaceRepository templates repository.TemplateRepository quota repository.QuotaRepository + aiPointUsage repository.AIPointUsageRepository supportedPlatformCount int cache sharedcache.Cache cacheGroup singleflight.Group @@ -29,11 +30,13 @@ func NewWorkspaceService( repo repository.WorkspaceRepository, templates repository.TemplateRepository, quota repository.QuotaRepository, + aiPointUsage repository.AIPointUsageRepository, ) *WorkspaceService { return &WorkspaceService{ repo: repo, templates: templates, quota: quota, + aiPointUsage: aiPointUsage, supportedPlatformCount: 5, } } @@ -105,18 +108,80 @@ func (s *WorkspaceService) QuotaSummary(ctx context.Context) (*domain.QuotaSumma } return nil, response.ErrInternal(50010, "query_failed", "failed to get quota balance") } + aiStatus, err := s.quota.GetAIQuotaStatus(loadCtx, actor.TenantID, time.Now().UTC()) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + aiStatus = &repository.AIQuotaStatus{} + } else { + return nil, response.ErrInternal(50010, "query_failed", "failed to get ai points balance") + } + } return &domain.QuotaSummary{ - PlanCode: status.PlanCode, - PlanName: status.PlanName, - TotalQuota: status.Total, - UsedQuota: status.Used, - Balance: status.Balance, - ResetAt: status.ResetAt, + PlanCode: status.PlanCode, + PlanName: status.PlanName, + TotalQuota: status.Total, + UsedQuota: status.Used, + Balance: status.Balance, + ResetAt: status.ResetAt, + AIPointsTotal: aiStatus.Total, + AIPointsUsed: aiStatus.Used, + AIPointsBalance: aiStatus.Balance, + AIPointBaseChars: aiStatus.BaseChars, + AIPointsResetAt: aiStatus.ResetAt, }, nil }) } +func (s *WorkspaceService) AIPointUsageLogs(ctx context.Context, page, pageSize int) (*domain.AIPointUsageListResponse, error) { + actor := auth.MustActor(ctx) + if page < 1 { + page = 1 + } + if pageSize < 1 || pageSize > 100 { + pageSize = 20 + } + offset := (page - 1) * pageSize + + rows, err := s.aiPointUsage.List(ctx, repository.AIPointUsageListParams{ + TenantID: actor.TenantID, + Limit: pageSize, + Offset: offset, + }) + if err != nil { + return nil, response.ErrInternal(50010, "query_failed", "failed to list ai point usage") + } + total, err := s.aiPointUsage.Count(ctx, actor.TenantID) + if err != nil { + return nil, response.ErrInternal(50010, "query_failed", "failed to count ai point usage") + } + + items := make([]domain.AIPointUsageLog, 0, len(rows)) + for _, row := range rows { + items = append(items, domain.AIPointUsageLog{ + ID: row.ID, + UsageType: row.UsageType, + ResourceType: row.ResourceType, + ResourceID: row.ResourceID, + ResourceUID: row.ResourceUID, + RequestChars: row.RequestChars, + BaseChars: row.BaseChars, + Points: row.Points, + Status: row.Status, + Model: row.Model, + ErrorMessage: row.ErrorMessage, + CompletedAt: row.CompletedAt, + CreatedAt: row.CreatedAt, + }) + } + return &domain.AIPointUsageListResponse{ + Items: items, + Total: total, + Page: page, + PageSize: pageSize, + }, nil +} + type TemplateCard struct { ID int64 `json:"id"` Scope string `json:"scope"` diff --git a/server/internal/tenant/domain/workspace.go b/server/internal/tenant/domain/workspace.go index 7329e1a..d14e4dd 100644 --- a/server/internal/tenant/domain/workspace.go +++ b/server/internal/tenant/domain/workspace.go @@ -23,12 +23,17 @@ type RecentArticle struct { } type QuotaSummary struct { - PlanCode string `json:"plan_code"` - PlanName string `json:"plan_name"` - TotalQuota int `json:"total_quota"` - UsedQuota int `json:"used_quota"` - Balance int `json:"balance"` - ResetAt *time.Time `json:"reset_at"` + PlanCode string `json:"plan_code"` + PlanName string `json:"plan_name"` + TotalQuota int `json:"total_quota"` + UsedQuota int `json:"used_quota"` + Balance int `json:"balance"` + ResetAt *time.Time `json:"reset_at"` + AIPointsTotal int `json:"ai_points_total"` + AIPointsUsed int `json:"ai_points_used"` + AIPointsBalance int `json:"ai_points_balance"` + AIPointBaseChars int `json:"ai_point_base_chars"` + AIPointsResetAt *time.Time `json:"ai_points_reset_at"` } type KolWorkspaceCard struct { @@ -41,3 +46,26 @@ type KolWorkspaceCard struct { PackageCover *string `json:"package_cover"` KolDisplayName string `json:"kol_display_name"` } + +type AIPointUsageLog struct { + ID int64 `json:"id"` + UsageType string `json:"usage_type"` + ResourceType *string `json:"resource_type"` + ResourceID *int64 `json:"resource_id"` + ResourceUID *string `json:"resource_uid"` + RequestChars int `json:"request_chars"` + BaseChars int `json:"base_chars"` + Points int `json:"points"` + Status string `json:"status"` + Model *string `json:"model"` + ErrorMessage *string `json:"error_message"` + CompletedAt *time.Time `json:"completed_at"` + CreatedAt time.Time `json:"created_at"` +} + +type AIPointUsageListResponse struct { + Items []AIPointUsageLog `json:"items"` + Total int64 `json:"total"` + Page int `json:"page"` + PageSize int `json:"page_size"` +} diff --git a/server/internal/tenant/repository/ai_point_usage_repo.go b/server/internal/tenant/repository/ai_point_usage_repo.go new file mode 100644 index 0000000..a3e28b7 --- /dev/null +++ b/server/internal/tenant/repository/ai_point_usage_repo.go @@ -0,0 +1,254 @@ +package repository + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5/pgtype" + + "github.com/geo-platform/tenant-api/internal/tenant/repository/generated" +) + +type AIPointUsageLog struct { + ID int64 + TenantID int64 + OperatorID *int64 + QuotaReservationID *int64 + UsageType string + ResourceType *string + ResourceID *int64 + ResourceUID *string + RequestChars int + BaseChars int + Points int + Status string + Model *string + MetadataJSON []byte + ErrorMessage *string + CompletedAt *time.Time + CreatedAt time.Time + UpdatedAt time.Time +} + +type CreateAIPointUsageLogInput struct { + TenantID int64 + OperatorID int64 + QuotaReservationID int64 + UsageType string + ResourceType *string + ResourceID *int64 + ResourceUID *string + RequestChars int + BaseChars int + Points int + MetadataJSON []byte +} + +type AIPointUsageListParams struct { + TenantID int64 + Limit int + Offset int +} + +type AIPointUsageRepository interface { + Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error) + MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error + MarkRefunded(ctx context.Context, tenantID, id int64, errorMessage string) error + FindPendingByResource(ctx context.Context, tenantID int64, resourceType string, resourceID *int64, resourceUID *string) (*AIPointUsageLog, error) + List(ctx context.Context, params AIPointUsageListParams) ([]AIPointUsageLog, error) + Count(ctx context.Context, tenantID int64) (int64, error) +} + +type aiPointUsageRepository struct { + db generated.DBTX +} + +func NewAIPointUsageRepository(db generated.DBTX) AIPointUsageRepository { + return &aiPointUsageRepository{db: db} +} + +func (r *aiPointUsageRepository) Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error) { + metadataJSON := input.MetadataJSON + if len(metadataJSON) == 0 { + metadataJSON = []byte("{}") + } + var operatorID *int64 + if input.OperatorID > 0 { + operatorID = &input.OperatorID + } + var reservationID *int64 + if input.QuotaReservationID > 0 { + reservationID = &input.QuotaReservationID + } + + var id int64 + err := r.db.QueryRow(ctx, ` +INSERT INTO ai_point_usage_logs ( + tenant_id, operator_id, quota_reservation_id, usage_type, + resource_type, resource_id, resource_uid, + request_chars, base_chars, points, metadata_json +) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb) +RETURNING id +`, input.TenantID, operatorID, reservationID, input.UsageType, + input.ResourceType, input.ResourceID, input.ResourceUID, + input.RequestChars, input.BaseChars, input.Points, metadataJSON, + ).Scan(&id) + return id, err +} + +func (r *aiPointUsageRepository) MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error { + _, err := r.db.Exec(ctx, ` +UPDATE ai_point_usage_logs +SET status = 'completed', + model = $1, + error_message = NULL, + completed_at = COALESCE(completed_at, NOW()), + updated_at = NOW() +WHERE id = $2 + AND tenant_id = $3 + AND status = 'pending' +`, model, id, tenantID) + return err +} + +func (r *aiPointUsageRepository) MarkRefunded(ctx context.Context, tenantID, id int64, errorMessage string) error { + _, err := r.db.Exec(ctx, ` +UPDATE ai_point_usage_logs +SET status = 'refunded', + error_message = $1, + completed_at = COALESCE(completed_at, NOW()), + updated_at = NOW() +WHERE id = $2 + AND tenant_id = $3 + AND status = 'pending' +`, errorMessage, id, tenantID) + return err +} + +func (r *aiPointUsageRepository) FindPendingByResource( + ctx context.Context, + tenantID int64, + resourceType string, + resourceID *int64, + resourceUID *string, +) (*AIPointUsageLog, error) { + row := r.db.QueryRow(ctx, ` +SELECT id, tenant_id, operator_id, quota_reservation_id, usage_type, + resource_type, resource_id, resource_uid, request_chars, base_chars, + points, status, model, metadata_json, error_message, completed_at, + created_at, updated_at +FROM ai_point_usage_logs +WHERE tenant_id = $1 + AND resource_type = $2 + AND ($3::bigint IS NULL OR resource_id = $3) + AND ($4::varchar IS NULL OR resource_uid = $4) + AND status = 'pending' +ORDER BY created_at DESC, id DESC +LIMIT 1 +`, tenantID, resourceType, resourceID, resourceUID) + + var item AIPointUsageLog + err := scanAIPointUsageLog(row, &item) + if err != nil { + return nil, err + } + return &item, nil +} + +func (r *aiPointUsageRepository) List(ctx context.Context, params AIPointUsageListParams) ([]AIPointUsageLog, error) { + limit := params.Limit + if limit <= 0 || limit > 100 { + limit = 50 + } + offset := params.Offset + if offset < 0 { + offset = 0 + } + + rows, err := r.db.Query(ctx, ` +SELECT id, tenant_id, operator_id, quota_reservation_id, usage_type, + resource_type, resource_id, resource_uid, request_chars, base_chars, + points, status, model, metadata_json, error_message, completed_at, + created_at, updated_at +FROM ai_point_usage_logs +WHERE tenant_id = $1 +ORDER BY created_at DESC, id DESC +LIMIT $2 OFFSET $3 +`, params.TenantID, limit, offset) + if err != nil { + return nil, err + } + defer rows.Close() + + items := make([]AIPointUsageLog, 0, limit) + for rows.Next() { + var item AIPointUsageLog + if err := scanAIPointUsageLog(rows, &item); err != nil { + return nil, err + } + items = append(items, item) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +func (r *aiPointUsageRepository) Count(ctx context.Context, tenantID int64) (int64, error) { + var total int64 + err := r.db.QueryRow(ctx, ` + SELECT COUNT(*) + FROM ai_point_usage_logs + WHERE tenant_id = $1 + `, tenantID).Scan(&total) + return total, err +} + +type aiPointUsageScanner interface { + Scan(dest ...any) error +} + +func scanAIPointUsageLog(scanner aiPointUsageScanner, item *AIPointUsageLog) error { + var operatorID pgtype.Int8 + var quotaReservationID pgtype.Int8 + var resourceType pgtype.Text + var resourceID pgtype.Int8 + var resourceUID pgtype.Text + var model pgtype.Text + var errorMessage pgtype.Text + var completedAt pgtype.Timestamptz + + if err := scanner.Scan( + &item.ID, + &item.TenantID, + &operatorID, + "aReservationID, + &item.UsageType, + &resourceType, + &resourceID, + &resourceUID, + &item.RequestChars, + &item.BaseChars, + &item.Points, + &item.Status, + &model, + &item.MetadataJSON, + &errorMessage, + &completedAt, + &item.CreatedAt, + &item.UpdatedAt, + ); err != nil { + return err + } + + item.OperatorID = nullableInt64(operatorID) + item.QuotaReservationID = nullableInt64(quotaReservationID) + item.ResourceType = nullableText(resourceType) + item.ResourceID = nullableInt64(resourceID) + item.ResourceUID = nullableText(resourceUID) + item.Model = nullableText(model) + item.ErrorMessage = nullableText(errorMessage) + item.CompletedAt = optionalTime(completedAt) + return nil +} diff --git a/server/internal/tenant/repository/quota_repo.go b/server/internal/tenant/repository/quota_repo.go index 899ceb9..2fc60a5 100644 --- a/server/internal/tenant/repository/quota_repo.go +++ b/server/internal/tenant/repository/quota_repo.go @@ -19,6 +19,16 @@ type ArticleQuotaStatus struct { ResetAt *time.Time } +type AIQuotaStatus struct { + PlanCode string + PlanName string + Total int + Used int + Balance int + BaseChars int + ResetAt *time.Time +} + type QuotaLedgerInput struct { TenantID int64 OperatorID int64 @@ -43,6 +53,7 @@ type QuotaReservationInput struct { type QuotaRepository interface { GetCurrentBalance(ctx context.Context, tenantID int64, quotaType string) (int, error) GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error) + GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error) InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error) CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error) UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error @@ -78,6 +89,58 @@ func (r *quotaRepository) GetCurrentBalance(ctx context.Context, tenantID int64, return status.Balance, nil } +func (r *quotaRepository) GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error) { + access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now) + if err != nil { + return nil, err + } + if access == nil { + return nil, pgx.ErrNoRows + } + + baseChars := access.AIPointBaseChars() + if !access.HasActiveAccess(now) { + return &AIQuotaStatus{ + PlanCode: access.PlanCode, + PlanName: access.PlanName, + BaseChars: baseChars, + }, nil + } + + total := access.AIPointsMonthlyLimit() + if total <= 0 { + return &AIQuotaStatus{ + PlanCode: access.PlanCode, + PlanName: access.PlanName, + BaseChars: baseChars, + }, nil + } + + windowStart, windowEnd, resetAt := access.AIPointsWindow(now) + if windowEnd.IsZero() { + return nil, errors.New("ai points quota window end is missing") + } + + used, err := CountEffectiveQuotaReservations(ctx, r.db, tenantID, "ai_points", windowStart, windowEnd) + if err != nil { + return nil, err + } + balance := total - used + if balance < 0 { + balance = 0 + } + + return &AIQuotaStatus{ + PlanCode: access.PlanCode, + PlanName: access.PlanName, + Total: total, + Used: used, + Balance: balance, + BaseChars: baseChars, + ResetAt: resetAt, + }, nil +} + func (r *quotaRepository) GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error) { access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now) if err != nil { diff --git a/server/internal/tenant/transport/article_handler.go b/server/internal/tenant/transport/article_handler.go index 96b6214..dde1b64 100644 --- a/server/internal/tenant/transport/article_handler.go +++ b/server/internal/tenant/transport/article_handler.go @@ -48,7 +48,7 @@ func NewArticleHandler(a *bootstrap.App) *ArticleHandler { a.DB, a.LLM, a.Config.LLM.MaxOutputTokens, - ), + ).WithCache(a.Cache), promptGenerateSvc: app.NewPromptRuleGenerationService( a.DB, a.LLM, @@ -344,7 +344,7 @@ func (h *ArticleHandler) OptimizeSelectionStream(c *gin.Context) { var streamed strings.Builder var streamWriteErr error - result, err := h.selectionOptimize.OptimizeSelection(streamCtx, req, func(delta string) { + result, err := h.selectionOptimize.OptimizeSelection(streamCtx, id, req, func(delta string) { if delta == "" || streamWriteErr != nil { return } diff --git a/server/internal/tenant/transport/kol_manage_handler.go b/server/internal/tenant/transport/kol_manage_handler.go index 12411cc..e2f4d56 100644 --- a/server/internal/tenant/transport/kol_manage_handler.go +++ b/server/internal/tenant/transport/kol_manage_handler.go @@ -42,7 +42,7 @@ func NewKolManageHandler(a *bootstrap.App, assets *AssetHandler) *KolManageHandl profileSvc: profileSvc, pkgSvc: app.NewKolPackageService(profileSvc, a.KolPackages, a.KolPrompts).WithCache(a.Cache), promptSvc: app.NewKolPromptService(a.DB, profileSvc, a.KolPackages, a.KolPrompts, a.KolSubscriptions, a.KolPromptAsset).WithCache(a.Cache), - assistSvc: app.NewKolAssistService(profileSvc, a.KolAssists, a.LLM, a.RabbitMQ, a.Config.LLM, a.Logger), + assistSvc: app.NewKolAssistService(a.DB, profileSvc, a.KolAssists, a.LLM, a.RabbitMQ, a.Config.LLM, a.Logger).WithCache(a.Cache), assets: assets, } } diff --git a/server/internal/tenant/transport/router.go b/server/internal/tenant/transport/router.go index ecdd5a8..6dcad9a 100644 --- a/server/internal/tenant/transport/router.go +++ b/server/internal/tenant/transport/router.go @@ -80,6 +80,7 @@ func RegisterRoutes(app *bootstrap.App) { workspace.GET("/overview", wsHandler.Overview) workspace.GET("/recent-articles", wsHandler.RecentArticles) workspace.GET("/quota-summary", wsHandler.QuotaSummary) + workspace.GET("/ai-point-usage", wsHandler.AIPointUsageLogs) workspace.GET("/template-cards", wsHandler.TemplateCards) workspace.GET("/kol-cards", wsHandler.KolCards) diff --git a/server/internal/tenant/transport/workspace_handler.go b/server/internal/tenant/transport/workspace_handler.go index 6ad75f6..9674bda 100644 --- a/server/internal/tenant/transport/workspace_handler.go +++ b/server/internal/tenant/transport/workspace_handler.go @@ -22,6 +22,7 @@ func NewWorkspaceHandler(a *bootstrap.App) *WorkspaceHandler { a.Cache, ), repository.NewQuotaRepository(a.DB), + repository.NewAIPointUsageRepository(a.DB), ).WithCache(a.Cache), } } @@ -53,6 +54,30 @@ func (h *WorkspaceHandler) QuotaSummary(c *gin.Context) { response.Success(c, data) } +func (h *WorkspaceHandler) AIPointUsageLogs(c *gin.Context) { + page := parseIntQuery(c.Query("page"), 1) + pageSize := parseIntQuery(c.Query("page_size"), 20) + if c.Query("page") == "" && c.Query("page_size") == "" && c.Query("limit") != "" { + limit := parseIntQuery(c.Query("limit"), 20) + offset := parseIntQuery(c.Query("offset"), 0) + if limit > 0 { + pageSize = limit + page = offset/limit + 1 + } + } + + data, err := h.svc.AIPointUsageLogs( + c.Request.Context(), + page, + pageSize, + ) + if err != nil { + response.Error(c, err) + return + } + response.Success(c, data) +} + func (h *WorkspaceHandler) TemplateCards(c *gin.Context) { data, err := h.svc.TemplateCards(c.Request.Context()) if err != nil { diff --git a/server/internal/worker/assist/kol_assist_worker.go b/server/internal/worker/assist/kol_assist_worker.go index 12685cb..1f6ff9d 100644 --- a/server/internal/worker/assist/kol_assist_worker.go +++ b/server/internal/worker/assist/kol_assist_worker.go @@ -8,9 +8,11 @@ import ( "os" "time" + "github.com/jackc/pgx/v5/pgxpool" amqp "github.com/rabbitmq/amqp091-go" "go.uber.org/zap" + 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/shared/messaging/rabbitmq" tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" @@ -20,10 +22,12 @@ import ( var errKolAssistConsumerClosed = errors.New("kol assist consumer channel closed") type KolAssistWorker struct { + pool *pgxpool.Pool rabbitMQ *rabbitmq.Client repo repository.KolAssistRepository llm llm.Client logger *zap.Logger + cache sharedcache.Cache retryInterval time.Duration processTimeout time.Duration consumerPrefix string @@ -31,10 +35,12 @@ type KolAssistWorker struct { } func NewKolAssistWorker( + pool *pgxpool.Pool, rabbitMQClient *rabbitmq.Client, repo repository.KolAssistRepository, llmClient llm.Client, logger *zap.Logger, + cache sharedcache.Cache, workerConcurrency int, processTimeout time.Duration, ) *KolAssistWorker { @@ -46,10 +52,12 @@ func NewKolAssistWorker( } return &KolAssistWorker{ + pool: pool, rabbitMQ: rabbitMQClient, repo: repo, llm: llmClient, logger: logger, + cache: cache, retryInterval: 5 * time.Second, processTimeout: processTimeout, consumerPrefix: fmt.Sprintf("kol-assist-%d", os.Getpid()), @@ -73,7 +81,19 @@ func (w *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJo if err != nil { return fmt.Errorf("load kol assist task: %w", err) } - if task == nil || task.Status == "completed" || task.Status == "failed" { + if task == nil { + return nil + } + switch task.Status { + case "completed": + w.completeAIPoints(ctx, job, "") + return nil + case "failed": + errorMessage := "kol assist task failed" + if task.ErrorMessage != nil && *task.ErrorMessage != "" { + errorMessage = *task.ErrorMessage + } + w.refundAIPoints(ctx, job, errorMessage) return nil } @@ -83,6 +103,7 @@ func (w *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJo if err := w.processTask(ctx, job, task); err != nil { _ = w.repo.MarkFailed(context.Background(), job.TenantID, job.TaskID, err.Error()) + w.refundAIPoints(context.Background(), job, err.Error()) if w.logger != nil { w.logger.Warn("kol assist task failed", zap.Error(err), @@ -215,5 +236,30 @@ func (w *KolAssistWorker) processTask(ctx context.Context, job tenantapp.KolAssi if err := w.repo.MarkCompleted(ctx, job.TenantID, job.TaskID, resultJSON); err != nil { return fmt.Errorf("mark kol assist completed: %w", err) } + w.completeAIPoints(ctx, job, result.Model) return nil } + +func (w *KolAssistWorker) completeAIPoints(ctx context.Context, job tenantapp.KolAssistJob, model string) { + resourceType := "kol_assist_task" + resourceUID := job.TaskID + if err := tenantapp.CompletePendingAIPointsByResource(ctx, w.pool, job.TenantID, resourceType, nil, &resourceUID, model); err != nil && w.logger != nil { + w.logger.Warn("kol assist ai points confirmation failed", + zap.Error(err), + zap.String("task_id", job.TaskID), + zap.Int64("tenant_id", job.TenantID), + ) + } +} + +func (w *KolAssistWorker) refundAIPoints(ctx context.Context, job tenantapp.KolAssistJob, errorMessage string) { + resourceType := "kol_assist_task" + resourceUID := job.TaskID + if err := tenantapp.RefundPendingAIPointsByResource(ctx, w.pool, w.cache, job.TenantID, resourceType, nil, &resourceUID, errorMessage); err != nil && w.logger != nil { + w.logger.Warn("kol assist ai points refund failed", + zap.Error(err), + zap.String("task_id", job.TaskID), + zap.Int64("tenant_id", job.TenantID), + ) + } +} diff --git a/server/migrations/20260428110000_create_ai_point_usage_logs.down.sql b/server/migrations/20260428110000_create_ai_point_usage_logs.down.sql new file mode 100644 index 0000000..d922a80 --- /dev/null +++ b/server/migrations/20260428110000_create_ai_point_usage_logs.down.sql @@ -0,0 +1,3 @@ +DROP INDEX IF EXISTS idx_quota_reservations_tenant_type_created_status; + +DROP TABLE IF EXISTS ai_point_usage_logs; diff --git a/server/migrations/20260428110000_create_ai_point_usage_logs.up.sql b/server/migrations/20260428110000_create_ai_point_usage_logs.up.sql new file mode 100644 index 0000000..fa35676 --- /dev/null +++ b/server/migrations/20260428110000_create_ai_point_usage_logs.up.sql @@ -0,0 +1,37 @@ +CREATE TABLE ai_point_usage_logs ( + id BIGSERIAL PRIMARY KEY, + tenant_id BIGINT NOT NULL REFERENCES tenants(id), + operator_id BIGINT REFERENCES users(id), + quota_reservation_id BIGINT REFERENCES quota_reservations(id), + usage_type VARCHAR(64) NOT NULL, + resource_type VARCHAR(64), + resource_id BIGINT, + resource_uid VARCHAR(128), + request_chars INT NOT NULL DEFAULT 0, + base_chars INT NOT NULL DEFAULT 1000, + points INT NOT NULL, + status VARCHAR(20) NOT NULL DEFAULT 'pending', + model VARCHAR(128), + metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb, + error_message TEXT, + completed_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT chk_ai_point_usage_status + CHECK (status IN ('pending', 'completed', 'refunded', 'failed')) +); + +CREATE INDEX idx_ai_point_usage_tenant_created + ON ai_point_usage_logs(tenant_id, created_at DESC); + +CREATE INDEX idx_ai_point_usage_tenant_operator_created + ON ai_point_usage_logs(tenant_id, operator_id, created_at DESC); + +CREATE INDEX idx_ai_point_usage_reservation + ON ai_point_usage_logs(quota_reservation_id); + +CREATE INDEX idx_ai_point_usage_resource + ON ai_point_usage_logs(tenant_id, resource_type, resource_id, resource_uid); + +CREATE INDEX idx_quota_reservations_tenant_type_created_status + ON quota_reservations(tenant_id, quota_type, created_at, status);