| 1 | package tui |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | |
| 6 | tea "github.com/charmbracelet/bubbletea" |
| 7 | "sourcecraft.dev/bigbes/claudio/cache" |
| 8 | ) |
| 9 | |
| 10 | var _ ModelFetcher = copilotFetcher{} |
| 11 | var _ ModelFetcher = openrouterFetcher{} |
| 12 | var _ ModelFetcher = nvidiaFetcher{} |
| 13 | var _ ModelFetcher = deepseekFetcher{} |
| 14 | var _ ModelFetcher = mistralFetcher{} |
| 15 | var _ ModelFetcher = kimiFetcher{} |
| 16 | var _ ModelFetcher = lmstudioFetcher{} |
| 17 | var _ ModelFetcher = ollamaFetcher{} |
| 18 | |
| 19 | // CategorizedResult holds the output of a model fetcher's categorization step. |
| 20 | type CategorizedResult struct { |
| 21 | Families []string |
| 22 | Items map[string][]PickerItem |
| 23 | PopularProviders []string |
| 24 | PopularProviderModels map[string][]PickerItem |
| 25 | AllModelIDs map[string]bool |
| 26 | } |
| 27 | |
| 28 | // ModelFetcher abstracts per-provider model fetching, filtering, and categorization. |
| 29 | type ModelFetcher interface { |
| 30 | // ProviderKey returns the cache/provider key (e.g. "copilot", "openrouter"). |
| 31 | ProviderKey() string |
| 32 | // FetchAndCategorize handles cache check, API fetch, cache store, |
| 33 | // filtering, and categorization. Returns the final result for the picker. |
| 34 | FetchAndCategorize(cdb *cache.DB) (CategorizedResult, error) |
| 35 | } |
| 36 | |
| 37 | // recentCategory is the label used for recently-used models prepended to the picker. |
| 38 | const recentCategory = "Recent" |
| 39 | |
| 40 | // startFetch launches an async model fetch using the given fetcher. |
| 41 | func (m Model) startFetch(f ModelFetcher) (tea.Model, tea.Cmd) { |
| 42 | m.phase = phaseLoadingModels |
| 43 | m.pickerProvider = f.ProviderKey() |
| 44 | cdb := m.cache |
| 45 | return m, func() tea.Msg { |
| 46 | result, err := f.FetchAndCategorize(cdb) |
| 47 | if err != nil { |
| 48 | return pickerModelsMsg{err: err} |
| 49 | } |
| 50 | return pickerModelsMsg{ |
| 51 | families: result.Families, |
| 52 | models: result.Items, |
| 53 | popularProviders: result.PopularProviders, |
| 54 | popularProviderModels: result.PopularProviderModels, |
| 55 | allModelIDs: result.AllModelIDs, |
| 56 | } |
| 57 | } |
| 58 | } |
| 59 | |
| 60 | // saveCacheModels marshals models and stores them in the cache. |
| 61 | func saveCacheModels(cdb *cache.DB, providerKey string, models any) { |
| 62 | if cdb == nil { |
| 63 | return |
| 64 | } |
| 65 | if data, err := json.Marshal(models); err == nil { |
| 66 | cdb.SetModels(providerKey, data) |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | // prependRecentModels prepends a "Recent" category from used_models_v2. |
| 71 | func prependRecentModels(cdb *cache.DB, providerKey string, result *CategorizedResult) { |
| 72 | if cdb == nil { |
| 73 | return |
| 74 | } |
| 75 | used, err := cdb.UsedModelsForProvider(providerKey) |
| 76 | if err != nil || len(used) == 0 { |
| 77 | return |
| 78 | } |
| 79 | var usedItems []PickerItem |
| 80 | for _, u := range used { |
| 81 | usedItems = append(usedItems, PickerItem{ID: u.ID, Name: u.Name}) |
| 82 | } |
| 83 | result.Items[recentCategory] = usedItems |
| 84 | result.Families = append([]string{recentCategory}, result.Families...) |
| 85 | } |
| 86 | |
| 87 | |