From f69edc2218ef010a3a4fffb7d77c1975c482ab7b Mon Sep 17 00:00:00 2001 From: liangxu Date: Thu, 9 Jul 2026 17:04:37 +0800 Subject: [PATCH] feat(tenant): bill AI points by output characters Previously AI point cost was computed from request (input) characters at reservation time. Switch to reserving a single point up front and settling the final cost from the generated output length on completion. - Reserve 1 point (or FixedPoints) instead of pricing on request chars - On completion, recompute points from output chars vs base chars, update the reservation amounts, and post a ledger delta for the top-up/refund; invalidate the workspace quota summary cache - MarkCompleted now persists final request_chars/base_chars/points - Use ceil division in calculateAIPointCost so an exact base boundary charges one point instead of rolling over - Thread output text through all CompleteAIPoints callers (article selection, KOL assist, question expansion, template assist, compliance judge) and return the settled reservation - Add unit tests plus an integration test gated on TEST_DATABASE_URL - Update the user manual: billing is by output characters Co-Authored-By: Claude Opus 4.8 (1M context) --- server/internal/tenant/app/ai_point_usage.go | 134 ++++++++++-- .../app/ai_point_usage_integration_test.go | 200 ++++++++++++++++++ .../tenant/app/ai_point_usage_test.go | 165 ++++++++++++++- .../app/article_selection_optimize_service.go | 2 +- .../internal/tenant/app/kol_assist_runtime.go | 2 +- .../tenant/app/question_expansion_service.go | 10 +- server/internal/tenant/app/template_assist.go | 12 +- server/internal/tenant/compliance/service.go | 76 ++++++- .../tenant/compliance/service_test.go | 37 ++++ .../tenant/repository/ai_point_usage_repo.go | 13 +- .../worker/assist/kol_assist_worker.go | 8 +- user-manual/11-AI点数与套餐.md | 4 +- 12 files changed, 612 insertions(+), 51 deletions(-) create mode 100644 server/internal/tenant/app/ai_point_usage_integration_test.go diff --git a/server/internal/tenant/app/ai_point_usage.go b/server/internal/tenant/app/ai_point_usage.go index 81ed454..2aacaf8 100644 --- a/server/internal/tenant/app/ai_point_usage.go +++ b/server/internal/tenant/app/ai_point_usage.go @@ -44,11 +44,20 @@ type AIPointReserveInput struct { type AIPointReservation struct { ReservationID int64 UsageLogID int64 + OperatorID int64 Points int RequestChars int BaseChars int } +type aiPointOutputSettlement struct { + OutputChars int + BaseChars int + FinalPoints int + ReservedPoints int + LedgerDelta int +} + func ReserveAIPoints( ctx context.Context, pool *pgxpool.Pool, @@ -85,8 +94,8 @@ func ReserveAIPoints( } requestChars := countAIPointChars(input.MeteredText) - points := calculateAIPointCost(requestChars, status.BaseChars) - if input.FixedPoints > 0 { + points := 1 + if input.FixedPoints > points { points = input.FixedPoints } if status.Total <= 0 || status.Balance < points { @@ -160,6 +169,7 @@ func ReserveAIPoints( return &AIPointReservation{ ReservationID: reservationID, UsageLogID: usageID, + OperatorID: input.OperatorID, Points: points, RequestChars: requestChars, BaseChars: status.BaseChars, @@ -169,36 +179,79 @@ func ReserveAIPoints( func CompleteAIPoints( ctx context.Context, pool *pgxpool.Pool, + cache sharedcache.Cache, tenantID int64, reservation AIPointReservation, model string, -) error { + meteredOutputText string, +) (*AIPointReservation, error) { + completed := reservation if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 { - return nil + return &completed, nil } tx, err := pool.Begin(ctx) if err != nil { - return err + return nil, err } defer rollbackGenerationTx(tx) - pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID) - if err != nil || !pending { - return err + if err := lockTenantAIPoints(ctx, tx, tenantID); err != nil { + return nil, err } - if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil { - return err + pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID) + if err != nil { + return nil, err + } + if !pending { + return &completed, nil + } + + settlement := calculateAIPointOutputSettlement(reservation, meteredOutputText) + + completed.RequestChars = settlement.OutputChars + completed.BaseChars = settlement.BaseChars + completed.Points = settlement.FinalPoints + + quotaRepo := repository.NewQuotaRepository(tx) + adjusted, err := completeAIPointReservationWithFinalAmountIfPending(ctx, tx, tenantID, reservation.ReservationID, settlement.FinalPoints) + if err != nil { + return nil, err + } + if adjusted && settlement.LedgerDelta != 0 { + currentBalance, err := quotaRepo.GetCurrentBalance(ctx, tenantID, aiPointsQuotaType) + if err != nil { + return nil, err + } + reason := "ai_points_output_settle" + referenceType := "ai_point_usage" + usageID := reservation.UsageLogID + if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ + TenantID: tenantID, + OperatorID: reservation.OperatorID, + QuotaType: aiPointsQuotaType, + Delta: settlement.LedgerDelta, + BalanceAfter: currentBalance + settlement.LedgerDelta, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &usageID, + }); err != nil { + return nil, 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 + if err := repository.NewAIPointUsageRepository(tx).MarkCompleted(ctx, tenantID, reservation.UsageLogID, settlement.OutputChars, settlement.BaseChars, settlement.FinalPoints, modelPtr); err != nil { + return nil, err } - return tx.Commit(ctx) + if err := tx.Commit(ctx); err != nil { + return nil, err + } + deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(tenantID)) + return &completed, nil } func RefundAIPoints( @@ -267,17 +320,19 @@ func RefundAIPoints( func CompletePendingAIPointsByResource( ctx context.Context, pool *pgxpool.Pool, + cache sharedcache.Cache, tenantID int64, resourceType string, resourceID *int64, resourceUID *string, model string, -) error { + meteredOutputText string, +) (*AIPointReservation, error) { reservation, err := findPendingAIPointReservation(ctx, pool, tenantID, resourceType, resourceID, resourceUID) if err != nil || reservation == nil { - return err + return reservation, err } - return CompleteAIPoints(ctx, pool, tenantID, *reservation, model) + return CompleteAIPoints(ctx, pool, cache, tenantID, *reservation, model, meteredOutputText) } func RefundPendingAIPointsByResource( @@ -337,6 +392,7 @@ func findPendingAIPointReservationWithOperator( return &AIPointReservation{ ReservationID: *log.QuotaReservationID, UsageLogID: log.ID, + OperatorID: operatorID, Points: log.Points, RequestChars: log.RequestChars, BaseChars: log.BaseChars, @@ -354,7 +410,27 @@ func calculateAIPointCost(chars, baseChars int) int { if chars <= 0 { return 1 } - return chars/baseChars + 1 + return (chars + baseChars - 1) / baseChars +} + +func calculateAIPointOutputSettlement(reservation AIPointReservation, meteredOutputText string) aiPointOutputSettlement { + baseChars := reservation.BaseChars + if baseChars <= 0 { + baseChars = 1000 + } + outputChars := countAIPointChars(meteredOutputText) + finalPoints := calculateAIPointCost(outputChars, baseChars) + reservedPoints := reservation.Points + if reservedPoints <= 0 { + reservedPoints = 1 + } + return aiPointOutputSettlement{ + OutputChars: outputChars, + BaseChars: baseChars, + FinalPoints: finalPoints, + ReservedPoints: reservedPoints, + LedgerDelta: reservedPoints - finalPoints, + } } func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error { @@ -379,6 +455,30 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI return status == "pending", nil } +func completeAIPointReservationWithFinalAmountIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64, finalPoints int) (bool, error) { + if finalPoints <= 0 { + finalPoints = 1 + } + tag, err := tx.Exec(ctx, ` + UPDATE quota_reservations + SET status = 'confirmed', + reserved_amount = $4, + consumed_amount = $4, + updated_at = NOW() + WHERE id = $1 + AND tenant_id = $2 + AND quota_type = $3 + AND status = 'pending' + `, reservationID, tenantID, aiPointsQuotaType, finalPoints) + if err != nil { + return false, err + } + if tag.RowsAffected() > 0 { + return true, nil + } + return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID) +} + func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations diff --git a/server/internal/tenant/app/ai_point_usage_integration_test.go b/server/internal/tenant/app/ai_point_usage_integration_test.go new file mode 100644 index 0000000..d16ee73 --- /dev/null +++ b/server/internal/tenant/app/ai_point_usage_integration_test.go @@ -0,0 +1,200 @@ +//go:build integration + +package app + +import ( + "context" + "fmt" + "os" + "strings" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/geo-platform/tenant-api/internal/tenant/repository" +) + +func TestAIPointsOutputSettlementRealDatabase(t *testing.T) { + dsn := os.Getenv("TEST_DATABASE_URL") + if dsn == "" { + t.Skip("set TEST_DATABASE_URL to run the real database AI point settlement test") + } + + ctx := context.Background() + pool, err := pgxpool.New(ctx, dsn) + if err != nil { + t.Fatalf("connect database: %v", err) + } + t.Cleanup(pool.Close) + + tenantID, userID := seedAIPointSettlementTenant(t, ctx, pool, 10, 1000) + + reservation, err := ReserveAIPoints(ctx, pool, nil, AIPointReserveInput{ + TenantID: tenantID, + OperatorID: userID, + UsageType: "test_output_billing", + MeteredText: "cheap input", + }) + if err != nil { + t.Fatalf("reserve ai points: %v", err) + } + if reservation.Points != 1 { + t.Fatalf("reserved points = %d, want minimum 1", reservation.Points) + } + + output := strings.Repeat("贵", 2500) + completed, err := CompleteAIPoints(ctx, pool, nil, tenantID, *reservation, "test-model", output) + if err != nil { + t.Fatalf("complete ai points: %v", err) + } + if completed.Points != 3 { + t.Fatalf("completed points = %d, want 3", completed.Points) + } + if completed.RequestChars != 2500 { + t.Fatalf("completed request/output chars = %d, want 2500", completed.RequestChars) + } + + var usageStatus, model string + var usageChars, usagePoints int + if err := pool.QueryRow(ctx, ` +SELECT status, request_chars, points, model +FROM ai_point_usage_logs +WHERE id = $1 AND tenant_id = $2 +`, reservation.UsageLogID, tenantID).Scan(&usageStatus, &usageChars, &usagePoints, &model); err != nil { + t.Fatalf("query usage log: %v", err) + } + if usageStatus != "completed" || usageChars != 2500 || usagePoints != 3 || model != "test-model" { + t.Fatalf("usage log = status %q chars %d points %d model %q; want completed 2500 3 test-model", usageStatus, usageChars, usagePoints, model) + } + + var reservationStatus string + var reservedAmount, consumedAmount int + if err := pool.QueryRow(ctx, ` +SELECT status, reserved_amount, consumed_amount +FROM quota_reservations +WHERE id = $1 AND tenant_id = $2 +`, reservation.ReservationID, tenantID).Scan(&reservationStatus, &reservedAmount, &consumedAmount); err != nil { + t.Fatalf("query quota reservation: %v", err) + } + if reservationStatus != "confirmed" || reservedAmount != 3 || consumedAmount != 3 { + t.Fatalf("reservation = status %q reserved %d consumed %d; want confirmed 3 3", reservationStatus, reservedAmount, consumedAmount) + } + + rows, err := pool.Query(ctx, ` +SELECT delta, balance_after, reason +FROM tenant_quota_ledgers +WHERE tenant_id = $1 + AND quota_type = 'ai_points' + AND reference_id = $2 +ORDER BY id +`, tenantID, reservation.UsageLogID) + if err != nil { + t.Fatalf("query quota ledger: %v", err) + } + defer rows.Close() + + var deltas []int + var balances []int + var reasons []string + for rows.Next() { + var delta, balance int + var reason string + if err := rows.Scan(&delta, &balance, &reason); err != nil { + t.Fatalf("scan quota ledger: %v", err) + } + deltas = append(deltas, delta) + balances = append(balances, balance) + reasons = append(reasons, reason) + } + if err := rows.Err(); err != nil { + t.Fatalf("iterate quota ledger: %v", err) + } + if fmt.Sprint(deltas) != "[-1 -2]" || fmt.Sprint(balances) != "[9 7]" || fmt.Sprint(reasons) != "[ai_points_reserve ai_points_output_settle]" { + t.Fatalf("ledger deltas=%v balances=%v reasons=%v; want [-1 -2], [9 7], reserve+settle", deltas, balances, reasons) + } + + status, err := repository.NewQuotaRepository(pool).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC()) + if err != nil { + t.Fatalf("query quota status: %v", err) + } + if status.Used != 3 || status.Balance != 7 { + t.Fatalf("quota status used=%d balance=%d; want 3/7", status.Used, status.Balance) + } +} + +func seedAIPointSettlementTenant(t *testing.T, ctx context.Context, pool *pgxpool.Pool, aiPointsMonthly int, baseChars int) (int64, int64) { + t.Helper() + + suffix := fmt.Sprintf("%d", time.Now().UnixNano()) + var tenantID int64 + if err := pool.QueryRow(ctx, ` +INSERT INTO tenants (name, status) +VALUES ($1, 'active') +RETURNING id +`, "ai-point-output-test-"+suffix).Scan(&tenantID); err != nil { + t.Fatalf("insert tenant: %v", err) + } + + var userID int64 + if err := pool.QueryRow(ctx, ` +INSERT INTO users (phone, password_hash, name, status) +VALUES ($1, 'test', 'AI Point Output Test', 'active') +RETURNING id +`, "199"+suffix[len(suffix)-8:]).Scan(&userID); err != nil { + t.Fatalf("insert user: %v", err) + } + + if _, err := pool.Exec(ctx, ` +INSERT INTO tenant_memberships (tenant_id, user_id, tenant_role) +VALUES ($1, $2, 'owner') +`, tenantID, userID); err != nil { + t.Fatalf("insert tenant membership: %v", err) + } + + var planID int64 + if err := pool.QueryRow(ctx, ` +INSERT INTO plans (plan_code, name, quota_policy_json, status) +VALUES ($1, 'AI Point Output Plan', jsonb_build_object( + 'article_generation', 0, + 'article_quota_cycle', 'monthly', + 'ai_points_monthly', $2::int, + 'ai_point_base_chars', $3::int, + 'image_storage_bytes', 0, + 'brand_limit', 1 +), 'active') +RETURNING id +`, "ai-point-output-"+suffix, aiPointsMonthly, baseChars).Scan(&planID); err != nil { + t.Fatalf("insert plan: %v", err) + } + + if _, err := pool.Exec(ctx, ` +INSERT INTO tenant_plan_subscriptions (tenant_id, plan_id, start_at, end_at, status) +VALUES ($1, $2, NOW() - INTERVAL '1 hour', NOW() + INTERVAL '24 hours', 'active') +`, tenantID, planID); err != nil { + t.Fatalf("insert tenant plan subscription: %v", err) + } + + t.Cleanup(func() { + cleanupCtx := context.Background() + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM ai_point_usage_logs WHERE tenant_id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_quota_ledgers WHERE tenant_id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM quota_reservations WHERE tenant_id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_plan_subscriptions WHERE tenant_id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM workspace_memberships WHERE tenant_id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM workspaces WHERE tenant_id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_memberships WHERE tenant_id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM users WHERE id = $1`, userID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenants WHERE id = $1`, tenantID) + mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM plans WHERE id = $1`, planID) + }) + + return tenantID, userID +} + +func mustCleanupAIPointSettlementTestData(t *testing.T, ctx context.Context, pool *pgxpool.Pool, sql string, args ...any) { + t.Helper() + if _, err := pool.Exec(ctx, sql, args...); err != nil { + t.Errorf("cleanup AI point settlement test data: %v; sql=%s", err, sql) + } +} diff --git a/server/internal/tenant/app/ai_point_usage_test.go b/server/internal/tenant/app/ai_point_usage_test.go index c44221b..21402bb 100644 --- a/server/internal/tenant/app/ai_point_usage_test.go +++ b/server/internal/tenant/app/ai_point_usage_test.go @@ -1,6 +1,14 @@ package app -import "testing" +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) func TestCalculateAIPointCost(t *testing.T) { t.Parallel() @@ -13,12 +21,12 @@ func TestCalculateAIPointCost(t *testing.T) { }{ {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: "first base boundary", chars: 1000, baseChars: 1000, want: 1}, {name: "over first base", chars: 1001, baseChars: 1000, want: 2}, - {name: "second base boundary", chars: 2000, baseChars: 1000, want: 3}, + {name: "second base boundary", chars: 2000, baseChars: 1000, want: 2}, {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}, + {name: "default base chars", chars: 1000, baseChars: 0, want: 1}, + {name: "custom base chars", chars: 1500, baseChars: 500, want: 3}, } for _, tt := range tests { @@ -31,3 +39,150 @@ func TestCalculateAIPointCost(t *testing.T) { }) } } + +func TestCalculateAIPointOutputSettlementUsesOutputText(t *testing.T) { + t.Parallel() + + reservation := AIPointReservation{ + Points: 1, + RequestChars: 8, + BaseChars: 1000, + } + + got := calculateAIPointOutputSettlement(reservation, strings.Repeat("贵", 2500)) + + if got.OutputChars != 2500 { + t.Fatalf("OutputChars = %d, want 2500", got.OutputChars) + } + if got.FinalPoints != 3 { + t.Fatalf("FinalPoints = %d, want 3", got.FinalPoints) + } + if got.LedgerDelta != -2 { + t.Fatalf("LedgerDelta = %d, want -2 for output-side top-up", got.LedgerDelta) + } +} + +func TestCalculateAIPointOutputSettlementBoundaries(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + outputChars int + wantPoints int + }{ + {name: "minimum empty output still one point", outputChars: 0, wantPoints: 1}, + {name: "below base", outputChars: 999, wantPoints: 1}, + {name: "exactly base", outputChars: 1000, wantPoints: 1}, + {name: "one over base", outputChars: 1001, wantPoints: 2}, + {name: "exactly second base", outputChars: 2000, wantPoints: 2}, + {name: "one over second base", outputChars: 2001, wantPoints: 3}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := calculateAIPointOutputSettlement( + AIPointReservation{Points: 1, RequestChars: 50000, BaseChars: 1000}, + strings.Repeat("字", tt.outputChars), + ) + if got.FinalPoints != tt.wantPoints { + t.Fatalf("FinalPoints = %d, want %d", got.FinalPoints, tt.wantPoints) + } + }) + } +} + +func TestCalculateAIPointOutputSettlementRefundsOverReservation(t *testing.T) { + t.Parallel() + + got := calculateAIPointOutputSettlement( + AIPointReservation{Points: 5, RequestChars: 5000, BaseChars: 1000}, + strings.Repeat("字", 700), + ) + + if got.FinalPoints != 1 { + t.Fatalf("FinalPoints = %d, want 1", got.FinalPoints) + } + if got.LedgerDelta != 4 { + t.Fatalf("LedgerDelta = %d, want +4 for output-side refund", got.LedgerDelta) + } +} + +func TestCalculateAIPointOutputSettlementUsesDefaultBaseChars(t *testing.T) { + t.Parallel() + + got := calculateAIPointOutputSettlement( + AIPointReservation{Points: 1, BaseChars: 0}, + strings.Repeat("字", 1001), + ) + + if got.BaseChars != 1000 { + t.Fatalf("BaseChars = %d, want default 1000", got.BaseChars) + } + if got.FinalPoints != 2 { + t.Fatalf("FinalPoints = %d, want 2", got.FinalPoints) + } + if got.LedgerDelta != -1 { + t.Fatalf("LedgerDelta = %d, want -1", got.LedgerDelta) + } +} + +func TestCompleteAIPointReservationWithFinalAmountUpdatesReservationAmounts(t *testing.T) { + tx := &captureAIPointTx{} + + adjusted, err := completeAIPointReservationWithFinalAmountIfPending(context.Background(), tx, 10, 20, 3) + if err != nil { + t.Fatalf("complete reservation: %v", err) + } + if !adjusted { + t.Fatalf("adjusted = false, want true") + } + if !strings.Contains(tx.sql, "reserved_amount = $4") { + t.Fatalf("reservation SQL must set final reserved_amount; query:\n%s", tx.sql) + } + if !strings.Contains(tx.sql, "consumed_amount = $4") { + t.Fatalf("reservation SQL must set final consumed_amount; query:\n%s", tx.sql) + } + if len(tx.args) != 4 { + t.Fatalf("args length = %d, want 4", len(tx.args)) + } + if tx.args[0] != int64(20) || tx.args[1] != int64(10) || tx.args[2] != aiPointsQuotaType || tx.args[3] != 3 { + t.Fatalf("args = %#v, want reservation, tenant, quota type, final points", tx.args) + } +} + +type captureAIPointTx struct { + sql string + args []any +} + +func (tx *captureAIPointTx) Begin(context.Context) (pgx.Tx, error) { return nil, errors.New("unused") } +func (tx *captureAIPointTx) Commit(context.Context) error { return nil } +func (tx *captureAIPointTx) Rollback(context.Context) error { return nil } +func (tx *captureAIPointTx) CopyFrom(context.Context, pgx.Identifier, []string, pgx.CopyFromSource) (int64, error) { + return 0, errors.New("unused") +} +func (tx *captureAIPointTx) SendBatch(context.Context, *pgx.Batch) pgx.BatchResults { + return nil +} +func (tx *captureAIPointTx) LargeObjects() pgx.LargeObjects { return pgx.LargeObjects{} } +func (tx *captureAIPointTx) Prepare(context.Context, string, string) (*pgconn.StatementDescription, error) { + return nil, errors.New("unused") +} +func (tx *captureAIPointTx) Exec(_ context.Context, sql string, args ...any) (pgconn.CommandTag, error) { + tx.sql = sql + tx.args = args + return pgconn.NewCommandTag("UPDATE 1"), nil +} +func (tx *captureAIPointTx) Query(context.Context, string, ...any) (pgx.Rows, error) { + return nil, errors.New("unused") +} +func (tx *captureAIPointTx) QueryRow(context.Context, string, ...any) pgx.Row { + return failingAIPointRow{} +} +func (tx *captureAIPointTx) Conn() *pgx.Conn { return nil } + +type failingAIPointRow struct{} + +func (failingAIPointRow) Scan(...any) error { return errors.New("unused") } diff --git a/server/internal/tenant/app/article_selection_optimize_service.go b/server/internal/tenant/app/article_selection_optimize_service.go index 72f47cb..6127acc 100644 --- a/server/internal/tenant/app/article_selection_optimize_service.go +++ b/server/internal/tenant/app/article_selection_optimize_service.go @@ -135,7 +135,7 @@ func (s *ArticleSelectionOptimizeService) OptimizeSelection( } completeCtx, cancel := newGenerationCleanupContext() - if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, result.Model); err != nil { + if _, err := CompleteAIPoints(completeCtx, s.pool, s.cache, actor.TenantID, *reservation, result.Model, content); err != nil { cancel() return nil, response.ErrInternal(50019, "article_selection_optimize_failed", "failed to confirm ai points usage") } diff --git a/server/internal/tenant/app/kol_assist_runtime.go b/server/internal/tenant/app/kol_assist_runtime.go index 6fc7710..9021f5a 100644 --- a/server/internal/tenant/app/kol_assist_runtime.go +++ b/server/internal/tenant/app/kol_assist_runtime.go @@ -77,7 +77,7 @@ func (s *KolAssistService) Run( } completeCtx, cancel := newGenerationCleanupContext() - if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, generated.Model); err != nil { + if _, err := CompleteAIPoints(completeCtx, s.pool, s.cache, actor.TenantID, *reservation, generated.Model, generated.Content); err != nil { cancel() return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", "failed to confirm ai points usage") } diff --git a/server/internal/tenant/app/question_expansion_service.go b/server/internal/tenant/app/question_expansion_service.go index 7ff5e48..08e04c6 100644 --- a/server/internal/tenant/app/question_expansion_service.go +++ b/server/internal/tenant/app/question_expansion_service.go @@ -317,7 +317,8 @@ func (s *QuestionExpansionService) GenerateByAIDistill(ctx context.Context, bran candidates = append(candidates, buildCandidate(text, QuestionSourceAIDistill, false, existing, brandCtx)) } - if err := CompleteAIPoints(context.Background(), s.pool, actor.TenantID, *reservation, result.Model); err != nil { + completed, err := CompleteAIPoints(context.Background(), s.pool, s.cache, actor.TenantID, *reservation, result.Model, result.Content) + if err != nil { return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to confirm ai points") } @@ -329,7 +330,7 @@ func (s *QuestionExpansionService) GenerateByAIDistill(ctx context.Context, bran return &QuestionCandidateResult{ Candidates: candidates, - AIPointsCharged: reservation.Points, + AIPointsCharged: completed.Points, }, nil } @@ -413,10 +414,11 @@ func (s *QuestionExpansionService) FillQuestionCombination(ctx context.Context, } filled := sanitizeQuestionCombinationFillResult(payload, fallback) - if err := CompleteAIPoints(context.Background(), s.pool, actor.TenantID, *reservation, result.Model); err != nil { + completed, err := CompleteAIPoints(context.Background(), s.pool, s.cache, actor.TenantID, *reservation, result.Model, result.Content) + if err != nil { return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to confirm ai points") } - filled.AIPointsCharged = reservation.Points + filled.AIPointsCharged = completed.Points if s.cache != nil { cached := filled cached.AIPointsCharged = 0 diff --git a/server/internal/tenant/app/template_assist.go b/server/internal/tenant/app/template_assist.go index 37ed275..0a0838a 100644 --- a/server/internal/tenant/app/template_assist.go +++ b/server/internal/tenant/app/template_assist.go @@ -430,7 +430,7 @@ func (s *TemplateService) IsAssistTaskTerminal(ctx context.Context, job AssistJo } switch record.Status { case "completed": - s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, "") + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, "", string(record.ResultJSON)) return true, nil case "failed": if isTemplateAssistAIPointTask(job.TaskType) { @@ -479,7 +479,7 @@ func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repositor s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err) return } - s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model) + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model, result.Content) } func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) { @@ -518,7 +518,7 @@ func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository. s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err) return } - s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model) + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model, result.Content) } func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) { @@ -562,7 +562,7 @@ func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repositor s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err) return } - s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model) + s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model, result.Content) } func (s *TemplateService) failAssist(ctx context.Context, taskID, taskType string, tenantID int64, err error) { @@ -619,10 +619,10 @@ func (s *TemplateService) refundReservedAIPoints(ctx context.Context, tenantID, _ = RefundAIPoints(ctx, s.pool, s.cache, tenantID, operatorID, *reservation, errorMessage) } -func (s *TemplateService) completeTemplateAssistAIPoints(ctx context.Context, tenantID int64, taskID string, model string) { +func (s *TemplateService) completeTemplateAssistAIPoints(ctx context.Context, tenantID int64, taskID string, model string, meteredOutputText string) { resourceType := "template_assist_task" resourceUID := taskID - _ = CompletePendingAIPointsByResource(ctx, s.pool, tenantID, resourceType, nil, &resourceUID, model) + _, _ = CompletePendingAIPointsByResource(ctx, s.pool, s.cache, tenantID, resourceType, nil, &resourceUID, model, meteredOutputText) } func isTemplateAssistAIPointTask(taskType string) bool { diff --git a/server/internal/tenant/compliance/service.go b/server/internal/tenant/compliance/service.go index 10a80b8..5d366a3 100644 --- a/server/internal/tenant/compliance/service.go +++ b/server/internal/tenant/compliance/service.go @@ -1139,7 +1139,7 @@ func (s *Service) ProcessReviewJob(ctx context.Context, msg ReviewJobMessage) er _ = s.refundComplianceLLMJudgePoints(context.Background(), job, reservation, err) return s.failReviewJob(ctx, job.ID, cfg.ReviewMaxAttempts, err) } - if err := s.completeComplianceLLMJudgePoints(ctx, job, reservation, result.Model); err != nil && s.logger != nil { + if err := s.completeComplianceLLMJudgePoints(ctx, job, reservation, result.Model, result.Content); err != nil && s.logger != nil { s.logger.Warn("compliance llm billing completion failed", zap.Error(err), zap.Int64("tenant_id", job.TenantID), @@ -1348,7 +1348,7 @@ func (s *Service) reserveComplianceLLMJudgePoints(ctx context.Context, job revie return nil, err } requestChars := countAIPointChars(meteredText) - points := calculateAIPointCost(requestChars, status.BaseChars) + points := 1 if status.Total <= 0 || status.Balance < points { return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are insufficient") } @@ -1416,7 +1416,7 @@ func (s *Service) reserveComplianceLLMJudgePoints(ctx context.Context, job revie }, nil } -func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job reviewJob, reservation *aiPointReservation, model string) error { +func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job reviewJob, reservation *aiPointReservation, model string, meteredOutputText string) error { if s == nil || s.pool == nil || reservation == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 { return nil } @@ -1425,19 +1425,59 @@ func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job revi return err } defer rollbackTx(tx) + if err := lockTenantAIPoints(ctx, tx, job.TenantID); err != nil { + return err + } pending, err := lockPendingAIPointUsage(ctx, tx, job.TenantID, reservation.UsageLogID) if err != nil || !pending { return err } - if err := confirmAIPointReservationIfPending(ctx, tx, job.TenantID, reservation.ReservationID); err != nil { + + baseChars := reservation.BaseChars + if baseChars <= 0 { + baseChars = 1000 + } + outputChars := countAIPointChars(meteredOutputText) + finalPoints := calculateAIPointCost(outputChars, baseChars) + reservedPoints := reservation.Points + if reservedPoints <= 0 { + reservedPoints = 1 + } + + quotaRepo := tenantrepo.NewQuotaRepository(tx) + adjusted, err := completeAIPointReservationWithFinalAmountIfPending(ctx, tx, job.TenantID, reservation.ReservationID, finalPoints) + if err != nil { return err } + pointDelta := finalPoints - reservedPoints + if adjusted && pointDelta != 0 { + currentBalance, err := quotaRepo.GetCurrentBalance(ctx, job.TenantID, aiPointsQuotaType) + if err != nil { + return err + } + ledgerDelta := -pointDelta + reason := "ai_points_output_settle" + referenceType := "ai_point_usage" + usageID := reservation.UsageLogID + if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{ + TenantID: job.TenantID, + OperatorID: job.CheckedBy, + QuotaType: aiPointsQuotaType, + Delta: ledgerDelta, + BalanceAfter: currentBalance + ledgerDelta, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &usageID, + }); err != nil { + return err + } + } var modelPtr *string if strings.TrimSpace(model) != "" { normalized := strings.TrimSpace(model) modelPtr = &normalized } - if err := tenantrepo.NewAIPointUsageRepository(tx).MarkCompleted(ctx, job.TenantID, reservation.UsageLogID, modelPtr); err != nil { + if err := tenantrepo.NewAIPointUsageRepository(tx).MarkCompleted(ctx, job.TenantID, reservation.UsageLogID, outputChars, baseChars, finalPoints, modelPtr); err != nil { return err } return tx.Commit(ctx) @@ -1502,7 +1542,7 @@ func calculateAIPointCost(chars, baseChars int) int { if chars <= 0 { return 1 } - return chars/baseChars + 1 + return (chars + baseChars - 1) / baseChars } func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error { @@ -1527,6 +1567,30 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI return status == "pending", nil } +func completeAIPointReservationWithFinalAmountIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64, finalPoints int) (bool, error) { + if finalPoints <= 0 { + finalPoints = 1 + } + tag, err := tx.Exec(ctx, ` + UPDATE quota_reservations + SET status = 'confirmed', + reserved_amount = $4, + consumed_amount = $4, + updated_at = NOW() + WHERE id = $1 + AND tenant_id = $2 + AND quota_type = $3 + AND status = 'pending' + `, reservationID, tenantID, aiPointsQuotaType, finalPoints) + if err != nil { + return false, err + } + if tag.RowsAffected() > 0 { + return true, nil + } + return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID) +} + func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations diff --git a/server/internal/tenant/compliance/service_test.go b/server/internal/tenant/compliance/service_test.go index 6f8e940..1481ed7 100644 --- a/server/internal/tenant/compliance/service_test.go +++ b/server/internal/tenant/compliance/service_test.go @@ -1,6 +1,7 @@ package compliance import ( + "strings" "testing" sharedcompliance "github.com/geo-platform/tenant-api/internal/shared/compliance" @@ -34,3 +35,39 @@ func TestDecidePassesWithoutHits(t *testing.T) { t.Fatalf("decide without hits = %s, want %s", got, sharedcompliance.GateDecisionPass) } } + +func TestComplianceAIPointCostUsesOutputBoundary(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + chars int + baseChars int + want int + }{ + {name: "empty output minimum", chars: 0, baseChars: 1000, want: 1}, + {name: "single char", chars: 1, baseChars: 1000, want: 1}, + {name: "exactly one base", chars: 1000, baseChars: 1000, want: 1}, + {name: "one over base", chars: 1001, baseChars: 1000, want: 2}, + {name: "default base", chars: 2000, baseChars: 0, want: 2}, + } + + 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) + } + }) + } +} + +func TestComplianceAIPointCharsTrimWhitespace(t *testing.T) { + t.Parallel() + + got := countAIPointChars(" " + strings.Repeat("字", 1001) + "\n") + if got != 1001 { + t.Fatalf("countAIPointChars = %d, want 1001", got) + } +} diff --git a/server/internal/tenant/repository/ai_point_usage_repo.go b/server/internal/tenant/repository/ai_point_usage_repo.go index a3e28b7..ab1fed6 100644 --- a/server/internal/tenant/repository/ai_point_usage_repo.go +++ b/server/internal/tenant/repository/ai_point_usage_repo.go @@ -52,7 +52,7 @@ type AIPointUsageListParams struct { type AIPointUsageRepository interface { Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error) - MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error + MarkCompleted(ctx context.Context, tenantID, id int64, requestChars, baseChars, points int, 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) @@ -97,18 +97,21 @@ RETURNING id return id, err } -func (r *aiPointUsageRepository) MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error { +func (r *aiPointUsageRepository) MarkCompleted(ctx context.Context, tenantID, id int64, requestChars, baseChars, points int, model *string) error { _, err := r.db.Exec(ctx, ` UPDATE ai_point_usage_logs SET status = 'completed', model = $1, + request_chars = $2, + base_chars = $3, + points = $4, error_message = NULL, completed_at = COALESCE(completed_at, NOW()), updated_at = NOW() -WHERE id = $2 - AND tenant_id = $3 +WHERE id = $5 + AND tenant_id = $6 AND status = 'pending' -`, model, id, tenantID) +`, model, requestChars, baseChars, points, id, tenantID) return err } diff --git a/server/internal/worker/assist/kol_assist_worker.go b/server/internal/worker/assist/kol_assist_worker.go index 0ecd5a6..a03fd8c 100644 --- a/server/internal/worker/assist/kol_assist_worker.go +++ b/server/internal/worker/assist/kol_assist_worker.go @@ -93,7 +93,7 @@ func (w *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJo } switch task.Status { case "completed": - w.completeAIPoints(ctx, job, "") + w.completeAIPoints(ctx, job, "", string(task.ResultJSON)) return nil case "failed": errorMessage := "kol assist task failed" @@ -262,14 +262,14 @@ 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) + w.completeAIPoints(ctx, job, result.Model, result.Content) return nil } -func (w *KolAssistWorker) completeAIPoints(ctx context.Context, job tenantapp.KolAssistJob, model string) { +func (w *KolAssistWorker) completeAIPoints(ctx context.Context, job tenantapp.KolAssistJob, model string, meteredOutputText 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 { + if _, err := tenantapp.CompletePendingAIPointsByResource(ctx, w.pool, w.cache, job.TenantID, resourceType, nil, &resourceUID, model, meteredOutputText); err != nil && w.logger != nil { w.logger.Warn("kol assist ai points confirmation failed", zap.Error(err), zap.String("task_id", job.TaskID), diff --git a/user-manual/11-AI点数与套餐.md b/user-manual/11-AI点数与套餐.md index 153d980..7561dea 100644 --- a/user-manual/11-AI点数与套餐.md +++ b/user-manual/11-AI点数与套餐.md @@ -26,7 +26,7 @@ ### 计费基础 -- 每个 AI 调用按"请求字数 / 基准字数"换算成点数 +- 每个 AI 调用按"输出字数 / 基准字数"换算成点数 - 套餐越高，**基准字数** 越大（同样请求消耗的点数越少） - 流水里能看到「{字数} 字 / 基准 {基准}」的换算说明 @@ -80,7 +80,7 @@ AI 任务在提交时会**先预占**对应点数，任务完成后才正式确 - **使用场景**：哪类操作触发的扣点 - **点数**：本次消耗（负数表示扣除，正数表示退回） -- **计费字数**：实际请求字数 +- **计费字数**：实际输出字数 - **创建时间**、**状态** 如果有"莫名其妙的扣点"：