root
01fa2b8309
Desktop Client Build / Resolve Build Metadata (push) Successful in 28s
Frontend CI / Frontend (push) Successful in 3m57s
Backend CI / Backend (push) Successful in 17m3s
Desktop Client Build / Build Desktop Client (push) Successful in 26m8s
Desktop Client Build / Publish Client Artifacts to NAS (push) Successful in 51s
437 lines
12 KiB
Go
437 lines
12 KiB
Go
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 := confirmAIPointReservationIfPending(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
|
|
}
|
|
|
|
refunded, err := refundAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if refunded {
|
|
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 := 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 confirmAIPointReservationIfPending(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
|
|
}
|
|
return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
|
|
}
|
|
|
|
func refundAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) (bool, 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 false, err
|
|
}
|
|
if tag.RowsAffected() > 0 {
|
|
return true, nil
|
|
}
|
|
return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID)
|
|
}
|
|
|
|
func ignoreAIPointReservationReset(ctx context.Context, tx pgx.Tx, updated bool, tenantID, reservationID int64) error {
|
|
if updated {
|
|
return nil
|
|
}
|
|
var status string
|
|
err := tx.QueryRow(ctx, `
|
|
SELECT status
|
|
FROM quota_reservations
|
|
WHERE id = $1
|
|
AND tenant_id = $2
|
|
AND quota_type = $3
|
|
`, reservationID, tenantID, aiPointsQuotaType).Scan(&status)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if status == "reset" {
|
|
return nil
|
|
}
|
|
return errors.New("ai points reservation is not pending")
|
|
}
|