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:
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user