Model Picker — Go source

// model-picker — filter + rank the AI model catalog.
//
// Language: Go (1.21+, standard library only)
// Source:   CosmoDev polyglot showcase port of the Model Picker tool, ported
//           from src/lib/modelPicker.ts (the canonical TypeScript
//           implementation).
// Live:     https://dev.cosmolabs.org/tools/model-picker
// License:  display source — part of CosmoDev's polyglot tool pages.
//
// Design goals:
//   - Pure + deterministic; never panics.
//   - Functionally equivalent to the TS reference: same inputs -> same outputs.
//   - Self-contained: stdlib only (no external Go modules).
//
// Port notes: the TS lib imports `listModels` from src/lib/ai/models.ts and
// lets `models` default to the bundled pricing snapshot
// (src/data/ai-models.json). A dependency-free port cannot load that file, so
// the filter step is inlined here and `models` is an explicit slice parameter.
// Go has no default arguments; the TS default `sort = 'price'` is expressed by
// SortPrice being the zero-value-friendly first constant — callers wanting the
// default pass SortPrice. sort.SliceStable is stable, matching
// Array.prototype.sort's stability guarantee in the TS original — rows
// comparing equal keep their input order.
//
// Every sort places nil/missing values last and breaks ties on id ascending,
// so output order is fully deterministic for a given model list.

package modelpicker

import (
	"slices"
	"sort"
	"strings"
)

// Tier is a model's qualitative class. Mirrors Tier in src/lib/ai/models.ts.
type Tier string

const (
	TierFlagship Tier = "flagship"
	TierBalanced Tier = "balanced"
	TierFast     Tier = "fast"
	TierBudget   Tier = "budget"
)

// Model is one catalog entry. Pointer fields are nil when the upstream value
// is null (unpriced / unknown), mirroring `number | null` in the TS lib.
type Model struct {
	ID              string   `json:"id"`
	Name            string   `json:"name"`
	Vendor          string   `json:"vendor"`
	Family          string   `json:"family"`
	Tier            Tier     `json:"tier"`
	ContextWindow   int      `json:"contextWindow"`
	MaxOutput       int      `json:"maxOutput"`
	InputPerM       *float64 `json:"inputPerM"`    // USD per 1M input tokens; nil = unpriced
	OutputPerM      *float64 `json:"outputPerM"`   // USD per 1M output tokens; nil = unpriced
	CacheReadPerM   *float64 `json:"cacheReadPerM"`
	CacheWritePerM  *float64 `json:"cacheWritePerM"`
	Modalities      []string `json:"modalities"`
	KnowledgeCutoff *string  `json:"knowledgeCutoff"`
	Released        *string  `json:"released"` // ISO date; nil = unknown
	OpenWeights     bool     `json:"openWeights"`
	Reasoning       bool     `json:"reasoning"`
	ToolCall        bool     `json:"toolCall"`
}

// Filter mirrors ModelFilter: nil pointer = filter not applied (TS undefined).
type Filter struct {
	Vendor       *string
	Tier         *Tier
	MinContext   *int    // minimum usable context window in tokens
	MaxInputPerM *float64 // maximum input price (USD per 1M); nil-priced models are skipped
	Modality     *string // required input modality, e.g. "image"
	Search       *string // case-insensitive substring match on id, name, and vendor
}

// SortKey selects the ranking. String values match the TS SortKey union
// ("price" | "context" | "tokensPerDollar" | "released"); SortPrice is the
// TS default sort.
type SortKey string

const (
	SortPrice           SortKey = "price"
	SortContext         SortKey = "context"
	SortTokensPerDollar SortKey = "tokensPerDollar"
	SortReleased        SortKey = "released"
)

// Preset couples a curated filter with the sort that makes that filter
// useful, mirroring the TS PRESETS record.
type Preset struct {
	Filter Filter
	Sort   SortKey
}

func intPtr(v int) *int                 { return &v }
func float64Ptr(v float64) *float64     { return &v }
func stringPtr(v string) *string        { return &v }
func tierPtr(v Tier) *Tier              { return &v }

// Presets mirrors the TS PRESETS record: long-context (minContext 500k,
// context sort), cheap-bulk (maxInputPerM 1, price sort), flagship
// (flagship tier, tokensPerDollar sort).
var Presets = map[string]Preset{
	"long-context": {Filter: Filter{MinContext: intPtr(500_000)}, Sort: SortContext},
	"cheap-bulk":   {Filter: Filter{MaxInputPerM: float64Ptr(1)}, Sort: SortPrice},
	"flagship":     {Filter: Filter{Tier: tierPtr(TierFlagship)}, Sort: SortTokensPerDollar},
}

// TokensPerDollar returns output tokens per USD (1e6 / outputPerM). The bool
// is false for unpriced models, mirroring the null return in the TS lib.
func TokensPerDollar(m Model) (float64, bool) {
	if m.OutputPerM == nil {
		return 0, false
	}
	return 1_000_000 / *m.OutputPerM, true
}

// ListModels filters the given list. Inline port of listModels() from
// src/lib/ai/models.ts: nil prices never satisfy MaxInputPerM, and search
// lowercases both sides before the substring check.
func ListModels(f *Filter, models []Model) []Model {
	var q string
	if f != nil && f.Search != nil {
		q = strings.ToLower(*f.Search)
	}
	out := make([]Model, 0, len(models))
	for _, m := range models {
		if f != nil {
			if f.Vendor != nil && m.Vendor != *f.Vendor {
				continue
			}
			if f.Tier != nil && m.Tier != *f.Tier {
				continue
			}
			if f.MinContext != nil && m.ContextWindow < *f.MinContext {
				continue
			}
			if f.MaxInputPerM != nil && (m.InputPerM == nil || *m.InputPerM > *f.MaxInputPerM) {
				continue
			}
			if f.Modality != nil && !slices.Contains(m.Modalities, *f.Modality) {
				continue
			}
			if f.Search != nil &&
				!strings.Contains(strings.ToLower(m.ID+" "+m.Name+" "+m.Vendor), q) {
				continue
			}
		}
		out = append(out, m)
	}
	return out
}

// byID is the shared stable tie-break for every sort: id ascending.
func byID(a, b Model) int {
	switch {
	case a.ID < b.ID:
		return -1
	case a.ID > b.ID:
		return 1
	default:
		return 0
	}
}

// comparePrice: input price ascending, unpriced last, id asc on ties.
func comparePrice(a, b Model) int {
	switch {
	case a.InputPerM == nil && b.InputPerM == nil:
		return byID(a, b)
	case a.InputPerM == nil:
		return 1
	case b.InputPerM == nil:
		return -1
	case *a.InputPerM != *b.InputPerM:
		if *a.InputPerM < *b.InputPerM {
			return -1
		}
		return 1
	default:
		return byID(a, b)
	}
}

// compareContext: context window descending, id asc on ties.
func compareContext(a, b Model) int {
	switch {
	case a.ContextWindow > b.ContextWindow:
		return -1
	case a.ContextWindow < b.ContextWindow:
		return 1
	default:
		return byID(a, b)
	}
}

// compareTokensPerDollar: 1e6/output price descending, unpriced last,
// id asc on ties.
func compareTokensPerDollar(a, b Model) int {
	ta, aOk := TokensPerDollar(a)
	tb, bOk := TokensPerDollar(b)
	switch {
	case !aOk && !bOk:
		return byID(a, b)
	case !aOk:
		return 1
	case !bOk:
		return -1
	case ta != tb:
		if ta > tb {
			return -1
		}
		return 1
	default:
		return byID(a, b)
	}
}

// compareReleased: newest first (ISO dates compare correctly as strings),
// unknown last, id asc on ties.
func compareReleased(a, b Model) int {
	switch {
	case a.Released == nil && b.Released == nil:
		return byID(a, b)
	case a.Released == nil:
		return 1
	case b.Released == nil:
		return -1
	case *a.Released != *b.Released:
		if *a.Released > *b.Released {
			return -1
		}
		return 1
	default:
		return byID(a, b)
	}
}

// comparators maps each sort key to its comparator: value order first,
// nil last, id asc on ties. Mirrors the TS COMPARATORS record.
var comparators = map[SortKey]func(a, b Model) int{
	SortPrice:           comparePrice,
	SortContext:         compareContext,
	SortTokensPerDollar: compareTokensPerDollar,
	SortReleased:        compareReleased,
}

// PickModels filters the given list, then sorts it. See comparators for the
// exact order. The TS original indexes a Record keyed by the sort union, where
// an unknown key is a type error; here an unknown SortKey string value falls
// back to the default price sort instead of the accidental lexicographic
// order a nil comparator would produce.
func PickModels(f *Filter, key SortKey, models []Model) []Model {
	ranked := ListModels(f, models)
	cmpFn, ok := comparators[key]
	if !ok {
		cmpFn = comparePrice
	}
	sort.SliceStable(ranked, func(i, j int) bool { return cmpFn(ranked[i], ranked[j]) < 0 })
	return ranked
}