306 lines
8.7 KiB
Go
306 lines
8.7 KiB
Go
package aiaudit
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
// Recovering cost for products processed before capture existed.
|
|
//
|
|
// processed_products.gpt_response keeps the provider's usage block for some rows
|
|
// (model + prompt/completion tokens). That is enough to reconstruct spend per
|
|
// company and day, so history is not simply lost. It is explicitly best effort:
|
|
//
|
|
// - Only rows that still carry a usage block are recoverable. gpt_response is
|
|
// overwritten on every reprocess and absent on older/failed rows.
|
|
// - There is no per-user attribution: processed_products.user_id is null on these
|
|
// rows, so backfilled spend lands under "Unattributed" in the by-user view.
|
|
// - The day is processed_products.updated_at, which is the LAST write to the row,
|
|
// not necessarily when the call was made.
|
|
// - Cached-token discounts cannot be recovered (the field did not exist), so a
|
|
// backfilled call is priced as if none of its input was cached — an
|
|
// over-estimate, never an under-estimate.
|
|
//
|
|
// Rows are written with source='backfill' so they never collide with live capture
|
|
// and a re-run replaces rather than doubles them.
|
|
|
|
// BackfillResult reports what a backfill pass recovered.
|
|
type BackfillResult struct {
|
|
ProductsScanned int `json:"products_scanned"`
|
|
ProductsUsed int `json:"products_used"`
|
|
ProductsNoUsage int `json:"products_without_usage"`
|
|
CallsRecovered int `json:"calls_recovered"`
|
|
PromptTokens int64 `json:"prompt_tokens"`
|
|
CompletionTokens int64 `json:"completion_tokens"`
|
|
CostMicros int64 `json:"cost_micros"`
|
|
RowsWritten int `json:"rows_written"`
|
|
RowsDeleted int64 `json:"rows_deleted"`
|
|
Cutoff string `json:"cutoff"`
|
|
DryRun bool `json:"dry_run"`
|
|
}
|
|
|
|
// CostUSD renders the recovered spend for CLI output.
|
|
func (r BackfillResult) CostUSD() float64 {
|
|
return float64(r.CostMicros) / float64(MicrosPerUSD)
|
|
}
|
|
|
|
type backfillKey struct {
|
|
day time.Time
|
|
companyID uuid.UUID
|
|
model string
|
|
role string
|
|
}
|
|
|
|
type backfillAgg struct {
|
|
calls int
|
|
promptTokens int64
|
|
completionTokens int64
|
|
totalTokens int64
|
|
costMicros int64
|
|
}
|
|
|
|
// LiveCaptureStart returns the earliest day live capture recorded, or zero when
|
|
// nothing has been captured yet. The backfill uses it as a default cutoff so a
|
|
// product processed after capture began is never counted twice.
|
|
func LiveCaptureStart(ctx context.Context, pool *pgxpool.Pool) (time.Time, error) {
|
|
if pool == nil {
|
|
return time.Time{}, nil
|
|
}
|
|
var day *time.Time
|
|
err := pool.QueryRow(ctx, `
|
|
SELECT min(day) FROM ai_usage_daily WHERE source = $1`, SourceLive).Scan(&day)
|
|
if err != nil {
|
|
return time.Time{}, err
|
|
}
|
|
if day == nil {
|
|
return time.Time{}, nil
|
|
}
|
|
return *day, nil
|
|
}
|
|
|
|
// BackfillFromProcessedProducts reconstructs cost history from stored provider
|
|
// responses for products last written before cutoff. Passing a zero cutoff scans
|
|
// everything, which is only safe when live capture has recorded nothing yet.
|
|
func BackfillFromProcessedProducts(ctx context.Context, pool *pgxpool.Pool, cutoff time.Time, dryRun bool) (BackfillResult, error) {
|
|
out := BackfillResult{DryRun: dryRun}
|
|
if pool == nil {
|
|
return out, fmt.Errorf("nil pool")
|
|
}
|
|
if cutoff.IsZero() {
|
|
// Far future = no bound.
|
|
cutoff = time.Now().UTC().AddDate(100, 0, 0)
|
|
}
|
|
out.Cutoff = cutoff.Format("2006-01-02")
|
|
|
|
prices, err := loadPrices(ctx, pool)
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
|
|
rows, err := pool.Query(ctx, `
|
|
SELECT company_id, updated_at, COALESCE(gpt_response, '{}'::jsonb)
|
|
FROM processed_products
|
|
WHERE gpt_response IS NOT NULL
|
|
AND updated_at < $1
|
|
ORDER BY updated_at`, cutoff)
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
agg := map[backfillKey]*backfillAgg{}
|
|
for rows.Next() {
|
|
var companyID uuid.UUID
|
|
var updatedAt time.Time
|
|
var raw []byte
|
|
if err := rows.Scan(&companyID, &updatedAt, &raw); err != nil {
|
|
return out, err
|
|
}
|
|
out.ProductsScanned++
|
|
calls := extractUsageCalls(raw)
|
|
if len(calls) == 0 {
|
|
out.ProductsNoUsage++
|
|
continue
|
|
}
|
|
out.ProductsUsed++
|
|
day := updatedAt.UTC().Truncate(24 * time.Hour)
|
|
for _, c := range calls {
|
|
price := priceFromMap(prices, c.Model)
|
|
// Cached tokens are unknown for history — price all input as fresh.
|
|
cost := CostMicros(price, c.PromptTokens, 0, c.CompletionTokens)
|
|
k := backfillKey{day: day, companyID: companyID, model: NormalizeModel(c.Model), role: c.Role}
|
|
a := agg[k]
|
|
if a == nil {
|
|
a = &backfillAgg{}
|
|
agg[k] = a
|
|
}
|
|
a.calls++
|
|
a.promptTokens += int64(c.PromptTokens)
|
|
a.completionTokens += int64(c.CompletionTokens)
|
|
a.totalTokens += int64(c.TotalTokens)
|
|
a.costMicros += cost
|
|
|
|
out.CallsRecovered++
|
|
out.PromptTokens += int64(c.PromptTokens)
|
|
out.CompletionTokens += int64(c.CompletionTokens)
|
|
out.CostMicros += cost
|
|
}
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return out, err
|
|
}
|
|
out.RowsWritten = len(agg)
|
|
if dryRun {
|
|
return out, nil
|
|
}
|
|
|
|
tx, err := pool.Begin(ctx)
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
|
|
// Replace rather than add: a re-run must not double the recovered history.
|
|
tag, err := tx.Exec(ctx, `DELETE FROM ai_usage_daily WHERE source = $1`, SourceBackfill)
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
out.RowsDeleted = tag.RowsAffected()
|
|
|
|
for k, a := range agg {
|
|
if _, err := tx.Exec(ctx, `
|
|
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 ($1, $2, '00000000-0000-0000-0000-000000000000', $3, $4, $5,
|
|
$6, 0, $7, 0, $8, $9, $10, now())`,
|
|
k.day, k.companyID, k.model, k.role, SourceBackfill,
|
|
a.calls, a.promptTokens, a.completionTokens, a.totalTokens, a.costMicros); err != nil {
|
|
return out, err
|
|
}
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return out, err
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func loadPrices(ctx context.Context, pool *pgxpool.Pool) (map[string]ModelPrice, error) {
|
|
out := map[string]ModelPrice{}
|
|
rows, err := 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 {
|
|
return out, err
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
var model string
|
|
var in, cached, o int64
|
|
if err := rows.Scan(&model, &in, &cached, &o); err != nil {
|
|
return out, err
|
|
}
|
|
out[NormalizeModel(model)] = ModelPrice{InputPerMTok: in, CachedInputPerMTok: cached, OutputPerMTok: o}
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func priceFromMap(prices map[string]ModelPrice, model string) ModelPrice {
|
|
key := NormalizeModel(model)
|
|
if p, ok := prices[key]; ok {
|
|
return p
|
|
}
|
|
if p, ok := DefaultModelPrice(key); ok {
|
|
return p
|
|
}
|
|
return ModelPrice{}
|
|
}
|
|
|
|
// recoveredCall is one provider response found inside gpt_response.
|
|
type recoveredCall struct {
|
|
Model string
|
|
Role string
|
|
PromptTokens int
|
|
CompletionTokens int
|
|
TotalTokens int
|
|
}
|
|
|
|
// extractUsageCalls walks a stored gpt_response and returns every provider usage
|
|
// block it holds. The shape varies by pipeline version, so this looks for any
|
|
// object with a "usage" child rather than assuming one fixed path — older rows
|
|
// nest it differently and a rigid path would silently recover nothing.
|
|
func extractUsageCalls(raw []byte) []recoveredCall {
|
|
var doc any
|
|
if err := json.Unmarshal(raw, &doc); err != nil {
|
|
return nil
|
|
}
|
|
var out []recoveredCall
|
|
walkUsage(doc, "", &out)
|
|
return out
|
|
}
|
|
|
|
func walkUsage(node any, stepHint string, out *[]recoveredCall) {
|
|
switch v := node.(type) {
|
|
case map[string]any:
|
|
if s, ok := v["step"].(string); ok && s != "" {
|
|
stepHint = s
|
|
}
|
|
if usage, ok := v["usage"].(map[string]any); ok {
|
|
call := recoveredCall{
|
|
Model: stringField(v, "model"),
|
|
Role: roleFromStep(stepHint),
|
|
PromptTokens: intField(usage, "prompt_tokens"),
|
|
CompletionTokens: intField(usage, "completion_tokens"),
|
|
TotalTokens: intField(usage, "total_tokens"),
|
|
}
|
|
if call.PromptTokens > 0 || call.CompletionTokens > 0 {
|
|
*out = append(*out, call)
|
|
}
|
|
}
|
|
for _, child := range v {
|
|
walkUsage(child, stepHint, out)
|
|
}
|
|
case []any:
|
|
for _, child := range v {
|
|
walkUsage(child, stepHint, out)
|
|
}
|
|
}
|
|
}
|
|
|
|
func roleFromStep(step string) string {
|
|
switch step {
|
|
case "categorize":
|
|
return RoleCategorize
|
|
case "ai_enhance", "":
|
|
return RoleProcessing
|
|
default:
|
|
return RoleProcessing
|
|
}
|
|
}
|
|
|
|
func stringField(m map[string]any, key string) string {
|
|
if s, ok := m[key].(string); ok {
|
|
return s
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func intField(m map[string]any, key string) int {
|
|
switch v := m[key].(type) {
|
|
case float64:
|
|
return int(v)
|
|
case int:
|
|
return v
|
|
case json.Number:
|
|
n, _ := v.Int64()
|
|
return int(n)
|
|
}
|
|
return 0
|
|
}
|