feat(ai): centralize durable usage
Add a content-free usage ledger with legacy backfill. Reserve Workspace and durable Strategy/CV work before execution so deleted history or duplicate admission cannot reset limits.
This commit is contained in:
cesnimda
@@ -0,0 +1,155 @@
|
||||
using JobTrackerApi.Data;
|
||||
using JobTrackerApi.Models;
|
||||
using Microsoft.EntityFrameworkCore;
|
||||
|
||||
namespace JobTrackerApi.Services;
|
||||
|
||||
public sealed record AiUsageTotals(int Calls, long InputCharacters, long OutputCharacters, long EstimatedTokens);
|
||||
public sealed record AiUsageReservation(AiUsageRecord Record, bool Created);
|
||||
|
||||
public sealed class AiUsageLimitException(string code, string message) : Exception(message)
|
||||
{
|
||||
public string Code { get; } = code;
|
||||
}
|
||||
|
||||
public sealed class AiUsageMeter(JobTrackerContext db, TimeProvider timeProvider)
|
||||
{
|
||||
private static readonly SemaphoreSlim Gate = new(1, 1);
|
||||
|
||||
public static (int InputCharacters, int EstimatedTokens) ReservationFor(string taskType) => taskType switch
|
||||
{
|
||||
StrategySnapshotService.TaskType => (48_000, 12_000),
|
||||
CvProcessingQueue.TaskType => (64_000, 16_000),
|
||||
_ => (16_000, 4_000),
|
||||
};
|
||||
|
||||
public async Task<AiUsageTotals> CurrentMonthAsync(string ownerUserId, CancellationToken cancellationToken)
|
||||
=> await SinceAsync(ownerUserId, MonthStart(timeProvider.GetUtcNow()), cancellationToken);
|
||||
|
||||
public async Task<AiUsageTotals> AllTimeAsync(string ownerUserId, CancellationToken cancellationToken)
|
||||
=> await SinceAsync(ownerUserId, null, cancellationToken);
|
||||
|
||||
public async Task EnsureCanReserveAsync(
|
||||
string ownerUserId,
|
||||
AccountEntitlements entitlements,
|
||||
int calls,
|
||||
int estimatedTokens,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
var used = await CurrentMonthAsync(ownerUserId, cancellationToken);
|
||||
EnsureWithinLimit(used, entitlements, calls, estimatedTokens);
|
||||
}
|
||||
|
||||
public async Task<AiUsageReservation> ReserveAsync(
|
||||
string ownerUserId,
|
||||
AccountEntitlements entitlements,
|
||||
string sourceType,
|
||||
string sourceId,
|
||||
string taskType,
|
||||
int inputCharacters,
|
||||
int estimatedTokens,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
Validate(sourceType, sourceId, taskType, inputCharacters, estimatedTokens);
|
||||
await Gate.WaitAsync(cancellationToken);
|
||||
try
|
||||
{
|
||||
var existing = await db.AiUsageRecords.FirstOrDefaultAsync(
|
||||
item => item.SourceType == sourceType && item.SourceId == sourceId,
|
||||
cancellationToken);
|
||||
if (existing is not null) return new AiUsageReservation(existing, false);
|
||||
|
||||
await EnsureCanReserveAsync(ownerUserId, entitlements, 1, estimatedTokens, cancellationToken);
|
||||
var record = NewRecord(ownerUserId, sourceType, sourceId, taskType, inputCharacters, estimatedTokens, timeProvider.GetUtcNow());
|
||||
db.AiUsageRecords.Add(record);
|
||||
await db.SaveChangesAsync(cancellationToken);
|
||||
return new AiUsageReservation(record, true);
|
||||
}
|
||||
finally
|
||||
{
|
||||
Gate.Release();
|
||||
}
|
||||
}
|
||||
|
||||
public async Task FinalizeAsync(long id, int inputCharacters, int outputCharacters, CancellationToken cancellationToken)
|
||||
{
|
||||
if (inputCharacters < 0 || outputCharacters < 0) throw new ArgumentOutOfRangeException();
|
||||
var estimatedTokens = (inputCharacters + outputCharacters + 3) / 4;
|
||||
await db.AiUsageRecords.Where(item => item.Id == id).ExecuteUpdateAsync(setters => setters
|
||||
.SetProperty(item => item.InputCharacterCount, inputCharacters)
|
||||
.SetProperty(item => item.OutputCharacterCount, outputCharacters)
|
||||
.SetProperty(item => item.EstimatedTokenCount, estimatedTokens), cancellationToken);
|
||||
}
|
||||
|
||||
public async Task ReleaseAsync(long id, CancellationToken cancellationToken)
|
||||
=> await db.AiUsageRecords.Where(item => item.Id == id).ExecuteDeleteAsync(cancellationToken);
|
||||
|
||||
public static AiUsageRecord NewOperationRecord(
|
||||
string ownerUserId,
|
||||
Guid operationId,
|
||||
string taskType,
|
||||
int inputCharacters,
|
||||
int estimatedTokens,
|
||||
DateTimeOffset createdAtUtc)
|
||||
=> NewRecord(ownerUserId, "operation", operationId.ToString("D"), taskType, inputCharacters, estimatedTokens, createdAtUtc);
|
||||
|
||||
private async Task<AiUsageTotals> SinceAsync(string ownerUserId, DateTimeOffset? since, CancellationToken cancellationToken)
|
||||
{
|
||||
var query = db.AiUsageRecords.Where(item => item.OwnerUserId == ownerUserId);
|
||||
if (db.Database.IsSqlite())
|
||||
{
|
||||
var rows = await query.AsNoTracking().ToListAsync(cancellationToken);
|
||||
if (since is not null) rows = rows.Where(item => item.CreatedAtUtc >= since.Value).ToList();
|
||||
return Sum(rows);
|
||||
}
|
||||
|
||||
if (since is not null) query = query.Where(item => item.CreatedAtUtc >= since.Value);
|
||||
var totals = await query.GroupBy(_ => 1).Select(group => new AiUsageTotals(
|
||||
group.Sum(item => item.CallCount),
|
||||
group.Sum(item => (long)item.InputCharacterCount),
|
||||
group.Sum(item => (long)item.OutputCharacterCount),
|
||||
group.Sum(item => (long)item.EstimatedTokenCount))).FirstOrDefaultAsync(cancellationToken);
|
||||
return totals ?? new AiUsageTotals(0, 0, 0, 0);
|
||||
}
|
||||
|
||||
private static AiUsageTotals Sum(IEnumerable<AiUsageRecord> records) => new(
|
||||
records.Sum(item => item.CallCount),
|
||||
records.Sum(item => (long)item.InputCharacterCount),
|
||||
records.Sum(item => (long)item.OutputCharacterCount),
|
||||
records.Sum(item => (long)item.EstimatedTokenCount));
|
||||
|
||||
private static void EnsureWithinLimit(AiUsageTotals used, AccountEntitlements entitlements, int calls, int tokens)
|
||||
{
|
||||
if (used.Calls + calls > entitlements.MonthlyAiCalls)
|
||||
throw new AiUsageLimitException("monthly_ai_calls_exhausted", $"Monthly AI limit reached ({entitlements.MonthlyAiCalls} generations). Try again next month.");
|
||||
if (used.EstimatedTokens + tokens > entitlements.MonthlyAiTokens)
|
||||
throw new AiUsageLimitException("monthly_ai_tokens_exhausted", $"Monthly AI cost limit reached ({entitlements.MonthlyAiTokens:N0} estimated tokens). Try again next month.");
|
||||
}
|
||||
|
||||
private static AiUsageRecord NewRecord(string ownerUserId, string sourceType, string sourceId, string taskType, int inputCharacters, int estimatedTokens, DateTimeOffset createdAtUtc)
|
||||
{
|
||||
Validate(sourceType, sourceId, taskType, inputCharacters, estimatedTokens);
|
||||
return new AiUsageRecord
|
||||
{
|
||||
OwnerUserId = ownerUserId,
|
||||
SourceType = sourceType,
|
||||
SourceId = sourceId,
|
||||
TaskType = taskType,
|
||||
InputCharacterCount = inputCharacters,
|
||||
EstimatedTokenCount = estimatedTokens,
|
||||
CreatedAtUtc = createdAtUtc,
|
||||
};
|
||||
}
|
||||
|
||||
private static void Validate(string sourceType, string sourceId, string taskType, int inputCharacters, int estimatedTokens)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(sourceType) || sourceType.Length > 32) throw new ArgumentOutOfRangeException(nameof(sourceType));
|
||||
if (string.IsNullOrWhiteSpace(sourceId) || sourceId.Length > 64) throw new ArgumentOutOfRangeException(nameof(sourceId));
|
||||
if (string.IsNullOrWhiteSpace(taskType) || taskType.Length > 64) throw new ArgumentOutOfRangeException(nameof(taskType));
|
||||
if (inputCharacters < 0) throw new ArgumentOutOfRangeException(nameof(inputCharacters));
|
||||
if (estimatedTokens < 0) throw new ArgumentOutOfRangeException(nameof(estimatedTokens));
|
||||
}
|
||||
|
||||
private static DateTimeOffset MonthStart(DateTimeOffset value)
|
||||
=> new(value.Year, value.Month, 1, 0, 0, 0, TimeSpan.Zero);
|
||||
}
|
||||
Reference in New Issue
Block a user