1package ollama
2
3import (
4 "context"
5 "fmt"
6 "strings"
7)
8
9// ProviderModel represents a model in the provider format
10type ProviderModel struct {
11 ID string
12 Model string
13 CostPer1MIn float64
14 CostPer1MOut float64
15 CostPer1MInCached float64
16 CostPer1MOutCached float64
17 ContextWindow int64
18 DefaultMaxTokens int64
19 CanReason bool
20 HasReasoningEffort bool
21 SupportsImages bool
22}
23
24// Provider represents an Ollama provider
25type Provider struct {
26 Name string
27 ID string
28 Models []ProviderModel
29}
30
31// GetProvider returns a Provider for Ollama
32func GetProvider(ctx context.Context) (*Provider, error) {
33 if err := EnsureRunning(ctx); err != nil {
34 return nil, fmt.Errorf("failed to ensure Ollama is running: %w", err)
35 }
36
37 models, err := GetModels(ctx)
38 if err != nil {
39 return nil, fmt.Errorf("failed to get models: %w", err)
40 }
41
42 providerModels := make([]ProviderModel, len(models))
43 for i, model := range models {
44 family := extractModelFamily(model.Name)
45 providerModels[i] = ProviderModel{
46 ID: model.Name,
47 Model: model.Name,
48 CostPer1MIn: 0, // Local models have no cost
49 CostPer1MOut: 0,
50 CostPer1MInCached: 0,
51 CostPer1MOutCached: 0,
52 ContextWindow: getContextWindow(family),
53 DefaultMaxTokens: 4096,
54 CanReason: false,
55 HasReasoningEffort: false,
56 SupportsImages: supportsImages(family),
57 }
58 }
59
60 return &Provider{
61 Name: "Ollama",
62 ID: "ollama",
63 Models: providerModels,
64 }, nil
65}
66
67// extractModelFamily extracts the model family from a model name
68func extractModelFamily(modelName string) string {
69 // Extract the family from model names like "llama3.2:3b" -> "llama"
70 parts := strings.Split(modelName, ":")
71 if len(parts) > 0 {
72 name := strings.ToLower(parts[0])
73
74 // Handle various model families in specific order
75 switch {
76 case strings.Contains(name, "llama-vision"):
77 return "llama-vision"
78 case strings.Contains(name, "codellama"):
79 return "codellama"
80 case strings.Contains(name, "llava"):
81 return "llava"
82 case strings.Contains(name, "llama"):
83 return "llama"
84 case strings.Contains(name, "mistral"):
85 return "mistral"
86 case strings.Contains(name, "gemma"):
87 return "gemma"
88 case strings.Contains(name, "qwen"):
89 return "qwen"
90 case strings.Contains(name, "phi"):
91 return "phi"
92 case strings.Contains(name, "vision"):
93 return "llama-vision"
94 }
95 }
96 return "unknown"
97}
98
99// getContextWindow returns an estimated context window based on model family
100func getContextWindow(family string) int64 {
101 switch family {
102 case "llama":
103 return 131072 // Llama 3.x context window
104 case "mistral":
105 return 32768
106 case "gemma":
107 return 8192
108 case "qwen":
109 return 131072
110 case "phi":
111 return 131072
112 case "codellama":
113 return 16384
114 default:
115 return 8192 // Conservative default
116 }
117}
118
119// supportsImages returns whether a model family supports image inputs
120func supportsImages(family string) bool {
121 switch family {
122 case "llama-vision", "llava":
123 return true
124 default:
125 return false
126 }
127}