From 084d34ba6c2febea2833dd8d66c98e6a074845c0 Mon Sep 17 00:00:00 2001 From: David Gageot Date: Tue, 12 May 2026 12:14:46 +0200 Subject: [PATCH] fix: rewrite vertex model garden provider to publisher for capability lookups --- pkg/model/provider/vertexai/modelgarden.go | 52 +++++- .../provider/vertexai/modelgarden_test.go | 160 ++++++++++++++++++ 2 files changed, 210 insertions(+), 2 deletions(-) diff --git a/pkg/model/provider/vertexai/modelgarden.go b/pkg/model/provider/vertexai/modelgarden.go index e0a40bf3a..8d4488360 100644 --- a/pkg/model/provider/vertexai/modelgarden.go +++ b/pkg/model/provider/vertexai/modelgarden.go @@ -77,15 +77,63 @@ func IsModelGardenConfig(cfg *latest.ModelConfig) bool { // // - publisher: anthropic → Anthropic-native `:streamRawPredict` endpoint. // - other publishers → Vertex AI's OpenAI-compatible `/chat/completions`. +// +// The underlying client is constructed with a Provider field rewritten so +// that downstream capability lookups against models.dev resolve correctly. +// Without this rewrite, attachments and other capability-gated features are +// silently dropped because `/` IDs like +// `google/claude-sonnet-4-20250514` do not exist in the models.dev database. +// See modelsDevProvider for the exact mapping. func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Client, error) { project, location, err := resolveProjectLocation(ctx, cfg, env) if err != nil { return nil, err } + rewritten := withModelsDevProvider(cfg) if strings.EqualFold(publisher(cfg), "anthropic") { - return anthropic.NewVertexClient(ctx, cfg, env, project, location, opts...) + return anthropic.NewVertexClient(ctx, rewritten, env, project, location, opts...) + } + return newOpenAIClient(ctx, rewritten, env, project, location, opts...) +} + +// withModelsDevProvider returns a deep copy of cfg with its Provider field +// rewritten so that `Provider + "/" + Model` becomes a valid models.dev key. +// The original cfg is never mutated so callers — and the wider runtime — +// keep their "google" provider view. +func withModelsDevProvider(cfg *latest.ModelConfig) *latest.ModelConfig { + out := cfg.Clone() + if p := modelsDevProvider(cfg); p != "" { + out.Provider = p + } + return out +} + +// modelsDevProvider returns the models.dev provider key that hosts the model +// described by cfg, or "" if the publisher is unset. +// +// On Vertex AI Model Garden the model-name format depends on the publisher: +// +// - Anthropic Claude is referenced by a bare model name (no prefix), e.g. +// `claude-sonnet-4-20250514`. models.dev keys these under the top-level +// `anthropic` provider. +// - Other publishers are referenced via Vertex AI's OpenAI-compatible API +// with a `/` prefix, e.g. +// `meta/llama-4-maverick-17b-128e-instruct-maas`. models.dev keys these +// under `google-vertex` with the publisher prefix kept in the model name. +// +// Returning the publisher name unchanged for non-Anthropic models would +// produce double-prefixed IDs like `meta/meta/llama-4-...` that do not exist +// in models.dev. +func modelsDevProvider(cfg *latest.ModelConfig) string { + p := strings.ToLower(strings.TrimSpace(publisher(cfg))) + switch p { + case "": + return "" + case "anthropic": + return "anthropic" + default: + return "google-vertex" } - return newOpenAIClient(ctx, cfg, env, project, location, opts...) } // publisher returns the provider_opts.publisher string, or "" if unset. diff --git a/pkg/model/provider/vertexai/modelgarden_test.go b/pkg/model/provider/vertexai/modelgarden_test.go index 5e9a3b878..be4f7ea66 100644 --- a/pkg/model/provider/vertexai/modelgarden_test.go +++ b/pkg/model/provider/vertexai/modelgarden_test.go @@ -4,6 +4,9 @@ import ( "strings" "testing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/environment" ) @@ -72,6 +75,163 @@ func TestIsModelGardenConfig(t *testing.T) { } } +func TestWithModelsDevProvider(t *testing.T) { + tests := []struct { + name string + cfg *latest.ModelConfig + wantProvider string + }{ + { + name: "anthropic publisher rewrites Provider to anthropic", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "claude-sonnet-4-20250514", + ProviderOpts: map[string]any{"publisher": "anthropic"}, + }, + wantProvider: "anthropic", + }, + { + name: "meta publisher rewrites Provider to google-vertex", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "meta/llama-4-maverick-17b-128e-instruct-maas", + ProviderOpts: map[string]any{"publisher": "meta"}, + }, + wantProvider: "google-vertex", + }, + { + name: "mistral publisher rewrites Provider to google-vertex", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "mistral/mistral-large-2411", + ProviderOpts: map[string]any{"publisher": "mistral"}, + }, + wantProvider: "google-vertex", + }, + { + name: "anthropic publisher is case-folded", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "claude-sonnet-4-20250514", + ProviderOpts: map[string]any{"publisher": "Anthropic"}, + }, + wantProvider: "anthropic", + }, + { + name: "publisher with surrounding whitespace is trimmed", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "claude-sonnet-4-20250514", + ProviderOpts: map[string]any{"publisher": " anthropic "}, + }, + wantProvider: "anthropic", + }, + { + name: "missing publisher leaves Provider untouched", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "gemini-2.5-flash", + }, + wantProvider: "google", + }, + { + name: "empty publisher leaves Provider untouched", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "gemini-2.5-flash", + ProviderOpts: map[string]any{"publisher": ""}, + }, + wantProvider: "google", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + originalProvider := tt.cfg.Provider + originalPublisher, _ := tt.cfg.ProviderOpts["publisher"].(string) + + got := withModelsDevProvider(tt.cfg) + require.NotNil(t, got) + require.NotSame(t, tt.cfg, got, "withModelsDevProvider must return a copy") + assert.Equal(t, tt.wantProvider, got.Provider) + + // The original config must be unchanged so the wider runtime keeps its + // "google" provider view (used for routing, telemetry, ...). + assert.Equal(t, originalProvider, tt.cfg.Provider, "input cfg.Provider was mutated") + gotPublisher, _ := tt.cfg.ProviderOpts["publisher"].(string) + assert.Equal(t, originalPublisher, gotPublisher, "input cfg.ProviderOpts[publisher] was mutated") + }) + } +} + +// TestWithModelsDevProvider_CapabilityLookupIDs pins the bug fix for issue +// #2740: the underlying provider client must compute its capability-lookup ID +// in a form that exists in the models.dev database. Concretely: +// +// - Anthropic Claude on Vertex AI: model name is bare (no prefix) and +// models.dev keys it under "anthropic/". +// - Other publishers on Vertex AI: model name carries a "/" +// prefix and models.dev keys it under "google-vertex//". +// +// The test guards against the regression where rewriting Provider to the +// publisher name (e.g. "meta") produced double-prefixed IDs like +// "meta/meta/llama-..." that do not exist in models.dev. +func TestWithModelsDevProvider_CapabilityLookupIDs(t *testing.T) { + tests := []struct { + name string + cfg *latest.ModelConfig + wantID string + }{ + { + name: "anthropic claude", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "claude-sonnet-4-20250514", + ProviderOpts: map[string]any{ + "project": "my-project", + "location": "us-east5", + "publisher": "anthropic", + }, + }, + wantID: "anthropic/claude-sonnet-4-20250514", + }, + { + name: "meta llama keeps the meta/ prefix exactly once", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "meta/llama-4-maverick-17b-128e-instruct-maas", + ProviderOpts: map[string]any{ + "project": "my-project", + "location": "us-central1", + "publisher": "meta", + }, + }, + wantID: "google-vertex/meta/llama-4-maverick-17b-128e-instruct-maas", + }, + { + name: "mistral keeps its mistral/ prefix exactly once", + cfg: &latest.ModelConfig{ + Provider: "google", + Model: "mistral/mistral-large-2411", + ProviderOpts: map[string]any{ + "project": "my-project", + "location": "us-central1", + "publisher": "mistral", + }, + }, + wantID: "google-vertex/mistral/mistral-large-2411", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rewritten := withModelsDevProvider(tt.cfg) + gotID := rewritten.Provider + "/" + rewritten.DisplayOrModel() + assert.Equal(t, tt.wantID, gotID) + }) + } +} + func TestPublisher(t *testing.T) { tests := []struct { name string