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) <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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") }
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user