ModelRouter/internal/billing/compile.go

package billing

import (
	"fmt"
	"strings"

	"llmsproxy/internal/config"
)

// Compile turns one billing profile into the prices table the billing plugin
// expects, resolving URL rules against the gateway's actual sources.
//
// WHY URL RULES NEED RESOLVING AT ALL: the operator declares pricing by URL
// because that is what a provider's price list is keyed on, and because several
// sources can point at the same URL. The plugin, however, looks up by
// `payload.source` (a source NAME) and by model — it has no idea what URL a
// request went to. So the URL match happens here, at compile time, where the
// config's name->base_url mapping is known, and the result is expressed in the
// dimensions the plugin already supports.
//
// Profiles are the "let the user choose" axis: the same URL can appear in
// several profiles and switching recomputes this table, so a gateway can be
// repriced without editing the plugin.
func Compile(profile *config.BillingProfile, sources []config.Source) (map[string]interface{}, error) {
	return CompileOpts(profile, sources, false)
}

// CompileOpts is Compile with an explicit policy for rules that match no
// source.
//
// The default is STRICT and that is correct for startup: a profile with a
// typo'd URL prices nothing, every request on it is recorded as unpriced, and
// the bill silently comes out wrong. Failing the load is the right response.
//
// The rules editor needs the opposite. An operator editing a half-finished
// profile — adding a rule before the source that will use it exists, or
// renaming a URL while typing — must be able to save and then see the warning,
// not be blocked by an error they can only resolve by guessing. Lenient mode
// keeps the offending rule in the config and drops it from the compiled table,
// so it is preserved as text and clearly reported, without pretending it
// prices anything.
func CompileOpts(profile *config.BillingProfile, sources []config.Source, lenient bool) (map[string]interface{}, error) {
	if profile == nil {
		return nil, fmt.Errorf("no billing profile")
	}
	prices := map[string]interface{}{
		"currency": "USD",
		"default":  map[string]interface{}{"prompt": 0.0, "completion": 0.0, "per_request": 0.0},
		"sources":  map[string]interface{}{},
		"models":   map[string]interface{}{},
		"keys":     map[string]interface{}{},
	}
	if profile.Default != "" {
		prices["default_mode"] = profile.Default
	}

	// Currency is a profile-level statement; a rule may override it.
	for i := range profile.Rules {
		if c := profile.Rules[i].Currency; c != "" {
			prices["currency"] = c
			break
		}
	}

	matched := map[string]bool{}
	for i := range profile.Rules {
		rule := &profile.Rules[i]
		targets := matchSources(rule.URL, sources)
		if len(targets) == 0 && rule.URL != "*" {
			// A rule for a URL no source uses is almost always a typo (or a
			// source that was removed). Failing loudly beats a profile that
			// silently prices nothing — the whole reason this is a config file
			// instead of a hand-written JSON blob.
			if lenient {
				continue // kept in the config text, reported as a warning
			}
			return nil, fmt.Errorf("rule url %q matches no configured source base_url", rule.URL)
		}
		for _, srcName := range targets {
			if matched[srcName] {
				// First rule wins. Two rules matching one source is ambiguous,
				// and silently letting the later one win makes the file's
				// meaning depend on ordering the operator cannot see.
				continue
			}
			matched[srcName] = true
			if err := applyRule(prices, rule, srcName, sources); err != nil {
				return nil, err
			}
		}
	}

	// Sources no rule matched fall through to the profile default, which the
	// plugin applies via prices.default. Recording them explicitly means the
	// UI can say "this source is unpriced" instead of leaving the operator to
	// infer it from a zero.
	var unmatched []string
	for _, src := range sources {
		if !matched[src.Name] {
			unmatched = append(unmatched, src.Name)
		}
	}
	if len(unmatched) > 0 {
		prices["unmatched_sources"] = unmatched
	}
	return prices, nil
}

// matchSources returns the source names whose base_url matches pattern.
// "*" matches every source (used as a catch-all default rule).
func matchSources(pattern string, sources []config.Source) []string {
	var out []string
	if pattern == "*" {
		for _, s := range sources {
			out = append(out, s.Name)
		}
		return out
	}
	want := normalizeURL(pattern)
	for _, s := range sources {
		if normalizeURL(s.BaseURL) == want {
			out = append(out, s.Name)
		}
	}
	return out
}

// normalizeURL compares URLs the way an operator expects: trailing slashes and
// case in the host are differences the provider's price list does not care
// about, and requiring an exact byte match would make the config brittle.
func normalizeURL(u string) string {
	u = strings.TrimSpace(u)
	u = strings.TrimRight(u, "/")
	return strings.ToLower(u)
}

// applyRule writes one rule's pricing for one source into the prices table.
func applyRule(prices map[string]interface{}, rule *config.BillingRule, srcName string, sources []config.Source) error {
	switch rule.Mode {
	case "free":
		// Explicitly priced at zero. This is NOT the same as unpriced: a source
		// the operator says is free must not appear in the unpriced warnings,
		// or those warnings become noise and stop being read.
		prices["sources"].(map[string]interface{})[srcName] = map[string]interface{}{
			"prompt": 0.0, "completion": 0.0, "per_request": 0.0,
		}
		return nil

	case "subscription":
		// A fixed monthly commitment: the per-request MARGINAL cost is zero, and
		// the flat fee is reported separately. Spreading a monthly fee across
		// requests would invent a per-request number the provider never charges,
		// and it would change every time traffic did.
		prices["sources"].(map[string]interface{})[srcName] = map[string]interface{}{
			"prompt": 0.0, "completion": 0.0, "per_request": 0.0,
		}
		fixed, _ := prices["subscriptions"].(map[string]interface{})
		if fixed == nil {
			fixed = map[string]interface{}{}
			prices["subscriptions"] = fixed
		}
		cur := "USD"
		if rule.Currency != "" {
			cur = rule.Currency
		}
		fixed[srcName] = map[string]interface{}{
			"monthly": rule.Subscription, "currency": cur,
		}
		return nil

	case "unpriced":
		// Deliberately left out of `sources` so the plugin's unpriced_models /
		// unpriced_reqs counters catch it. That is the point: a subscription
		// plan whose credits cannot be converted to tokens must be VISIBLE as
		// unbilled, not quietly estimated.
		return nil

	case "token":
		if len(rule.Models) == 0 {
			return fmt.Errorf("rule for url %q: mode token requires at least one model", rule.URL)
		}
		models := prices["models"].(map[string]interface{})
		for modelID, t := range rule.Models {
			entry := map[string]interface{}{}
			prompt, err := usdPerM(t.Prompt)
			if err != nil {
				return fmt.Errorf("model %q prompt: %w", modelID, err)
			}
			completion, err := usdPerM(t.Completion)
			if err != nil {
				return fmt.Errorf("model %q completion: %w", modelID, err)
			}
			entry["prompt"] = prompt
			entry["completion"] = completion
			if t.CacheDiscount != nil {
				entry["cache_discount"] = *t.CacheDiscount
			}
			// reasoning is optional and stays ABSENT (not nil, not 0) when the
			// rule does not set it, which is what makes the billing plugin skip
			// the repricing branch entirely. Writing an explicit 0 here would
			// mean "thinking is free" and silently under-bill every request
			// from a provider that bills it like ordinary output.
			if strings.TrimSpace(t.Reasoning) != "" {
				rp, err := usdPerM(t.Reasoning)
				if err != nil {
					return fmt.Errorf("model %q reasoning: %w", modelID, err)
				}
				entry["reasoning_price"] = rp
			}
			if lc := t.LongContext; lc != nil {
				if lc.Above <= 0 {
					return fmt.Errorf("model %q long_context: above must be > 0", modelID)
				}
				above, err := longContextRate(lc, prompt)
				if err != nil {
					return fmt.Errorf("model %q long_context: %w", modelID, err)
				}
				entry["long_context"] = map[string]interface{}{
					"above": lc.Above,
					// Marginal, not blended. The plugin charges the tokens past the
					// threshold at this rate and the rest at the base rate, so the
					// threshold composes with a cache discount (which applies to the
					// cached portion) instead of forcing the plugin to guess how
					// the two interact.
					"prompt": above,
				}
			}
			if rule.Peak != nil {
				entry["peak"] = compilePeak(rule.Peak)
			}
			models[modelID] = entry
		}
		// A specific-URL token rule needs a source entry so a model served
		// from THIS url but absent from `rule.Models` is still recognised as
		// belonging to a priced provider.
		//
		// A WILDCARD rule ("*") must NOT get one. "*" is a model catalogue — it
		// says "these model ids have these prices, wherever they come from" —
		// and marking every source priced at {0,0} would make priceFor() set
		// priced=true for every model the catalogue does NOT list. Paid
		// sources (commandcode, alittokenplan, …) would then record $0 instead
		// of showing up in unpriced_reqs, which is precisely the silent
		// zero-billing failure this whole plugin is built to make visible.
		// The model lookup p.models[payload.model] already marks a listed
		// model priced on its own, so the wildcard loses nothing.
		if rule.URL != "*" {
			if _, ok := prices["sources"].(map[string]interface{})[srcName]; !ok {
				prices["sources"].(map[string]interface{})[srcName] = map[string]interface{}{
					"prompt": 0.0, "completion": 0.0, "per_request": 0.0,
				}
			}
		}
		return nil
	}
	return fmt.Errorf("unknown mode %q", rule.Mode)
}

// usdPerM converts a USD-per-million string into the per-token rate the plugin
// expects. Delegates to config.ParseFloatUSDPerM so the DSL validator and the
// compiler agree byte-for-byte on what a valid price is — a value that passes
// Validate() but fails here (or vice versa) would be the worst kind of drift.
func usdPerM(s string) (float64, error) {
	return config.ParseFloatUSDPerM(s)
}

// compilePeak renders the peak window in the shape the plugin reads:
// { multiplier, windows = { { days = {...}, hours = { {lo,hi}, ... } } } }.
//
// The plugin reads `days` / `hours` pairs. The first version of the deployed
// price table used { start, end, weekdays } — a shape nothing reads — so peak
// traffic was billed at off-peak rates with no error anywhere. Compiling from
// typed config fields removes the chance of writing the wrong key names by hand.
func compilePeak(p *config.BillingPeak) map[string]interface{} {
	win := map[string]interface{}{}
	if len(p.Weekdays) > 0 {
		days := make([]interface{}, 0, len(p.Weekdays))
		for _, d := range p.Weekdays {
			days = append(days, d)
		}
		win["days"] = days
	}
	if len(p.Hours) > 0 {
		hours := make([]interface{}, 0, len(p.Hours))
		for _, h := range p.Hours {
			hours = append(hours, []interface{}{h[0], h[1]})
		}
		win["hours"] = hours
	}
	out := map[string]interface{}{"multiplier": p.Multiplier}
	if len(win) > 0 {
		out["windows"] = []interface{}{win}
	}
	return out
}

// longContextRate resolves the marginal per-token rate for the tokens ABOVE a
// threshold, given the model's base per-million prompt price.
//
// Two ways to express it are supported because providers publish both: an
// absolute rate (OpenAI quotes \$3.00/M past 200k) and a multiple (DeepSeek
// quotes exactly 2x past 128k). Neither may be missing: a step with no rate is a
// silent no-op, and an operator who wrote one meant something.
func longContextRate(lc *config.BillingLongContext, basePerToken float64) (float64, error) {
	hasPrompt := strings.TrimSpace(lc.Prompt) != ""
	hasMult := lc.Multiplier > 0
	switch {
	case hasPrompt && hasMult:
		p, err := usdPerM(lc.Prompt)
		if err != nil {
			return 0, err
		}
		// Both given: the absolute rate wins, because it is the one the provider
		// actually quotes. Silently multiplying instead would produce a number
		// the operator never wrote.
		return p, nil
	case hasPrompt:
		return usdPerM(lc.Prompt)
	case hasMult:
		return basePerToken * lc.Multiplier, nil
	default:
		return 0, fmt.Errorf("needs either prompt or multiplier")
	}
}