LLM Cost Calculator — Go source

// Package llmcost is the Go twin of CosmoDev's src/lib/llmCost.ts (dual
// source: the web lib is TypeScript, the CLI lib is Go — kept in lock-step).
// All rates flow from the model snapshot accessor (cosmodev/aimodels) — never
// hardcoded here; custom rates are the one exception, supplied by the caller
// as a CostRates value. Pure + deterministic, never panics. The table-driven
// tests in llm-cost-calculator_test.go share vectors with
// src/lib/llmCost.test.ts so the two implementations are held to the same
// contract.
package llmcost

import (
	"fmt"
	"math"
	"slices"
	"strconv"
	"strings"

	"cosmodev/aimodels"
)

// BatchDiscount is the multiplier applied to Batch API pricing (the standard
// 50% discount). Mirrors BATCH_DISCOUNT in src/lib/llmCost.ts.
const BatchDiscount = 0.5

// CostRates is a per-million-token price pair. Pointer fields are nil when
// the upstream value is null (unpriced / unknown), mirroring `number | null`.
type CostRates struct {
	InputPerM  *float64
	OutputPerM *float64
}

// CostInput describes one workload. Requests is a *int so the zero value
// means "1 request" (the TS default `requests ?? 1`) while an explicit 0 is
// still honored verbatim — same pointer-for-omitted pattern as slugify's
// Options.Separator.
type CostInput struct {
	InputTokens  int  // input tokens per request
	OutputTokens int  // output tokens per request
	Requests     *int // nil → 1
	Batch        bool // apply BatchDiscount
}

// ModelCost is one comparison row: a model's id, its rates, the workload cost
// (nil when unpriced), and output tokens per USD (nil-safe).
type ModelCost struct {
	ID              string
	InputPerM       *float64
	OutputPerM      *float64
	Cost            *float64
	TokensPerDollar *float64
}

// CostFor returns the cost in USD for a workload, or nil when either rate is
// unpriced. It is the Go twin of costFor() in src/lib/llmCost.ts and must
// agree with it on every shared vector:
// ((inputTokens/1e6)·inputPerM + (outputTokens/1e6)·outputPerM) × requests
// × BatchDiscount when batch.
func CostFor(rates CostRates, input CostInput) *float64 {
	if rates.InputPerM == nil || rates.OutputPerM == nil {
		return nil
	}
	requests := 1
	if input.Requests != nil {
		requests = *input.Requests
	}
	discount := 1.0
	if input.Batch {
		discount = BatchDiscount
	}
	base := (float64(input.InputTokens)/1e6)*(*rates.InputPerM) +
		(float64(input.OutputTokens)/1e6)*(*rates.OutputPerM)
	cost := base * float64(requests) * discount
	return &cost
}

// RatesFor projects a model onto its CostRates pair, mirroring ratesFor().
func RatesFor(m *aimodels.Model) CostRates {
	return CostRates{InputPerM: m.InputPerM, OutputPerM: m.OutputPerM}
}

// -- Cache-aware monthly estimator (FEAT-070; twin of the TS extension) --

// CacheRates is all four pricing lanes; cache fields nil when the model has
// no caching, mirroring CacheRates in src/lib/llmCost.ts.
type CacheRates struct {
	InputPerM      *float64
	OutputPerM     *float64
	CacheReadPerM  *float64
	CacheWritePerM *float64
}

// MonthlyInput is a monthly workload at scale. Pointer fields mirror the TS
// optionals: CacheableFraction nil → 0.9; CacheWriteTokensPerMonth nil →
// derive from this month's misses.
type MonthlyInput struct {
	InputTokensPerMonth  int64
	OutputTokensPerMonth int64
	CacheHit             float64
	CacheableFraction    *float64
	CacheWritePerMonth   *int64
	Batch                bool
}

// MonthlyCost is the four-lane monthly estimate. Lane pointers are nil when
// unpriced (input/output) or on the degrade path (cache wanted, rates
// absent); CacheUnavailable flags that degrade. BlendedPerM is
// total / (input+output in millions), nil when total is nil.
type MonthlyCostResult struct {
	Total            *float64
	UncachedInput    *float64
	CacheReads       *float64
	CacheWrites      *float64
	Output           *float64
	CacheUnavailable bool
	BlendedPerM      *float64
}

// CacheRatesFor projects a model onto all four lanes, mirroring cacheRatesFor().
func CacheRatesFor(m *aimodels.Model) CacheRates {
	return CacheRates{
		InputPerM:      m.InputPerM,
		OutputPerM:     m.OutputPerM,
		CacheReadPerM:  m.CacheReadPerM,
		CacheWritePerM: m.CacheWritePerM,
	}
}

// MonthlyCost prices a monthly workload with cache economics. Caching is
// active iff CacheHit > 0 AND both cache rates are priced; otherwise the
// estimate degrades to plain no-cache math with CacheUnavailable set.
// Twin of monthlyCost() in src/lib/llmCost.ts — must agree on every shared
// vector in the test files.
func MonthlyCost(rates CacheRates, w MonthlyInput) MonthlyCostResult {
	unpriced := rates.InputPerM == nil || rates.OutputPerM == nil
	cacheable := 0.9
	if w.CacheableFraction != nil {
		cacheable = *w.CacheableFraction
	}
	wantedCache := w.CacheHit > 0
	hasCacheRates := rates.CacheReadPerM != nil && rates.CacheWritePerM != nil
	active := wantedCache && hasCacheRates
	hit := 0.0
	if active {
		hit = w.CacheHit
	}

	if unpriced {
		return MonthlyCostResult{CacheUnavailable: wantedCache && !hasCacheRates}
	}

	inT := float64(w.InputTokensPerMonth)
	missTokens := inT * (1 - hit)
	writeTokens := 0.0
	if w.CacheWritePerMonth != nil {
		writeTokens = float64(*w.CacheWritePerMonth)
	} else if active {
		writeTokens = missTokens * cacheable
	}

	uncachedInput := (missTokens / 1e6) * *rates.InputPerM
	// Lane semantics: real 0 when caching is simply off; nil only on the
	// degrade path (cache wanted, rates absent — the price is unknowable).
	var cacheReads, cacheWrites *float64
	degrade := wantedCache && !hasCacheRates
	if active {
		cr := (inT * hit / 1e6) * *rates.CacheReadPerM
		cw := (writeTokens / 1e6) * *rates.CacheWritePerM
		cacheReads, cacheWrites = &cr, &cw
	} else if degrade {
		cacheReads, cacheWrites = nil, nil
	} else {
		zero := 0.0
		cacheReads, cacheWrites = &zero, &zero
	}
	output := (float64(w.OutputTokensPerMonth) / 1e6) * *rates.OutputPerM

	lanes := uncachedInput + output
	if cacheReads != nil {
		lanes += *cacheReads
	}
	if cacheWrites != nil {
		lanes += *cacheWrites
	}
	if w.Batch {
		lanes *= BatchDiscount
	}
	totalMTok := (inT + float64(w.OutputTokensPerMonth)) / 1e6

	out := MonthlyCostResult{
		Total:            &lanes,
		UncachedInput:    &uncachedInput,
		CacheReads:       cacheReads,
		CacheWrites:      cacheWrites,
		Output:           &output,
		CacheUnavailable: wantedCache && !hasCacheRates,
	}
	if totalMTok > 0 {
		b := lanes / totalMTok
		out.BlendedPerM = &b
	}
	return out
}

// ParseTokensPerMonth parses "2B"/"1.5b"/"500M"/"10m"/"2K"/"2000"/"1_000_000"
// (case-insensitive suffix, underscores allowed) into a token count. Twin of
// parseTokensPerMonth(); returns an error on anything else.
func ParseTokensPerMonth(s string) (int64, error) {
	t := strings.TrimSpace(s)
	lower := strings.ToLower(t)
	mult := 1.0
	switch {
	case strings.HasSuffix(lower, "b"):
		mult, t = 1e9, lower[:len(lower)-1]
	case strings.HasSuffix(lower, "m"):
		mult, t = 1e6, lower[:len(lower)-1]
	case strings.HasSuffix(lower, "k"):
		mult, t = 1e3, lower[:len(lower)-1]
	}
	t = strings.ReplaceAll(t, "_", "")
	if t == "" {
		return 0, fmt.Errorf("unparseable tokens/month: %q", s)
	}
	v, err := strconv.ParseFloat(t, 64)
	if err != nil || v < 0 {
		return 0, fmt.Errorf("unparseable tokens/month: %q", s)
	}
	return int64(math.Round(v * mult)), nil
}

// WorkloadPreset is an archetype default for realistic cache modeling.
type WorkloadPreset struct {
	ID                string
	Label             string
	CacheHit          float64
	CacheableFraction float64
	Rationale         string
}

// WorkloadPresets mirrors WORKLOAD_PRESETS in src/lib/llmCost.ts — planning
// numbers from published prompt-caching writeups, refined as data accumulates.
var WorkloadPresets = []WorkloadPreset{
	{ID: "agent-coding", Label: "Agent / coding", CacheHit: 0.75, CacheableFraction: 0.95,
		Rationale: "Large stable system prompt plus repo context, re-sent every turn — the highest-hit profile."},
	{ID: "chat", Label: "Chat / assistant", CacheHit: 0.30, CacheableFraction: 0.60,
		Rationale: "Short varied conversations; only the system prompt and early turns stay hot."},
	{ID: "rag", Label: "RAG / retrieval", CacheHit: 0.50, CacheableFraction: 0.80,
		Rationale: "Preamble caches well but each query pulls fresh retrieved chunks."},
	{ID: "batch-summarize", Label: "Batch summarize", CacheHit: 0.60, CacheableFraction: 0.50,
		Rationale: "Shared instruction prefix across many unique documents."},
}

// CompareModels costs every requested model for one workload and sorts the
// rows: cost asc, nil-cost last, ties (equal cost or nil-nil) by id asc.
// Unknown ids are dropped. A nil models slice means "use the live snapshot"
// (via aimodels.AllModels), mirroring the TS default parameter. It is the Go
// twin of compareModels() in src/lib/llmCost.ts.
func CompareModels(ids []string, input CostInput, models []aimodels.Model) []ModelCost {
	all := models
	if all == nil {
		var err error
		if all, err = aimodels.AllModels(); err != nil {
			return nil
		}
	}
	rows := make([]ModelCost, 0, len(ids))
	for _, id := range ids {
		for i := range all {
			if all[i].ID != id {
				continue
			}
			m := &all[i]
			var tpd *float64
			if v, ok := aimodels.TokensPerDollar(m); ok {
				tpd = &v
			}
			rows = append(rows, ModelCost{
				ID:              id,
				InputPerM:       m.InputPerM,
				OutputPerM:      m.OutputPerM,
				Cost:            CostFor(RatesFor(m), input),
				TokensPerDollar: tpd,
			})
			break
		}
	}
	slices.SortStableFunc(rows, func(a, b ModelCost) int {
		if a.Cost == nil || b.Cost == nil {
			if a.Cost == nil && b.Cost == nil {
				return strings.Compare(a.ID, b.ID)
			}
			if a.Cost == nil {
				return 1
			}
			return -1
		}
		if *a.Cost == *b.Cost {
			return strings.Compare(a.ID, b.ID)
		}
		if *a.Cost < *b.Cost {
			return -1
		}
		return 1
	})
	return rows
}