cpa-plugin/internal/repository/sqlite_pricing.go

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
}