6a7d5c5 · crush · git.iko.soy

git.iko.soy/crush

chore: add default models handling and tests

commit6a7d5c50e23faf7f27672262d08ccf515a4da00e authorKujtim Hoxha <kujtimii.h@gmail.com> dateSat 5 Jul 2025, 11:11 (1 year ago) parente157024

4 files changed+562−37

compare with main

unifiedsplit

pkg/config/config.go+11 pkg/config/load.go+161−37 pkg/config/load_test.go+387 pkg/env/env.go+3

pkg/config/config.go+11

@@ -195,3 +195,14 @@ func (c *Config) EnabledProviders() []ProviderConfig {

195195func (c *Config) IsConfigured() bool {

196196 return len(c.EnabledProviders()) > 0

197197}

198+

199+func (c *Config) GetModel(provider, model string) *provider.Model {

200+ if providerConfig, ok := c.Providers[provider]; ok {

201+ for _, m := range providerConfig.Models {

202+ if m.ID == model {

203+ return &m

204+ }

205+ }

206+ }

207+ return nil

208+}

pkg/config/load.go+161−37

@@ -77,15 +77,10 @@ func Load(workingDir string, debug bool) (*Config, error) {

7777 return cfg, nil

7878 }

7979

80− // largeModel, ok := cfg.Models[SelectedModelTypeLarge]

81− // if !ok {

82− // // set default

83− // }

84− // smallModel, ok := cfg.Models[SelectedModelTypeSmall]

85− // if !ok {

86− // // set default

87− // }

88− //

80+ if err := cfg.configureSelectedModels(providers); err != nil {

81+ return nil, fmt.Errorf("failed to configure selected models: %w", err)

82+ }

83+

8984 return cfg, nil

9085}

9186

@@ -236,34 +231,6 @@ func (cfg *Config) configureProviders(env env.Env, resolver VariableResolver, kn

236231 return nil

237232}

238233

239−func hasVertexCredentials(env env.Env) bool {

240− useVertex := env.Get("GOOGLE_GENAI_USE_VERTEXAI") == "true"

241− hasProject := env.Get("GOOGLE_CLOUD_PROJECT") != ""

242− hasLocation := env.Get("GOOGLE_CLOUD_LOCATION") != ""

243− return useVertex && hasProject && hasLocation

244−}

245−

246−func hasAWSCredentials(env env.Env) bool {

247− if env.Get("AWS_ACCESS_KEY_ID") != "" && env.Get("AWS_SECRET_ACCESS_KEY") != "" {

248− return true

249− }

250−

251− if env.Get("AWS_PROFILE") != "" || env.Get("AWS_DEFAULT_PROFILE") != "" {

252− return true

253− }

254−

255− if env.Get("AWS_REGION") != "" || env.Get("AWS_DEFAULT_REGION") != "" {

256− return true

257− }

258−

259− if env.Get("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI") != "" ||

260− env.Get("AWS_CONTAINER_CREDENTIALS_FULL_URI") != "" {

261− return true

262− }

263−

264− return false

265−}

266−

267234func (cfg *Config) setDefaults(workingDir string) {

268235 cfg.workingDir = workingDir

269236 if cfg.Options == nil {

@@ -297,6 +264,135 @@ func (cfg *Config) setDefaults(workingDir string) {

297264 cfg.Options.ContextPaths = slices.Compact(cfg.Options.ContextPaths)

298265}

299266

267+func (cfg *Config) defaultModelSelection(knownProviders []provider.Provider) (largeModel SelectedModel, smallModel SelectedModel, err error) {

268+ if len(knownProviders) == 0 && len(cfg.Providers) == 0 {

269+ err = fmt.Errorf("no providers configured, please configure at least one provider")

270+ return

271+ }

272+

273+ // Use the first provider enabled based on the known providers order

274+ // if no provider found that is known use the first provider configured

275+ for _, p := range knownProviders {

276+ providerConfig, ok := cfg.Providers[string(p.ID)]

277+ if !ok || providerConfig.Disable {

278+ continue

279+ }

280+ defaultLargeModel := cfg.GetModel(string(p.ID), p.DefaultLargeModelID)

281+ if defaultLargeModel == nil {

282+ err = fmt.Errorf("default large model %s not found for provider %s", p.DefaultLargeModelID, p.ID)

283+ return

284+ }

285+ largeModel = SelectedModel{

286+ Provider: string(p.ID),

287+ Model: defaultLargeModel.ID,

288+ MaxTokens: defaultLargeModel.DefaultMaxTokens,

289+ }

290+

291+ defaultSmallModel := cfg.GetModel(string(p.ID), p.DefaultSmallModelID)

292+ if defaultSmallModel == nil {

293+ err = fmt.Errorf("default small model %s not found for provider %s", p.DefaultSmallModelID, p.ID)

294+ }

295+ smallModel = SelectedModel{

296+ Provider: string(p.ID),

297+ Model: defaultSmallModel.ID,

298+ MaxTokens: defaultSmallModel.DefaultMaxTokens,

299+ }

300+ return

301+ }

302+

303+ enabledProviders := cfg.EnabledProviders()

304+ slices.SortFunc(enabledProviders, func(a, b ProviderConfig) int {

305+ return strings.Compare(a.ID, b.ID)

306+ })

307+

308+ if len(enabledProviders) == 0 {

309+ err = fmt.Errorf("no providers configured, please configure at least one provider")

310+ return

311+ }

312+

313+ providerConfig := enabledProviders[0]

314+ if len(providerConfig.Models) == 0 {

315+ err = fmt.Errorf("provider %s has no models configured", providerConfig.ID)

316+ return

317+ }

318+ defaultLargeModel := cfg.GetModel(providerConfig.ID, providerConfig.Models[0].ID)

319+ largeModel = SelectedModel{

320+ Provider: providerConfig.ID,

321+ Model: defaultLargeModel.ID,

322+ MaxTokens: defaultLargeModel.DefaultMaxTokens,

323+ }

324+ defaultSmallModel := cfg.GetModel(providerConfig.ID, providerConfig.Models[0].ID)

325+ smallModel = SelectedModel{

326+ Provider: providerConfig.ID,

327+ Model: defaultSmallModel.ID,

328+ MaxTokens: defaultSmallModel.DefaultMaxTokens,

329+ }

330+ return

331+}

332+

333+func (cfg *Config) configureSelectedModels(knownProviders []provider.Provider) error {

334+ large, small, err := cfg.defaultModelSelection(knownProviders)

335+ if err != nil {

336+ return fmt.Errorf("failed to select default models: %w", err)

337+ }

338+

339+ largeModelSelected, largeModelConfigured := cfg.Models[SelectedModelTypeLarge]

340+ if largeModelConfigured {

341+ if largeModelSelected.Model != "" {

342+ large.Model = largeModelSelected.Model

343+ }

344+ if largeModelSelected.Provider != "" {

345+ large.Provider = largeModelSelected.Provider

346+ }

347+ model := cfg.GetModel(large.Provider, large.Model)

348+ if model == nil {

349+ return fmt.Errorf("large model %s not found for provider %s", large.Model, large.Provider)

350+ }

351+ if largeModelSelected.MaxTokens > 0 {

352+ large.MaxTokens = largeModelSelected.MaxTokens

353+ } else {

354+ large.MaxTokens = model.DefaultMaxTokens

355+ }

356+ large.ReasoningEffort = largeModelSelected.ReasoningEffort

357+ large.Think = largeModelSelected.Think

358+

359+ }

360+ smallModelSelected, smallModelConfigured := cfg.Models[SelectedModelTypeSmall]

361+ if smallModelConfigured {

362+ if smallModelSelected.Model != "" {

363+ small.Model = smallModelSelected.Model

364+ }

365+ if smallModelSelected.Provider != "" {

366+ small.Provider = smallModelSelected.Provider

367+ }

368+

369+ model := cfg.GetModel(small.Provider, small.Model)

370+ if model == nil {

371+ return fmt.Errorf("large model %s not found for provider %s", large.Model, large.Provider)

372+ }

373+ if smallModelSelected.MaxTokens > 0 {

374+ small.MaxTokens = smallModelSelected.MaxTokens

375+ } else {

376+ small.MaxTokens = model.DefaultMaxTokens

377+ }

378+ small.ReasoningEffort = smallModelSelected.ReasoningEffort

379+ small.Think = smallModelSelected.Think

380+ }

381+

382+ // validate the selected models

383+ largeModel := cfg.GetModel(large.Provider, large.Model)

384+ if largeModel == nil {

385+ return fmt.Errorf("large model %s not found for provider %s", large.Model, large.Provider)

386+ }

387+ smallModel := cfg.GetModel(small.Provider, small.Model)

388+ if smallModel == nil {

389+ return fmt.Errorf("small model %s not found for provider %s", small.Model, small.Provider)

390+ }

391+ cfg.Models[SelectedModelTypeLarge] = large

392+ cfg.Models[SelectedModelTypeSmall] = small

393+ return nil

394+}

395+

300396func loadFromConfigPaths(configPaths []string) (*Config, error) {

301397 var configs []io.Reader

302398

@@ -329,6 +425,34 @@ func loadFromReaders(readers []io.Reader) (*Config, error) {

329425 return LoadReader(merged)

330426}

331427

428+func hasVertexCredentials(env env.Env) bool {

429+ useVertex := env.Get("GOOGLE_GENAI_USE_VERTEXAI") == "true"

430+ hasProject := env.Get("GOOGLE_CLOUD_PROJECT") != ""

431+ hasLocation := env.Get("GOOGLE_CLOUD_LOCATION") != ""

432+ return useVertex && hasProject && hasLocation

433+}

434+

435+func hasAWSCredentials(env env.Env) bool {

436+ if env.Get("AWS_ACCESS_KEY_ID") != "" && env.Get("AWS_SECRET_ACCESS_KEY") != "" {

437+ return true

438+ }

439+

440+ if env.Get("AWS_PROFILE") != "" || env.Get("AWS_DEFAULT_PROFILE") != "" {

441+ return true

442+ }

443+

444+ if env.Get("AWS_REGION") != "" || env.Get("AWS_DEFAULT_REGION") != "" {

445+ return true

446+ }

447+

448+ if env.Get("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI") != "" ||

449+ env.Get("AWS_CONTAINER_CREDENTIALS_FULL_URI") != "" {

450+ return true

451+ }

452+

453+ return false

454+}

455+

332456func globalConfig() string {

333457 xdgConfigHome := os.Getenv("XDG_CONFIG_HOME")

334458 if xdgConfigHome != "" {

pkg/config/load_test.go+387

@@ -760,3 +760,390 @@ func TestConfig_configureProvidersEnhancedCredentialValidation(t *testing.T) {

760760 assert.True(t, exists)

761761 })

762762}

763+

764+func TestConfig_defaultModelSelection(t *testing.T) {

765+ t.Run("default behavior uses the default models for given provider", func(t *testing.T) {

766+ knownProviders := []provider.Provider{

767+ {

768+ ID: "openai",

769+ APIKey: "abc",

770+ DefaultLargeModelID: "large-model",

771+ DefaultSmallModelID: "small-model",

772+ Models: []provider.Model{

773+ {

774+ ID: "large-model",

775+ DefaultMaxTokens: 1000,

776+ },

777+ {

778+ ID: "small-model",

779+ DefaultMaxTokens: 500,

780+ },

781+ },

782+ },

783+ }

784+

785+ cfg := &Config{}

786+ cfg.setDefaults("/tmp")

787+ env := env.NewFromMap(map[string]string{})

788+ resolver := NewEnvironmentVariableResolver(env)

789+ err := cfg.configureProviders(env, resolver, knownProviders)

790+ assert.NoError(t, err)

791+

792+ large, small, err := cfg.defaultModelSelection(knownProviders)

793+ assert.NoError(t, err)

794+ assert.Equal(t, "large-model", large.Model)

795+ assert.Equal(t, "openai", large.Provider)

796+ assert.Equal(t, int64(1000), large.MaxTokens)

797+ assert.Equal(t, "small-model", small.Model)

798+ assert.Equal(t, "openai", small.Provider)

799+ assert.Equal(t, int64(500), small.MaxTokens)

800+ })

801+ t.Run("should error if no providers configured", func(t *testing.T) {

802+ knownProviders := []provider.Provider{

803+ {

804+ ID: "openai",

805+ APIKey: "$MISSING_KEY",

806+ DefaultLargeModelID: "large-model",

807+ DefaultSmallModelID: "small-model",

808+ Models: []provider.Model{

809+ {

810+ ID: "large-model",

811+ DefaultMaxTokens: 1000,

812+ },

813+ {

814+ ID: "small-model",

815+ DefaultMaxTokens: 500,

816+ },

817+ },

818+ },

819+ }

820+

821+ cfg := &Config{}

822+ cfg.setDefaults("/tmp")

823+ env := env.NewFromMap(map[string]string{})

824+ resolver := NewEnvironmentVariableResolver(env)

825+ err := cfg.configureProviders(env, resolver, knownProviders)

826+ assert.NoError(t, err)

827+

828+ _, _, err = cfg.defaultModelSelection(knownProviders)

829+ assert.Error(t, err)

830+ })

831+ t.Run("should error if model is missing", func(t *testing.T) {

832+ knownProviders := []provider.Provider{

833+ {

834+ ID: "openai",

835+ APIKey: "abc",

836+ DefaultLargeModelID: "large-model",

837+ DefaultSmallModelID: "small-model",

838+ Models: []provider.Model{

839+ {

840+ ID: "not-large-model",

841+ DefaultMaxTokens: 1000,

842+ },

843+ {

844+ ID: "small-model",

845+ DefaultMaxTokens: 500,

846+ },

847+ },

848+ },

849+ }

850+

851+ cfg := &Config{}

852+ cfg.setDefaults("/tmp")

853+ env := env.NewFromMap(map[string]string{})

854+ resolver := NewEnvironmentVariableResolver(env)

855+ err := cfg.configureProviders(env, resolver, knownProviders)

856+ assert.NoError(t, err)

857+ _, _, err = cfg.defaultModelSelection(knownProviders)

858+ assert.Error(t, err)

859+ })

860+

861+ t.Run("should configure the default models with a custom provider", func(t *testing.T) {

862+ knownProviders := []provider.Provider{

863+ {

864+ ID: "openai",

865+ APIKey: "$MISSING", // will not be included in the config

866+ DefaultLargeModelID: "large-model",

867+ DefaultSmallModelID: "small-model",

868+ Models: []provider.Model{

869+ {

870+ ID: "not-large-model",

871+ DefaultMaxTokens: 1000,

872+ },

873+ {

874+ ID: "small-model",

875+ DefaultMaxTokens: 500,

876+ },

877+ },

878+ },

879+ }

880+

881+ cfg := &Config{

882+ Providers: map[string]ProviderConfig{

883+ "custom": {

884+ APIKey: "test-key",

885+ BaseURL: "https://api.custom.com/v1",

886+ Models: []provider.Model{

887+ {

888+ ID: "model",

889+ DefaultMaxTokens: 600,

890+ },

891+ },

892+ },

893+ },

894+ }

895+ cfg.setDefaults("/tmp")

896+ env := env.NewFromMap(map[string]string{})

897+ resolver := NewEnvironmentVariableResolver(env)

898+ err := cfg.configureProviders(env, resolver, knownProviders)

899+ assert.NoError(t, err)

900+ large, small, err := cfg.defaultModelSelection(knownProviders)

901+ assert.NoError(t, err)

902+ assert.Equal(t, "model", large.Model)

903+ assert.Equal(t, "custom", large.Provider)

904+ assert.Equal(t, int64(600), large.MaxTokens)

905+ assert.Equal(t, "model", small.Model)

906+ assert.Equal(t, "custom", small.Provider)

907+ assert.Equal(t, int64(600), small.MaxTokens)

908+ })

909+

910+ t.Run("should fail if no model configured", func(t *testing.T) {

911+ knownProviders := []provider.Provider{

912+ {

913+ ID: "openai",

914+ APIKey: "$MISSING", // will not be included in the config

915+ DefaultLargeModelID: "large-model",

916+ DefaultSmallModelID: "small-model",

917+ Models: []provider.Model{

918+ {

919+ ID: "not-large-model",

920+ DefaultMaxTokens: 1000,

921+ },

922+ {

923+ ID: "small-model",

924+ DefaultMaxTokens: 500,

925+ },

926+ },

927+ },

928+ }

929+

930+ cfg := &Config{

931+ Providers: map[string]ProviderConfig{

932+ "custom": {

933+ APIKey: "test-key",

934+ BaseURL: "https://api.custom.com/v1",

935+ Models: []provider.Model{},

936+ },

937+ },

938+ }

939+ cfg.setDefaults("/tmp")

940+ env := env.NewFromMap(map[string]string{})

941+ resolver := NewEnvironmentVariableResolver(env)

942+ err := cfg.configureProviders(env, resolver, knownProviders)

943+ assert.NoError(t, err)

944+ _, _, err = cfg.defaultModelSelection(knownProviders)

945+ assert.Error(t, err)

946+ })

947+ t.Run("should use the default provider first", func(t *testing.T) {

948+ knownProviders := []provider.Provider{

949+ {

950+ ID: "openai",

951+ APIKey: "set",

952+ DefaultLargeModelID: "large-model",

953+ DefaultSmallModelID: "small-model",

954+ Models: []provider.Model{

955+ {

956+ ID: "large-model",

957+ DefaultMaxTokens: 1000,

958+ },

959+ {

960+ ID: "small-model",

961+ DefaultMaxTokens: 500,

962+ },

963+ },

964+ },

965+ }

966+

967+ cfg := &Config{

968+ Providers: map[string]ProviderConfig{

969+ "custom": {

970+ APIKey: "test-key",

971+ BaseURL: "https://api.custom.com/v1",

972+ Models: []provider.Model{

973+ {

974+ ID: "large-model",

975+ DefaultMaxTokens: 1000,

976+ },

977+ },

978+ },

979+ },

980+ }

981+ cfg.setDefaults("/tmp")

982+ env := env.NewFromMap(map[string]string{})

983+ resolver := NewEnvironmentVariableResolver(env)

984+ err := cfg.configureProviders(env, resolver, knownProviders)

985+ assert.NoError(t, err)

986+ large, small, err := cfg.defaultModelSelection(knownProviders)

987+ assert.NoError(t, err)

988+ assert.Equal(t, "large-model", large.Model)

989+ assert.Equal(t, "openai", large.Provider)

990+ assert.Equal(t, int64(1000), large.MaxTokens)

991+ assert.Equal(t, "small-model", small.Model)

992+ assert.Equal(t, "openai", small.Provider)

993+ assert.Equal(t, int64(500), small.MaxTokens)

994+ })

995+}

996+

997+func TestConfig_configureSelectedModels(t *testing.T) {

998+ t.Run("should override defaults", func(t *testing.T) {

999+ knownProviders := []provider.Provider{

1000+ {

1001+ ID: "openai",

1002+ APIKey: "abc",

1003+ DefaultLargeModelID: "large-model",

1004+ DefaultSmallModelID: "small-model",

1005+ Models: []provider.Model{

1006+ {

1007+ ID: "larger-model",

1008+ DefaultMaxTokens: 2000,

1009+ },

1010+ {

1011+ ID: "large-model",

1012+ DefaultMaxTokens: 1000,

1013+ },

1014+ {

1015+ ID: "small-model",

1016+ DefaultMaxTokens: 500,

1017+ },

1018+ },

1019+ },

1020+ }

1021+

1022+ cfg := &Config{

1023+ Models: map[SelectedModelType]SelectedModel{

1024+ "large": {

1025+ Model: "larger-model",

1026+ },

1027+ },

1028+ }

1029+ cfg.setDefaults("/tmp")

1030+ env := env.NewFromMap(map[string]string{})

1031+ resolver := NewEnvironmentVariableResolver(env)

1032+ err := cfg.configureProviders(env, resolver, knownProviders)

1033+ assert.NoError(t, err)

1034+

1035+ err = cfg.configureSelectedModels(knownProviders)

1036+ assert.NoError(t, err)

1037+ large := cfg.Models[SelectedModelTypeLarge]

1038+ small := cfg.Models[SelectedModelTypeSmall]

1039+ assert.Equal(t, "larger-model", large.Model)

1040+ assert.Equal(t, "openai", large.Provider)

1041+ assert.Equal(t, int64(2000), large.MaxTokens)

1042+ assert.Equal(t, "small-model", small.Model)

1043+ assert.Equal(t, "openai", small.Provider)

1044+ assert.Equal(t, int64(500), small.MaxTokens)

1045+ })

1046+ t.Run("should be possible to use multiple providers", func(t *testing.T) {

1047+ knownProviders := []provider.Provider{

1048+ {

1049+ ID: "openai",

1050+ APIKey: "abc",

1051+ DefaultLargeModelID: "large-model",

1052+ DefaultSmallModelID: "small-model",

1053+ Models: []provider.Model{

1054+ {

1055+ ID: "large-model",

1056+ DefaultMaxTokens: 1000,

1057+ },

1058+ {

1059+ ID: "small-model",

1060+ DefaultMaxTokens: 500,

1061+ },

1062+ },

1063+ },

1064+ {

1065+ ID: "anthropic",

1066+ APIKey: "abc",

1067+ DefaultLargeModelID: "a-large-model",

1068+ DefaultSmallModelID: "a-small-model",

1069+ Models: []provider.Model{

1070+ {

1071+ ID: "a-large-model",

1072+ DefaultMaxTokens: 1000,

1073+ },

1074+ {

1075+ ID: "a-small-model",

1076+ DefaultMaxTokens: 200,

1077+ },

1078+ },

1079+ },

1080+ }

1081+

1082+ cfg := &Config{

1083+ Models: map[SelectedModelType]SelectedModel{

1084+ "small": {

1085+ Model: "a-small-model",

1086+ Provider: "anthropic",

1087+ MaxTokens: 300,

1088+ },

1089+ },

1090+ }

1091+ cfg.setDefaults("/tmp")

1092+ env := env.NewFromMap(map[string]string{})

1093+ resolver := NewEnvironmentVariableResolver(env)

1094+ err := cfg.configureProviders(env, resolver, knownProviders)

1095+ assert.NoError(t, err)

1096+

1097+ err = cfg.configureSelectedModels(knownProviders)

1098+ assert.NoError(t, err)

1099+ large := cfg.Models[SelectedModelTypeLarge]

1100+ small := cfg.Models[SelectedModelTypeSmall]

1101+ assert.Equal(t, "large-model", large.Model)

1102+ assert.Equal(t, "openai", large.Provider)

1103+ assert.Equal(t, int64(1000), large.MaxTokens)

1104+ assert.Equal(t, "a-small-model", small.Model)

1105+ assert.Equal(t, "anthropic", small.Provider)

1106+ assert.Equal(t, int64(300), small.MaxTokens)

1107+ })

1108+

1109+ t.Run("should override the max tokens only", func(t *testing.T) {

1110+ knownProviders := []provider.Provider{

1111+ {

1112+ ID: "openai",

1113+ APIKey: "abc",

1114+ DefaultLargeModelID: "large-model",

1115+ DefaultSmallModelID: "small-model",

1116+ Models: []provider.Model{

1117+ {

1118+ ID: "large-model",

1119+ DefaultMaxTokens: 1000,

1120+ },

1121+ {

1122+ ID: "small-model",

1123+ DefaultMaxTokens: 500,

1124+ },

1125+ },

1126+ },

1127+ }

1128+

1129+ cfg := &Config{

1130+ Models: map[SelectedModelType]SelectedModel{

1131+ "large": {

1132+ MaxTokens: 100,

1133+ },

1134+ },

1135+ }

1136+ cfg.setDefaults("/tmp")

1137+ env := env.NewFromMap(map[string]string{})

1138+ resolver := NewEnvironmentVariableResolver(env)

1139+ err := cfg.configureProviders(env, resolver, knownProviders)

1140+ assert.NoError(t, err)

1141+

1142+ err = cfg.configureSelectedModels(knownProviders)

1143+ assert.NoError(t, err)

1144+ large := cfg.Models[SelectedModelTypeLarge]

1145+ assert.Equal(t, "large-model", large.Model)

1146+ assert.Equal(t, "openai", large.Provider)

1147+ assert.Equal(t, int64(100), large.MaxTokens)

1148+ })

1149+}

pkg/env/env.go+3

@@ -51,5 +51,8 @@ func (m *mapEnv) Env() []string {

5151}

5252

5353func NewFromMap(m map[string]string) Env {

54+ if m == nil {

55+ m = make(map[string]string)

56+ }

5457 return &mapEnv{m: m}

5558}