Skip to content

Commit 8dfb25c

Browse files
authored
feat(catalog): centralize model policy helpers (#56)
* feat(catalog): centralize model policy helpers * feat(catalog): expose tier model preferences
1 parent 6706e73 commit 8dfb25c

2 files changed

Lines changed: 505 additions & 0 deletions

File tree

catalog/model_policy.go

Lines changed: 356 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,356 @@
1+
package catalog
2+
3+
import (
4+
"sort"
5+
"strings"
6+
7+
"github.com/GrayCodeAI/eyrie/catalog/registry"
8+
)
9+
10+
// ModelCostTier is a relative cost band for model selection.
11+
type ModelCostTier int
12+
13+
const (
14+
// CostTierCheap is appropriate for short, low-risk, or summarization work.
15+
CostTierCheap ModelCostTier = iota
16+
// CostTierMid is the default balanced cost band.
17+
CostTierMid
18+
// CostTierExpensive is appropriate for complex planning or generation work.
19+
CostTierExpensive
20+
)
21+
22+
// ModelRole identifies the purpose of a model in a multi-model workflow.
23+
type ModelRole string
24+
25+
const (
26+
ModelRolePlanner ModelRole = "planner"
27+
ModelRoleCoder ModelRole = "coder"
28+
ModelRoleReviewer ModelRole = "reviewer"
29+
ModelRoleCommit ModelRole = "commit"
30+
)
31+
32+
// ModelRoleAssignments maps workflow roles to concrete model IDs.
33+
type ModelRoleAssignments struct {
34+
Planner string `json:"planner,omitempty"`
35+
Coder string `json:"coder,omitempty"`
36+
Reviewer string `json:"reviewer,omitempty"`
37+
Commit string `json:"commit,omitempty"`
38+
}
39+
40+
// ProviderDefaultModelV1 returns the provider's first catalog model.
41+
func ProviderDefaultModelV1(compiled *CompiledCatalogV1, provider, fallback string) string {
42+
models := ModelEntriesForProvider(compiled, provider)
43+
if len(models) == 0 {
44+
return fallback
45+
}
46+
if model := strings.TrimSpace(models[0].ID); model != "" {
47+
return model
48+
}
49+
return fallback
50+
}
51+
52+
// PreferredProviderModelV1 returns the preferred model for provider and tier.
53+
func PreferredProviderModelV1(compiled *CompiledCatalogV1, provider string, tier ModelTier, fallback string) string {
54+
switch tier {
55+
case TierHaiku:
56+
return CheapestModelForProviderV1(compiled, provider, fallback)
57+
case TierOpus:
58+
return MostExpensiveModelForProviderV1(compiled, provider, fallback)
59+
default:
60+
return middleModelForProviderV1(compiled, provider, fallback)
61+
}
62+
}
63+
64+
// PreferredModelsForTierV1 returns unique preferred models for a tier, starting
65+
// with primaryProvider and then following the registry chat preference order.
66+
func PreferredModelsForTierV1(compiled *CompiledCatalogV1, primaryProvider string, tier ModelTier, limit int) []string {
67+
seenProviders := map[string]bool{}
68+
seenModels := map[string]bool{}
69+
providers := make([]string, 0, len(registry.ChatProviderPreferenceOrder())+1)
70+
if primary := CanonicalProviderID(strings.TrimSpace(primaryProvider)); primary != "" {
71+
providers = append(providers, primary)
72+
}
73+
providers = append(providers, registry.ChatProviderPreferenceOrder()...)
74+
providers = append(providers, AllModelProvidersV1(compiled)...)
75+
76+
models := make([]string, 0, len(providers))
77+
for _, provider := range providers {
78+
provider = CanonicalProviderID(strings.TrimSpace(provider))
79+
if provider == "" || seenProviders[provider] {
80+
continue
81+
}
82+
seenProviders[provider] = true
83+
model := PreferredProviderModelV1(compiled, provider, tier, "")
84+
if model == "" || seenModels[model] {
85+
continue
86+
}
87+
seenModels[model] = true
88+
models = append(models, model)
89+
if limit > 0 && len(models) >= limit {
90+
return models
91+
}
92+
}
93+
return models
94+
}
95+
96+
// CheapestModelForProviderV1 returns the lowest known input-priced model for a provider.
97+
func CheapestModelForProviderV1(compiled *CompiledCatalogV1, provider, fallback string) string {
98+
models := ModelEntriesForProvider(compiled, provider)
99+
if len(models) == 0 {
100+
return fallback
101+
}
102+
best := models[0]
103+
foundPriced := false
104+
for _, model := range models {
105+
if strings.TrimSpace(model.ID) == "" {
106+
continue
107+
}
108+
if model.InputPricePer1M <= 0 {
109+
continue
110+
}
111+
if !foundPriced || model.InputPricePer1M < best.InputPricePer1M {
112+
best = model
113+
foundPriced = true
114+
}
115+
}
116+
if id := strings.TrimSpace(best.ID); id != "" {
117+
return id
118+
}
119+
return fallback
120+
}
121+
122+
// MostExpensiveModelForProviderV1 returns the highest known input-priced model for a provider.
123+
func MostExpensiveModelForProviderV1(compiled *CompiledCatalogV1, provider, fallback string) string {
124+
models := ModelEntriesForProvider(compiled, provider)
125+
if len(models) == 0 {
126+
return fallback
127+
}
128+
best := models[0]
129+
for _, model := range models[1:] {
130+
if strings.TrimSpace(model.ID) == "" {
131+
continue
132+
}
133+
if model.InputPricePer1M > best.InputPricePer1M {
134+
best = model
135+
}
136+
}
137+
if id := strings.TrimSpace(best.ID); id != "" {
138+
return id
139+
}
140+
return fallback
141+
}
142+
143+
// ModelCostTierOf resolves a model's cost tier from catalog family and pricing data.
144+
func ModelCostTierOf(compiled *CompiledCatalogV1, modelName string) ModelCostTier {
145+
if tier, ok := costTierFromCatalogFamily(compiled, modelName); ok {
146+
return mapModelTierToCostTier(tier)
147+
}
148+
if tier, ok := costTierFromCatalogPricing(compiled, modelName); ok {
149+
return tier
150+
}
151+
return costTierFromModelName(modelName)
152+
}
153+
154+
// ProviderForModelV1 returns the canonical owner provider for modelName.
155+
func ProviderForModelV1(compiled *CompiledCatalogV1, modelName string) string {
156+
if compiled == nil {
157+
return ""
158+
}
159+
canonical := strings.TrimSpace(modelName)
160+
if resolved, ok := compiled.CanonicalModelForAliasOrID(canonical); ok {
161+
canonical = resolved
162+
}
163+
model := compiled.ModelsByID[canonical]
164+
return CanonicalProviderID(model.ProviderID)
165+
}
166+
167+
// AllModelProvidersV1 lists canonical model owner providers in the catalog.
168+
func AllModelProvidersV1(compiled *CompiledCatalogV1) []string {
169+
if compiled == nil {
170+
return nil
171+
}
172+
seen := map[string]bool{}
173+
providers := make([]string, 0, len(compiled.ModelsByID))
174+
for _, model := range compiled.ModelsByID {
175+
provider := CanonicalProviderID(model.ProviderID)
176+
if provider == "" || seen[provider] {
177+
continue
178+
}
179+
seen[provider] = true
180+
providers = append(providers, provider)
181+
}
182+
sort.Strings(providers)
183+
return providers
184+
}
185+
186+
// DefaultModelRolesV1 uses the primary model for interactive roles and the cheapest
187+
// same-provider model for commit/summarization work when catalog data is available.
188+
func DefaultModelRolesV1(compiled *CompiledCatalogV1, primaryModel string) ModelRoleAssignments {
189+
primaryModel = strings.TrimSpace(primaryModel)
190+
commit := primaryModel
191+
if provider := ProviderForModelV1(compiled, primaryModel); provider != "" {
192+
commit = CheapestModelForProviderV1(compiled, provider, primaryModel)
193+
}
194+
return ModelRoleAssignments{
195+
Planner: primaryModel,
196+
Coder: primaryModel,
197+
Reviewer: primaryModel,
198+
Commit: commit,
199+
}
200+
}
201+
202+
// ModelForRoleV1 resolves a role assignment with coder then catalog fallback.
203+
func ModelForRoleV1(compiled *CompiledCatalogV1, roles ModelRoleAssignments, role ModelRole) string {
204+
var model string
205+
switch role {
206+
case ModelRolePlanner:
207+
model = roles.Planner
208+
case ModelRoleCoder:
209+
model = roles.Coder
210+
case ModelRoleReviewer:
211+
model = roles.Reviewer
212+
case ModelRoleCommit:
213+
model = roles.Commit
214+
}
215+
if strings.TrimSpace(model) != "" {
216+
return model
217+
}
218+
if strings.TrimSpace(roles.Coder) != "" {
219+
return roles.Coder
220+
}
221+
return PrimaryModelV1(compiled)
222+
}
223+
224+
// PrimaryModelV1 returns a stable best-effort model from chat-preferred providers.
225+
func PrimaryModelV1(compiled *CompiledCatalogV1) string {
226+
for _, provider := range registry.ChatProviderPreferenceOrder() {
227+
if model := ProviderDefaultModelV1(compiled, provider, ""); model != "" {
228+
return model
229+
}
230+
}
231+
for _, provider := range AllModelProvidersV1(compiled) {
232+
if model := ProviderDefaultModelV1(compiled, provider, ""); model != "" {
233+
return model
234+
}
235+
}
236+
return ""
237+
}
238+
239+
func middleModelForProviderV1(compiled *CompiledCatalogV1, provider, fallback string) string {
240+
models := ModelEntriesForProvider(compiled, provider)
241+
if len(models) == 0 {
242+
return fallback
243+
}
244+
priced := make([]ModelCatalogEntry, 0, len(models))
245+
for _, model := range models {
246+
if strings.TrimSpace(model.ID) != "" && model.InputPricePer1M > 0 {
247+
priced = append(priced, model)
248+
}
249+
}
250+
if len(priced) == 0 {
251+
return ProviderDefaultModelV1(compiled, provider, fallback)
252+
}
253+
sort.SliceStable(priced, func(i, j int) bool {
254+
if priced[i].InputPricePer1M == priced[j].InputPricePer1M {
255+
return priced[i].ID < priced[j].ID
256+
}
257+
return priced[i].InputPricePer1M < priced[j].InputPricePer1M
258+
})
259+
return priced[len(priced)/2].ID
260+
}
261+
262+
func costTierFromCatalogFamily(compiled *CompiledCatalogV1, modelName string) (ModelTier, bool) {
263+
if compiled == nil {
264+
return "", false
265+
}
266+
canonical := strings.TrimSpace(modelName)
267+
if resolved, ok := compiled.CanonicalModelForAliasOrID(canonical); ok {
268+
canonical = resolved
269+
}
270+
model := compiled.ModelsByID[canonical]
271+
if model.ID == "" {
272+
return "", false
273+
}
274+
switch strings.ToLower(strings.TrimSpace(model.Family)) {
275+
case "haiku", "cheap", "lite", "flash", "mini":
276+
return TierHaiku, true
277+
case "opus", "pro", "max", "heavy", "ultra":
278+
return TierOpus, true
279+
case "sonnet", "standard", "balanced", "medium":
280+
return TierSonnet, true
281+
default:
282+
return "", false
283+
}
284+
}
285+
286+
func costTierFromCatalogPricing(compiled *CompiledCatalogV1, modelName string) (ModelCostTier, bool) {
287+
if compiled == nil {
288+
return 0, false
289+
}
290+
canonical := strings.TrimSpace(modelName)
291+
if resolved, ok := compiled.CanonicalModelForAliasOrID(canonical); ok {
292+
canonical = resolved
293+
}
294+
model := compiled.ModelsByID[canonical]
295+
if model.ID == "" {
296+
return 0, false
297+
}
298+
offering := firstOfferingForModel(compiled, canonical)
299+
price := offering.Pricing.RatesPer1M["input_tokens"]
300+
if price <= 0 {
301+
return 0, false
302+
}
303+
models := ModelEntriesForProvider(compiled, model.ProviderID)
304+
prices := make([]float64, 0, len(models))
305+
seen := map[float64]bool{}
306+
for _, candidate := range models {
307+
if candidate.InputPricePer1M <= 0 || seen[candidate.InputPricePer1M] {
308+
continue
309+
}
310+
seen[candidate.InputPricePer1M] = true
311+
prices = append(prices, candidate.InputPricePer1M)
312+
}
313+
if len(prices) < 2 {
314+
return 0, false
315+
}
316+
sort.Float64s(prices)
317+
switch {
318+
case price <= prices[0]:
319+
return CostTierCheap, true
320+
case price >= prices[len(prices)-1]:
321+
return CostTierExpensive, true
322+
default:
323+
return CostTierMid, true
324+
}
325+
}
326+
327+
func mapModelTierToCostTier(tier ModelTier) ModelCostTier {
328+
switch tier {
329+
case TierHaiku:
330+
return CostTierCheap
331+
case TierOpus:
332+
return CostTierExpensive
333+
default:
334+
return CostTierMid
335+
}
336+
}
337+
338+
func costTierFromModelName(modelName string) ModelCostTier {
339+
lower := strings.ToLower(strings.TrimSpace(modelName))
340+
for _, pattern := range cheapModelNamePatterns {
341+
if strings.Contains(lower, pattern) {
342+
return CostTierCheap
343+
}
344+
}
345+
for _, pattern := range expensiveModelNamePatterns {
346+
if strings.Contains(lower, pattern) {
347+
return CostTierExpensive
348+
}
349+
}
350+
return CostTierMid
351+
}
352+
353+
var (
354+
cheapModelNamePatterns = []string{"haiku", "mini", "flash", "lite", "nano", "micro", "small", "tiny"}
355+
expensiveModelNamePatterns = []string{"opus", "pro", "max", "ultra", "heavy", "large", "o1", "o3"}
356+
)

0 commit comments

Comments
 (0)