git.iko.soy/crush
on main help sign in
feat(enricher): add omlx enricher
89b680dKieran Klukas
internal/discover/omlx.gonew file+70
compare with main
unifiedsplit
@@ -0,0 +1,70 @@
1+package discover
2+
3+import (
4+ "context"
5+ "encoding/json"
6+ "net/http"
7+
8+ "charm.land/catwalk/pkg/catwalk"
9+)
10+
11+func init() {
12+ RegisterEnricher("omlx", &omlxEnricher{})
13+}
14+
15+// omlxModelsStatusResponse mirrors the response from oMLX's
16+// GET /v1/models/status endpoint, which returns model metadata
17+// including max_context_window and max_tokens.
18+type omlxModelsStatusResponse struct {
19+ Models []omlxModelStatus `json:"models"`
20+}
21+
22+// omlxModelStatus is a single entry from /v1/models/status.
23+type omlxModelStatus struct {
24+ ID string `json:"id"`
25+ MaxContextWindow *int64 `json:"max_context_window"`
26+ MaxTokens *int64 `json:"max_tokens"`
27+}
28+
29+// omlxEnricher fetches model metadata from oMLX's /v1/models/status
30+// endpoint and populates context window and max tokens on discovered
31+// models.
32+type omlxEnricher struct{}
33+
34+func (e *omlxEnricher) EnrichModels(ctx context.Context, cfg Config, resolver Resolver, models []catwalk.Model) ([]catwalk.Model, error) {
35+ resp, err := doRequest(ctx, http.MethodGet, cfg.BaseURL, "/v1/models/status", cfg.APIKey, cfg.ExtraHeaders, resolver, nil)
36+ if err != nil {
37+ return models, nil
38+ }
39+ defer resp.Body.Close()
40+
41+ if resp.StatusCode != http.StatusOK {
42+ return models, nil
43+ }
44+
45+ var statusResp omlxModelsStatusResponse
46+ if err := json.NewDecoder(resp.Body).Decode(&statusResp); err != nil {
47+ return models, nil
48+ }
49+
50+ // Index by model ID for O(1) lookup.
51+ metaByID := make(map[string]omlxModelStatus, len(statusResp.Models))
52+ for _, m := range statusResp.Models {
53+ metaByID[m.ID] = m
54+ }
55+
56+ for i := range models {
57+ meta, ok := metaByID[models[i].ID]
58+ if !ok {
59+ continue
60+ }
61+ if models[i].ContextWindow == 0 && meta.MaxContextWindow != nil {
62+ models[i].ContextWindow = *meta.MaxContextWindow
63+ }
64+ if models[i].DefaultMaxTokens == 0 && meta.MaxTokens != nil {
65+ models[i].DefaultMaxTokens = *meta.MaxTokens
66+ }
67+ }
68+
69+ return models, nil
70+}
internal/discover/omlx_test.gonew file+90
@@ -0,0 +1,90 @@
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 TestOmlxEnricher(t *testing.T) {
15+ t.Parallel()
16+
17+ t.Run("populates context window and max tokens from /v1/models/status", func(t *testing.T) {
18+ t.Parallel()
19+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
20+ require.Equal(t, "/v1/models/status", r.URL.Path)
21+ w.Header().Set("Content-Type", "application/json")
22+ json.NewEncoder(w).Encode(omlxModelsStatusResponse{
23+ Models: []omlxModelStatus{
24+ {ID: "qwen3:latest", MaxContextWindow: ptr(int64(32768)), MaxTokens: ptr(int64(16384))},
25+ {ID: "llama3:latest", MaxContextWindow: ptr(int64(8192)), MaxTokens: ptr(int64(4096))},
26+ },
27+ })
28+ }))
29+ defer srv.Close()
30+
31+ cfg := Config{ID: "test-omlx", BaseURL: srv.URL}
32+ models := []catwalk.Model{
33+ {ID: "qwen3:latest", Name: "qwen3:latest"},
34+ {ID: "llama3:latest", Name: "llama3:latest"},
35+ {ID: "unknown:latest", Name: "unknown:latest"},
36+ }
37+
38+ e := &omlxEnricher{}
39+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
40+ require.NoError(t, err)
41+ require.Len(t, result, 3)
42+ require.Equal(t, int64(32768), result[0].ContextWindow)
43+ require.Equal(t, int64(16384), result[0].DefaultMaxTokens)
44+ require.Equal(t, int64(8192), result[1].ContextWindow)
45+ require.Equal(t, int64(4096), result[1].DefaultMaxTokens)
46+ require.Equal(t, int64(0), result[2].ContextWindow)
47+ })
48+
49+ t.Run("preserves existing non-zero values", func(t *testing.T) {
50+ t.Parallel()
51+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
52+ w.Header().Set("Content-Type", "application/json")
53+ json.NewEncoder(w).Encode(omlxModelsStatusResponse{
54+ Models: []omlxModelStatus{
55+ {ID: "m1", MaxContextWindow: ptr(int64(32768)), MaxTokens: ptr(int64(16384))},
56+ },
57+ })
58+ }))
59+ defer srv.Close()
60+
61+ cfg := Config{ID: "test-omlx", BaseURL: srv.URL}
62+ models := []catwalk.Model{
63+ {ID: "m1", ContextWindow: 65536, DefaultMaxTokens: 8192},
64+ }
65+
66+ e := &omlxEnricher{}
67+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
68+ require.NoError(t, err)
69+ require.Equal(t, int64(65536), result[0].ContextWindow)
70+ require.Equal(t, int64(8192), result[0].DefaultMaxTokens)
71+ })
72+
73+ t.Run("returns models unchanged on HTTP error", func(t *testing.T) {
74+ t.Parallel()
75+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
76+ w.WriteHeader(http.StatusNotFound)
77+ }))
78+ defer srv.Close()
79+
80+ cfg := Config{ID: "test-omlx", BaseURL: srv.URL}
81+ models := []catwalk.Model{{ID: "m1"}}
82+
83+ e := &omlxEnricher{}
84+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
85+ require.NoError(t, err)
86+ require.Len(t, result, 1)
87+ })
88+}
89+
90+func ptr[T any](v T) *T { return &v }