| 1 | package ollama |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "io" |
| 8 | "net/http" |
| 9 | "regexp" |
| 10 | "sort" |
| 11 | "strings" |
| 12 | "time" |
| 13 | ) |
| 14 | |
| 15 | // OllamaModel represents a model returned by the Ollama /api/tags endpoint. |
| 16 | type OllamaModel struct { |
| 17 | Name string `json:"name"` |
| 18 | Model string `json:"model"` |
| 19 | ModifiedAt string `json:"modified_at"` |
| 20 | Size int64 `json:"size"` |
| 21 | Details OllamaDetails `json:"details"` |
| 22 | ContextLength int `json:"context_length,omitempty"` |
| 23 | Capabilities []string `json:"capabilities,omitempty"` |
| 24 | } |
| 25 | |
| 26 | // OllamaDetails contains model metadata. |
| 27 | type OllamaDetails struct { |
| 28 | Family string `json:"family"` |
| 29 | Families []string `json:"families"` |
| 30 | ParameterSize string `json:"parameter_size"` |
| 31 | QuantizationLvl string `json:"quantization_level"` |
| 32 | } |
| 33 | |
| 34 | // DisplayName returns a human-friendly label for the model. |
| 35 | func (m OllamaModel) DisplayName() string { |
| 36 | name := m.Name |
| 37 | // Strip ":latest" suffix for cleaner display. |
| 38 | name = strings.TrimSuffix(name, ":latest") |
| 39 | ps := m.Details.ParameterSize |
| 40 | ql := m.Details.QuantizationLvl |
| 41 | switch { |
| 42 | case ps != "" && ql != "": |
| 43 | return name + " (" + ps + ", " + ql + ")" |
| 44 | case ps != "": |
| 45 | return name + " (" + ps + ")" |
| 46 | case ql != "": |
| 47 | return name + " (" + ql + ")" |
| 48 | } |
| 49 | return name |
| 50 | } |
| 51 | |
| 52 | // Family returns the model family for grouping. |
| 53 | func (m OllamaModel) Family() string { |
| 54 | if m.Details.Family != "" { |
| 55 | return m.Details.Family |
| 56 | } |
| 57 | // Fall back to base model name (before ':'). |
| 58 | if i := strings.Index(m.Name, ":"); i > 0 { |
| 59 | return m.Name[:i] |
| 60 | } |
| 61 | return m.Name |
| 62 | } |
| 63 | |
| 64 | // ListModels fetches available models from an Ollama instance. |
| 65 | // baseURL is e.g. "http://localhost:11434" or "https://ollama.com". |
| 66 | // apiKey may be empty for local instances. |
| 67 | func ListModels(baseURL, apiKey string) ([]OllamaModel, error) { |
| 68 | url := strings.TrimRight(baseURL, "/") + "/api/tags" |
| 69 | |
| 70 | req, err := http.NewRequest("GET", url, nil) |
| 71 | if err != nil { |
| 72 | return nil, err |
| 73 | } |
| 74 | if apiKey != "" { |
| 75 | req.Header.Set("Authorization", "Bearer "+apiKey) |
| 76 | } |
| 77 | req.Header.Set("Accept", "application/json") |
| 78 | |
| 79 | resp, err := http.DefaultClient.Do(req) |
| 80 | if err != nil { |
| 81 | return nil, fmt.Errorf("fetching models: %w", err) |
| 82 | } |
| 83 | defer resp.Body.Close() |
| 84 | |
| 85 | if resp.StatusCode != http.StatusOK { |
| 86 | return nil, fmt.Errorf("models endpoint returned %s", resp.Status) |
| 87 | } |
| 88 | |
| 89 | var result struct { |
| 90 | Models []OllamaModel `json:"models"` |
| 91 | } |
| 92 | if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { |
| 93 | return nil, fmt.Errorf("parsing models: %w", err) |
| 94 | } |
| 95 | |
| 96 | return result.Models, nil |
| 97 | } |
| 98 | |
| 99 | // ModelShowInfo holds the fields extracted from /api/show. |
| 100 | type ModelShowInfo struct { |
| 101 | ContextLength int |
| 102 | Capabilities []string |
| 103 | } |
| 104 | |
| 105 | // ShowModel calls /api/show to get detailed info for a single model. |
| 106 | func ShowModel(baseURL, apiKey, modelName string) (ModelShowInfo, error) { |
| 107 | url := strings.TrimRight(baseURL, "/") + "/api/show" |
| 108 | |
| 109 | body, _ := json.Marshal(map[string]string{"model": modelName}) |
| 110 | req, err := http.NewRequest("POST", url, bytes.NewReader(body)) |
| 111 | if err != nil { |
| 112 | return ModelShowInfo{}, err |
| 113 | } |
| 114 | if apiKey != "" { |
| 115 | req.Header.Set("Authorization", "Bearer "+apiKey) |
| 116 | } |
| 117 | req.Header.Set("Content-Type", "application/json") |
| 118 | req.Header.Set("Accept", "application/json") |
| 119 | |
| 120 | resp, err := http.DefaultClient.Do(req) |
| 121 | if err != nil { |
| 122 | return ModelShowInfo{}, fmt.Errorf("show model %s: %w", modelName, err) |
| 123 | } |
| 124 | defer resp.Body.Close() |
| 125 | |
| 126 | if resp.StatusCode != http.StatusOK { |
| 127 | return ModelShowInfo{}, fmt.Errorf("show model %s: %s", modelName, resp.Status) |
| 128 | } |
| 129 | |
| 130 | var result struct { |
| 131 | ModelInfo map[string]json.RawMessage `json:"model_info"` |
| 132 | Capabilities []string `json:"capabilities"` |
| 133 | } |
| 134 | if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { |
| 135 | return ModelShowInfo{}, fmt.Errorf("parsing show %s: %w", modelName, err) |
| 136 | } |
| 137 | |
| 138 | info := ModelShowInfo{Capabilities: result.Capabilities} |
| 139 | |
| 140 | // Read general.architecture, then look up "<arch>.context_length". |
| 141 | if archRaw, ok := result.ModelInfo["general.architecture"]; ok { |
| 142 | var arch string |
| 143 | if json.Unmarshal(archRaw, &arch) == nil && arch != "" { |
| 144 | if raw, ok := result.ModelInfo[arch+".context_length"]; ok { |
| 145 | var cl int |
| 146 | if json.Unmarshal(raw, &cl) == nil && cl > 0 { |
| 147 | info.ContextLength = cl |
| 148 | } |
| 149 | } |
| 150 | } |
| 151 | } |
| 152 | return info, nil |
| 153 | } |
| 154 | |
| 155 | // ShowCache provides get/set for per-model /api/show info. |
| 156 | type ShowCache interface { |
| 157 | GetOllamaModelInfo(provider, model string) (modifiedAt string, contextLength int, capabilities string, ok bool) |
| 158 | SetOllamaModelInfo(provider, model, modifiedAt string, contextLength int, capabilities string) error |
| 159 | } |
| 160 | |
| 161 | // FetchModelDetails enriches models with context_length and capabilities from /api/show. |
| 162 | // It uses the cache to avoid re-fetching models whose ModifiedAt hasn't changed. |
| 163 | func FetchModelDetails(models []OllamaModel, baseURL, apiKey, provKey string, sc ShowCache) { |
| 164 | for i := range models { |
| 165 | m := &models[i] |
| 166 | |
| 167 | // Check cache first: if modified_at matches, use cached value. |
| 168 | if sc != nil { |
| 169 | if cachedModAt, cl, caps, ok := sc.GetOllamaModelInfo(provKey, m.Name); ok && cachedModAt == m.ModifiedAt { |
| 170 | m.ContextLength = cl |
| 171 | if caps != "" { |
| 172 | _ = json.Unmarshal([]byte(caps), &m.Capabilities) |
| 173 | } |
| 174 | continue |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | info, err := ShowModel(baseURL, apiKey, m.Name) |
| 179 | if err != nil { |
| 180 | continue |
| 181 | } |
| 182 | m.ContextLength = info.ContextLength |
| 183 | m.Capabilities = info.Capabilities |
| 184 | |
| 185 | if sc != nil { |
| 186 | capsJSON := "" |
| 187 | if len(info.Capabilities) > 0 { |
| 188 | if data, err := json.Marshal(info.Capabilities); err == nil { |
| 189 | capsJSON = string(data) |
| 190 | } |
| 191 | } |
| 192 | _ = sc.SetOllamaModelInfo(provKey, m.Name, m.ModifiedAt, info.ContextLength, capsJSON) |
| 193 | } |
| 194 | } |
| 195 | } |
| 196 | |
| 197 | // CategorizeModels groups models by family and returns ordered family names |
| 198 | // plus the grouped map. |
| 199 | func CategorizeModels(models []OllamaModel) ([]string, map[string][]OllamaModel) { |
| 200 | grouped := make(map[string][]OllamaModel) |
| 201 | var familyOrder []string |
| 202 | seen := make(map[string]bool) |
| 203 | |
| 204 | // Sort by name for consistent ordering within families. |
| 205 | sorted := make([]OllamaModel, len(models)) |
| 206 | copy(sorted, models) |
| 207 | sort.Slice(sorted, func(i, j int) bool { |
| 208 | return sorted[i].Name < sorted[j].Name |
| 209 | }) |
| 210 | |
| 211 | for _, m := range sorted { |
| 212 | fam := m.Family() |
| 213 | if !seen[fam] { |
| 214 | seen[fam] = true |
| 215 | familyOrder = append(familyOrder, fam) |
| 216 | } |
| 217 | grouped[fam] = append(grouped[fam], m) |
| 218 | } |
| 219 | |
| 220 | return familyOrder, grouped |
| 221 | } |
| 222 | |
| 223 | // Provider returns the provider/namespace for a model. |
| 224 | // For namespaced models ("namespace/model:tag") it returns the namespace; |
| 225 | // otherwise it falls back to the model family. |
| 226 | func (m OllamaModel) Provider() string { |
| 227 | if i := strings.Index(m.Name, "/"); i > 0 { |
| 228 | return m.Name[:i] |
| 229 | } |
| 230 | return m.Family() |
| 231 | } |
| 232 | |
| 233 | // knownProviders is a set of known provider prefixes. |
| 234 | var knownProviders = map[string]bool{ |
| 235 | "qwen": true, |
| 236 | "llama": true, |
| 237 | "gemma": true, |
| 238 | "phi": true, |
| 239 | "mistral": true, |
| 240 | "deepseek": true, |
| 241 | "command-r": true, |
| 242 | "starcoder": true, |
| 243 | "glm": true, |
| 244 | "kimi": true, |
| 245 | "minimax": true, |
| 246 | "ministral": true, |
| 247 | "devstral": true, |
| 248 | "gemini": true, |
| 249 | "nemotron": true, |
| 250 | "gpt-oss": true, |
| 251 | } |
| 252 | |
| 253 | // providerDisplayNames maps raw provider keys to human-friendly names. |
| 254 | var providerDisplayNames = map[string]string{ |
| 255 | "qwen": "Qwen", |
| 256 | "llama": "Llama", |
| 257 | "gemma": "Google", |
| 258 | "phi": "Phi", |
| 259 | "mistral": "Mistral", |
| 260 | "deepseek": "DeepSeek", |
| 261 | "command-r": "Command R", |
| 262 | "starcoder": "StarCoder", |
| 263 | "glm": "GLM", |
| 264 | "kimi": "Kimi", |
| 265 | "minimax": "MiniMax", |
| 266 | "ministral": "Mistral", |
| 267 | "devstral": "Mistral", |
| 268 | "gemini": "Google", |
| 269 | "nemotron": "NVIDIA", |
| 270 | "gpt-oss": "OpenAI", |
| 271 | } |
| 272 | |
| 273 | // extractModelBase extracts the base provider key from a full model name. |
| 274 | // It tries to match against knownProviders first (longest match wins), |
| 275 | // then falls back to extracting the first name segment. |
| 276 | // E.g., "deepseek-v3.1:671b" -> "deepseek", "gpt-oss:20b" -> "gpt-oss", |
| 277 | // "qwen3.5:397b" -> "qwen", "glm-5" -> "glm" |
| 278 | func extractModelBase(modelName string) string { |
| 279 | // Remove tag suffix (:something) |
| 280 | if idx := strings.Index(modelName, ":"); idx > 0 { |
| 281 | modelName = modelName[:idx] |
| 282 | } |
| 283 | |
| 284 | // Try matching against known providers (longest match wins). |
| 285 | // A match requires the provider key to be followed by EOF, a digit, |
| 286 | // a dot, or a hyphen (to avoid partial-word matches). |
| 287 | bestMatch := "" |
| 288 | for key := range knownProviders { |
| 289 | if len(key) <= len(bestMatch) { |
| 290 | continue |
| 291 | } |
| 292 | if !strings.HasPrefix(modelName, key) { |
| 293 | continue |
| 294 | } |
| 295 | rest := modelName[len(key):] |
| 296 | if rest == "" || rest[0] == '-' || rest[0] == '.' || (rest[0] >= '0' && rest[0] <= '9') { |
| 297 | bestMatch = key |
| 298 | } |
| 299 | } |
| 300 | if bestMatch != "" { |
| 301 | return bestMatch |
| 302 | } |
| 303 | |
| 304 | // Fallback: first segment split by '-', with trailing digits stripped. |
| 305 | parts := strings.Split(modelName, "-") |
| 306 | if len(parts) > 0 { |
| 307 | base := parts[0] |
| 308 | for len(base) > 0 && base[len(base)-1] >= '0' && base[len(base)-1] <= '9' { |
| 309 | base = base[:len(base)-1] |
| 310 | } |
| 311 | if base != "" { |
| 312 | return base |
| 313 | } |
| 314 | } |
| 315 | return modelName |
| 316 | } |
| 317 | |
| 318 | // parseModifiedAt parses the modified_at timestamp and returns a time.Time. |
| 319 | // Returns zero time if parsing fails. |
| 320 | func parseModifiedAt(s string) time.Time { |
| 321 | if s == "" { |
| 322 | return time.Time{} |
| 323 | } |
| 324 | // Try RFC3339 format first |
| 325 | t, err := time.Parse(time.RFC3339, s) |
| 326 | if err == nil { |
| 327 | return t |
| 328 | } |
| 329 | // Try alternative common formats |
| 330 | formats := []string{ |
| 331 | "2006-01-02T15:04:05Z", |
| 332 | "2006-01-02T15:04:05.999Z", |
| 333 | "2006-01-02T15:04:05Z07:00", |
| 334 | "2006-01-02", |
| 335 | } |
| 336 | for _, format := range formats { |
| 337 | t, err = time.Parse(format, s) |
| 338 | if err == nil { |
| 339 | return t |
| 340 | } |
| 341 | } |
| 342 | return time.Time{} |
| 343 | } |
| 344 | |
| 345 | func providerDisplayName(key string) string { |
| 346 | if name, ok := providerDisplayNames[key]; ok { |
| 347 | return name |
| 348 | } |
| 349 | if knownProviders[key] { |
| 350 | if len(key) > 0 { |
| 351 | return strings.ToUpper(key[:1]) + key[1:] |
| 352 | } |
| 353 | return key |
| 354 | } |
| 355 | return "Other" |
| 356 | } |
| 357 | |
| 358 | // SeparatorCategory is a special category name rendered as a divider line. |
| 359 | const SeparatorCategory = "───" |
| 360 | |
| 361 | var libraryLinkRe = regexp.MustCompile(`href="/library/([^"]+)"`) |
| 362 | |
| 363 | // FetchPopularNames scrapes ollama.com/search?c=cloud for popular model names. |
| 364 | // Returns base names (e.g. "qwen3.5", "glm-5") in popularity order. |
| 365 | func FetchPopularNames() ([]string, error) { |
| 366 | resp, err := http.Get("https://ollama.com/search?c=cloud") |
| 367 | if err != nil { |
| 368 | return nil, err |
| 369 | } |
| 370 | defer resp.Body.Close() |
| 371 | |
| 372 | if resp.StatusCode != http.StatusOK { |
| 373 | return nil, fmt.Errorf("ollama search returned %s", resp.Status) |
| 374 | } |
| 375 | |
| 376 | body, err := io.ReadAll(resp.Body) |
| 377 | if err != nil { |
| 378 | return nil, err |
| 379 | } |
| 380 | |
| 381 | matches := libraryLinkRe.FindAllSubmatch(body, -1) |
| 382 | seen := make(map[string]bool) |
| 383 | var names []string |
| 384 | for _, m := range matches { |
| 385 | name := string(m[1]) |
| 386 | if !seen[name] { |
| 387 | seen[name] = true |
| 388 | names = append(names, name) |
| 389 | } |
| 390 | } |
| 391 | return names, nil |
| 392 | } |
| 393 | |
| 394 | // CategorizeByProvider groups models by provider for cloud instances. |
| 395 | // Models are grouped by company name extracted from the model name prefix. |
| 396 | // popularNames (from FetchPopularNames) determines the "Popular" category; |
| 397 | // models whose base name (before ":") matches a popular name are included, |
| 398 | // preserving the popularity order. |
| 399 | // The returned order is: Popular, separator, then provider groups alphabetically. |
| 400 | // Within each company, models are sorted by modified_at (newest first). |
| 401 | func CategorizeByProvider(models []OllamaModel, popularNames []string) ([]string, map[string][]OllamaModel) { |
| 402 | grouped := make(map[string][]OllamaModel) |
| 403 | |
| 404 | // Group by company |
| 405 | for _, m := range models { |
| 406 | company := providerDisplayName(extractModelBase(m.Name)) |
| 407 | grouped[company] = append(grouped[company], m) |
| 408 | } |
| 409 | |
| 410 | // Sort models within each company by modified_at (newest first) |
| 411 | for company := range grouped { |
| 412 | sort.Slice(grouped[company], func(i, j int) bool { |
| 413 | ti := parseModifiedAt(grouped[company][i].ModifiedAt) |
| 414 | tj := parseModifiedAt(grouped[company][j].ModifiedAt) |
| 415 | return ti.After(tj) |
| 416 | }) |
| 417 | } |
| 418 | |
| 419 | // Build "Popular" from popularNames by matching against available models. |
| 420 | // For each popular base name, include all model variants (tags) sorted by modified_at. |
| 421 | if len(popularNames) > 0 { |
| 422 | // Index models by base name (before ":"). |
| 423 | byBase := make(map[string][]OllamaModel) |
| 424 | for _, m := range models { |
| 425 | base := m.Name |
| 426 | if i := strings.Index(base, ":"); i > 0 { |
| 427 | base = base[:i] |
| 428 | } |
| 429 | byBase[base] = append(byBase[base], m) |
| 430 | } |
| 431 | var popular []OllamaModel |
| 432 | for _, name := range popularNames { |
| 433 | if ms, ok := byBase[name]; ok { |
| 434 | sort.Slice(ms, func(i, j int) bool { |
| 435 | ti := parseModifiedAt(ms[i].ModifiedAt) |
| 436 | tj := parseModifiedAt(ms[j].ModifiedAt) |
| 437 | return ti.After(tj) |
| 438 | }) |
| 439 | popular = append(popular, ms...) |
| 440 | } |
| 441 | } |
| 442 | if len(popular) > 0 { |
| 443 | grouped["Popular"] = popular |
| 444 | } |
| 445 | } |
| 446 | |
| 447 | // Extract companies and sort alphabetically (but "Other" at the end) |
| 448 | var companies []string |
| 449 | for company := range grouped { |
| 450 | if company == "Popular" { |
| 451 | continue |
| 452 | } |
| 453 | companies = append(companies, company) |
| 454 | } |
| 455 | sort.Slice(companies, func(i, j int) bool { |
| 456 | if companies[i] == "Other" { |
| 457 | return false |
| 458 | } |
| 459 | if companies[j] == "Other" { |
| 460 | return true |
| 461 | } |
| 462 | return companies[i] < companies[j] |
| 463 | }) |
| 464 | |
| 465 | var order []string |
| 466 | if _, ok := grouped["Popular"]; ok { |
| 467 | order = append(order, "Popular", SeparatorCategory) |
| 468 | } |
| 469 | order = append(order, companies...) |
| 470 | |
| 471 | return order, grouped |
| 472 | } |
| 473 | |