fix

2026-08-23 22:52:13 +02:00
parent 8983cfc8a1
commit 296dc3c841
32 changed files with 2025 additions and 52 deletions
@@ -0,0 +1,305 @@
+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
+}