1package models
2
3import (
4 "fmt"
5 "slices"
6 "strings"
7
8 tea "github.com/charmbracelet/bubbletea/v2"
9 "github.com/charmbracelet/catwalk/pkg/catwalk"
10 "github.com/charmbracelet/crush/internal/config"
11 "github.com/charmbracelet/crush/internal/tui/exp/list"
12 "github.com/charmbracelet/crush/internal/tui/styles"
13 "github.com/charmbracelet/crush/internal/tui/util"
14)
15
16type listModel = list.FilterableGroupList[list.CompletionItem[ModelOption]]
17
18type ModelListComponent struct {
19 list listModel
20 modelType int
21 providers []catwalk.Provider
22}
23
24func NewModelListComponent(keyMap list.KeyMap, inputPlaceholder string, shouldResize bool) *ModelListComponent {
25 t := styles.CurrentTheme()
26 inputStyle := t.S().Base.PaddingLeft(1).PaddingBottom(1)
27 options := []list.ListOption{
28 list.WithKeyMap(keyMap),
29 list.WithWrapNavigation(),
30 }
31 if shouldResize {
32 options = append(options, list.WithResizeByList())
33 }
34 modelList := list.NewFilterableGroupedList(
35 []list.Group[list.CompletionItem[ModelOption]]{},
36 list.WithFilterInputStyle(inputStyle),
37 list.WithFilterPlaceholder(inputPlaceholder),
38 list.WithFilterListOptions(
39 options...,
40 ),
41 )
42
43 return &ModelListComponent{
44 list: modelList,
45 modelType: LargeModelType,
46 }
47}
48
49func (m *ModelListComponent) Init() tea.Cmd {
50 var cmds []tea.Cmd
51 if len(m.providers) == 0 {
52 cfg := config.Get()
53 providers, err := config.Providers(cfg)
54 filteredProviders := []catwalk.Provider{}
55 for _, p := range providers {
56 hasAPIKeyEnv := strings.HasPrefix(p.APIKey, "$")
57 if hasAPIKeyEnv && p.ID != catwalk.InferenceProviderAzure {
58 filteredProviders = append(filteredProviders, p)
59 }
60 }
61
62 m.providers = filteredProviders
63 if err != nil {
64 cmds = append(cmds, util.ReportError(err))
65 }
66 }
67 cmds = append(cmds, m.list.Init(), m.SetModelType(m.modelType))
68 return tea.Batch(cmds...)
69}
70
71func (m *ModelListComponent) Update(msg tea.Msg) (*ModelListComponent, tea.Cmd) {
72 u, cmd := m.list.Update(msg)
73 m.list = u.(listModel)
74 return m, cmd
75}
76
77func (m *ModelListComponent) View() string {
78 return m.list.View()
79}
80
81func (m *ModelListComponent) Cursor() *tea.Cursor {
82 return m.list.Cursor()
83}
84
85func (m *ModelListComponent) SetSize(width, height int) tea.Cmd {
86 return m.list.SetSize(width, height)
87}
88
89func (m *ModelListComponent) SelectedModel() *ModelOption {
90 s := m.list.SelectedItem()
91 if s == nil {
92 return nil
93 }
94 sv := *s
95 model := sv.Value()
96 return &model
97}
98
99func (m *ModelListComponent) SetModelType(modelType int) tea.Cmd {
100 t := styles.CurrentTheme()
101 m.modelType = modelType
102
103 var groups []list.Group[list.CompletionItem[ModelOption]]
104 // first none section
105 selectedItemID := ""
106
107 cfg := config.Get()
108 var currentModel config.SelectedModel
109 if m.modelType == LargeModelType {
110 currentModel = cfg.Models[config.SelectedModelTypeLarge]
111 } else {
112 currentModel = cfg.Models[config.SelectedModelTypeSmall]
113 }
114
115 configuredIcon := t.S().Base.Foreground(t.Success).Render(styles.CheckIcon)
116 configured := fmt.Sprintf("%s %s", configuredIcon, t.S().Subtle.Render("Configured"))
117
118 // Create a map to track which providers we've already added
119 addedProviders := make(map[string]bool)
120
121 // First, add any configured providers that are not in the known providers list
122 // These should appear at the top of the list
123 knownProviders, err := config.Providers(cfg)
124 if err != nil {
125 return util.ReportError(err)
126 }
127 for providerID, providerConfig := range cfg.Providers.Seq2() {
128 if providerConfig.Disable {
129 continue
130 }
131
132 // Check if this provider is not in the known providers list
133 if !slices.ContainsFunc(knownProviders, func(p catwalk.Provider) bool { return p.ID == catwalk.InferenceProvider(providerID) }) ||
134 !slices.ContainsFunc(m.providers, func(p catwalk.Provider) bool { return p.ID == catwalk.InferenceProvider(providerID) }) {
135 // Convert config provider to provider.Provider format
136 configProvider := catwalk.Provider{
137 Name: providerConfig.Name,
138 ID: catwalk.InferenceProvider(providerID),
139 Models: make([]catwalk.Model, len(providerConfig.Models)),
140 }
141
142 // Convert models
143 for i, model := range providerConfig.Models {
144 configProvider.Models[i] = catwalk.Model{
145 ID: model.ID,
146 Name: model.Name,
147 CostPer1MIn: model.CostPer1MIn,
148 CostPer1MOut: model.CostPer1MOut,
149 CostPer1MInCached: model.CostPer1MInCached,
150 CostPer1MOutCached: model.CostPer1MOutCached,
151 ContextWindow: model.ContextWindow,
152 DefaultMaxTokens: model.DefaultMaxTokens,
153 CanReason: model.CanReason,
154 HasReasoningEffort: model.HasReasoningEffort,
155 DefaultReasoningEffort: model.DefaultReasoningEffort,
156 SupportsImages: model.SupportsImages,
157 }
158 }
159
160 // Add this unknown provider to the list
161 name := configProvider.Name
162 if name == "" {
163 name = string(configProvider.ID)
164 }
165 section := list.NewItemSection(name)
166 section.SetInfo(configured)
167 group := list.Group[list.CompletionItem[ModelOption]]{
168 Section: section,
169 }
170 for _, model := range configProvider.Models {
171 item := list.NewCompletionItem(model.Name, ModelOption{
172 Provider: configProvider,
173 Model: model,
174 },
175 list.WithCompletionID(
176 fmt.Sprintf("%s:%s", providerConfig.ID, model.ID),
177 ),
178 )
179
180 group.Items = append(group.Items, item)
181 if model.ID == currentModel.Model && string(configProvider.ID) == currentModel.Provider {
182 selectedItemID = item.ID()
183 }
184 }
185 groups = append(groups, group)
186
187 addedProviders[providerID] = true
188 }
189 }
190
191 // Then add the known providers from the predefined list
192 for _, provider := range m.providers {
193 // Skip if we already added this provider as an unknown provider
194 if addedProviders[string(provider.ID)] {
195 continue
196 }
197
198 // Check if this provider is configured and not disabled
199 if providerConfig, exists := cfg.Providers.Get(string(provider.ID)); exists && providerConfig.Disable {
200 continue
201 }
202
203 name := provider.Name
204 if name == "" {
205 name = string(provider.ID)
206 }
207
208 section := list.NewItemSection(name)
209 if _, ok := cfg.Providers.Get(string(provider.ID)); ok {
210 section.SetInfo(configured)
211 }
212 group := list.Group[list.CompletionItem[ModelOption]]{
213 Section: section,
214 }
215 for _, model := range provider.Models {
216 item := list.NewCompletionItem(model.Name, ModelOption{
217 Provider: provider,
218 Model: model,
219 },
220 list.WithCompletionID(
221 fmt.Sprintf("%s:%s", provider.ID, model.ID),
222 ),
223 )
224 group.Items = append(group.Items, item)
225 if model.ID == currentModel.Model && string(provider.ID) == currentModel.Provider {
226 selectedItemID = item.ID()
227 }
228 }
229 groups = append(groups, group)
230 }
231
232 var cmds []tea.Cmd
233
234 cmd := m.list.SetGroups(groups)
235
236 if cmd != nil {
237 cmds = append(cmds, cmd)
238 }
239 cmd = m.list.SetSelected(selectedItemID)
240 if cmd != nil {
241 cmds = append(cmds, cmd)
242 }
243
244 return tea.Sequence(cmds...)
245}
246
247// GetModelType returns the current model type
248func (m *ModelListComponent) GetModelType() int {
249 return m.modelType
250}
251
252func (m *ModelListComponent) SetInputPlaceholder(placeholder string) {
253 m.list.SetInputPlaceholder(placeholder)
254}