Skip to content

Commit 9212d6f

Browse files
molty3000molty3000
committed
feat: dynamic model context discovery via GET /models endpoint on startup
1 parent 681257e commit 9212d6f

3 files changed

Lines changed: 297 additions & 3 deletions

File tree

‎internal/llm/models.go‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
// Package llm provides an OpenAI-compatible HTTP client using only stdlib.
2+
package llm
3+
4+
import (
5+
"context"
6+
"encoding/json"
7+
"io"
8+
"net/http"
9+
"strings"
10+
"sync"
11+
"time"
12+
)
13+
14+
// modelCache caches the full model list per endpoint so multiple model
15+
// lookups from the same provider share a single API call.
16+
var (
17+
modelCache map[string]map[string]int // key: "baseURL|apiKey" → modelID → contextLength
18+
modelCacheMu sync.RWMutex
19+
)
20+
21+
func init() {
22+
modelCache = make(map[string]map[string]int)
23+
}
24+
25+
// cacheKey returns a unique key for the (baseURL, apiKey) pair.
26+
func cacheKey(baseURL, apiKey string) string {
27+
return baseURL + "|" + apiKey
28+
}
29+
30+
// rawModel is a single model entry from the /models endpoint.
31+
// Different providers use different field names for context length.
32+
type rawModel struct {
33+
ID string `json:"id"`
34+
ContextLength int `json:"context_length"` // OpenRouter, Together, some providers
35+
MaxContext int `json:"max_context"` // Fallback field name
36+
MaxInput int `json:"max_input_tokens"` // Common alternative
37+
}
38+
39+
// modelsResponse is the top-level response from GET /models.
40+
type modelsResponse struct {
41+
Data []rawModel `json:"data"`
42+
Models []rawModel `json:"models"` // Some providers use this wrapper
43+
}
44+
45+
// ResetModelCache clears the model discovery cache. Used in tests.
46+
func ResetModelCache() {
47+
modelCacheMu.Lock()
48+
modelCache = make(map[string]map[string]int)
49+
modelCacheMu.Unlock()
50+
}
51+
52+
// DiscoverModelContext queries the /models endpoint of the configured base URL
53+
// to discover the context window for the given model. Returns 0 if the
54+
// endpoint doesn't support model attribute discovery or the model isn't found.
55+
//
56+
// Results are cached per (baseURL, apiKey) so multiple agents using the same
57+
// provider share a single API call. The full model list is cached so
58+
// different model lookups from the same endpoint don't re-query.
59+
//
60+
// Call this at startup before creating the engine. The HTTP call uses a 5s
61+
// timeout and never blocks startup for more than that.
62+
func DiscoverModelContext(baseURL, apiKey, model string) int {
63+
cacheK := cacheKey(baseURL, apiKey)
64+
65+
// Check cache first
66+
modelCacheMu.RLock()
67+
if models, ok := modelCache[cacheK]; ok {
68+
if val, ok2 := models[model]; ok2 {
69+
modelCacheMu.RUnlock()
70+
return val
71+
}
72+
// Model not in cached list — don't re-query
73+
modelCacheMu.RUnlock()
74+
return 0
75+
}
76+
modelCacheMu.RUnlock()
77+
78+
// Query the API
79+
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
80+
defer cancel()
81+
82+
url := strings.TrimRight(baseURL, "/") + "/models"
83+
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
84+
if err != nil {
85+
return 0
86+
}
87+
if apiKey != "" {
88+
req.Header.Set("Authorization", "Bearer "+apiKey)
89+
}
90+
req.Header.Set("Accept", "application/json")
91+
92+
client := &http.Client{Timeout: 5 * time.Second}
93+
resp, err := client.Do(req)
94+
if err != nil {
95+
return 0
96+
}
97+
defer resp.Body.Close()
98+
99+
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) // 1MB max
100+
if err != nil || resp.StatusCode != http.StatusOK {
101+
return 0
102+
}
103+
104+
var parsed modelsResponse
105+
if err := json.Unmarshal(body, &parsed); err != nil {
106+
return 0
107+
}
108+
109+
// Search both possible array fields
110+
models := parsed.Data
111+
if len(models) == 0 {
112+
models = parsed.Models
113+
}
114+
115+
// Build a lookup map from the model list
116+
lookup := make(map[string]int, len(models))
117+
for _, m := range models {
118+
val := m.ContextLength
119+
if val == 0 {
120+
val = m.MaxContext
121+
}
122+
if val == 0 {
123+
val = m.MaxInput
124+
}
125+
lookup[m.ID] = val
126+
}
127+
128+
// Cache the full list
129+
modelCacheMu.Lock()
130+
modelCache[cacheK] = lookup
131+
modelCacheMu.Unlock()
132+
133+
return lookup[model]
134+
}

‎internal/llm/models_test.go‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
package llm
2+
3+
import (
4+
"encoding/json"
5+
"net/http"
6+
"net/http/httptest"
7+
"testing"
8+
)
9+
10+
func TestDiscoverModelContext_OpenRouterFormat(t *testing.T) {
11+
ResetModelCache()
12+
13+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
14+
if r.URL.Path != "/models" {
15+
t.Errorf("expected /models, got %s", r.URL.Path)
16+
}
17+
resp := modelsResponse{
18+
Data: []rawModel{
19+
{ID: "deepseek-v4-flash", ContextLength: 131072},
20+
{ID: "deepseek-v4-pro", ContextLength: 1048576},
21+
},
22+
}
23+
json.NewEncoder(w).Encode(resp)
24+
}))
25+
defer srv.Close()
26+
27+
ctx := DiscoverModelContext(srv.URL, "test-key", "deepseek-v4-flash")
28+
if ctx != 131072 {
29+
t.Errorf("expected 131072, got %d", ctx)
30+
}
31+
32+
// Second model
33+
ctx = DiscoverModelContext(srv.URL, "test-key", "deepseek-v4-pro")
34+
if ctx != 1048576 {
35+
t.Errorf("expected 1048576, got %d", ctx)
36+
}
37+
}
38+
39+
func TestDiscoverModelContext_UnknownModel(t *testing.T) {
40+
ResetModelCache()
41+
42+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
43+
resp := modelsResponse{
44+
Data: []rawModel{
45+
{ID: "gpt-4o", ContextLength: 128000},
46+
},
47+
}
48+
json.NewEncoder(w).Encode(resp)
49+
}))
50+
defer srv.Close()
51+
52+
ctx := DiscoverModelContext(srv.URL, "test-key", "unknown-model")
53+
if ctx != 0 {
54+
t.Errorf("expected 0 for unknown model, got %d", ctx)
55+
}
56+
}
57+
58+
func TestDiscoverModelContext_ServerError(t *testing.T) {
59+
ResetModelCache()
60+
61+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
62+
w.WriteHeader(http.StatusInternalServerError)
63+
}))
64+
defer srv.Close()
65+
66+
ctx := DiscoverModelContext(srv.URL, "test-key", "any-model")
67+
if ctx != 0 {
68+
t.Errorf("expected 0 on server error, got %d", ctx)
69+
}
70+
}
71+
72+
func TestDiscoverModelContext_Timeout(t *testing.T) {
73+
ResetModelCache()
74+
75+
// A port that's very unlikely to be listening — tests connection refused handling
76+
ctx := DiscoverModelContext("http://127.0.0.1:1", "test-key", "any-model")
77+
if ctx != 0 {
78+
t.Errorf("expected 0 on connection error, got %d", ctx)
79+
}
80+
}
81+
82+
func TestDiscoverModelContext_Cached(t *testing.T) {
83+
ResetModelCache()
84+
85+
callCount := 0
86+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
87+
callCount++
88+
resp := modelsResponse{
89+
Data: []rawModel{
90+
{ID: "flash", ContextLength: 131072},
91+
},
92+
}
93+
json.NewEncoder(w).Encode(resp)
94+
}))
95+
defer srv.Close()
96+
97+
// First call hits the server
98+
ctx1 := DiscoverModelContext(srv.URL, "test-key", "flash")
99+
if ctx1 != 131072 {
100+
t.Errorf("expected 131072, got %d", ctx1)
101+
}
102+
103+
// Second call should be cached
104+
ctx2 := DiscoverModelContext(srv.URL, "test-key", "flash")
105+
if ctx2 != 131072 {
106+
t.Errorf("expected 131072, got %d", ctx2)
107+
}
108+
109+
if callCount != 1 {
110+
t.Errorf("expected 1 server call (cached), got %d", callCount)
111+
}
112+
}
113+
114+
func TestDiscoverModelContext_MaxContextFallback(t *testing.T) {
115+
ResetModelCache()
116+
117+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
118+
resp := modelsResponse{
119+
Data: []rawModel{
120+
{ID: "my-model", MaxContext: 64000},
121+
},
122+
}
123+
json.NewEncoder(w).Encode(resp)
124+
}))
125+
defer srv.Close()
126+
127+
ctx := DiscoverModelContext(srv.URL, "test-key", "my-model")
128+
if ctx != 64000 {
129+
t.Errorf("expected 64000 from max_context fallback, got %d", ctx)
130+
}
131+
}
132+
133+
func TestDiscoverModelContext_ModelsArrayFallback(t *testing.T) {
134+
ResetModelCache()
135+
136+
// Some providers use "models" instead of "data"
137+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
138+
json.NewEncoder(w).Encode(map[string][]rawModel{
139+
"models": {
140+
{ID: "claude-sonnet-4", ContextLength: 200000},
141+
},
142+
})
143+
}))
144+
defer srv.Close()
145+
146+
ctx := DiscoverModelContext(srv.URL, "test-key", "claude-sonnet-4")
147+
if ctx != 200000 {
148+
t.Errorf("expected 200000, got %d", ctx)
149+
}
150+
}

‎odek.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -370,10 +370,20 @@ func New(cfg Config) (*Agent, error) {
370370
timeout = profile.Timeout
371371
}
372372

373-
// Resolve max context: profile value (0 = no limit)
373+
// Resolve max context: discovered API values > profile > 0 (no limit)
374374
maxContext := 0
375-
if profile := LookupProfile(cfg.Model); profile != nil && profile.MaxContext > 0 {
376-
maxContext = profile.MaxContext
375+
// Priority 1: dynamic discovery via GET /models endpoint
376+
if discovered := llm.DiscoverModelContext(cfg.BaseURL, cfg.APIKey, cfg.Model); discovered > 0 {
377+
maxContext = discovered
378+
}
379+
// Priority 2: static profile fallback (only if discovery returned nothing)
380+
if maxContext == 0 {
381+
if profile := LookupProfile(cfg.Model); profile != nil && profile.MaxContext > 0 {
382+
maxContext = profile.MaxContext
383+
}
384+
}
385+
if maxContext > 0 {
386+
log.Printf("odek: model %q context window: %d tokens", cfg.Model, maxContext)
377387
}
378388

379389
// Build tool registry from external Tool interface

0 commit comments

Comments
 (0)