package repository import ( "context" "fmt" "time" "cpa-ext/internal/pricing" ) // ListPrices returns every configured exact-model policy. func (r *SQLiteUsageRepository) ListPrices(ctx context.Context) ([]pricing.Policy, error) { rows, err := r.db.QueryContext(ctx, ` SELECT model, input_rate_micros, cache_read_rate_micros, cache_write_rate_micros, output_rate_micros, long_context_enabled, long_context_threshold, long_context_comparison, long_input_rate_micros, long_cache_read_rate_micros, long_cache_write_rate_micros, long_output_rate_micros, fast_pricing_enabled, fast_multiplier_numerator, fast_multiplier_denominator FROM model_prices ORDER BY model`) if err != nil { return nil, fmt.Errorf("查询模型价格: %w", err) } defer rows.Close() var policies []pricing.Policy for rows.Next() { var policy pricing.Policy var longEnabled bool var long pricing.LongContext if err := rows.Scan( &policy.Model, &policy.Base.InputMicrosPer1M, &policy.Base.CacheReadMicrosPer1M, &policy.Base.CacheWriteMicrosPer1M, &policy.Base.OutputMicrosPer1M, &longEnabled, &long.ThresholdInputTokens, &long.Comparison, &long.Rates.InputMicrosPer1M, &long.Rates.CacheReadMicrosPer1M, &long.Rates.CacheWriteMicrosPer1M, &long.Rates.OutputMicrosPer1M, &policy.FastPricingEnabled, &policy.FastMultiplier.Numerator, &policy.FastMultiplier.Denominator, ); err != nil { return nil, fmt.Errorf("读取模型价格: %w", err) } if longEnabled { policy.LongContext = &long } policies = append(policies, policy) } if err := rows.Err(); err != nil { return nil, fmt.Errorf("遍历模型价格: %w", err) } return policies, nil } // UpsertPrice atomically replaces one model policy. func (r *SQLiteUsageRepository) UpsertPrice(ctx context.Context, policy pricing.Policy) error { if err := policy.Validate(); err != nil { return err } longEnabled := policy.LongContext != nil long := pricing.LongContext{Comparison: "gt"} if policy.LongContext != nil { long = *policy.LongContext } _, err := r.db.ExecContext(ctx, ` INSERT INTO model_prices ( model, input_rate_micros, cache_read_rate_micros, cache_write_rate_micros, output_rate_micros, long_context_enabled, long_context_threshold, long_context_comparison, long_input_rate_micros, long_cache_read_rate_micros, long_cache_write_rate_micros, long_output_rate_micros, fast_pricing_enabled, fast_multiplier_numerator, fast_multiplier_denominator, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(model) DO UPDATE SET input_rate_micros=excluded.input_rate_micros, cache_read_rate_micros=excluded.cache_read_rate_micros, cache_write_rate_micros=excluded.cache_write_rate_micros, output_rate_micros=excluded.output_rate_micros, long_context_enabled=excluded.long_context_enabled, long_context_threshold=excluded.long_context_threshold, long_context_comparison=excluded.long_context_comparison, long_input_rate_micros=excluded.long_input_rate_micros, long_cache_read_rate_micros=excluded.long_cache_read_rate_micros, long_cache_write_rate_micros=excluded.long_cache_write_rate_micros, long_output_rate_micros=excluded.long_output_rate_micros, fast_pricing_enabled=excluded.fast_pricing_enabled, fast_multiplier_numerator=excluded.fast_multiplier_numerator, fast_multiplier_denominator=excluded.fast_multiplier_denominator, updated_at=excluded.updated_at`, policy.Model, policy.Base.InputMicrosPer1M, policy.Base.CacheReadMicrosPer1M, policy.Base.CacheWriteMicrosPer1M, policy.Base.OutputMicrosPer1M, longEnabled, long.ThresholdInputTokens, long.Comparison, long.Rates.InputMicrosPer1M, long.Rates.CacheReadMicrosPer1M, long.Rates.CacheWriteMicrosPer1M, long.Rates.OutputMicrosPer1M, policy.FastPricingEnabled, policy.FastMultiplier.Numerator, policy.FastMultiplier.Denominator, time.Now().UTC().Format(time.RFC3339Nano), ) if err != nil { return fmt.Errorf("保存模型价格: %w", err) } return nil } // BackfillMissingCosts 只补算尚未定价的历史记录,已经保存的账单金额不会随价格修改而变化。 func (r *SQLiteUsageRepository) BackfillMissingCosts(ctx context.Context, policy pricing.Policy) (int64, error) { if err := policy.Validate(); err != nil { return 0, err } rows, err := r.db.QueryContext(ctx, ` SELECT id, input_tokens, cache_read_tokens, cache_write_tokens, output_tokens, service_tier, speed FROM usage_records WHERE model = ? AND cost_micros IS NULL`, policy.Model) if err != nil { return 0, fmt.Errorf("查询待补算用量: %w", err) } type pendingCost struct { id int64 inputTokens int64 cacheReadTokens int64 cacheWriteTokens int64 outputTokens int64 serviceTier, speed string } var pending []pendingCost for rows.Next() { var item pendingCost if err := rows.Scan(&item.id, &item.inputTokens, &item.cacheReadTokens, &item.cacheWriteTokens, &item.outputTokens, &item.serviceTier, &item.speed); err != nil { _ = rows.Close() return 0, fmt.Errorf("读取待补算用量: %w", err) } pending = append(pending, item) } if err := rows.Err(); err != nil { _ = rows.Close() return 0, fmt.Errorf("遍历待补算用量: %w", err) } if err := rows.Close(); err != nil { return 0, fmt.Errorf("关闭待补算查询: %w", err) } tx, err := r.db.BeginTx(ctx, nil) if err != nil { return 0, fmt.Errorf("开始补算事务: %w", err) } defer func() { _ = tx.Rollback() }() var updated int64 for _, item := range pending { result, calculateErr := pricing.Calculate(policy, pricing.Usage{ InputTokens: item.inputTokens, CacheReadTokens: item.cacheReadTokens, CacheWriteTokens: item.cacheWriteTokens, OutputTokens: item.outputTokens, ServiceTier: item.serviceTier, Speed: item.speed, }) if calculateErr != nil { continue } change, updateErr := tx.ExecContext(ctx, ` UPDATE usage_records SET cost_micros = ?, price_tier = ?, fast_requested = ?, fast_pricing_applied = ?, price_multiplier_numerator = ?, price_multiplier_denominator = ? WHERE id = ? AND cost_micros IS NULL`, result.CostMicros, result.PriceTier, result.FastRequested, result.FastApplied, result.MultiplierNumerator, result.MultiplierDenominator, item.id) if updateErr != nil { return 0, fmt.Errorf("补算用量价格: %w", updateErr) } count, countErr := change.RowsAffected() if countErr != nil { return 0, fmt.Errorf("读取补算数量: %w", countErr) } updated += count } if err := tx.Commit(); err != nil { return 0, fmt.Errorf("提交补算事务: %w", err) } return updated, nil } func (r *SQLiteUsageRepository) DeletePrice(ctx context.Context, model string) error { if _, err := r.db.ExecContext(ctx, `DELETE FROM model_prices WHERE model = ?`, model); err != nil { return fmt.Errorf("删除模型价格: %w", err) } return nil }