git.iko.soy/crush
on main help sign in
feat(enricher): add lmstudio enricher
3c4e654Kieran Klukas
internal/discover/lmstudio.gonew file+94
compare with main
unifiedsplit
@@ -0,0 +1,94 @@
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("lmstudio", &lmstudioEnricher{})
13+}
14+
15+// lmstudioModelsResponse mirrors the response from LM Studio's native
16+// GET /api/v1/models endpoint. Only the fields we care about are
17+// decoded.
18+type lmstudioModelsResponse struct {
19+ Data []lmstudioModelEntry `json:"data"`
20+}
21+
22+// lmstudioModelEntry is a single entry from /api/v1/models.
23+type lmstudioModelEntry struct {
24+ Key string `json:"key"`
25+ DisplayName string `json:"display_name"`
26+ MaxContextLength int64 `json:"max_context_length"`
27+ LoadedInstances []lmstudioInstance `json:"loaded_instances"`
28+}
29+
30+// lmstudioInstance is a currently loaded model instance with its
31+// runtime config.
32+type lmstudioInstance struct {
33+ Config lmstudioInstanceConfig `json:"config"`
34+}
35+
36+// lmstudioInstanceConfig holds per-instance runtime settings.
37+type lmstudioInstanceConfig struct {
38+ ContextLength int64 `json:"context_length"`
39+}
40+
41+// lmstudioEnricher fetches model metadata from LM Studio's native
42+// /api/v1/models endpoint and populates context window and display
43+// name on discovered models.
44+type lmstudioEnricher struct{}
45+
46+func (e *lmstudioEnricher) EnrichModels(ctx context.Context, cfg Config, resolver Resolver, models []catwalk.Model) ([]catwalk.Model, error) {
47+ resp, err := doRequest(ctx, http.MethodGet, cfg.BaseURL, "/api/v1/models", cfg.APIKey, cfg.ExtraHeaders, resolver, nil)
48+ if err != nil {
49+ return models, nil
50+ }
51+ defer resp.Body.Close()
52+
53+ if resp.StatusCode != http.StatusOK {
54+ return models, nil
55+ }
56+
57+ var modelsResp lmstudioModelsResponse
58+ if err := json.NewDecoder(resp.Body).Decode(&modelsResp); err != nil {
59+ return models, nil
60+ }
61+
62+ // Index by key for O(1) lookup.
63+ metaByKey := make(map[string]lmstudioModelEntry, len(modelsResp.Data))
64+ for _, m := range modelsResp.Data {
65+ metaByKey[m.Key] = m
66+ }
67+
68+ for i := range models {
69+ meta, ok := metaByKey[models[i].ID]
70+ if !ok {
71+ continue
72+ }
73+
74+ // Context window: prefer loaded instance config, fall back
75+ // to the model-level max.
76+ if models[i].ContextWindow == 0 {
77+ if len(meta.LoadedInstances) > 0 && meta.LoadedInstances[0].Config.ContextLength > 0 {
78+ models[i].ContextWindow = meta.LoadedInstances[0].Config.ContextLength
79+ } else if meta.MaxContextLength > 0 {
80+ models[i].ContextWindow = meta.MaxContextLength
81+ }
82+ }
83+
84+ // Display name if not already set by user.
85+ if models[i].Name == models[i].ID && meta.DisplayName != "" {
86+ models[i].Name = meta.DisplayName
87+ }
88+
89+ // TODO: populate vision/tool-use flags when catwalk.Model
90+ // gains dedicated fields for them.
91+ }
92+
93+ return models, nil
94+}
internal/discover/lmstudio_test.gonew file+150
@@ -0,0 +1,150 @@
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 TestLmstudioEnricher(t *testing.T) {
15+ t.Parallel()
16+
17+ t.Run("populates context window and name from /api/v1/models", 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/v1/models", r.URL.Path)
21+ w.Header().Set("Content-Type", "application/json")
22+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
23+ Data: []lmstudioModelEntry{
24+ {
25+ Key: "qwen2.5-7b-instruct",
26+ DisplayName: "Qwen 2.5 7B Instruct",
27+ MaxContextLength: 32768,
28+ },
29+ {
30+ Key: "llama-3.1-8b",
31+ DisplayName: "Llama 3.1 8B",
32+ MaxContextLength: 131072,
33+ },
34+ },
35+ })
36+ }))
37+ defer srv.Close()
38+
39+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
40+ models := []catwalk.Model{
41+ {ID: "qwen2.5-7b-instruct", Name: "qwen2.5-7b-instruct"},
42+ {ID: "llama-3.1-8b", Name: "llama-3.1-8b"},
43+ {ID: "unknown-model", Name: "unknown-model"},
44+ }
45+
46+ e := &lmstudioEnricher{}
47+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
48+ require.NoError(t, err)
49+ require.Len(t, result, 3)
50+ require.Equal(t, int64(32768), result[0].ContextWindow)
51+ require.Equal(t, "Qwen 2.5 7B Instruct", result[0].Name)
52+ require.Equal(t, int64(131072), result[1].ContextWindow)
53+ require.Equal(t, "Llama 3.1 8B", result[1].Name)
54+ require.Equal(t, int64(0), result[2].ContextWindow)
55+ require.Equal(t, "unknown-model", result[2].Name)
56+ })
57+
58+ t.Run("prefers loaded instance context length over model max", func(t *testing.T) {
59+ t.Parallel()
60+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
61+ w.Header().Set("Content-Type", "application/json")
62+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
63+ Data: []lmstudioModelEntry{
64+ {
65+ Key: "m1",
66+ MaxContextLength: 131072,
67+ LoadedInstances: []lmstudioInstance{
68+ {Config: lmstudioInstanceConfig{ContextLength: 8192}},
69+ },
70+ },
71+ },
72+ })
73+ }))
74+ defer srv.Close()
75+
76+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
77+ models := []catwalk.Model{{ID: "m1", Name: "m1"}}
78+
79+ e := &lmstudioEnricher{}
80+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
81+ require.NoError(t, err)
82+ require.Equal(t, int64(8192), result[0].ContextWindow)
83+ })
84+
85+ t.Run("preserves existing non-zero values", func(t *testing.T) {
86+ t.Parallel()
87+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
88+ w.Header().Set("Content-Type", "application/json")
89+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
90+ Data: []lmstudioModelEntry{
91+ {
92+ Key: "m1",
93+ DisplayName: "Should Not Override",
94+ MaxContextLength: 131072,
95+ },
96+ },
97+ })
98+ }))
99+ defer srv.Close()
100+
101+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
102+ models := []catwalk.Model{
103+ {ID: "m1", Name: "My Custom Name", ContextWindow: 65536},
104+ }
105+
106+ e := &lmstudioEnricher{}
107+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
108+ require.NoError(t, err)
109+ require.Equal(t, int64(65536), result[0].ContextWindow)
110+ require.Equal(t, "My Custom Name", result[0].Name)
111+ })
112+
113+ t.Run("returns models unchanged on HTTP error", func(t *testing.T) {
114+ t.Parallel()
115+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
116+ w.WriteHeader(http.StatusServiceUnavailable)
117+ }))
118+ defer srv.Close()
119+
120+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
121+ models := []catwalk.Model{{ID: "m1"}}
122+
123+ e := &lmstudioEnricher{}
124+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
125+ require.NoError(t, err)
126+ require.Len(t, result, 1)
127+ require.Equal(t, int64(0), result[0].ContextWindow)
128+ })
129+
130+ t.Run("does not override custom name with display name", func(t *testing.T) {
131+ t.Parallel()
132+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
133+ w.Header().Set("Content-Type", "application/json")
134+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
135+ Data: []lmstudioModelEntry{
136+ {Key: "m1", DisplayName: "API Name"},
137+ },
138+ })
139+ }))
140+ defer srv.Close()
141+
142+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
143+ models := []catwalk.Model{{ID: "m1", Name: "User Name"}}
144+
145+ e := &lmstudioEnricher{}
146+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
147+ require.NoError(t, err)
148+ require.Equal(t, "User Name", result[0].Name)
149+ })
150+}