git.iko.soy/crush
on main help sign in
feat(enricher): add litellm enricher
0f27905Kieran Klukas
internal/discover/litellm.gonew file+87
compare with main
unifiedsplit
@@ -0,0 +1,87 @@
1+package discover
2+
3+import (
4+ "context"
5+ "encoding/json"
6+ "net/http"
7+
8+ "charm.land/catwalk/pkg/catwalk"
9+)
10+
11+// litellmModelInfoResponse mirrors the response from LiteLLM's
12+// /model/info endpoint, which returns rich metadata including context
13+// windows, max tokens, and pricing.
14+type litellmModelInfoResponse struct {
15+ Data []litellmModelInfo `json:"data"`
16+}
17+
18+// litellmModelInfo is a single entry from /model/info.
19+type litellmModelInfo struct {
20+ ModelName string `json:"model_name"`
21+ ModelInfo litellmModelMeta `json:"model_info"`
22+}
23+
24+// litellmModelMeta holds the metadata fields we care about from
25+// LiteLLM's model_info block.
26+type litellmModelMeta struct {
27+ MaxInputTokens *int64 `json:"max_input_tokens"`
28+ MaxOutputTokens *int64 `json:"max_output_tokens"`
29+ InputCostPerToken *float64 `json:"input_cost_per_token"`
30+ OutputCostPerToken *float64 `json:"output_cost_per_token"`
31+ Mode string `json:"mode"`
32+}
33+
34+func init() {
35+ RegisterEnricher("litellm", &litellmEnricher{})
36+}
37+
38+// litellmEnricher fetches model metadata from LiteLLM's /model/info
39+// endpoint and populates context window, max tokens, and pricing on
40+// discovered models.
41+type litellmEnricher struct{}
42+
43+func (e *litellmEnricher) EnrichModels(ctx context.Context, cfg Config, resolver Resolver, models []catwalk.Model) ([]catwalk.Model, error) {
44+ resp, err := doRequest(ctx, http.MethodGet, cfg.BaseURL, "/model/info", cfg.APIKey, cfg.ExtraHeaders, resolver, nil)
45+ if err != nil {
46+ return models, nil
47+ }
48+ defer resp.Body.Close()
49+
50+ if resp.StatusCode != http.StatusOK {
51+ return models, nil
52+ }
53+
54+ var infoResp litellmModelInfoResponse
55+ if err := json.NewDecoder(resp.Body).Decode(&infoResp); err != nil {
56+ return models, nil
57+ }
58+
59+ // Index metadata by model name for O(1) lookup.
60+ metaByID := make(map[string]litellmModelMeta, len(infoResp.Data))
61+ for _, entry := range infoResp.Data {
62+ metaByID[entry.ModelName] = entry.ModelInfo
63+ }
64+
65+ // Apply metadata to discovered models, preserving existing
66+ // non-zero values (user overrides win).
67+ for i := range models {
68+ meta, ok := metaByID[models[i].ID]
69+ if !ok {
70+ continue
71+ }
72+ if models[i].ContextWindow == 0 && meta.MaxInputTokens != nil {
73+ models[i].ContextWindow = *meta.MaxInputTokens
74+ }
75+ if models[i].DefaultMaxTokens == 0 && meta.MaxOutputTokens != nil {
76+ models[i].DefaultMaxTokens = *meta.MaxOutputTokens
77+ }
78+ if models[i].CostPer1MIn == 0 && meta.InputCostPerToken != nil {
79+ models[i].CostPer1MIn = *meta.InputCostPerToken * 1_000_000
80+ }
81+ if models[i].CostPer1MOut == 0 && meta.OutputCostPerToken != nil {
82+ models[i].CostPer1MOut = *meta.OutputCostPerToken * 1_000_000
83+ }
84+ }
85+
86+ return models, nil
87+}
internal/discover/litellm_test.gonew file+171
@@ -0,0 +1,171 @@
1+package discover
2+
3+import (
4+ "context"
5+ "net/http"
6+ "net/http/httptest"
7+ "testing"
8+
9+ "charm.land/catwalk/pkg/catwalk"
10+ "github.com/stretchr/testify/require"
11+)
12+
13+func TestLitellmEnricher(t *testing.T) {
14+ t.Parallel()
15+
16+ t.Run("populates metadata from /model/info", func(t *testing.T) {
17+ t.Parallel()
18+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
19+ require.Equal(t, "/model/info", r.URL.Path)
20+ w.Header().Set("Content-Type", "application/json")
21+ _, _ = w.Write([]byte(`{
22+ "data": [
23+ {
24+ "model_name": "gpt-4o",
25+ "model_info": {
26+ "max_input_tokens": 128000,
27+ "max_output_tokens": 16384,
28+ "input_cost_per_token": 2.5e-06,
29+ "output_cost_per_token": 1e-05
30+ }
31+ },
32+ {
33+ "model_name": "claude-3-opus",
34+ "model_info": {
35+ "max_input_tokens": 200000,
36+ "max_output_tokens": 4096
37+ }
38+ }
39+ ]
40+ }`))
41+ }))
42+ defer srv.Close()
43+
44+ cfg := Config{
45+ ID: "test-litellm",
46+ BaseURL: srv.URL,
47+ APIKey: "test-key",
48+ }
49+ models := []catwalk.Model{
50+ {ID: "gpt-4o", Name: "gpt-4o"},
51+ {ID: "claude-3-opus", Name: "claude-3-opus"},
52+ {ID: "unknown-model", Name: "unknown-model"},
53+ }
54+
55+ e := &litellmEnricher{}
56+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
57+ require.NoError(t, err)
58+ require.Len(t, result, 3)
59+
60+ // gpt-4o should have all fields populated.
61+ require.Equal(t, int64(128000), result[0].ContextWindow)
62+ require.Equal(t, int64(16384), result[0].DefaultMaxTokens)
63+ require.InDelta(t, 2.5, result[0].CostPer1MIn, 0.01)
64+ require.InDelta(t, 10.0, result[0].CostPer1MOut, 0.01)
65+
66+ // claude-3-opus should have context window and max tokens.
67+ require.Equal(t, int64(200000), result[1].ContextWindow)
68+ require.Equal(t, int64(4096), result[1].DefaultMaxTokens)
69+
70+ // unknown-model should be unchanged.
71+ require.Equal(t, int64(0), result[2].ContextWindow)
72+ })
73+
74+ t.Run("preserves existing non-zero values", func(t *testing.T) {
75+ t.Parallel()
76+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
77+ w.Header().Set("Content-Type", "application/json")
78+ _, _ = w.Write([]byte(`{
79+ "data": [{
80+ "model_name": "gpt-4o",
81+ "model_info": {
82+ "max_input_tokens": 128000,
83+ "max_output_tokens": 16384
84+ }
85+ }]
86+ }`))
87+ }))
88+ defer srv.Close()
89+
90+ cfg := Config{
91+ ID: "test-litellm",
92+ BaseURL: srv.URL,
93+ }
94+ models := []catwalk.Model{
95+ {ID: "gpt-4o", Name: "GPT-4o Custom", ContextWindow: 200000, DefaultMaxTokens: 32768},
96+ }
97+
98+ e := &litellmEnricher{}
99+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
100+ require.NoError(t, err)
101+
102+ // User overrides should be preserved.
103+ require.Equal(t, int64(200000), result[0].ContextWindow)
104+ require.Equal(t, int64(32768), result[0].DefaultMaxTokens)
105+ require.Equal(t, "GPT-4o Custom", result[0].Name)
106+ })
107+
108+ t.Run("returns models unchanged on HTTP error", func(t *testing.T) {
109+ t.Parallel()
110+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
111+ w.WriteHeader(http.StatusInternalServerError)
112+ }))
113+ defer srv.Close()
114+
115+ cfg := Config{
116+ ID: "test-litellm",
117+ BaseURL: srv.URL,
118+ }
119+ models := []catwalk.Model{{ID: "m1"}}
120+
121+ e := &litellmEnricher{}
122+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
123+ require.NoError(t, err)
124+ require.Len(t, result, 1)
125+ require.Equal(t, "m1", result[0].ID)
126+ })
127+
128+ t.Run("sends auth and extra headers", func(t *testing.T) {
129+ t.Parallel()
130+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
131+ require.Equal(t, "Bearer my-key", r.Header.Get("Authorization"))
132+ require.Equal(t, "custom-val", r.Header.Get("X-Custom"))
133+ w.Header().Set("Content-Type", "application/json")
134+ _, _ = w.Write([]byte(`{"data": []}`))
135+ }))
136+ defer srv.Close()
137+
138+ cfg := Config{
139+ ID: "test-litellm",
140+ BaseURL: srv.URL,
141+ APIKey: "my-key",
142+ ExtraHeaders: map[string]string{
143+ "X-Custom": "custom-val",
144+ },
145+ }
146+
147+ e := &litellmEnricher{}
148+ _, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, nil)
149+ require.NoError(t, err)
150+ })
151+}
152+
153+func TestGetEnricher(t *testing.T) {
154+ t.Parallel()
155+
156+ require.NotNil(t, GetEnricher("litellm"))
157+ require.Nil(t, GetEnricher("openai-compat"))
158+ require.Nil(t, GetEnricher(""))
159+}
160+
161+func TestIsKnownCustomProvider(t *testing.T) {
162+ t.Parallel()
163+
164+ require.True(t, IsKnownCustomProvider("litellm"))
165+ require.True(t, IsKnownCustomProvider("ollama"))
166+ require.True(t, IsKnownCustomProvider("omlx"))
167+ require.True(t, IsKnownCustomProvider("lmstudio"))
168+ require.False(t, IsKnownCustomProvider("openai-compat"))
169+ require.False(t, IsKnownCustomProvider("anthropic"))
170+ require.False(t, IsKnownCustomProvider(""))
171+}