1package models
2
3import (
4 "fmt"
5 "slices"
6
7 tea "github.com/charmbracelet/bubbletea/v2"
8 "github.com/charmbracelet/crush/internal/config"
9 "github.com/charmbracelet/crush/internal/fur/provider"
10 "github.com/charmbracelet/crush/internal/tui/components/completions"
11 "github.com/charmbracelet/crush/internal/tui/components/core/list"
12 "github.com/charmbracelet/crush/internal/tui/components/dialogs/commands"
13 "github.com/charmbracelet/crush/internal/tui/styles"
14 "github.com/charmbracelet/crush/internal/tui/util"
15 "github.com/charmbracelet/lipgloss/v2"
16)
17
18type ModelListComponent struct {
19 list list.ListModel
20 modelType int
21 providers []provider.Provider
22}
23
24func NewModelListComponent(keyMap list.KeyMap, inputStyle lipgloss.Style, inputPlaceholder string) *ModelListComponent {
25 modelList := list.New(
26 list.WithFilterable(true),
27 list.WithKeyMap(keyMap),
28 list.WithInputStyle(inputStyle),
29 list.WithFilterPlaceholder(inputPlaceholder),
30 list.WithWrapNavigation(true),
31 )
32
33 return &ModelListComponent{
34 list: modelList,
35 modelType: LargeModelType,
36 }
37}
38
39func (m *ModelListComponent) Init() tea.Cmd {
40 var cmds []tea.Cmd
41 if len(m.providers) == 0 {
42 providers, err := config.Providers()
43 m.providers = providers
44 if err != nil {
45 cmds = append(cmds, util.ReportError(err))
46 }
47 }
48 cmds = append(cmds, m.list.Init(), m.SetModelType(m.modelType))
49 return tea.Batch(cmds...)
50}
51
52func (m *ModelListComponent) Update(msg tea.Msg) (*ModelListComponent, tea.Cmd) {
53 u, cmd := m.list.Update(msg)
54 m.list = u.(list.ListModel)
55 return m, cmd
56}
57
58func (m *ModelListComponent) View() string {
59 return m.list.View()
60}
61
62func (m *ModelListComponent) Cursor() *tea.Cursor {
63 return m.list.Cursor()
64}
65
66func (m *ModelListComponent) SetSize(width, height int) tea.Cmd {
67 return m.list.SetSize(width, height)
68}
69
70func (m *ModelListComponent) Items() []util.Model {
71 return m.list.Items()
72}
73
74func (m *ModelListComponent) SelectedIndex() int {
75 return m.list.SelectedIndex()
76}
77
78func (m *ModelListComponent) SetModelType(modelType int) tea.Cmd {
79 t := styles.CurrentTheme()
80 m.modelType = modelType
81
82 modelItems := []util.Model{}
83 selectIndex := 0
84
85 cfg := config.Get()
86 var currentModel config.SelectedModel
87 if m.modelType == LargeModelType {
88 currentModel = cfg.Models[config.SelectedModelTypeLarge]
89 } else {
90 currentModel = cfg.Models[config.SelectedModelTypeSmall]
91 }
92
93 configuredIcon := t.S().Base.Foreground(t.Success).Render(styles.CheckIcon)
94 configured := fmt.Sprintf("%s %s", configuredIcon, t.S().Subtle.Render("Configured"))
95
96 // Create a map to track which providers we've already added
97 addedProviders := make(map[string]bool)
98
99 // First, add any configured providers that are not in the known providers list
100 // These should appear at the top of the list
101 knownProviders := provider.KnownProviders()
102 for providerID, providerConfig := range cfg.Providers {
103 if providerConfig.Disable {
104 continue
105 }
106
107 // Check if this provider is not in the known providers list
108 if !slices.Contains(knownProviders, provider.InferenceProvider(providerID)) {
109 // Convert config provider to provider.Provider format
110 configProvider := provider.Provider{
111 Name: providerConfig.Name,
112 ID: provider.InferenceProvider(providerID),
113 Models: make([]provider.Model, len(providerConfig.Models)),
114 }
115
116 // Convert models
117 for i, model := range providerConfig.Models {
118 configProvider.Models[i] = provider.Model{
119 ID: model.ID,
120 Model: model.Model,
121 CostPer1MIn: model.CostPer1MIn,
122 CostPer1MOut: model.CostPer1MOut,
123 CostPer1MInCached: model.CostPer1MInCached,
124 CostPer1MOutCached: model.CostPer1MOutCached,
125 ContextWindow: model.ContextWindow,
126 DefaultMaxTokens: model.DefaultMaxTokens,
127 CanReason: model.CanReason,
128 HasReasoningEffort: model.HasReasoningEffort,
129 DefaultReasoningEffort: model.DefaultReasoningEffort,
130 SupportsImages: model.SupportsImages,
131 }
132 }
133
134 // Add this unknown provider to the list
135 name := configProvider.Name
136 if name == "" {
137 name = string(configProvider.ID)
138 }
139 section := commands.NewItemSection(name)
140 section.SetInfo(configured)
141 modelItems = append(modelItems, section)
142 for _, model := range configProvider.Models {
143 modelItems = append(modelItems, completions.NewCompletionItem(model.Model, ModelOption{
144 Provider: configProvider,
145 Model: model,
146 }))
147 if model.ID == currentModel.Model && string(configProvider.ID) == currentModel.Provider {
148 selectIndex = len(modelItems) - 1 // Set the selected index to the current model
149 }
150 }
151 addedProviders[providerID] = true
152 }
153 }
154
155 // Then add the known providers from the predefined list
156 for _, provider := range m.providers {
157 // Skip if we already added this provider as an unknown provider
158 if addedProviders[string(provider.ID)] {
159 continue
160 }
161
162 // Check if this provider is configured and not disabled
163 if providerConfig, exists := cfg.Providers[string(provider.ID)]; exists && providerConfig.Disable {
164 continue
165 }
166
167 name := provider.Name
168 if name == "" {
169 name = string(provider.ID)
170 }
171
172 section := commands.NewItemSection(name)
173 if _, ok := cfg.Providers[string(provider.ID)]; ok {
174 section.SetInfo(configured)
175 }
176 modelItems = append(modelItems, section)
177 for _, model := range provider.Models {
178 modelItems = append(modelItems, completions.NewCompletionItem(model.Model, ModelOption{
179 Provider: provider,
180 Model: model,
181 }))
182 if model.ID == currentModel.Model && string(provider.ID) == currentModel.Provider {
183 selectIndex = len(modelItems) - 1 // Set the selected index to the current model
184 }
185 }
186 }
187
188 return tea.Sequence(m.list.SetItems(modelItems), m.list.SetSelected(selectIndex))
189}
190
191// GetModelType returns the current model type
192func (m *ModelListComponent) GetModelType() int {
193 return m.modelType
194}
195
196func (m *ModelListComponent) SetInputPlaceholder(placeholder string) {
197 m.list.SetFilterPlaceholder(placeholder)
198}
199
200func (m *ModelListComponent) SetProviders(providers []provider.Provider) {
201 m.providers = providers
202}