billing.go

package llmproxy

// CostInfo contains pricing information for a model.
type CostInfo struct {
	// Input is the cost per 1M input tokens in USD.
	Input float64
	// Output is the cost per 1M output tokens in USD.
	Output float64
	// CacheRead is the cost per 1M cached input tokens (optional).
	CacheRead float64
	// CacheWrite is the cost per 1M cache write tokens (optional, Anthropic).
	CacheWrite float64
	// Unit describes the billing unit for the cost values. Token-based models use
	// "per_million_tokens"; non-token modalities may use provider-specific units.
	Unit string
}

// CostLookup is a function that returns the cost for a given provider and model.
// It should return the pricing info or false if the model is not found.
//
// The lookup function allows the pricing data to be managed externally,
// such as downloading from models.dev or using a custom pricing database.
type CostLookup func(provider string, model string) (CostInfo, bool)

// BillingResult contains the calculated cost for a request.
type BillingResult struct {
	// Provider is the provider name.
	Provider string
	// Model is the model identifier.
	Model string
	// PromptTokens is the number of input tokens.
	PromptTokens int
	// CompletionTokens is the number of output tokens.
	CompletionTokens int
	// CachedTokens is the number of prompt tokens served from cache.
	CachedTokens int
	// TotalTokens is the sum of prompt and completion tokens.
	TotalTokens int
	// Unit is the billing unit used for this calculation.
	Unit string
	// InputQuantity is the input quantity for non-token units.
	InputQuantity float64
	// OutputQuantity is the output quantity for non-token units.
	OutputQuantity float64
	// InputCost is the calculated input cost in USD (non-cached prompt tokens).
	InputCost float64
	// CachedInputCost is the cost for cached prompt tokens in USD.
	CachedInputCost float64
	// OutputCost is the calculated output cost in USD.
	OutputCost float64
	// TotalCost is the sum of all costs in USD.
	TotalCost float64
}

// CalculateCost computes the billing result from cost info, token usage, and cache usage.
// Cached tokens are billed at the CacheRead rate (if available), and non-cached prompt
// tokens are billed at the full Input rate.
func CalculateCost(provider, model string, costInfo CostInfo, promptTokens, completionTokens int, cacheUsage *CacheUsage) BillingResult {
	return CalculateCostWithMeteredUsage(provider, model, costInfo, promptTokens, completionTokens, cacheUsage, MeteredUsage{})
}

func CalculateCostWithMeteredUsage(provider, model string, costInfo CostInfo, promptTokens, completionTokens int, cacheUsage *CacheUsage, meteredUsage MeteredUsage) BillingResult {
	unit := costInfo.Unit
	if unit == "" {
		unit = "per_million_tokens"
	}
	if costInfo.Unit != "" && costInfo.Unit != "per_million_tokens" {
		inputQuantity, outputQuantity := billingQuantitiesForUnit(unit, meteredUsage)
		inputCost, outputCost := meteredCostForUnit(costInfo, unit, inputQuantity, outputQuantity)
		return BillingResult{
			Provider:         provider,
			Model:            model,
			PromptTokens:     promptTokens,
			CompletionTokens: completionTokens,
			TotalTokens:      promptTokens + completionTokens,
			Unit:             unit,
			InputQuantity:    inputQuantity,
			OutputQuantity:   outputQuantity,
			InputCost:        inputCost,
			OutputCost:       outputCost,
			TotalCost:        inputCost + outputCost,
		}
	}

	cachedTokens := 0
	if cacheUsage != nil {
		// Providers populate only one of these fields — OpenAI/Fireworks/Bedrock
		// set CachedTokens while Anthropic sets CacheReadInputTokens. We sum them
		// so the same code path works for any provider. The clamp below guards
		// against overcounting if a future provider ever sets both fields.
		cachedTokens = cacheUsage.CachedTokens + cacheUsage.CacheReadInputTokens
	}

	// Providers report prompt tokens differently:
	//   - OpenAI/Fireworks/Bedrock: promptTokens INCLUDES cached tokens
	//     → non-cached = promptTokens - cachedTokens
	//   - Anthropic: input_tokens EXCLUDES cached tokens (only new tokens)
	//     → non-cached = promptTokens (as reported), cached is additional
	//
	// We detect the style by comparing: if cached > prompt, the provider
	// must be reporting non-cached only (Anthropic style).
	var nonCachedTokens int
	if cachedTokens > promptTokens {
		// Anthropic style: promptTokens = non-cached only, cached is separate
		nonCachedTokens = promptTokens
		// Adjust promptTokens to reflect the true total for the BillingResult
		promptTokens = promptTokens + cachedTokens
	} else {
		// OpenAI style: promptTokens includes cached
		nonCachedTokens = promptTokens - cachedTokens
	}

	// Non-cached prompt tokens at full input rate
	inputCost := costInfo.Input * float64(nonCachedTokens) / 1_000_000

	// Cached tokens at cache read rate (falls back to full input rate if no cache pricing)
	var cachedInputCost float64
	if cachedTokens > 0 {
		cacheRate := costInfo.CacheRead
		if cacheRate <= 0 {
			cacheRate = costInfo.Input // fallback to full rate
		}
		cachedInputCost = cacheRate * float64(cachedTokens) / 1_000_000
	}

	outputCost := costInfo.Output * float64(completionTokens) / 1_000_000

	return BillingResult{
		Provider:         provider,
		Model:            model,
		PromptTokens:     promptTokens,
		CompletionTokens: completionTokens,
		CachedTokens:     cachedTokens,
		TotalTokens:      promptTokens + completionTokens,
		Unit:             unit,
		InputQuantity:    float64(promptTokens),
		OutputQuantity:   float64(completionTokens),
		InputCost:        inputCost,
		CachedInputCost:  cachedInputCost,
		OutputCost:       outputCost,
		TotalCost:        inputCost + cachedInputCost + outputCost,
	}
}

func billingQuantitiesForUnit(unit string, usage MeteredUsage) (float64, float64) {
	switch unit {
	case "per_million_characters":
		return float64(usage.InputCharacters), float64(usage.OutputCharacters)
	case "per_minute_audio":
		return usage.InputAudioSeconds / 60, usage.OutputAudioSeconds / 60
	case "per_second_720p_video", "per_second_720p_1080p_video":
		return 0, usage.OutputVideoSeconds
	default:
		return 0, 0
	}
}

func meteredCostForUnit(costInfo CostInfo, unit string, inputQuantity float64, outputQuantity float64) (float64, float64) {
	switch unit {
	case "per_million_characters":
		return costInfo.Input * inputQuantity / 1_000_000, costInfo.Output * outputQuantity / 1_000_000
	case "per_minute_audio":
		return costInfo.Input * inputQuantity, costInfo.Output * outputQuantity
	case "per_second_720p_video", "per_second_720p_1080p_video":
		return costInfo.Input * inputQuantity, costInfo.Output * outputQuantity
	default:
		return 0, 0
	}
}