fixup! Add pi support

This commit is contained in:
2026-07-22 11:00:04 +01:00
parent b5e3bcd301
commit 62debbfcda
5 changed files with 107 additions and 11 deletions
+10 -10
View File
@@ -31,16 +31,16 @@ func GatewayProviderKey(gatewayURL string) string {
} }
// crushModel builds a 10-field model entry required by the crush schema. // crushModel builds a 10-field model entry required by the crush schema.
func crushModel(id string) map[string]any { func crushModel(m canonical.Model) map[string]any {
return map[string]any{ return map[string]any{
"id": id, "id": m.ID,
"name": id, "name": m.ID,
"cost_per_1m_in": 0, "cost_per_1m_in": 0,
"cost_per_1m_out": 0, "cost_per_1m_out": 0,
"cost_per_1m_in_cached": 0, "cost_per_1m_in_cached": 0,
"cost_per_1m_out_cached": 0, "cost_per_1m_out_cached": 0,
"context_window": 200000, "context_window": m.ContextLimit(),
"default_max_tokens": 8192, "default_max_tokens": m.OutputLimit(),
"can_reason": false, "can_reason": false,
"supports_attachments": true, "supports_attachments": true,
} }
@@ -56,10 +56,10 @@ func ProvidersAsCrushMap(p *canonical.Profile) map[string]any {
} }
models := make([]map[string]any, 0) models := make([]map[string]any, 0)
for _, m := range p.Models { for _, m := range p.Models {
models = append(models, crushModel(m.ID)) models = append(models, crushModel(m))
} }
if len(models) == 0 && p.Gateway.DefaultModel != "" { if len(models) == 0 && p.Gateway.DefaultModel != "" {
models = append(models, crushModel(p.Gateway.DefaultModel)) models = append(models, crushModel(canonical.Model{ID: p.Gateway.DefaultModel}))
} }
out[GatewayProviderKey(p.Gateway.URL)] = map[string]any{ out[GatewayProviderKey(p.Gateway.URL)] = map[string]any{
"type": "openai-compat", "type": "openai-compat",
@@ -101,7 +101,7 @@ func ProvidersAsMap(p *canonical.Profile) map[string]any {
} }
models[m.ID] = map[string]any{ models[m.ID] = map[string]any{
"name": displayName, "name": displayName,
"limit": map[string]any{"context": 200000, "output": 8192}, "limit": map[string]any{"context": m.ContextLimit(), "output": m.OutputLimit()},
} }
} }
entry := map[string]any{ entry := map[string]any{
@@ -172,7 +172,7 @@ func GooseCustomProviderFile(p *canonical.Profile) ([]byte, string) {
for _, m := range p.Models { for _, m := range p.Models {
models = append(models, map[string]any{ models = append(models, map[string]any{
"name": m.ID, "name": m.ID,
"context_limit": 200000, "context_limit": m.ContextLimit(),
}) })
} }
if len(models) == 0 && p.Gateway.DefaultModel != "" { if len(models) == 0 && p.Gateway.DefaultModel != "" {
@@ -183,7 +183,7 @@ func GooseCustomProviderFile(p *canonical.Profile) ([]byte, string) {
} }
models = append(models, map[string]any{ models = append(models, map[string]any{
"name": modelID, "name": modelID,
"context_limit": 200000, "context_limit": canonical.DefaultContextLimit,
}) })
} }
+50
View File
@@ -1,6 +1,7 @@
package common package common
import ( import (
"encoding/json"
"testing" "testing"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
@@ -64,3 +65,52 @@ func TestParseMCPFromJSONMissingKey(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
assert.Nil(t, servers) assert.Nil(t, servers)
} }
func TestProvidersAsMapHonorsPerModelLimits(t *testing.T) {
p := &canonical.Profile{
Gateway: canonical.Gateway{URL: "https://gw.example.com/v1", Token: "t"},
Models: []canonical.Model{
{ID: "prov/big", Context: 262144, Output: 65536},
{ID: "prov/default"},
},
}
models := ProvidersAsMap(p)[GatewayProviderKey(p.Gateway.URL)].(map[string]any)["models"].(map[string]any)
big := models["prov/big"].(map[string]any)["limit"].(map[string]any)
assert.Equal(t, 262144, big["context"], "explicit context override should win")
assert.Equal(t, 65536, big["output"], "explicit output override should win")
def := models["prov/default"].(map[string]any)["limit"].(map[string]any)
assert.Equal(t, 200000, def["context"], "unset model falls back to default context")
assert.Equal(t, 8192, def["output"], "unset model falls back to default output")
}
func TestProvidersAsCrushMapHonorsPerModelLimits(t *testing.T) {
p := &canonical.Profile{
Gateway: canonical.Gateway{URL: "https://gw.example.com/v1", Token: "t"},
Models: []canonical.Model{
{ID: "prov/big", Context: 262144, Output: 65536},
{ID: "prov/default"},
},
}
models := ProvidersAsCrushMap(p)[GatewayProviderKey(p.Gateway.URL)].(map[string]any)["models"].([]map[string]any)
assert.Equal(t, 262144, models[0]["context_window"], "explicit context override should win")
assert.Equal(t, 65536, models[0]["default_max_tokens"], "explicit output override should win")
assert.Equal(t, 200000, models[1]["context_window"], "unset model falls back to default context")
assert.Equal(t, 8192, models[1]["default_max_tokens"], "unset model falls back to default output")
}
func TestGooseCustomProviderFileHonorsPerModelContext(t *testing.T) {
p := &canonical.Profile{
Gateway: canonical.Gateway{URL: "https://gw.example.com/v1", Token: "t"},
Models: []canonical.Model{{ID: "big", Context: 262144}, {ID: "default"}},
}
body, _ := GooseCustomProviderFile(p)
var entry map[string]any
require.NoError(t, json.Unmarshal(body, &entry))
models := entry["models"].([]any)
assert.Equal(t, float64(262144), models[0].(map[string]any)["context_limit"], "explicit context override should win")
assert.Equal(t, float64(200000), models[1].(map[string]any)["context_limit"], "unset model falls back to default context")
}
+1 -1
View File
@@ -122,7 +122,7 @@ func zedLanguageModels(p *canonical.Profile) map[string]any {
for _, m := range p.Models { for _, m := range p.Models {
me := map[string]any{ me := map[string]any{
"name": m.ID, "name": m.ID,
"max_tokens": 200000, "max_tokens": m.ContextLimit(),
"capabilities": map[string]any{ "capabilities": map[string]any{
"tools": true, "tools": true,
"images": false, "images": false,
+19
View File
@@ -63,6 +63,25 @@ func TestZedRenderProducesExpectedTargets(t *testing.T) {
assert.NotContains(t, parsed, "context_servers") assert.NotContains(t, parsed, "context_servers")
} }
func TestZedRenderHonorsPerModelContext(t *testing.T) {
home := t.TempDir()
b := bundle()
b.Profile.Models = []canonical.Model{{ID: "big", Context: 262144}}
fs, err := New(WithHome(home)).Render(b)
require.NoError(t, err)
seen := map[string]adapter.File{}
fs.ForEach(func(f adapter.File) { seen[f.Dest] = f })
cfgDest := filepath.Join(home, ".config", "zed", "settings.json")
require.Contains(t, seen, cfgDest)
var parsed map[string]any
require.NoError(t, json.Unmarshal(seen[cfgDest].Content, &parsed))
models := parsed["language_models"].(map[string]any)["openai_compatible"].(map[string]any)["hs-llmgw-example-com"].(map[string]any)["available_models"].([]any)
assert.Equal(t, float64(262144), models[0].(map[string]any)["max_tokens"], "explicit context override should win")
}
func TestZedRenderMergesExistingKeys(t *testing.T) { func TestZedRenderMergesExistingKeys(t *testing.T) {
home := t.TempDir() home := t.TempDir()
base := filepath.Join(home, ".config", "zed") base := filepath.Join(home, ".config", "zed")
+27
View File
@@ -35,6 +35,33 @@ type Upstream struct {
type Model struct { type Model struct {
ID string `yaml:"id"` ID string `yaml:"id"`
Alias string `yaml:"alias,omitempty"` Alias string `yaml:"alias,omitempty"`
// Context and Output are optional per-model limit overrides (tokens). When
// unset (0), harnesses fall back to DefaultContextLimit / DefaultOutputLimit.
Context int `yaml:"context,omitempty"`
Output int `yaml:"output,omitempty"`
}
// Default per-model limits, applied by harnesses that surface context/output
// limits when a Model doesn't specify its own.
const (
DefaultContextLimit = 200000
DefaultOutputLimit = 8192
)
// ContextLimit returns the model's context window, or DefaultContextLimit if unset.
func (m Model) ContextLimit() int {
if m.Context > 0 {
return m.Context
}
return DefaultContextLimit
}
// OutputLimit returns the model's max output tokens, or DefaultOutputLimit if unset.
func (m Model) OutputLimit() int {
if m.Output > 0 {
return m.Output
}
return DefaultOutputLimit
} }
type MCPRegistry struct { type MCPRegistry struct {