@@ -1,8 +1,12 @@
11package models
22
33import (
4+ "fmt"
5+ "time"
6+
47 "github.com/charmbracelet/bubbles/v2/help"
58 "github.com/charmbracelet/bubbles/v2/key"
9+ "github.com/charmbracelet/bubbles/v2/spinner"
610 tea "github.com/charmbracelet/bubbletea/v2"
711 "github.com/charmbracelet/crush/internal/config"
812 "github.com/charmbracelet/crush/internal/fur/provider"
@@ -56,6 +60,14 @@ type modelDialogCmp struct {
5660 modelList *ModelListComponent
5761 keyMap KeyMap
5862 help help.Model
63+
64+ // API key state
65+ needsAPIKey bool
66+ apiKeyInput *APIKeyInput
67+ selectedModel *ModelOption
68+ selectedModelType config.SelectedModelType
69+ isAPIKeyValid bool
70+ apiKeyValue string
5971}
6072
6173func NewModelDialogCmp() ModelDialog {
@@ -75,19 +87,22 @@ func NewModelDialogCmp() ModelDialog {
7587 t := styles.CurrentTheme()
7688 inputStyle := t.S().Base.Padding(0, 1, 0, 1)
7789 modelList := NewModelListComponent(listKeyMap, inputStyle, "Choose a model for large, complex tasks")
90+ apiKeyInput := NewAPIKeyInput()
91+ apiKeyInput.SetShowTitle(false)
7892 help := help.New()
7993 help.Styles = t.S().Help
8094
8195 return &modelDialogCmp{
82− modelList: modelList,
83− width: defaultWidth,
84− keyMap: DefaultKeyMap(),
85− help: help,
96+ modelList: modelList,
97+ apiKeyInput: apiKeyInput,
98+ width: defaultWidth,
99+ keyMap: DefaultKeyMap(),
100+ help: help,
86101 }
87102}
88103
89104func (m *modelDialogCmp) Init() tea.Cmd {
90− return m.modelList.Init()
105+ return tea.Batch(m.modelList.Init(), m.apiKeyInput.Init())
91106}
92107
93108func (m *modelDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
@@ -95,10 +110,58 @@ func (m *modelDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
95110 case tea.WindowSizeMsg:
96111 m.wWidth = msg.Width
97112 m.wHeight = msg.Height
113+ m.apiKeyInput.SetWidth(m.width - 2)
114+ m.help.Width = m.width - 2
98115 return m, m.modelList.SetSize(m.listWidth(), m.listHeight())
116+ case APIKeyStateChangeMsg:
117+ u, cmd := m.apiKeyInput.Update(msg)
118+ m.apiKeyInput = u.(*APIKeyInput)
119+ return m, cmd
99120 case tea.KeyPressMsg:
100121 switch {
101122 case key.Matches(msg, m.keyMap.Select):
123+ if m.isAPIKeyValid {
124+ return m, m.saveAPIKeyAndContinue(m.apiKeyValue)
125+ }
126+ if m.needsAPIKey {
127+ // Handle API key submission
128+ m.apiKeyValue = m.apiKeyInput.Value()
129+ provider, err := m.getProvider(m.selectedModel.Provider.ID)
130+ if err != nil || provider == nil {
131+ return m, util.ReportError(fmt.Errorf("provider %s not found", m.selectedModel.Provider.ID))
132+ }
133+ providerConfig := config.ProviderConfig{
134+ ID: string(m.selectedModel.Provider.ID),
135+ Name: m.selectedModel.Provider.Name,
136+ APIKey: m.apiKeyValue,
137+ Type: provider.Type,
138+ BaseURL: provider.APIEndpoint,
139+ }
140+ return m, tea.Sequence(
141+ util.CmdHandler(APIKeyStateChangeMsg{
142+ State: APIKeyInputStateVerifying,
143+ }),
144+ func() tea.Msg {
145+ start := time.Now()
146+ err := providerConfig.TestConnection(config.Get().Resolver())
147+ // intentionally wait for at least 750ms to make sure the user sees the spinner
148+ elapsed := time.Since(start)
149+ if elapsed < 750*time.Millisecond {
150+ time.Sleep(750*time.Millisecond - elapsed)
151+ }
152+ if err == nil {
153+ m.isAPIKeyValid = true
154+ return APIKeyStateChangeMsg{
155+ State: APIKeyInputStateVerified,
156+ }
157+ }
158+ return APIKeyStateChangeMsg{
159+ State: APIKeyInputStateError,
160+ }
161+ },
162+ )
163+ }
164+ // Normal model selection
102165 selectedItemInx := m.modelList.SelectedIndex()
103166 if selectedItemInx == list.NoSelection {
104167 return m, nil
@@ -113,17 +176,32 @@ func (m *modelDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
113176 modelType = config.SelectedModelTypeSmall
114177 }
115178
116− return m, tea.Sequence(
117− util.CmdHandler(dialogs.CloseDialogMsg{}),
118− util.CmdHandler(ModelSelectedMsg{
119− Model: config.SelectedModel{
120− Model: selectedItem.Model.ID,
121− Provider: string(selectedItem.Provider.ID),
122− },
123− ModelType: modelType,
124− }),
125− )
179+ // Check if provider is configured
180+ if m.isProviderConfigured(string(selectedItem.Provider.ID)) {
181+ return m, tea.Sequence(
182+ util.CmdHandler(dialogs.CloseDialogMsg{}),
183+ util.CmdHandler(ModelSelectedMsg{
184+ Model: config.SelectedModel{
185+ Model: selectedItem.Model.ID,
186+ Provider: string(selectedItem.Provider.ID),
187+ },
188+ ModelType: modelType,
189+ }),
190+ )
191+ } else {
192+ // Provider not configured, show API key input
193+ m.needsAPIKey = true
194+ m.selectedModel = &selectedItem
195+ m.selectedModelType = modelType
196+ m.apiKeyInput.SetProviderName(selectedItem.Provider.Name)
197+ return m, nil
198+ }
126199 case key.Matches(msg, m.keyMap.Tab):
200+ if m.needsAPIKey {
201+ u, cmd := m.apiKeyInput.Update(msg)
202+ m.apiKeyInput = u.(*APIKeyInput)
203+ return m, cmd
204+ }
127205 if m.modelList.GetModelType() == LargeModelType {
128206 m.modelList.SetInputPlaceholder(smallModelInputPlaceholder)
129207 return m, m.modelList.SetModelType(SmallModelType)
@@ -132,18 +210,68 @@ func (m *modelDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
132210 return m, m.modelList.SetModelType(LargeModelType)
133211 }
134212 case key.Matches(msg, m.keyMap.Close):
213+ if m.needsAPIKey {
214+ if m.isAPIKeyValid {
215+ return m, nil
216+ }
217+ // Go back to model selection
218+ m.needsAPIKey = false
219+ m.selectedModel = nil
220+ m.isAPIKeyValid = false
221+ m.apiKeyValue = ""
222+ m.apiKeyInput.Reset()
223+ return m, nil
224+ }
135225 return m, util.CmdHandler(dialogs.CloseDialogMsg{})
136226 default:
137− u, cmd := m.modelList.Update(msg)
138− m.modelList = u
227+ if m.needsAPIKey {
228+ u, cmd := m.apiKeyInput.Update(msg)
229+ m.apiKeyInput = u.(*APIKeyInput)
230+ return m, cmd
231+ } else {
232+ u, cmd := m.modelList.Update(msg)
233+ m.modelList = u
234+ return m, cmd
235+ }
236+ }
237+ case tea.PasteMsg:
238+ if m.needsAPIKey {
239+ u, cmd := m.apiKeyInput.Update(msg)
240+ m.apiKeyInput = u.(*APIKeyInput)
241+ return m, cmd
242+ } else {
243+ var cmd tea.Cmd
244+ m.modelList, cmd = m.modelList.Update(msg)
139245 return m, cmd
140246 }
247+ case spinner.TickMsg:
248+ u, cmd := m.apiKeyInput.Update(msg)
249+ m.apiKeyInput = u.(*APIKeyInput)
250+ return m, cmd
141251 }
142252 return m, nil
143253}
144254
145255func (m *modelDialogCmp) View() string {
146256 t := styles.CurrentTheme()
257+
258+ if m.needsAPIKey {
259+ // Show API key input
260+ m.keyMap.isAPIKeyHelp = true
261+ m.keyMap.isAPIKeyValid = m.isAPIKeyValid
262+ apiKeyView := m.apiKeyInput.View()
263+ apiKeyView = t.S().Base.Width(m.width - 3).Height(lipgloss.Height(apiKeyView)).PaddingLeft(1).Render(apiKeyView)
264+ content := lipgloss.JoinVertical(
265+ lipgloss.Left,
266+ t.S().Base.Padding(0, 1, 1, 1).Render(core.Title(m.apiKeyInput.GetTitle(), m.width-4)),
267+ apiKeyView,
268+ "",
269+ t.S().Base.Width(m.width-2).PaddingLeft(1).AlignHorizontal(lipgloss.Left).Render(m.help.View(m.keyMap)),
270+ )
271+ return m.style().Render(content)
272+ }
273+
274+ // Show model selection
147275 listView := m.modelList.View()
148276 radio := m.modelTypeRadio()
149277 content := lipgloss.JoinVertical(
@@ -157,10 +285,18 @@ func (m *modelDialogCmp) View() string {
157285}
158286
159287func (m *modelDialogCmp) Cursor() *tea.Cursor {
160− cursor := m.modelList.Cursor()
161− if cursor != nil {
162− cursor = m.moveCursor(cursor)
163− return cursor
288+ if m.needsAPIKey {
289+ cursor := m.apiKeyInput.Cursor()
290+ if cursor != nil {
291+ cursor = m.moveCursor(cursor)
292+ return cursor
293+ }
294+ } else {
295+ cursor := m.modelList.Cursor()
296+ if cursor != nil {
297+ cursor = m.moveCursor(cursor)
298+ return cursor
299+ }
164300 }
165301 return nil
166302}
@@ -192,9 +328,15 @@ func (m *modelDialogCmp) Position() (int, int) {
192328
193329func (m *modelDialogCmp) moveCursor(cursor *tea.Cursor) *tea.Cursor {
194330 row, col := m.Position()
195− offset := row + 3 // Border + title
196− cursor.Y += offset
197− cursor.X = cursor.X + col + 2
331+ if m.needsAPIKey {
332+ offset := row + 3 // Border + title + API key input offset
333+ cursor.Y += offset
334+ cursor.X = cursor.X + col + 2
335+ } else {
336+ offset := row + 3 // Border + title
337+ cursor.Y += offset
338+ cursor.X = cursor.X + col + 2
339+ }
198340 return cursor
199341}
200342
@@ -212,3 +354,49 @@ func (m *modelDialogCmp) modelTypeRadio() string {
212354 }
213355 return t.S().Base.Foreground(t.FgHalfMuted).Render(iconUnselected + " " + choices[0] + " " + iconSelected + " " + choices[1])
214356}
357+
358+func (m *modelDialogCmp) isProviderConfigured(providerID string) bool {
359+ cfg := config.Get()
360+ if _, ok := cfg.Providers[providerID]; ok {
361+ return true
362+ }
363+ return false
364+}
365+
366+func (m *modelDialogCmp) getProvider(providerID provider.InferenceProvider) (*provider.Provider, error) {
367+ providers, err := config.Providers()
368+ if err != nil {
369+ return nil, err
370+ }
371+ for _, p := range providers {
372+ if p.ID == providerID {
373+ return &p, nil
374+ }
375+ }
376+ return nil, nil
377+}
378+
379+func (m *modelDialogCmp) saveAPIKeyAndContinue(apiKey string) tea.Cmd {
380+ if m.selectedModel == nil {
381+ return util.ReportError(fmt.Errorf("no model selected"))
382+ }
383+
384+ cfg := config.Get()
385+ err := cfg.SetProviderAPIKey(string(m.selectedModel.Provider.ID), apiKey)
386+ if err != nil {
387+ return util.ReportError(fmt.Errorf("failed to save API key: %w", err))
388+ }
389+
390+ // Reset API key state and continue with model selection
391+ selectedModel := *m.selectedModel
392+ return tea.Sequence(
393+ util.CmdHandler(dialogs.CloseDialogMsg{}),
394+ util.CmdHandler(ModelSelectedMsg{
395+ Model: config.SelectedModel{
396+ Model: selectedModel.Model.ID,
397+ Provider: string(selectedModel.Provider.ID),
398+ },
399+ ModelType: m.selectedModelType,
400+ }),
401+ )
402+}