git.iko.soy/crush
on main help sign in
feat(enricher): add ollama enricher
9299ad4Kieran Klukas
internal/discover/ollama.gonew file+105
compare with main
unifiedsplit
@@ -0,0 +1,105 @@
1+package discover
2+
3+import (
4+ "context"
5+ "encoding/json"
6+ "net/http"
7+ "strings"
8+ "sync"
9+
10+ "charm.land/catwalk/pkg/catwalk"
11+)
12+
13+func init() {
14+ RegisterEnricher("ollama", &ollamaEnricher{})
15+}
16+
17+// ollamaShowResponse mirrors the response from Ollama's POST /api/show
18+// endpoint. Only the fields we care about are decoded.
19+type ollamaShowResponse struct {
20+ ModelInfo map[string]any `json:"model_info"`
21+}
22+
23+// ollamaEnricher fetches model metadata from Ollama's /api/show
24+// endpoint and populates context window on discovered models.
25+type ollamaEnricher struct{}
26+
27+func (e *ollamaEnricher) EnrichModels(ctx context.Context, cfg Config, resolver Resolver, models []catwalk.Model) ([]catwalk.Model, error) {
28+ // Collect indices that need enrichment.
29+ var needEnrichment []int
30+ for i := range models {
31+ if models[i].ContextWindow == 0 {
32+ needEnrichment = append(needEnrichment, i)
33+ }
34+ }
35+ if len(needEnrichment) == 0 {
36+ return models, nil
37+ }
38+
39+ // Fetch metadata concurrently with bounded parallelism.
40+ type result struct {
41+ index int
42+ contextLength int64
43+ }
44+
45+ results := make([]result, len(needEnrichment))
46+ var wg sync.WaitGroup
47+ sem := make(chan struct{}, 5) // Max 5 concurrent requests.
48+
49+ for ri, idx := range needEnrichment {
50+ wg.Go(func() {
51+ sem <- struct{}{}
52+ defer func() { <-sem }()
53+
54+ resp, err := doRequest(ctx, http.MethodPost, cfg.BaseURL, "/api/show",
55+ cfg.APIKey, cfg.ExtraHeaders, resolver,
56+ map[string]string{"model": models[idx].ID})
57+ if err != nil {
58+ return
59+ }
60+ defer resp.Body.Close()
61+
62+ var showResp ollamaShowResponse
63+ if err := json.NewDecoder(resp.Body).Decode(&showResp); err != nil {
64+ return
65+ }
66+
67+ if cl := extractContextLength(showResp.ModelInfo); cl > 0 {
68+ results[ri] = result{index: idx, contextLength: cl}
69+ }
70+ })
71+ }
72+ wg.Wait()
73+
74+ for _, r := range results {
75+ if r.contextLength > 0 {
76+ models[r.index].ContextWindow = r.contextLength
77+ }
78+ }
79+
80+ return models, nil
81+}
82+
83+// extractContextLength finds the context_length value in Ollama's
84+// model_info map. The key is architecture-specific (e.g.
85+// "llama.context_length", "qwen2.context_length"), so we scan for any
86+// key ending in ".context_length".
87+func extractContextLength(info map[string]any) int64 {
88+ for k, v := range info {
89+ if !strings.HasSuffix(k, ".context_length") {
90+ continue
91+ }
92+ switch n := v.(type) {
93+ case float64:
94+ return int64(n)
95+ case int64:
96+ return n
97+ case json.Number:
98+ i, err := n.Int64()
99+ if err == nil {
100+ return i
101+ }
102+ }
103+ }
104+ return 0
105+}
internal/discover/ollama_test.gonew file+129
@@ -0,0 +1,129 @@
1+package discover
2+
3+import (
4+ "context"
5+ "encoding/json"
6+ "net/http"
7+ "net/http/httptest"
8+ "testing"
9+
10+ "charm.land/catwalk/pkg/catwalk"
11+ "github.com/stretchr/testify/require"
12+)
13+
14+func TestOllamaEnricher(t *testing.T) {
15+ t.Parallel()
16+
17+ t.Run("populates context window from /api/show", func(t *testing.T) {
18+ t.Parallel()
19+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
20+ require.Equal(t, "/api/show", r.URL.Path)
21+ require.Equal(t, http.MethodPost, r.Method)
22+
23+ var req map[string]string
24+ json.NewDecoder(r.Body).Decode(&req)
25+
26+ w.Header().Set("Content-Type", "application/json")
27+ switch req["model"] {
28+ case "llama3:latest":
29+ json.NewEncoder(w).Encode(ollamaShowResponse{
30+ ModelInfo: map[string]any{
31+ "llama.context_length": float64(8192),
32+ },
33+ })
34+ case "qwen2:latest":
35+ json.NewEncoder(w).Encode(ollamaShowResponse{
36+ ModelInfo: map[string]any{
37+ "qwen2.context_length": float64(32768),
38+ },
39+ })
40+ default:
41+ json.NewEncoder(w).Encode(ollamaShowResponse{})
42+ }
43+ }))
44+ defer srv.Close()
45+
46+ cfg := Config{ID: "test-ollama", BaseURL: srv.URL}
47+ models := []catwalk.Model{
48+ {ID: "llama3:latest", Name: "llama3:latest"},
49+ {ID: "qwen2:latest", Name: "qwen2:latest"},
50+ {ID: "unknown:latest", Name: "unknown:latest"},
51+ }
52+
53+ e := &ollamaEnricher{}
54+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
55+ require.NoError(t, err)
56+ require.Len(t, result, 3)
57+ require.Equal(t, int64(8192), result[0].ContextWindow)
58+ require.Equal(t, int64(32768), result[1].ContextWindow)
59+ require.Equal(t, int64(0), result[2].ContextWindow)
60+ })
61+
62+ t.Run("preserves existing context window", func(t *testing.T) {
63+ t.Parallel()
64+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
65+ w.Header().Set("Content-Type", "application/json")
66+ json.NewEncoder(w).Encode(ollamaShowResponse{
67+ ModelInfo: map[string]any{
68+ "llama.context_length": float64(8192),
69+ },
70+ })
71+ }))
72+ defer srv.Close()
73+
74+ cfg := Config{ID: "test-ollama", BaseURL: srv.URL}
75+ models := []catwalk.Model{
76+ {ID: "llama3:latest", ContextWindow: 16384},
77+ }
78+
79+ e := &ollamaEnricher{}
80+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
81+ require.NoError(t, err)
82+ require.Equal(t, int64(16384), result[0].ContextWindow)
83+ })
84+
85+ t.Run("skips /api/show call when context window already set", func(t *testing.T) {
86+ t.Parallel()
87+ calls := 0
88+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
89+ calls++
90+ w.Header().Set("Content-Type", "application/json")
91+ json.NewEncoder(w).Encode(ollamaShowResponse{})
92+ }))
93+ defer srv.Close()
94+
95+ cfg := Config{ID: "test-ollama", BaseURL: srv.URL}
96+ models := []catwalk.Model{
97+ {ID: "m1", ContextWindow: 4096},
98+ {ID: "m2", ContextWindow: 8192},
99+ }
100+
101+ e := &ollamaEnricher{}
102+ _, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
103+ require.NoError(t, err)
104+ require.Equal(t, 0, calls)
105+ })
106+}
107+
108+func TestExtractContextLength(t *testing.T) {
109+ t.Parallel()
110+
111+ tests := []struct {
112+ name string
113+ info map[string]any
114+ expected int64
115+ }{
116+ {"llama key", map[string]any{"llama.context_length": float64(8192)}, 8192},
117+ {"qwen2 key", map[string]any{"qwen2.context_length": float64(32768)}, 32768},
118+ {"no context key", map[string]any{"llama.block_count": float64(32)}, 0},
119+ {"empty map", map[string]any{}, 0},
120+ {"nil map", nil, 0},
121+ }
122+
123+ for _, tt := range tests {
124+ t.Run(tt.name, func(t *testing.T) {
125+ t.Parallel()
126+ require.Equal(t, tt.expected, extractContextLength(tt.info))
127+ })
128+ }
129+}