feat(server): meter AI point usage across tenant AI features

Adds an ai_point_usage_logs ledger and a reserve/refund/complete pipeline
that charges AI points for article selection optimize, template analyze
/title/outline, and KOL prompt generate/optimize. Pending reservations
are reconciled when the kol-assist worker and template-assist tasks
finish or fail, so points refund automatically on errors. Workspace now
exposes the AI quota status and a paginated usage ledger.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-28 11:33:47 +08:00
parent 794a2d89db
commit e307a048d1
18 changed files with 1270 additions and 27 deletions
+2
View File
@@ -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)
@@ -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
}
@@ -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)
}
})
}
}
@@ -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,
@@ -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
}
@@ -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"))
}
+135 -11
View File
@@ -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,
@@ -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"`
+34 -6
View File
@@ -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"`
}
@@ -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,
&quotaReservationID,
&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
}
@@ -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 {
@@ -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
}
@@ -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,
}
}
@@ -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)
@@ -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 {
@@ -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),
)
}
}
@@ -0,0 +1,3 @@
DROP INDEX IF EXISTS idx_quota_reservations_tenant_type_created_status;
DROP TABLE IF EXISTS ai_point_usage_logs;
@@ -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);