249 lines
7.6 KiB
Go
249 lines
7.6 KiB
Go
|
package aiaudit
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/google/uuid"
|
||
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Cost accounting. Kept apart from the prompt capture above on purpose:
|
||
|
|
// ai_call_logs is a 7-day debugging buffer, while spend has to answer "what did
|
||
|
|
// this tenant cost me last quarter" for years. Usage is therefore rolled up per
|
||
|
|
// day / company / user / model / role as each call completes.
|
||
|
|
//
|
||
|
|
// Money is micro-USD integers end to end. Summing millions of fractional-cent
|
||
|
|
// calls in float drifts; integers add up exactly.
|
||
|
|
|
||
|
|
// UsageRetentionDays bounds the rollup. Three years plus a leap day.
|
||
|
|
const UsageRetentionDays = 3*365 + 1
|
||
|
|
|
||
|
|
// MicrosPerUSD converts the stored integer unit to dollars.
|
||
|
|
const MicrosPerUSD = 1_000_000
|
||
|
|
|
||
|
|
// ModelPrice is micro-USD per 1,000,000 tokens.
|
||
|
|
type ModelPrice struct {
|
||
|
|
InputPerMTok int64
|
||
|
|
CachedInputPerMTok int64
|
||
|
|
OutputPerMTok int64
|
||
|
|
}
|
||
|
|
|
||
|
|
// defaultModelPrices is the fallback when ai_model_prices has no row for a model
|
||
|
|
// (a provider the admin pointed at without adding a price). Keeping the shipped
|
||
|
|
// GPT-5.6 rates here means a fresh database still costs correctly before the
|
||
|
|
// migration seed is edited. Rates: OpenAI list price, 2026-07-30.
|
||
|
|
var defaultModelPrices = map[string]ModelPrice{
|
||
|
|
"gpt-5.6-luna": {InputPerMTok: 200_000, CachedInputPerMTok: 20_000, OutputPerMTok: 1_200_000},
|
||
|
|
"gpt-5.6-terra": {InputPerMTok: 2_000_000, CachedInputPerMTok: 200_000, OutputPerMTok: 12_000_000},
|
||
|
|
"gpt-5.6-sol": {InputPerMTok: 5_000_000, CachedInputPerMTok: 500_000, OutputPerMTok: 30_000_000},
|
||
|
|
}
|
||
|
|
|
||
|
|
// DefaultModelPrice returns the shipped price for a model, ok=false when unknown.
|
||
|
|
func DefaultModelPrice(model string) (ModelPrice, bool) {
|
||
|
|
p, ok := defaultModelPrices[NormalizeModel(model)]
|
||
|
|
return p, ok
|
||
|
|
}
|
||
|
|
|
||
|
|
// NormalizeModel folds provider prefixes and dated suffixes onto the price key, so
|
||
|
|
// "openai/gpt-5.6-luna" and "gpt-5.6-luna-2026-07-30" both price as gpt-5.6-luna.
|
||
|
|
func NormalizeModel(model string) string {
|
||
|
|
m := strings.ToLower(strings.TrimSpace(model))
|
||
|
|
if i := strings.LastIndex(m, "/"); i >= 0 {
|
||
|
|
m = m[i+1:]
|
||
|
|
}
|
||
|
|
// Trim a trailing -YYYY-MM-DD snapshot tag.
|
||
|
|
if len(m) > 11 {
|
||
|
|
tail := m[len(m)-11:]
|
||
|
|
if tail[0] == '-' && isDigits(tail[1:5]) && tail[5] == '-' && isDigits(tail[6:8]) && tail[8] == '-' && isDigits(tail[9:11]) {
|
||
|
|
m = m[:len(m)-11]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return m
|
||
|
|
}
|
||
|
|
|
||
|
|
func isDigits(s string) bool {
|
||
|
|
for _, r := range s {
|
||
|
|
if r < '0' || r > '9' {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return len(s) > 0
|
||
|
|
}
|
||
|
|
|
||
|
|
// CostMicros prices one call. Cached prompt tokens are billed at the cached rate
|
||
|
|
// and subtracted from fresh input — they are a subset of prompt tokens, not extra.
|
||
|
|
func CostMicros(p ModelPrice, promptTokens, cachedPromptTokens, completionTokens int) int64 {
|
||
|
|
if cachedPromptTokens < 0 {
|
||
|
|
cachedPromptTokens = 0
|
||
|
|
}
|
||
|
|
if cachedPromptTokens > promptTokens {
|
||
|
|
cachedPromptTokens = promptTokens
|
||
|
|
}
|
||
|
|
fresh := int64(promptTokens - cachedPromptTokens)
|
||
|
|
cached := int64(cachedPromptTokens)
|
||
|
|
out := int64(completionTokens)
|
||
|
|
if out < 0 {
|
||
|
|
out = 0
|
||
|
|
}
|
||
|
|
if fresh < 0 {
|
||
|
|
fresh = 0
|
||
|
|
}
|
||
|
|
// Rounded division keeps small calls from always costing zero.
|
||
|
|
return divRound(fresh*p.InputPerMTok, 1_000_000) +
|
||
|
|
divRound(cached*p.CachedInputPerMTok, 1_000_000) +
|
||
|
|
divRound(out*p.OutputPerMTok, 1_000_000)
|
||
|
|
}
|
||
|
|
|
||
|
|
func divRound(n, d int64) int64 {
|
||
|
|
if d == 0 {
|
||
|
|
return 0
|
||
|
|
}
|
||
|
|
return (n + d/2) / d
|
||
|
|
}
|
||
|
|
|
||
|
|
// priceCache avoids a DB read per call; prices change at most a few times a year.
|
||
|
|
type priceCache struct {
|
||
|
|
mu sync.RWMutex
|
||
|
|
prices map[string]ModelPrice
|
||
|
|
loadedAt time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
const priceCacheTTL = 5 * time.Minute
|
||
|
|
|
||
|
|
func (r *Recorder) priceFor(ctx context.Context, model string) ModelPrice {
|
||
|
|
key := NormalizeModel(model)
|
||
|
|
if r != nil {
|
||
|
|
r.prices.mu.RLock()
|
||
|
|
fresh := time.Since(r.prices.loadedAt) < priceCacheTTL
|
||
|
|
p, ok := r.prices.prices[key]
|
||
|
|
r.prices.mu.RUnlock()
|
||
|
|
if fresh && ok {
|
||
|
|
return p
|
||
|
|
}
|
||
|
|
if !fresh {
|
||
|
|
r.reloadPrices(ctx)
|
||
|
|
r.prices.mu.RLock()
|
||
|
|
p, ok = r.prices.prices[key]
|
||
|
|
r.prices.mu.RUnlock()
|
||
|
|
if ok {
|
||
|
|
return p
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if p, ok := DefaultModelPrice(key); ok {
|
||
|
|
return p
|
||
|
|
}
|
||
|
|
// Unknown model: record usage with zero cost rather than guessing a rate.
|
||
|
|
return ModelPrice{}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *Recorder) reloadPrices(ctx context.Context) {
|
||
|
|
if r == nil || r.pool == nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
rows, err := r.pool.Query(ctx, `
|
||
|
|
SELECT model, input_micros_per_mtok, cached_input_micros_per_mtok, output_micros_per_mtok
|
||
|
|
FROM ai_model_prices`)
|
||
|
|
if err != nil {
|
||
|
|
// Keep whatever is cached; defaults still apply for misses.
|
||
|
|
r.prices.mu.Lock()
|
||
|
|
r.prices.loadedAt = time.Now()
|
||
|
|
r.prices.mu.Unlock()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
defer rows.Close()
|
||
|
|
loaded := map[string]ModelPrice{}
|
||
|
|
for rows.Next() {
|
||
|
|
var model string
|
||
|
|
var in, cached, out int64
|
||
|
|
if err := rows.Scan(&model, &in, &cached, &out); err != nil {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
loaded[NormalizeModel(model)] = ModelPrice{
|
||
|
|
InputPerMTok: in, CachedInputPerMTok: cached, OutputPerMTok: out,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
r.prices.mu.Lock()
|
||
|
|
r.prices.prices = loaded
|
||
|
|
r.prices.loadedAt = time.Now()
|
||
|
|
r.prices.mu.Unlock()
|
||
|
|
}
|
||
|
|
|
||
|
|
// SourceLive marks rows written by the recorder from a provider usage block.
|
||
|
|
// SourceBackfill marks history reconstructed from processed_products (see
|
||
|
|
// BackfillFromProcessedProducts) — no user attribution, best effort.
|
||
|
|
const (
|
||
|
|
SourceLive = "live"
|
||
|
|
SourceBackfill = "backfill"
|
||
|
|
)
|
||
|
|
|
||
|
|
const upsertUsageSQL = `
|
||
|
|
INSERT INTO ai_usage_daily (
|
||
|
|
day, company_id, user_id, model, role, source,
|
||
|
|
calls, failed_calls, prompt_tokens, cached_prompt_tokens, completion_tokens,
|
||
|
|
total_tokens, cost_micros, updated_at)
|
||
|
|
VALUES (
|
||
|
|
(now() AT TIME ZONE 'UTC')::date, $1, $2, $3, $4, 'live',
|
||
|
|
1, $5, $6, $7, $8,
|
||
|
|
$9, $10, now())
|
||
|
|
ON CONFLICT (day, company_id, user_id, model, role, source) DO UPDATE SET
|
||
|
|
calls = ai_usage_daily.calls + 1,
|
||
|
|
failed_calls = ai_usage_daily.failed_calls + EXCLUDED.failed_calls,
|
||
|
|
prompt_tokens = ai_usage_daily.prompt_tokens + EXCLUDED.prompt_tokens,
|
||
|
|
cached_prompt_tokens = ai_usage_daily.cached_prompt_tokens + EXCLUDED.cached_prompt_tokens,
|
||
|
|
completion_tokens = ai_usage_daily.completion_tokens + EXCLUDED.completion_tokens,
|
||
|
|
total_tokens = ai_usage_daily.total_tokens + EXCLUDED.total_tokens,
|
||
|
|
cost_micros = ai_usage_daily.cost_micros + EXCLUDED.cost_micros,
|
||
|
|
updated_at = now()`
|
||
|
|
|
||
|
|
// recordUsage rolls one call into the daily aggregate. Like Record it must never
|
||
|
|
// break the call it observes.
|
||
|
|
func (r *Recorder) recordUsage(ctx context.Context, c Call) {
|
||
|
|
if r == nil || r.pool == nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
price := r.priceFor(ctx, c.Model)
|
||
|
|
cost := CostMicros(price, c.PromptTokens, c.CachedPromptTokens, c.OutputTokens)
|
||
|
|
failed := int64(0)
|
||
|
|
if strings.TrimSpace(c.Error) != "" {
|
||
|
|
failed = 1
|
||
|
|
}
|
||
|
|
writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||
|
|
defer cancel()
|
||
|
|
_, err := r.pool.Exec(writeCtx, upsertUsageSQL,
|
||
|
|
orNilUUID(c.CompanyID), orNilUUID(c.UserID), NormalizeModel(c.Model), roleOrOther(c.Role),
|
||
|
|
failed, c.PromptTokens, c.CachedPromptTokens, c.OutputTokens,
|
||
|
|
c.TotalTokens, cost,
|
||
|
|
)
|
||
|
|
if err != nil {
|
||
|
|
r.noteFailure(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// orNilUUID keeps the primary key non-NULL (see the migration comment).
|
||
|
|
func orNilUUID(id uuid.UUID) uuid.UUID { return id }
|
||
|
|
|
||
|
|
func roleOrOther(role string) string {
|
||
|
|
if r := strings.TrimSpace(role); r != "" {
|
||
|
|
return r
|
||
|
|
}
|
||
|
|
return RoleOther
|
||
|
|
}
|
||
|
|
|
||
|
|
// CleanupExpiredUsage drops rollup rows past UsageRetentionDays.
|
||
|
|
func CleanupExpiredUsage(ctx context.Context, pool *pgxpool.Pool) (int64, error) {
|
||
|
|
if pool == nil {
|
||
|
|
return 0, nil
|
||
|
|
}
|
||
|
|
tag, err := pool.Exec(ctx, `
|
||
|
|
DELETE FROM ai_usage_daily
|
||
|
|
WHERE day < ((now() AT TIME ZONE 'UTC')::date - ($1::int))`, UsageRetentionDays)
|
||
|
|
if err != nil {
|
||
|
|
return 0, err
|
||
|
|
}
|
||
|
|
return tag.RowsAffected(), nil
|
||
|
|
}
|