From 7c92cad77ee670ece11fd38d63a372716631c6df Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 17:35:32 +0200 Subject: [PATCH 1/9] feat(providers): add vertex ai provider --- README.md | 14 +- cmd/gomodel/main.go | 2 + config/config.example.yaml | 15 +- config/config.go | 21 +- docs/advanced/config-yaml.mdx | 5 + docs/advanced/configuration.mdx | 22 ++ docs/guides/gemini.mdx | 93 ++++- go.mod | 2 + go.sum | 4 + internal/providers/config.go | 282 +++++++++++--- internal/providers/config_test.go | 205 +++++++++- internal/providers/factory.go | 1 + internal/providers/gemini/gemini.go | 405 ++++++++++++++++---- internal/providers/gemini/gemini_test.go | 228 ++++++++++- internal/providers/gemini/native.go | 25 ++ internal/providers/googleauth/googleauth.go | 112 ++++++ internal/providers/vertex/vertex.go | 377 ++++++++++++++++++ internal/providers/vertex/vertex_test.go | 152 ++++++++ 18 files changed, 1800 insertions(+), 165 deletions(-) create mode 100644 internal/providers/googleauth/googleauth.go create mode 100644 internal/providers/vertex/vertex.go create mode 100644 internal/providers/vertex/vertex_test.go diff --git a/README.md b/README.md index da6c4fde..6f594ea4 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,9 @@ docker run --rm -p 8080:8080 \ -e OPENAI_API_KEY="your-openai-key" \ -e ANTHROPIC_API_KEY="your-anthropic-key" \ -e GEMINI_API_KEY="your-gemini-key" \ + -e VERTEX_PROJECT="your-gcp-project" \ + -e VERTEX_LOCATION="us-central1" \ + -e VERTEX_AUTH_TYPE="gcp_adc" \ -e DEEPSEEK_API_KEY="your-deepseek-key" \ -e GROQ_API_KEY="your-groq-key" \ -e OPENROUTER_API_KEY="your-openrouter-key" \ @@ -90,7 +93,8 @@ Example model identifiers are illustrative and subject to change; consult provid | ------------- | ----------------------------------------------------------------- | ---------------------------------- | :--: | :----------: | :---: | :---: | :-----: | :------: | | OpenAI | `OPENAI_API_KEY` | `gpt-5.5` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | Anthropic | `ANTHROPIC_API_KEY` | `claude-sonnet-4-20250514` | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | -| Google Gemini | `GEMINI_API_KEY` | `gemini-2.5-flash` | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| Google Gemini | `GEMINI_API_KEY` | `gemini-2.5-flash` | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| Vertex AI | `VERTEX_PROJECT` + `VERTEX_LOCATION` | `google/gemini-2.5-flash` | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | | DeepSeek | `DEEPSEEK_API_KEY` | `deepseek-v4-pro` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | | Groq | `GROQ_API_KEY` | `llama-3.3-70b-versatile` | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | | OpenRouter | `OPENROUTER_API_KEY` | `google/gemini-2.5-flash` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | @@ -120,6 +124,10 @@ To register multiple instances of the same provider type without `config.yaml`, use suffixed env vars such as `OPENAI_EAST_API_KEY` and `OPENAI_EAST_BASE_URL`; add `OPENAI_EAST_MODELS` to configure that instance's model list. This registers provider `openai-east` with type `openai`. +Vertex AI uses the `VERTEX_*` prefix and registers provider `vertex` +with type `vertex`; suffixed variables such as `VERTEX_US_PROJECT` register +provider `vertex_us`. Vertex requires `VERTEX_PROJECT` and `VERTEX_LOCATION`; +`VERTEX_AUTH_TYPE` defaults to Application Default Credentials (`gcp_adc`). --- @@ -250,7 +258,9 @@ Key settings: | `ENABLE_PASSTHROUGH_ROUTES` | `true` | Enable provider-native passthrough routes under `/p/{provider}/...` | | `ALLOW_PASSTHROUGH_V1_ALIAS` | `true` | Allow `/p/{provider}/v1/...` aliases while keeping `/p/{provider}/...` canonical | | `ENABLED_PASSTHROUGH_PROVIDERS` | `openai,anthropic,openrouter,zai,vllm` | Comma-separated list of enabled passthrough providers | -| `USE_GOOGLE_GEMINI_NATIVE_API` | `true` | Use Gemini native `generateContent` for chat/responses; set `false` for Gemini's OpenAI-compatible API and image_url pass-through behavior | +| `GEMINI_API_MODE` | `native` | Gemini AI Studio upstream mode: `native` or `openai_compatible` | +| `VERTEX_API_MODE` | `native` | Vertex AI Gemini upstream mode: `native` or `openai_compatible` | +| `USE_GOOGLE_GEMINI_NATIVE_API` | `true` | Legacy global Gemini mode toggle used when per-provider `*_API_MODE` is unset | | `STORAGE_TYPE` | `sqlite` | Storage backend (`sqlite`, `postgresql`, `mongodb`) | | `METRICS_ENABLED` | `false` | Enable Prometheus metrics (experimental) | | `LOGGING_ENABLED` | `false` | Enable audit logging | diff --git a/cmd/gomodel/main.go b/cmd/gomodel/main.go index 18a9333e..789eee43 100644 --- a/cmd/gomodel/main.go +++ b/cmd/gomodel/main.go @@ -26,6 +26,7 @@ import ( "gomodel/internal/providers/openai" "gomodel/internal/providers/openrouter" "gomodel/internal/providers/oracle" + "gomodel/internal/providers/vertex" "gomodel/internal/providers/vllm" "gomodel/internal/providers/xai" "gomodel/internal/providers/zai" @@ -122,6 +123,7 @@ func main() { factory.Add(anthropic.Registration) factory.Add(deepseek.Registration) factory.Add(gemini.Registration) + factory.Add(vertex.Registration) factory.Add(groq.Registration) factory.Add(minimax.Registration) factory.Add(ollama.Registration) diff --git a/config/config.example.yaml b/config/config.example.yaml index 42575813..ace7acff 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -202,9 +202,22 @@ providers: type: gemini api_key: "..." # Chat/responses use Gemini's native generateContent API by default. - # Set USE_GOOGLE_GEMINI_NATIVE_API=false to use Gemini's OpenAI-compatible API instead. + # Set GEMINI_API_MODE=openai_compatible or api_mode: openai_compatible to use + # Gemini's OpenAI-compatible API instead. # Native mode accepts image data URLs but does not fetch remote image URLs yet. + vertex: + type: vertex + auth_type: gcp_adc # or gcp_service_account + vertex_project: "my-gcp-project" + vertex_location: "us-central1" + # service_account_file: "/var/run/secrets/google/service-account.json" + # service_account_json: "${VERTEX_SERVICE_ACCOUNT_JSON}" + # service_account_json_base64: "${VERTEX_SERVICE_ACCOUNT_JSON_BASE64}" + # api_mode: native # or openai_compatible + # models: + # - id: "google/gemini-2.5-flash" + xai: type: xai api_key: "..." diff --git a/config/config.go b/config/config.go index 9dd21a31..18d718e4 100644 --- a/config/config.go +++ b/config/config.go @@ -54,12 +54,21 @@ type LoadResult struct { // overrides, credential filtering, or resilience merging. Exported so the // providers package can resolve it into a fully-configured ProviderConfig. type RawProviderConfig struct { - Type string `yaml:"type"` - APIKey string `yaml:"api_key"` - BaseURL string `yaml:"base_url"` - APIVersion string `yaml:"api_version"` - Models []RawProviderModel `yaml:"models"` - Resilience *RawResilienceConfig `yaml:"resilience"` + Type string `yaml:"type"` + APIKey string `yaml:"api_key"` + BaseURL string `yaml:"base_url"` + APIVersion string `yaml:"api_version"` + Backend string `yaml:"backend"` + AuthType string `yaml:"auth_type"` + APIMode string `yaml:"api_mode"` + VertexProject string `yaml:"vertex_project"` + VertexLocation string `yaml:"vertex_location"` + ServiceAccountFile string `yaml:"service_account_file"` + ServiceAccountJSON string `yaml:"service_account_json"` + ServiceAccountJSONBase64 string `yaml:"service_account_json_base64"` + GCPScope string `yaml:"gcp_scope"` + Models []RawProviderModel `yaml:"models"` + Resilience *RawResilienceConfig `yaml:"resilience"` } // RawResilienceConfig holds optional per-provider resilience overrides from YAML. diff --git a/docs/advanced/config-yaml.mdx b/docs/advanced/config-yaml.mdx index 8fb05be9..9e85c5ad 100644 --- a/docs/advanced/config-yaml.mdx +++ b/docs/advanced/config-yaml.mdx @@ -25,6 +25,11 @@ For multiple provider instances, env vars support `OLLAMA_A_BASE_URL` registers `ollama-a`. Azure also supports `__API_VERSION`. +Vertex AI is the exception to the provider-name shape: use `VERTEX_*` +env vars to create `type: vertex` providers named `vertex`, `vertex_us`, +`vertex_eu`, and so on. Vertex requires `VERTEX_PROJECT` and +`VERTEX_LOCATION`; `VERTEX_AUTH_TYPE` defaults to `gcp_adc`. + Configured provider model lists can stay in env via `_MODELS`, for example `OPENROUTER_MODELS`, `ORACLE_MODELS`, `AZURE_MODELS`, or `VLLM_MODELS`. Set `CONFIGURED_PROVIDER_MODELS_MODE=fallback` (default) to use those lists only diff --git a/docs/advanced/configuration.mdx b/docs/advanced/configuration.mdx index d3376bd1..5b32d3b2 100644 --- a/docs/advanced/configuration.mdx +++ b/docs/advanced/configuration.mdx @@ -350,6 +350,18 @@ The same pattern works for every registered provider type: suffixed `BASE_URL` values because their endpoints are deployment- or region-specific. +Vertex AI uses the dedicated `VERTEX_*` prefix and registers provider +`vertex` with type `vertex`. Suffixed Vertex env vars keep underscores in the +provider name: + +```bash +export VERTEX_PROJECT="prod-ai" +export VERTEX_LOCATION="us-central1" # Registers "vertex", type "vertex" + +export VERTEX_US_PROJECT="prod-ai" +export VERTEX_US_LOCATION="us-central1" # Registers "vertex_us", type "vertex" +``` + ### YAML Provider Blocks For more control (custom names, per-provider resilience, or larger structured @@ -404,6 +416,16 @@ providers: models: - gemini-2.0-flash - gemini-1.5-pro + + # Add Vertex AI. Project and location are required. + vertex: + type: vertex + auth_type: gcp_adc + vertex_project: "prod-ai" + vertex_location: "us-central1" + api_mode: native + models: + - google/gemini-2.5-flash ``` diff --git a/docs/guides/gemini.mdx b/docs/guides/gemini.mdx index 85e18495..8789f901 100644 --- a/docs/guides/gemini.mdx +++ b/docs/guides/gemini.mdx @@ -9,12 +9,13 @@ GoModel routes Gemini chat and Responses API requests through Gemini's native Gemini's OpenAI-compatible API when you need compatibility behavior that the native adapter does not implement yet. -## Configure Gemini +## Configure AI Studio -Env-only configuration is enough: +Use the `GEMINI_*` prefix for Gemini API keys from AI Studio: ```bash export GEMINI_API_KEY="..." +export GEMINI_API_MODE="native" ``` Or in `config.yaml`: @@ -24,21 +25,80 @@ providers: gemini: type: gemini api_key: "${GEMINI_API_KEY}" + api_mode: native ``` +## Configure Vertex AI + +Use the `VERTEX_*` prefix for Gemini on Vertex AI. Vertex providers use the +dedicated `vertex` provider type: + +```bash +export VERTEX_PROJECT="my-gcp-project" +export VERTEX_LOCATION="us-central1" +export VERTEX_AUTH_TYPE="gcp_adc" +export VERTEX_API_MODE="native" +``` + +This registers provider `vertex`, type `vertex`. For multiple Vertex regions +or accounts, use suffixes: + +```bash +export VERTEX_US_PROJECT="prod-ai" +export VERTEX_US_LOCATION="us-central1" +export VERTEX_US_AUTH_TYPE="gcp_adc" + +export VERTEX_EU_PROJECT="prod-ai" +export VERTEX_EU_LOCATION="europe-west4" +export VERTEX_EU_AUTH_TYPE="gcp_service_account" +export VERTEX_EU_SERVICE_ACCOUNT_FILE="/secrets/vertex-eu.json" +``` + +These register providers `vertex_us` and `vertex_eu`. + +Or in `config.yaml`: + +```yaml +providers: + vertex: + type: vertex + auth_type: gcp_adc + vertex_project: my-gcp-project + vertex_location: us-central1 + api_mode: native +``` + +Vertex requires both project and location. `VERTEX_AUTH_TYPE` defaults to +Application Default Credentials (`gcp_adc`) when service-account fields are not +set. Use service-account JSON directly when ADC is not appropriate: + +```bash +export VERTEX_AUTH_TYPE="gcp_service_account" +export VERTEX_SERVICE_ACCOUNT_FILE="/secrets/service-account.json" +# or VERTEX_SERVICE_ACCOUNT_JSON / VERTEX_SERVICE_ACCOUNT_JSON_BASE64 +``` + +The default Vertex bases are derived from project and location: + +- OpenAI-compatible: `https://aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi` +- native Gemini: `https://aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/google` + ## Native versus OpenAI-compatible mode Gemini native mode is enabled by default: ```bash -export USE_GOOGLE_GEMINI_NATIVE_API=true +export GEMINI_API_MODE="native" +export VERTEX_API_MODE="native" ``` -Set it to `false` to route chat and Responses API requests through Gemini's -OpenAI-compatible `/chat/completions` endpoint: +Set the per-provider API mode to `openai_compatible` to route chat and +Responses API requests through the upstream OpenAI-compatible +`/chat/completions` endpoint: ```bash -export USE_GOOGLE_GEMINI_NATIVE_API=false +export GEMINI_API_MODE="openai_compatible" +export VERTEX_API_MODE="openai_compatible" ``` `GEMINI_BASE_URL` configures the Gemini base. GoModel keeps separate internal @@ -50,15 +110,17 @@ bases for native Gemini and the OpenAI-compatible API: When `GEMINI_BASE_URL` ends in `/openai`, GoModel uses that value for the OpenAI-compatible client and derives the native base by stripping `/openai`. Gemini embeddings, files, and batches still use the OpenAI-compatible surface. +Vertex embeddings use Vertex AI native prediction. Vertex does not expose +Gemini Files or OpenAI-compatible batch operations. ```bash export GEMINI_BASE_URL="https://generativelanguage.googleapis.com/v1beta/openai" ``` - `USE_GOOGLE_GEMINI_NATIVE_API` decides whether chat and Responses API calls - use native Gemini or the OpenAI-compatible API. `GEMINI_BASE_URL` only - configures the upstream base URLs. + `USE_GOOGLE_GEMINI_NATIVE_API` is still supported as a legacy global toggle + when per-provider `*_API_MODE` is unset. Prefer `GEMINI_API_MODE` and + `VERTEX_API_MODE` for new deployments. ## Image URL behavior @@ -84,8 +146,9 @@ mode. Google's native Gemini API supports inline image data and Files API references; for URL-hosted images, Google's examples fetch the URL first and send the bytes to `generateContent`. -Set `USE_GOOGLE_GEMINI_NATIVE_API=false` when you need GoModel to pass the -OpenAI-compatible `image_url` request shape through to Gemini's +Set `GEMINI_API_MODE=openai_compatible` or +`VERTEX_API_MODE=openai_compatible` when you need GoModel to pass the +OpenAI-compatible `image_url` request shape through to Google's OpenAI-compatible endpoint instead. Google documents image input for that endpoint using the OpenAI `image_url` field. @@ -95,17 +158,21 @@ Integrated: - chat completions and streaming - Responses API and streaming -- model listing through Gemini's native `/models` +- model listing through AI Studio `/models` and Vertex publisher models +- Vertex embeddings through native prediction - usage metadata normalization for native responses - tool calls and function-call results - inline image data via `data:` URLs in native mode -Not integrated in native mode yet: +Not integrated yet: - fetching remote `image_url` values - uploading remote images through the Gemini Files API before a chat request +- Vertex Files and batch prediction APIs References: - [Gemini image understanding](https://ai.google.dev/gemini-api/docs/image-understanding) - [Gemini OpenAI compatibility](https://ai.google.dev/gemini-api/docs/openai) +- [Vertex AI OpenAI compatibility](https://cloud.google.com/vertex-ai/generative-ai/docs/start/openai) +- [Gemini API in Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference) diff --git a/go.mod b/go.mod index fb82a9b9..626e0ce7 100644 --- a/go.mod +++ b/go.mod @@ -17,6 +17,7 @@ require ( github.com/swaggo/swag/v2 v2.0.0-rc5 github.com/tidwall/gjson v1.18.0 go.mongodb.org/mongo-driver/v2 v2.6.0 + golang.org/x/oauth2 v0.36.0 golang.org/x/sync v0.20.0 golang.org/x/term v0.42.0 gopkg.in/yaml.v3 v3.0.1 @@ -24,6 +25,7 @@ require ( ) require ( + cloud.google.com/go/compute/metadata v0.3.0 // indirect github.com/KyleBanks/depth v1.2.1 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/cpuguy83/go-md2man/v2 v2.0.5 // indirect diff --git a/go.sum b/go.sum index bc7a4803..08e781fb 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +cloud.google.com/go/compute/metadata v0.3.0 h1:Tz+eQXMEqDIKRsmY3cHTL6FVaynIjX2QxYC4trgAKZc= +cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k= github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc= github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE= github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro= @@ -161,6 +163,8 @@ golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= golang.org/x/net v0.50.0 h1:ucWh9eiCGyDR3vtzso0WMQinm2Dnt8cFMuQa9K33J60= golang.org/x/net v0.50.0/go.mod h1:UgoSli3F/pBgdJBHCTc+tp3gmrU4XswgGRgtnwWTfyM= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= diff --git a/internal/providers/config.go b/internal/providers/config.go index 28cbcc15..94f05f5a 100644 --- a/internal/providers/config.go +++ b/internal/providers/config.go @@ -14,11 +14,20 @@ import ( // ProviderConfig holds the fully resolved provider configuration after merging // global defaults with per-provider overrides. type ProviderConfig struct { - Type string - APIKey string - BaseURL string - APIVersion string - Models []string + Type string + APIKey string + BaseURL string + APIVersion string + Backend string + AuthType string + APIMode string + VertexProject string + VertexLocation string + ServiceAccountFile string + ServiceAccountJSON string + ServiceAccountJSONBase64 string + GCPScope string + Models []string // ModelMetadataOverrides holds operator-supplied metadata keyed by raw model // ID (as it appears in the provider's /models response). The registry merges // these onto remote-registry metadata after enrichment; non-zero fields here @@ -47,17 +56,19 @@ func applyProviderEnvVars(raw map[string]config.RawProviderConfig, discovery map for _, providerType := range sortedDiscoveryTypes(discovery) { spec := discovery[providerType] - envGroups := collectProviderEnvValues(providerType, spec, environ) + for _, source := range providerEnvSources(providerType, spec) { + envGroups := collectProviderEnvValues(source.Prefix, spec, environ) - if values, ok := envGroups[""]; ok { - applyUnsuffixedProviderEnvVars(result, providerType, spec, values) - } + if values, ok := envGroups[""]; ok { + applyUnsuffixedProviderEnvVars(result, providerType, spec, source, values) + } - for _, suffix := range sortedProviderEnvSuffixes(envGroups) { - if suffix == "" { - continue + for _, suffix := range sortedProviderEnvSuffixes(envGroups) { + if suffix == "" { + continue + } + applySuffixedProviderEnvVars(result, providerType, spec, source, suffix, envGroups[suffix]) } - applySuffixedProviderEnvVars(result, providerType, spec, suffix, envGroups[suffix]) } } @@ -71,25 +82,71 @@ const ( providerEnvFieldBaseURL providerEnvFieldAPIVersion providerEnvFieldModels + providerEnvFieldBackend + providerEnvFieldAuthType + providerEnvFieldAPIMode + providerEnvFieldVertexProject + providerEnvFieldVertexLocation + providerEnvFieldServiceAccountFile + providerEnvFieldServiceAccountJSON + providerEnvFieldServiceAccountJSONBase64 + providerEnvFieldGCPScope ) +type providerEnvSource struct { + Prefix string + DefaultName string + NameSeparator string + OverlayByType bool +} + type providerEnvValues struct { - APIKey string - BaseURL string - APIVersion string - Models []string + APIKey string + BaseURL string + APIVersion string + Backend string + AuthType string + APIMode string + VertexProject string + VertexLocation string + ServiceAccountFile string + ServiceAccountJSON string + ServiceAccountJSONBase64 string + GCPScope string + Models []string } func (v providerEnvValues) empty() bool { return strings.TrimSpace(v.APIKey) == "" && strings.TrimSpace(v.BaseURL) == "" && strings.TrimSpace(v.APIVersion) == "" && + strings.TrimSpace(v.Backend) == "" && + strings.TrimSpace(v.AuthType) == "" && + strings.TrimSpace(v.APIMode) == "" && + strings.TrimSpace(v.VertexProject) == "" && + strings.TrimSpace(v.VertexLocation) == "" && + strings.TrimSpace(v.ServiceAccountFile) == "" && + strings.TrimSpace(v.ServiceAccountJSON) == "" && + strings.TrimSpace(v.ServiceAccountJSONBase64) == "" && + strings.TrimSpace(v.GCPScope) == "" && len(v.Models) == 0 } -func collectProviderEnvValues(providerType string, spec DiscoveryConfig, environ []string) map[string]providerEnvValues { +func providerEnvSources(providerType string, spec DiscoveryConfig) []providerEnvSource { + separator := spec.NameSeparator + if separator == "" { + separator = "-" + } + return []providerEnvSource{{ + Prefix: envPrefix(providerType), + DefaultName: providerType, + NameSeparator: separator, + OverlayByType: true, + }} +} + +func collectProviderEnvValues(prefix string, spec DiscoveryConfig, environ []string) map[string]providerEnvValues { groups := make(map[string]providerEnvValues) - prefix := envPrefix(providerType) prefixWithSeparator := prefix + "_" for _, entry := range environ { @@ -113,6 +170,24 @@ func collectProviderEnvValues(providerType string, spec DiscoveryConfig, environ values.APIVersion = value case providerEnvFieldModels: values.Models = parseCSVEnvList(value) + case providerEnvFieldBackend: + values.Backend = value + case providerEnvFieldAuthType: + values.AuthType = value + case providerEnvFieldAPIMode: + values.APIMode = value + case providerEnvFieldVertexProject: + values.VertexProject = value + case providerEnvFieldVertexLocation: + values.VertexLocation = value + case providerEnvFieldServiceAccountFile: + values.ServiceAccountFile = value + case providerEnvFieldServiceAccountJSON: + values.ServiceAccountJSON = value + case providerEnvFieldServiceAccountJSONBase64: + values.ServiceAccountJSONBase64 = value + case providerEnvFieldGCPScope: + values.GCPScope = value } groups[suffix] = values } @@ -139,8 +214,19 @@ func parseProviderEnvKey(prefix, key string, spec DiscoveryConfig) (string, prov name string field providerEnvField }{ + {name: "SERVICE_ACCOUNT_JSON_BASE64", field: providerEnvFieldServiceAccountJSONBase64}, + {name: "SERVICE_ACCOUNT_JSON", field: providerEnvFieldServiceAccountJSON}, + {name: "SERVICE_ACCOUNT_FILE", field: providerEnvFieldServiceAccountFile}, + {name: "VERTEX_PROJECT", field: providerEnvFieldVertexProject}, + {name: "VERTEX_LOCATION", field: providerEnvFieldVertexLocation}, {name: "API_VERSION", field: providerEnvFieldAPIVersion}, {name: "BASE_URL", field: providerEnvFieldBaseURL}, + {name: "AUTH_TYPE", field: providerEnvFieldAuthType}, + {name: "API_MODE", field: providerEnvFieldAPIMode}, + {name: "PROJECT", field: providerEnvFieldVertexProject}, + {name: "LOCATION", field: providerEnvFieldVertexLocation}, + {name: "GCP_SCOPE", field: providerEnvFieldGCPScope}, + {name: "BACKEND", field: providerEnvFieldBackend}, {name: "API_KEY", field: providerEnvFieldAPIKey}, {name: "MODELS", field: providerEnvFieldModels}, } @@ -192,12 +278,12 @@ func sortedProviderEnvSuffixes(groups map[string]providerEnvValues) []string { return suffixes } -func applyUnsuffixedProviderEnvVars(result map[string]config.RawProviderConfig, providerType string, spec DiscoveryConfig, values providerEnvValues) { +func applyUnsuffixedProviderEnvVars(result map[string]config.RawProviderConfig, providerType string, spec DiscoveryConfig, source providerEnvSource, values providerEnvValues) { if values.empty() { return } - targetKey, matched, ambiguous := findEnvOverlayTarget(result, providerType) + targetKey, matched, ambiguous := findEnvOverlayTarget(result, providerType, source) if matched { result[targetKey] = overlayProviderEnvValues(result[targetKey], values, spec) return @@ -209,15 +295,15 @@ func applyUnsuffixedProviderEnvVars(result map[string]config.RawProviderConfig, return } - result[providerType] = values.rawConfig(providerType, spec) + result[source.DefaultName] = values.rawConfig(providerType, spec) } -func applySuffixedProviderEnvVars(result map[string]config.RawProviderConfig, providerType string, spec DiscoveryConfig, suffix string, values providerEnvValues) { +func applySuffixedProviderEnvVars(result map[string]config.RawProviderConfig, providerType string, spec DiscoveryConfig, source providerEnvSource, suffix string, values providerEnvValues) { if values.empty() { return } - targetKey := providerNameForEnvSuffix(providerType, suffix) + targetKey := providerNameForEnvSuffix(source, suffix) if targetKey == "" { return } @@ -238,12 +324,22 @@ func applySuffixedProviderEnvVars(result map[string]config.RawProviderConfig, pr } func (v providerEnvValues) rawConfig(providerType string, spec DiscoveryConfig) config.RawProviderConfig { + backend := v.Backend return config.RawProviderConfig{ - Type: providerType, - APIKey: v.APIKey, - BaseURL: v.resolvedBaseURL(spec), - APIVersion: v.APIVersion, - Models: rawProviderModelsFromIDs(v.Models), + Type: providerType, + APIKey: v.APIKey, + BaseURL: v.resolvedBaseURL(spec), + APIVersion: v.APIVersion, + Backend: backend, + AuthType: v.AuthType, + APIMode: v.APIMode, + VertexProject: v.VertexProject, + VertexLocation: v.VertexLocation, + ServiceAccountFile: v.ServiceAccountFile, + ServiceAccountJSON: v.ServiceAccountJSON, + ServiceAccountJSONBase64: v.ServiceAccountJSONBase64, + GCPScope: v.GCPScope, + Models: rawProviderModelsFromIDs(v.Models), } } @@ -267,25 +363,55 @@ func overlayProviderEnvValues(existing config.RawProviderConfig, values provider if values.APIVersion != "" { existing.APIVersion = values.APIVersion } + if values.Backend != "" { + existing.Backend = values.Backend + } + if values.AuthType != "" { + existing.AuthType = values.AuthType + } + if values.APIMode != "" { + existing.APIMode = values.APIMode + } + if values.VertexProject != "" { + existing.VertexProject = values.VertexProject + } + if values.VertexLocation != "" { + existing.VertexLocation = values.VertexLocation + } + if values.ServiceAccountFile != "" { + existing.ServiceAccountFile = values.ServiceAccountFile + } + if values.ServiceAccountJSON != "" { + existing.ServiceAccountJSON = values.ServiceAccountJSON + } + if values.ServiceAccountJSONBase64 != "" { + existing.ServiceAccountJSONBase64 = values.ServiceAccountJSONBase64 + } + if values.GCPScope != "" { + existing.GCPScope = values.GCPScope + } if len(values.Models) > 0 { existing.Models = rawProviderModelsFromIDs(values.Models) } return existing } -func providerNameForEnvSuffix(providerType, suffix string) string { - providerType = strings.TrimSpace(providerType) - suffixName := normalizeEnvSuffixForProviderName(suffix) +func providerNameForEnvSuffix(source providerEnvSource, suffix string) string { + baseName := strings.TrimSpace(source.DefaultName) + suffixName := normalizeEnvSuffixForProviderName(suffix, source.NameSeparator) if suffixName == "" { - return providerType + return baseName } - if providerType == "" { + if baseName == "" { return suffixName } - return providerType + "-" + suffixName + return baseName + source.NameSeparator + suffixName } -func normalizeEnvSuffixForProviderName(suffix string) string { +func normalizeEnvSuffixForProviderName(suffix, separator string) string { + if separator == "" { + separator = "-" + } var b strings.Builder lastHyphen := false for _, r := range strings.TrimSpace(suffix) { @@ -294,16 +420,19 @@ func normalizeEnvSuffixForProviderName(suffix string) string { b.WriteRune(unicode.ToLower(r)) lastHyphen = false case r == '_' && !lastHyphen: - b.WriteByte('-') + b.WriteString(separator) lastHyphen = true } } - return strings.Trim(b.String(), "-") + return strings.Trim(b.String(), separator) } -func findEnvOverlayTarget(raw map[string]config.RawProviderConfig, providerType string) (string, bool, bool) { - if existing, ok := raw[providerType]; ok && rawProviderMatchesType(existing, providerType) { - return providerType, true, false +func findEnvOverlayTarget(raw map[string]config.RawProviderConfig, providerType string, source providerEnvSource) (string, bool, bool) { + if existing, ok := raw[source.DefaultName]; ok && rawProviderMatchesType(existing, providerType) { + return source.DefaultName, true, false + } + if !source.OverlayByType { + return "", false, false } var matchedKey string @@ -326,6 +455,9 @@ func findEnvOverlayTarget(raw map[string]config.RawProviderConfig, providerType } func rawProviderMatchesType(cfg config.RawProviderConfig, providerType string) bool { + if strings.EqualFold(strings.TrimSpace(providerType), "vertex") { + return isVertexProviderConfig(cfg) + } return strings.TrimSpace(cfg.Type) == strings.TrimSpace(providerType) } @@ -412,10 +544,18 @@ func isUnresolvedEnvPlaceholder(value string) bool { func filterEmptyProviders(raw map[string]config.RawProviderConfig, discovery map[string]DiscoveryConfig) map[string]config.RawProviderConfig { result := make(map[string]config.RawProviderConfig, len(raw)) for name, p := range raw { - spec, known := discovery[strings.TrimSpace(p.Type)] + providerType := normalizeProviderType(p) + spec, known := discovery[providerType] if known && spec.RequireBaseURL && strings.TrimSpace(p.BaseURL) == "" { continue } + if isVertexProviderConfig(p) { + p.Type = providerType + if validVertexProviderConfig(p) { + result[name] = p + } + continue + } if known && spec.AllowAPIKeyless { result[name] = p continue @@ -427,6 +567,33 @@ func filterEmptyProviders(raw map[string]config.RawProviderConfig, discovery map return result } +func isVertexProviderConfig(p config.RawProviderConfig) bool { + return strings.EqualFold(strings.TrimSpace(p.Type), "vertex") || + (strings.EqualFold(strings.TrimSpace(p.Type), "gemini") && strings.EqualFold(strings.TrimSpace(p.Backend), "vertex")) +} + +func validVertexProviderConfig(p config.RawProviderConfig) bool { + if !hasResolvedProviderValue(p.VertexProject) || !hasResolvedProviderValue(p.VertexLocation) { + return false + } + authType := strings.ToLower(strings.TrimSpace(p.AuthType)) + switch authType { + case "", "gcp_adc", "adc", "google_adc": + return true + case "gcp_service_account", "service_account": + return hasResolvedProviderValue(p.ServiceAccountFile) || + hasResolvedProviderValue(p.ServiceAccountJSON) || + hasResolvedProviderValue(p.ServiceAccountJSONBase64) + default: + return false + } +} + +func hasResolvedProviderValue(value string) bool { + value = strings.TrimSpace(value) + return value != "" && !strings.Contains(value, "${") +} + // buildProviderConfigs merges each raw provider config with the global ResilienceConfig, // producing fully resolved ProviderConfig values. func buildProviderConfigs(raw map[string]config.RawProviderConfig, global config.ResilienceConfig) map[string]ProviderConfig { @@ -441,13 +608,22 @@ func buildProviderConfigs(raw map[string]config.RawProviderConfig, global config // Non-nil fields in the raw config override the global defaults. func buildProviderConfig(raw config.RawProviderConfig, global config.ResilienceConfig) ProviderConfig { resolved := ProviderConfig{ - Type: raw.Type, - APIKey: raw.APIKey, - BaseURL: raw.BaseURL, - APIVersion: raw.APIVersion, - Models: config.ProviderModelIDs(raw.Models), - ModelMetadataOverrides: config.ProviderModelMetadataOverrides(raw.Models), - Resilience: global, + Type: normalizeProviderType(raw), + APIKey: raw.APIKey, + BaseURL: raw.BaseURL, + APIVersion: raw.APIVersion, + Backend: raw.Backend, + AuthType: raw.AuthType, + APIMode: raw.APIMode, + VertexProject: raw.VertexProject, + VertexLocation: raw.VertexLocation, + ServiceAccountFile: raw.ServiceAccountFile, + ServiceAccountJSON: raw.ServiceAccountJSON, + ServiceAccountJSONBase64: raw.ServiceAccountJSONBase64, + GCPScope: raw.GCPScope, + Models: config.ProviderModelIDs(raw.Models), + ModelMetadataOverrides: config.ProviderModelMetadataOverrides(raw.Models), + Resilience: global, } if raw.Resilience == nil { @@ -487,6 +663,14 @@ func buildProviderConfig(raw config.RawProviderConfig, global config.ResilienceC return resolved } +func normalizeProviderType(raw config.RawProviderConfig) string { + providerType := strings.TrimSpace(raw.Type) + if strings.EqualFold(providerType, "gemini") && strings.EqualFold(strings.TrimSpace(raw.Backend), "vertex") { + return "vertex" + } + return providerType +} + // rawProviderModelsFromIDs wraps a plain string slice into RawProviderModel // entries. Used for env-var-sourced model lists where metadata is never present. func rawProviderModelsFromIDs(ids []string) []config.RawProviderModel { diff --git a/internal/providers/config_test.go b/internal/providers/config_test.go index 37f2736f..ce594a3f 100644 --- a/internal/providers/config_test.go +++ b/internal/providers/config_test.go @@ -27,6 +27,9 @@ var testDiscoveryConfigs = map[string]DiscoveryConfig{ "gemini": { DefaultBaseURL: "https://generativelanguage.googleapis.com/v1beta/openai", }, + "vertex": { + NameSeparator: "_", + }, "deepseek": { DefaultBaseURL: "https://api.deepseek.com", }, @@ -171,10 +174,17 @@ func TestBuildProviderConfig_ZeroValueOverride(t *testing.T) { func TestBuildProviderConfig_PreservesFields(t *testing.T) { raw := config.RawProviderConfig{ - Type: "openai", - APIKey: "sk-key", - BaseURL: "https://custom.endpoint.com", - Models: []config.RawProviderModel{{ID: "gpt-4"}, {ID: "gpt-3.5-turbo"}}, + Type: "gemini", + APIKey: "sk-key", + BaseURL: "https://custom.endpoint.com", + Backend: "vertex", + AuthType: "gcp_adc", + APIMode: "native", + VertexProject: "prod-ai", + VertexLocation: "us-central1", + ServiceAccountFile: "/secrets/vertex.json", + GCPScope: "scope-a", + Models: []config.RawProviderModel{{ID: "gpt-4"}, {ID: "gpt-3.5-turbo"}}, } got := buildProviderConfig(raw, globalResilience) @@ -184,11 +194,51 @@ func TestBuildProviderConfig_PreservesFields(t *testing.T) { if got.BaseURL != "https://custom.endpoint.com" { t.Errorf("BaseURL = %q, want https://custom.endpoint.com", got.BaseURL) } + if got.Backend != "vertex" { + t.Errorf("Backend = %q, want vertex", got.Backend) + } + if got.AuthType != "gcp_adc" { + t.Errorf("AuthType = %q, want gcp_adc", got.AuthType) + } + if got.APIMode != "native" { + t.Errorf("APIMode = %q, want native", got.APIMode) + } + if got.VertexProject != "prod-ai" { + t.Errorf("VertexProject = %q, want prod-ai", got.VertexProject) + } + if got.VertexLocation != "us-central1" { + t.Errorf("VertexLocation = %q, want us-central1", got.VertexLocation) + } + if got.ServiceAccountFile != "/secrets/vertex.json" { + t.Errorf("ServiceAccountFile = %q, want /secrets/vertex.json", got.ServiceAccountFile) + } + if got.GCPScope != "scope-a" { + t.Errorf("GCPScope = %q, want scope-a", got.GCPScope) + } if len(got.Models) != 2 || got.Models[0] != "gpt-4" { t.Errorf("Models = %v, want [gpt-4 gpt-3.5-turbo]", got.Models) } } +func TestBuildProviderConfig_NormalizesLegacyGeminiVertexType(t *testing.T) { + raw := config.RawProviderConfig{ + Type: "gemini", + Backend: "vertex", + AuthType: "gcp_adc", + VertexProject: "prod-ai", + VertexLocation: "us-central1", + } + + got := buildProviderConfig(raw, globalResilience) + + if got.Type != "vertex" { + t.Fatalf("Type = %q, want vertex", got.Type) + } + if got.Backend != "vertex" { + t.Fatalf("Backend = %q, want vertex", got.Backend) + } +} + // --- buildProviderConfigs --- func TestBuildProviderConfigs_MultipleProviders(t *testing.T) { @@ -467,6 +517,147 @@ func TestApplyProviderEnvVars_DiscoversZAIWithExplicitBaseURL(t *testing.T) { } } +func TestApplyProviderEnvVars_DiscoversVertexProviderFromEnvAlias(t *testing.T) { + t.Setenv("VERTEX_PROJECT", "prod-ai") + t.Setenv("VERTEX_LOCATION", "us-central1") + t.Setenv("VERTEX_AUTH_TYPE", "gcp_adc") + t.Setenv("VERTEX_API_MODE", "native") + t.Setenv("VERTEX_MODELS", "google/gemini-2.5-flash") + + got := applyProviderEnvVars(map[string]config.RawProviderConfig{}, testDiscoveryConfigs) + + p, exists := got["vertex"] + if !exists { + t.Fatal("expected vertex to be discovered from VERTEX_* env vars") + } + if p.Type != "vertex" { + t.Fatalf("Type = %q, want vertex", p.Type) + } + if p.VertexProject != "prod-ai" { + t.Fatalf("VertexProject = %q, want prod-ai", p.VertexProject) + } + if p.VertexLocation != "us-central1" { + t.Fatalf("VertexLocation = %q, want us-central1", p.VertexLocation) + } + if p.AuthType != "gcp_adc" { + t.Fatalf("AuthType = %q, want gcp_adc", p.AuthType) + } + if p.APIMode != "native" { + t.Fatalf("APIMode = %q, want native", p.APIMode) + } + if len(p.Models) != 1 || p.Models[0].ID != "google/gemini-2.5-flash" { + t.Fatalf("Models = %v, want [google/gemini-2.5-flash]", p.Models) + } +} + +func TestApplyProviderEnvVars_DiscoversSuffixedVertexProviderWithUnderscoreName(t *testing.T) { + t.Setenv("VERTEX_US_PROJECT", "prod-ai") + t.Setenv("VERTEX_US_LOCATION", "us-central1") + t.Setenv("VERTEX_US_AUTH_TYPE", "gcp_service_account") + t.Setenv("VERTEX_US_SERVICE_ACCOUNT_FILE", "/secrets/vertex.json") + + got := applyProviderEnvVars(map[string]config.RawProviderConfig{}, testDiscoveryConfigs) + + p, exists := got["vertex_us"] + if !exists { + t.Fatal("expected vertex_us to be discovered from VERTEX_US_* env vars") + } + if p.Type != "vertex" { + t.Fatalf("Type = %q, want vertex", p.Type) + } + if p.ServiceAccountFile != "/secrets/vertex.json" { + t.Fatalf("ServiceAccountFile = %q, want /secrets/vertex.json", p.ServiceAccountFile) + } +} + +func TestResolveProviders_FiltersVertexWithoutProjectOrLocation(t *testing.T) { + t.Setenv("VERTEX_PROJECT", "prod-ai") + t.Setenv("VERTEX_AUTH_TYPE", "gcp_adc") + + got, filteredRaw := resolveProviders(map[string]config.RawProviderConfig{}, globalResilience, testDiscoveryConfigs) + + if _, exists := got["vertex"]; exists { + t.Fatal("expected vertex without location to be filtered") + } + if _, exists := filteredRaw["vertex"]; exists { + t.Fatal("expected raw vertex without location to be filtered") + } +} + +func TestResolveProviders_KeepsVertexWithProjectAndLocationUsingDefaultADC(t *testing.T) { + t.Setenv("VERTEX_PROJECT", "prod-ai") + t.Setenv("VERTEX_LOCATION", "us-central1") + + got, filteredRaw := resolveProviders(map[string]config.RawProviderConfig{}, globalResilience, testDiscoveryConfigs) + + p, exists := got["vertex"] + if !exists { + t.Fatal("expected vertex with project and location to be resolved") + } + if p.Type != "vertex" { + t.Fatalf("Type = %q, want vertex", p.Type) + } + if _, exists := filteredRaw["vertex"]; !exists { + t.Fatal("expected raw vertex with project and location to be retained") + } +} + +func TestResolveProviders_FiltersVertexServiceAccountWithoutCredentials(t *testing.T) { + t.Setenv("VERTEX_PROJECT", "prod-ai") + t.Setenv("VERTEX_LOCATION", "us-central1") + t.Setenv("VERTEX_AUTH_TYPE", "gcp_service_account") + + got, filteredRaw := resolveProviders(map[string]config.RawProviderConfig{}, globalResilience, testDiscoveryConfigs) + + if _, exists := got["vertex"]; exists { + t.Fatal("expected vertex service account provider without credentials to be filtered") + } + if _, exists := filteredRaw["vertex"]; exists { + t.Fatal("expected raw vertex service account provider without credentials to be filtered") + } +} + +func TestResolveProviders_FiltersVertexWithUnresolvedProjectPlaceholder(t *testing.T) { + raw := map[string]config.RawProviderConfig{ + "vertex": { + Type: "vertex", + AuthType: "gcp_adc", + VertexProject: "${VERTEX_PROJECT}", + VertexLocation: "us-central1", + }, + } + + got, filteredRaw := resolveProviders(raw, globalResilience, testDiscoveryConfigs) + + if _, exists := got["vertex"]; exists { + t.Fatal("expected vertex with unresolved project placeholder to be filtered") + } + if _, exists := filteredRaw["vertex"]; exists { + t.Fatal("expected raw vertex with unresolved project placeholder to be filtered") + } +} + +func TestResolveProviders_FiltersVertexWithUnresolvedServiceAccountPlaceholder(t *testing.T) { + raw := map[string]config.RawProviderConfig{ + "vertex": { + Type: "vertex", + AuthType: "gcp_service_account", + VertexProject: "prod-ai", + VertexLocation: "us-central1", + ServiceAccountFile: "${VERTEX_SERVICE_ACCOUNT_FILE}", + }, + } + + got, filteredRaw := resolveProviders(raw, globalResilience, testDiscoveryConfigs) + + if _, exists := got["vertex"]; exists { + t.Fatal("expected vertex with unresolved service account placeholder to be filtered") + } + if _, exists := filteredRaw["vertex"]; exists { + t.Fatal("expected raw vertex with unresolved service account placeholder to be filtered") + } +} + func TestApplyProviderEnvVars_DiscoversVLLMFromBaseURLWithoutAPIKey(t *testing.T) { t.Setenv("VLLM_BASE_URL", "http://localhost:8000/v1") @@ -605,7 +796,11 @@ func TestApplyProviderEnvVars_DiscoversSuffixedProvidersForEveryRegisteredType(t got := applyProviderEnvVars(map[string]config.RawProviderConfig{}, testDiscoveryConfigs) for providerType, spec := range testDiscoveryConfigs { - name := providerType + "-east" + separator := spec.NameSeparator + if separator == "" { + separator = "-" + } + name := providerType + separator + "east" p, exists := got[name] if !exists { t.Fatalf("expected %s to be discovered from suffixed env vars", name) diff --git a/internal/providers/factory.go b/internal/providers/factory.go index efbf09e0..bbfa8931 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -28,6 +28,7 @@ type DiscoveryConfig struct { RequireBaseURL bool AllowAPIKeyless bool SupportsAPIVersion bool + NameSeparator string } // Registration contains metadata for registering a provider with the factory. diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 0603c1e5..c89aa55f 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -15,8 +15,10 @@ import ( "time" "gomodel/internal/core" + "gomodel/internal/httpclient" "gomodel/internal/llmclient" "gomodel/internal/providers" + "gomodel/internal/providers/googleauth" ) // Registration provides factory registration for the Gemini provider. @@ -32,48 +34,76 @@ const ( // Gemini provides an OpenAI-compatible endpoint defaultOpenAICompatibleBaseURL = "https://generativelanguage.googleapis.com/v1beta/openai" // Native Gemini API endpoint for generateContent and models listing - defaultModelsBaseURL = "https://generativelanguage.googleapis.com/v1beta" - useNativeAPIEnvVar = "USE_GOOGLE_GEMINI_NATIVE_API" + defaultModelsBaseURL = "https://generativelanguage.googleapis.com/v1beta" + useNativeAPIEnvVar = "USE_GOOGLE_GEMINI_NATIVE_API" + geminiBackendAIStudio = "aistudio" + geminiBackendVertex = "vertex" + geminiAuthTypeAPIKey = "api_key" + geminiAuthTypeGCPADC = "gcp_adc" + geminiAuthTypeServiceKey = "gcp_service_account" ) // Provider implements the core.Provider interface for Google Gemini type Provider struct { - client *llmclient.Client - nativeClient *llmclient.Client - httpClient *http.Client - hooks llmclient.Hooks - apiKey string - useNativeAPI bool - modelsURL string - modelsClientConf llmclient.Config + client *llmclient.Client + nativeClient *llmclient.Client + apiKey string + backend string + authType string + useNativeAPI bool + modelsURL string + configErr error } // New creates a new Gemini provider. func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { - baseURL, modelsURL := geminiBaseURLs(providerCfg.BaseURL) + return newProvider(providerCfg, opts, nil, false) +} + +// NewVertexWithHTTPClient creates a Vertex-configured Gemini provider using an +// already-authenticated HTTP client. It is used by the vertex package so Vertex +// owns Google auth while reusing Gemini request/response translation. +func NewVertexWithHTTPClient(providerCfg providers.ProviderConfig, opts providers.ProviderOptions, httpClient *http.Client) *Provider { + providerCfg.Backend = geminiBackendVertex + return newProvider(providerCfg, opts, httpClient, true) +} + +func newProvider(providerCfg providers.ProviderConfig, opts providers.ProviderOptions, httpClient *http.Client, preauthenticated bool) *Provider { + backend := normalizeGeminiBackend(providerCfg) + authType := normalizeGeminiAuthType(backend, providerCfg) + baseURL, modelsURL := geminiBaseURLs(providerCfg, backend) p := &Provider{ - httpClient: nil, apiKey: providerCfg.APIKey, - hooks: opts.Hooks, - useNativeAPI: useNativeAPIFromEnv(), + backend: backend, + authType: authType, + useNativeAPI: useNativeAPI(providerCfg.APIMode), modelsURL: modelsURL, - modelsClientConf: llmclient.Config{ - ProviderName: "gemini", - BaseURL: modelsURL, - Retry: opts.Resilience.Retry, - Hooks: opts.Hooks, - CircuitBreaker: opts.Resilience.CircuitBreaker, - }, + } + p.validateConfig(providerCfg) + if !preauthenticated { + httpClient = p.authHTTPClient(providerCfg, httpClient) + } + + clientProviderName := "gemini" + if backend == geminiBackendVertex { + clientProviderName = "vertex" } clientCfg := llmclient.Config{ - ProviderName: "gemini", + ProviderName: clientProviderName, BaseURL: baseURL, Retry: opts.Resilience.Retry, Hooks: opts.Hooks, CircuitBreaker: opts.Resilience.CircuitBreaker, } + modelsCfg := clientCfg + modelsCfg.BaseURL = modelsURL + if httpClient != nil { + p.client = llmclient.NewWithHTTPClient(httpClient, clientCfg, p.setHeaders) + p.nativeClient = llmclient.NewWithHTTPClient(httpClient, modelsCfg, p.setNativeHeaders) + return p + } p.client = llmclient.New(clientCfg, p.setHeaders) - p.nativeClient = llmclient.New(p.modelsClientConf, p.setNativeHeaders) + p.nativeClient = llmclient.New(modelsCfg, p.setNativeHeaders) return p } @@ -83,17 +113,17 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H if httpClient == nil { httpClient = http.DefaultClient } - baseURL, modelsURL := geminiBaseURLs("") + providerCfg := providers.ProviderConfig{APIKey: apiKey} + baseURL, modelsURL := geminiBaseURLs(providerCfg, geminiBackendAIStudio) p := &Provider{ - httpClient: httpClient, apiKey: apiKey, - hooks: hooks, + backend: geminiBackendAIStudio, + authType: geminiAuthTypeAPIKey, useNativeAPI: useNativeAPIFromEnv(), modelsURL: modelsURL, } modelsCfg := llmclient.DefaultConfig("gemini", modelsURL) modelsCfg.Hooks = hooks - p.modelsClientConf = modelsCfg cfg := llmclient.DefaultConfig("gemini", baseURL) cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) @@ -103,29 +133,70 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H // SetBaseURL allows configuring a custom base URL for the provider func (p *Provider) SetBaseURL(url string) { - baseURL, modelsURL := geminiBaseURLs(url) + baseURL, modelsURL := geminiBaseURLs(providers.ProviderConfig{BaseURL: url}, p.backend) p.client.SetBaseURL(baseURL) p.modelsURL = modelsURL - p.modelsClientConf.BaseURL = modelsURL if p.nativeClient != nil { p.nativeClient.SetBaseURL(modelsURL) } - p.useNativeAPI = useNativeAPIFromEnv() } // SetModelsURL allows configuring a custom models API base URL. // This is primarily useful for tests and local emulators. func (p *Provider) SetModelsURL(url string) { p.modelsURL = url - p.modelsClientConf.BaseURL = url if p.nativeClient != nil { p.nativeClient.SetBaseURL(url) } } +func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { + if p.backend == geminiBackendVertex && p.authType == geminiAuthTypeAPIKey { + p.configErr = fmt.Errorf("vertex Gemini requires gcp_adc or gcp_service_account auth") + return + } + if p.backend == geminiBackendVertex && (strings.TrimSpace(providerCfg.VertexProject) == "" || strings.TrimSpace(providerCfg.VertexLocation) == "") { + p.configErr = fmt.Errorf("vertex Gemini requires vertex_project and vertex_location") + return + } + if p.backend == geminiBackendAIStudio && p.authType == geminiAuthTypeAPIKey && strings.TrimSpace(providerCfg.APIKey) == "" { + p.configErr = fmt.Errorf("gemini API key is required") + } +} + +func (p *Provider) authHTTPClient(providerCfg providers.ProviderConfig, base *http.Client) *http.Client { + if p.configErr != nil || p.authType == geminiAuthTypeAPIKey { + return base + } + source, err := googleauth.TokenSource(context.Background(), googleauth.Config{ + AuthType: p.authType, + ServiceAccountFile: providerCfg.ServiceAccountFile, + ServiceAccountJSON: providerCfg.ServiceAccountJSON, + ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, + Scope: providerCfg.GCPScope, + }) + if err != nil { + p.configErr = err + return base + } + if base == nil { + base = httpclient.NewDefaultHTTPClient() + } + return googleauth.HTTPClient(base, source) +} + +func (p *Provider) ready() error { + if p.configErr == nil { + return nil + } + return core.NewProviderError("gemini", http.StatusBadGateway, "invalid Gemini provider configuration: "+p.configErr.Error(), p.configErr) +} + // setHeaders sets the required headers for Gemini API requests func (p *Provider) setHeaders(req *http.Request) { - req.Header.Set("Authorization", "Bearer "+p.apiKey) + if p.authType == geminiAuthTypeAPIKey { + req.Header.Set("Authorization", "Bearer "+p.apiKey) + } // Forward request ID if present in context for request tracing if requestID := core.GetRequestID(req.Context()); requestID != "" { @@ -135,13 +206,67 @@ func (p *Provider) setHeaders(req *http.Request) { // setNativeHeaders sets the required headers for Gemini native API requests. func (p *Provider) setNativeHeaders(req *http.Request) { - req.Header.Set("x-goog-api-key", p.apiKey) + if p.authType == geminiAuthTypeAPIKey { + req.Header.Set("x-goog-api-key", p.apiKey) + } if requestID := core.GetRequestID(req.Context()); requestID != "" { req.Header.Set("X-Request-Id", requestID) } } +func normalizeGeminiBackend(cfg providers.ProviderConfig) string { + backend := strings.ToLower(strings.TrimSpace(cfg.Backend)) + switch backend { + case geminiBackendVertex: + return geminiBackendVertex + case geminiBackendAIStudio, "ai_studio", "developer", "developer_api": + return geminiBackendAIStudio + } + if strings.TrimSpace(cfg.VertexProject) != "" || + strings.TrimSpace(cfg.VertexLocation) != "" || + strings.HasPrefix(strings.ToLower(strings.TrimSpace(cfg.AuthType)), "gcp") { + return geminiBackendVertex + } + return geminiBackendAIStudio +} + +func normalizeGeminiAuthType(backend string, cfg providers.ProviderConfig) string { + authType := strings.ToLower(strings.TrimSpace(cfg.AuthType)) + switch authType { + case "": + if backend == geminiBackendVertex { + return googleauth.NormalizeAuthType(authType, googleauth.HasServiceAccount(googleauth.Config{ + ServiceAccountFile: cfg.ServiceAccountFile, + ServiceAccountJSON: cfg.ServiceAccountJSON, + ServiceAccountJSONBase64: cfg.ServiceAccountJSONBase64, + })) + } + return geminiAuthTypeAPIKey + case "api_key", "key": + return geminiAuthTypeAPIKey + case "gcp_adc", "adc", "google_adc": + return geminiAuthTypeGCPADC + case "gcp_service_account", "service_account": + return geminiAuthTypeServiceKey + default: + return authType + } +} + +func useNativeAPI(apiMode string) bool { + switch strings.ToLower(strings.TrimSpace(apiMode)) { + case "native", "gemini_native", "generate_content": + return true + case "openai", "openai_compatible", "openai-compatible", "compat", "compatible": + return false + case "": + return useNativeAPIFromEnv() + default: + return useNativeAPIFromEnv() + } +} + func useNativeAPIFromEnv() bool { value, ok := os.LookupEnv(useNativeAPIEnvVar) if !ok || strings.TrimSpace(value) == "" { @@ -155,7 +280,11 @@ func useNativeAPIFromEnv() bool { } } -func geminiBaseURLs(configuredBaseURL string) (openAICompatibleBaseURL, nativeBaseURL string) { +func geminiBaseURLs(providerCfg providers.ProviderConfig, backend string) (openAICompatibleBaseURL, nativeBaseURL string) { + if backend == geminiBackendVertex { + return vertexBaseURLs(providerCfg) + } + configuredBaseURL := providerCfg.BaseURL baseURL := strings.TrimRight(strings.TrimSpace(configuredBaseURL), "/") if baseURL == "" { return defaultOpenAICompatibleBaseURL, defaultModelsBaseURL @@ -172,6 +301,47 @@ func geminiBaseURLs(configuredBaseURL string) (openAICompatibleBaseURL, nativeBa return baseURL, baseURL } +func vertexBaseURLs(providerCfg providers.ProviderConfig) (openAICompatibleBaseURL, nativeBaseURL string) { + baseURL := strings.TrimRight(strings.TrimSpace(providerCfg.BaseURL), "/") + if baseURL == "" { + project := strings.TrimSpace(providerCfg.VertexProject) + location := strings.TrimSpace(providerCfg.VertexLocation) + root := "https://aiplatform.googleapis.com/v1/projects/" + url.PathEscape(project) + "/locations/" + url.PathEscape(location) + return root + "/endpoints/openapi", root + "/publishers/google" + } + if nativeBaseURL, ok := vertexNativeBaseURLFromOpenAICompatibleBaseURL(baseURL); ok { + return baseURL, nativeBaseURL + } + if openAIBaseURL, ok := vertexOpenAICompatibleBaseURLFromNativeBaseURL(baseURL); ok { + return openAIBaseURL, baseURL + } + return baseURL, baseURL +} + +func vertexNativeBaseURLFromOpenAICompatibleBaseURL(baseURL string) (string, bool) { + const suffix = "/endpoints/openapi" + if !strings.HasSuffix(baseURL, suffix) { + return "", false + } + root := strings.TrimRight(strings.TrimSuffix(baseURL, suffix), "/") + if root == "" { + return "", false + } + return root + "/publishers/google", true +} + +func vertexOpenAICompatibleBaseURLFromNativeBaseURL(baseURL string) (string, bool) { + const suffix = "/publishers/google" + if !strings.HasSuffix(baseURL, suffix) { + return "", false + } + root := strings.TrimRight(strings.TrimSuffix(baseURL, suffix), "/") + if root == "" { + return "", false + } + return root + "/endpoints/openapi", true +} + func nativeBaseURLFromOpenAICompatibleBaseURL(baseURL string) (string, bool) { const suffix = "/openai" if !strings.HasSuffix(baseURL, suffix) { @@ -208,15 +378,53 @@ func adaptChatRequest(req *core.ChatRequest) (any, error) { return raw, nil } +func rewriteRequestModel(body any, model string) (any, error) { + if strings.TrimSpace(model) == "" { + return body, nil + } + encoded, err := json.Marshal(body) + if err != nil { + return nil, core.NewInvalidRequestError("failed to marshal gemini request: "+err.Error(), err) + } + var raw map[string]json.RawMessage + if err := json.Unmarshal(encoded, &raw); err != nil { + return nil, core.NewInvalidRequestError("failed to decode gemini request payload: "+err.Error(), err) + } + modelJSON, _ := json.Marshal(model) + raw["model"] = modelJSON + return raw, nil +} + +func (p *Provider) openAICompatibleChatBody(req *core.ChatRequest) (any, error) { + body, err := adaptChatRequest(req) + if err != nil { + return nil, err + } + if p.backend != geminiBackendVertex { + return body, nil + } + return rewriteRequestModel(body, vertexOpenAIModelID(req.Model)) +} + +func (p *Provider) openAICompatibleEmbeddingBody(req *core.EmbeddingRequest) (any, error) { + if p.backend != geminiBackendVertex { + return req, nil + } + return rewriteRequestModel(req, vertexOpenAIModelID(req.Model)) +} + // ChatCompletion sends a chat completion request to Gemini func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*core.ChatResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } if req == nil { return nil, core.NewInvalidRequestError("chat request is required", nil) } if p.useNativeAPI { return p.nativeChatCompletion(ctx, req) } - body, err := adaptChatRequest(req) + body, err := p.openAICompatibleChatBody(req) if err != nil { return nil, err } @@ -254,6 +462,9 @@ func (p *Provider) nativeChatCompletion(ctx context.Context, req *core.ChatReque // StreamChatCompletion returns a raw response body for streaming (caller must close) func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { + if err := p.ready(); err != nil { + return nil, err + } if req == nil { return nil, core.NewInvalidRequestError("chat request is required", nil) } @@ -261,7 +472,7 @@ func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatReque return p.nativeStreamChatCompletion(ctx, req) } streamReq := req.WithStreaming() - body, err := adaptChatRequest(streamReq) + body, err := p.openAICompatibleChatBody(streamReq) if err != nil { return nil, err } @@ -311,38 +522,25 @@ type geminiModel struct { // geminiModelsResponse represents the native Gemini models list response type geminiModelsResponse struct { - Models []geminiModel `json:"models"` + Models []geminiModel `json:"models"` + PublisherModels []geminiModel `json:"publisherModels"` } -// ListModels retrieves the list of available models from Gemini -func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { - // Use the native Gemini API to list models - // We need to create a separate client for the models endpoint since it uses a different URL - modelsCfg := p.modelsClientConf - modelsCfg.BaseURL = p.modelsURL - modelsCfg.Hooks = p.hooks - headers := func(req *http.Request) { - // Use header-based API key auth for models requests. - req.Header.Set("x-goog-api-key", p.apiKey) - - // Preserve request tracing across list-models requests. - requestID := req.Header.Get("X-Request-Id") - if requestID == "" { - requestID = core.GetRequestID(req.Context()) - } - if requestID != "" { - req.Header.Set("X-Request-Id", requestID) - } +func geminiModelSupportedMethods(modelID string, methods []string) (supportsGenerate, supportsEmbed bool) { + if len(methods) == 0 { + normalized := normalizeGeminiModelID(modelID) + return strings.HasPrefix(normalized, "gemini-"), strings.HasPrefix(normalized, "text-embedding-") } + return slices.Contains(methods, "generateContent") || slices.Contains(methods, "streamGenerateContent"), + slices.Contains(methods, "embedContent") +} - var modelsClient *llmclient.Client - if p.httpClient != nil { - modelsClient = llmclient.NewWithHTTPClient(p.httpClient, modelsCfg, headers) - } else { - modelsClient = llmclient.New(modelsCfg, headers) +// ListModels retrieves the list of available models from Gemini +func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { + if err := p.ready(); err != nil { + return nil, err } - - rawResp, err := modelsClient.DoRaw(ctx, llmclient.Request{ + rawResp, err := p.nativeClient.DoRaw(ctx, llmclient.Request{ Method: http.MethodGet, Endpoint: "/models", }) @@ -356,38 +554,31 @@ func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) // If the payload contains an explicit "models" field with an empty array, // return an empty list instead of falling through to fallback parsing. var nativeProbe struct { - Models json.RawMessage `json:"models"` + Models json.RawMessage `json:"models"` + PublisherModels json.RawMessage `json:"publisherModels"` } - if err := json.Unmarshal(rawResp.Body, &nativeProbe); err == nil && nativeProbe.Models != nil { + if err := json.Unmarshal(rawResp.Body, &nativeProbe); err == nil && (nativeProbe.Models != nil || nativeProbe.PublisherModels != nil) { var geminiResp geminiModelsResponse if err := json.Unmarshal(rawResp.Body, &geminiResp); err != nil { return nil, core.NewProviderError("gemini", http.StatusBadGateway, "failed to parse native Gemini models response", err) } - if len(geminiResp.Models) == 0 { + modelEntries := append(geminiResp.Models, geminiResp.PublisherModels...) + if len(modelEntries) == 0 { return &core.ModelsResponse{ Object: "list", Data: []core.Model{}, }, nil } - models := make([]core.Model, 0, len(geminiResp.Models)) + models := make([]core.Model, 0, len(modelEntries)) - for _, gm := range geminiResp.Models { - // Extract model ID from name (format: "models/gemini-...") - modelID := strings.TrimPrefix(gm.Name, "models/") + for _, gm := range modelEntries { + modelID := displayModelIDFromGemini(gm.Name, p.backend) // Only include models that support generateContent (chat/completion) - supportsGenerate := false - for _, method := range gm.SupportedMethods { - if method == "generateContent" || method == "streamGenerateContent" { - supportsGenerate = true - break - } - } + supportsGenerate, supportsEmbed := geminiModelSupportedMethods(modelID, gm.SupportedMethods) - supportsEmbed := slices.Contains(gm.SupportedMethods, "embedContent") - - isOpenAICompatModel := strings.HasPrefix(modelID, "gemini-") || strings.HasPrefix(modelID, "text-embedding-") + isOpenAICompatModel := isGeminiExposedModel(modelID) if (supportsGenerate || supportsEmbed) && isOpenAICompatModel { models = append(models, core.Model{ ID: modelID, @@ -409,8 +600,8 @@ func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) if err := json.Unmarshal(rawResp.Body, &openAIResp); err == nil && openAIResp.Object == "list" { models := make([]core.Model, 0, len(openAIResp.Data)) for _, m := range openAIResp.Data { - modelID := strings.TrimPrefix(m.ID, "models/") - isOpenAICompatModel := strings.HasPrefix(modelID, "gemini-") || strings.HasPrefix(modelID, "text-embedding-") + modelID := displayModelIDFromGemini(m.ID, p.backend) + isOpenAICompatModel := isGeminiExposedModel(modelID) if !isOpenAICompatModel { continue } @@ -436,16 +627,29 @@ func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) // Responses sends a Responses API request to Gemini (converted to chat format) func (p *Provider) Responses(ctx context.Context, req *core.ResponsesRequest) (*core.ResponsesResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } return providers.ResponsesViaChat(ctx, p, req) } // Embeddings sends an embeddings request to Gemini via its OpenAI-compatible endpoint func (p *Provider) Embeddings(ctx context.Context, req *core.EmbeddingRequest) (*core.EmbeddingResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } + if req == nil { + return nil, core.NewInvalidRequestError("embedding request is required", nil) + } + body, err := p.openAICompatibleEmbeddingBody(req) + if err != nil { + return nil, err + } var resp core.EmbeddingResponse - err := p.client.Do(ctx, llmclient.Request{ + err = p.client.Do(ctx, llmclient.Request{ Method: http.MethodPost, Endpoint: "/embeddings", - Body: req, + Body: body, }, &resp) if err != nil { return nil, err @@ -458,11 +662,17 @@ func (p *Provider) Embeddings(ctx context.Context, req *core.EmbeddingRequest) ( // StreamResponses returns a raw response body for streaming Responses API (caller must close) func (p *Provider) StreamResponses(ctx context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) { + if err := p.ready(); err != nil { + return nil, err + } return providers.StreamResponsesViaChat(ctx, p, req, "gemini") } // CreateBatch creates a native Gemini batch job through its OpenAI-compatible endpoint. func (p *Provider) CreateBatch(ctx context.Context, req *core.BatchRequest) (*core.BatchResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } var resp core.BatchResponse err := p.client.Do(ctx, llmclient.Request{ Method: http.MethodPost, @@ -480,6 +690,9 @@ func (p *Provider) CreateBatch(ctx context.Context, req *core.BatchRequest) (*co // GetBatch retrieves a native Gemini batch job. func (p *Provider) GetBatch(ctx context.Context, id string) (*core.BatchResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } var resp core.BatchResponse err := p.client.Do(ctx, llmclient.Request{ Method: http.MethodGet, @@ -496,6 +709,9 @@ func (p *Provider) GetBatch(ctx context.Context, id string) (*core.BatchResponse // ListBatches lists native Gemini batch jobs. func (p *Provider) ListBatches(ctx context.Context, limit int, after string) (*core.BatchListResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } values := url.Values{} if limit > 0 { values.Set("limit", strconv.Itoa(limit)) @@ -526,6 +742,9 @@ func (p *Provider) ListBatches(ctx context.Context, limit int, after string) (*c // CancelBatch cancels a native Gemini batch job. func (p *Provider) CancelBatch(ctx context.Context, id string) (*core.BatchResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } var resp core.BatchResponse err := p.client.Do(ctx, llmclient.Request{ Method: http.MethodPost, @@ -542,11 +761,17 @@ func (p *Provider) CancelBatch(ctx context.Context, id string) (*core.BatchRespo // GetBatchResults fetches Gemini batch results via the output file API. func (p *Provider) GetBatchResults(ctx context.Context, id string) (*core.BatchResultsResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } return providers.FetchBatchResultsFromOutputFile(ctx, p.client, "gemini", id) } // CreateFile uploads a file through Gemini's OpenAI-compatible /files API. func (p *Provider) CreateFile(ctx context.Context, req *core.FileCreateRequest) (*core.FileObject, error) { + if err := p.ready(); err != nil { + return nil, err + } resp, err := providers.CreateOpenAICompatibleFile(ctx, p.client, req) if err != nil { return nil, err @@ -557,6 +782,9 @@ func (p *Provider) CreateFile(ctx context.Context, req *core.FileCreateRequest) // ListFiles lists files through Gemini's OpenAI-compatible /files API. func (p *Provider) ListFiles(ctx context.Context, purpose string, limit int, after string) (*core.FileListResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } resp, err := providers.ListOpenAICompatibleFiles(ctx, p.client, purpose, limit, after) if err != nil { return nil, err @@ -569,6 +797,9 @@ func (p *Provider) ListFiles(ctx context.Context, purpose string, limit int, aft // GetFile retrieves one file object through Gemini's OpenAI-compatible /files API. func (p *Provider) GetFile(ctx context.Context, id string) (*core.FileObject, error) { + if err := p.ready(); err != nil { + return nil, err + } resp, err := providers.GetOpenAICompatibleFile(ctx, p.client, id) if err != nil { return nil, err @@ -579,10 +810,16 @@ func (p *Provider) GetFile(ctx context.Context, id string) (*core.FileObject, er // DeleteFile deletes a file object through Gemini's OpenAI-compatible /files API. func (p *Provider) DeleteFile(ctx context.Context, id string) (*core.FileDeleteResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } return providers.DeleteOpenAICompatibleFile(ctx, p.client, id) } // GetFileContent fetches raw file bytes through Gemini's /files/{id}/content API. func (p *Provider) GetFileContent(ctx context.Context, id string) (*core.FileContentResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } return providers.GetOpenAICompatibleFileContent(ctx, p.client, id) } diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index b8f17067..2c7f8d31 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -13,6 +13,9 @@ import ( "gomodel/internal/core" "gomodel/internal/llmclient" "gomodel/internal/providers" + "gomodel/internal/providers/googleauth" + + "golang.org/x/oauth2" ) func TestNew(t *testing.T) { @@ -79,7 +82,54 @@ func TestGeminiBaseURLs(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - gotCompat, gotNative := geminiBaseURLs(tt.configured) + gotCompat, gotNative := geminiBaseURLs(providers.ProviderConfig{BaseURL: tt.configured}, geminiBackendAIStudio) + if gotCompat != tt.wantCompat { + t.Fatalf("OpenAI-compatible base = %q, want %q", gotCompat, tt.wantCompat) + } + if gotNative != tt.wantNative { + t.Fatalf("native base = %q, want %q", gotNative, tt.wantNative) + } + }) + } +} + +func TestVertexBaseURLs(t *testing.T) { + tests := []struct { + name string + cfg providers.ProviderConfig + wantCompat string + wantNative string + }{ + { + name: "derives official vertex bases from project and location", + cfg: providers.ProviderConfig{ + VertexProject: "prod-ai", + VertexLocation: "us-central1", + }, + wantCompat: "https://aiplatform.googleapis.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi", + wantNative: "https://aiplatform.googleapis.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, + { + name: "custom OpenAI-compatible vertex URL derives native sibling", + cfg: providers.ProviderConfig{ + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi/", + }, + wantCompat: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi", + wantNative: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, + { + name: "custom native vertex URL derives OpenAI-compatible sibling", + cfg: providers.ProviderConfig{ + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google/", + }, + wantCompat: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi", + wantNative: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotCompat, gotNative := geminiBaseURLs(tt.cfg, geminiBackendVertex) if gotCompat != tt.wantCompat { t.Fatalf("OpenAI-compatible base = %q, want %q", gotCompat, tt.wantCompat) } @@ -90,6 +140,34 @@ func TestGeminiBaseURLs(t *testing.T) { } } +func TestVertexModelNormalization(t *testing.T) { + tests := []struct { + in string + wantNative string + wantOpenAI string + }{ + {in: "gemini-2.5-flash", wantNative: "gemini-2.5-flash", wantOpenAI: "google/gemini-2.5-flash"}, + {in: "models/gemini-2.5-flash", wantNative: "gemini-2.5-flash", wantOpenAI: "google/gemini-2.5-flash"}, + {in: "google/gemini-2.5-flash", wantNative: "gemini-2.5-flash", wantOpenAI: "google/gemini-2.5-flash"}, + { + in: "projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash", + wantNative: "gemini-2.5-flash", + wantOpenAI: "google/gemini-2.5-flash", + }, + } + + for _, tt := range tests { + t.Run(tt.in, func(t *testing.T) { + if got := normalizeGeminiModelID(tt.in); got != tt.wantNative { + t.Fatalf("normalizeGeminiModelID() = %q, want %q", got, tt.wantNative) + } + if got := vertexOpenAIModelID(tt.in); got != tt.wantOpenAI { + t.Fatalf("vertexOpenAIModelID() = %q, want %q", got, tt.wantOpenAI) + } + }) + } +} + func TestNew_CustomBaseURLDerivesNativeBaseURL(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "true") @@ -111,8 +189,8 @@ func TestNew_CustomBaseURLDerivesNativeBaseURL(t *testing.T) { if geminiProvider.modelsURL != "https://proxy.example.com/v1beta" { t.Fatalf("modelsURL = %q, want derived native base URL", geminiProvider.modelsURL) } - if geminiProvider.modelsClientConf.BaseURL != "https://proxy.example.com/v1beta" { - t.Fatalf("modelsClientConf.BaseURL = %q, want derived native base URL", geminiProvider.modelsClientConf.BaseURL) + if geminiProvider.nativeClient.BaseURL() != "https://proxy.example.com/v1beta" { + t.Fatalf("nativeClient.BaseURL() = %q, want derived native base URL", geminiProvider.nativeClient.BaseURL()) } } @@ -202,8 +280,8 @@ func TestSetBaseURLDerivesModelsURL(t *testing.T) { if provider.modelsURL != server.URL+"/v1beta" { t.Fatalf("modelsURL = %q, want derived native base URL", provider.modelsURL) } - if provider.modelsClientConf.BaseURL != server.URL+"/v1beta" { - t.Fatalf("modelsClientConf.BaseURL = %q, want derived native base URL", provider.modelsClientConf.BaseURL) + if provider.nativeClient.BaseURL() != server.URL+"/v1beta" { + t.Fatalf("nativeClient.BaseURL() = %q, want derived native base URL", provider.nativeClient.BaseURL()) } resp, err := provider.ListModels(context.Background()) @@ -218,6 +296,146 @@ func TestSetBaseURLDerivesModelsURL(t *testing.T) { } } +func TestVertexNativeChatUsesOAuthAuthorization(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash:generateContent" { + t.Errorf("Path = %q, want Vertex native generateContent endpoint", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer vertex-token" { + t.Errorf("Authorization = %q, want Bearer vertex-token", got) + } + if got := r.Header.Get("x-goog-api-key"); got != "" { + t.Errorf("x-goog-api-key = %q, want empty for Vertex OAuth", got) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "vertex-native", + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP" + }] + }`)) + })) + defer server.Close() + + p := newVertexTestProvider(server, true) + resp, err := p.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "google/gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp == nil || resp.ID != "vertex-native" { + t.Fatalf("response = %+v, want vertex-native", resp) + } +} + +func TestVertexOpenAICompatibleChatUsesOAuthAndGoogleModelPrefix(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "false") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/endpoints/openapi/chat/completions" { + t.Errorf("Path = %q, want Vertex OpenAI-compatible chat endpoint", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer vertex-token" { + t.Errorf("Authorization = %q, want Bearer vertex-token", got) + } + if got := r.Header.Get("x-goog-api-key"); got != "" { + t.Errorf("x-goog-api-key = %q, want empty for Vertex OAuth", got) + } + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode request: %v", err) + } + if got := payload["model"]; got != "google/gemini-2.5-flash" { + t.Fatalf("model = %#v, want google/gemini-2.5-flash", got) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "id": "vertex-openai", + "object": "chat.completion", + "created": 1677652288, + "model": "google/gemini-2.5-flash", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }] + }`)) + })) + defer server.Close() + + p := newVertexTestProvider(server, false) + resp, err := p.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp == nil || resp.ID != "vertex-openai" { + t.Fatalf("response = %+v, want vertex-openai", resp) + } +} + +func TestVertexListModelsAcceptsPublisherModelsResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models" { + t.Errorf("Path = %q, want Vertex publisher models endpoint", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer vertex-token" { + t.Errorf("Authorization = %q, want Bearer vertex-token", got) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "publisherModels": [ + {"name": "projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash"}, + {"name": "projects/prod-ai/locations/us-central1/publishers/google/models/text-embedding-005"}, + {"name": "projects/prod-ai/locations/us-central1/publishers/google/models/imagen-4.0"} + ] + }`)) + })) + defer server.Close() + + p := newVertexTestProvider(server, true) + resp, err := p.ListModels(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(resp.Data) != 2 { + t.Fatalf("models = %+v, want 2 Gemini-compatible models", resp.Data) + } + if resp.Data[0].ID != "google/gemini-2.5-flash" || resp.Data[1].ID != "google/text-embedding-005" { + t.Fatalf("models = %+v, want google/gemini-2.5-flash and google/text-embedding-005", resp.Data) + } +} + +func newVertexTestProvider(server *httptest.Server, native bool) *Provider { + tokenClient := googleauth.HTTPClient(server.Client(), oauth2.StaticTokenSource(&oauth2.Token{ + AccessToken: "vertex-token", + TokenType: "Bearer", + })) + p := &Provider{ + backend: geminiBackendVertex, + authType: geminiAuthTypeGCPADC, + useNativeAPI: native, + } + openAIBaseURL := server.URL + "/v1/projects/prod-ai/locations/us-central1/endpoints/openapi" + nativeBaseURL := server.URL + "/v1/projects/prod-ai/locations/us-central1/publishers/google" + openAICfg := llmclient.DefaultConfig("gemini", openAIBaseURL) + nativeCfg := llmclient.DefaultConfig("gemini", nativeBaseURL) + p.client = llmclient.NewWithHTTPClient(tokenClient, openAICfg, p.setHeaders) + p.nativeClient = llmclient.NewWithHTTPClient(tokenClient, nativeCfg, p.setNativeHeaders) + return p +} + func TestChatCompletion(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "false") diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go index 0a5a5a6c..daa9d563 100644 --- a/internal/providers/gemini/native.go +++ b/internal/providers/gemini/native.go @@ -811,10 +811,35 @@ func nativeStreamEndpoint(model string) string { func normalizeGeminiModelID(model string) string { model = strings.TrimSpace(model) + if idx := strings.LastIndex(model, "/models/"); idx >= 0 { + model = model[idx+len("/models/"):] + } model = strings.TrimPrefix(model, "models/") + model = strings.TrimPrefix(model, "google/") return model } +func vertexOpenAIModelID(model string) string { + model = normalizeGeminiModelID(model) + if model == "" { + return "" + } + return "google/" + model +} + +func displayModelIDFromGemini(model, backend string) string { + model = normalizeGeminiModelID(model) + if backend == geminiBackendVertex && model != "" { + return "google/" + model + } + return model +} + +func isGeminiExposedModel(modelID string) bool { + modelID = normalizeGeminiModelID(modelID) + return strings.HasPrefix(modelID, "gemini-") || strings.HasPrefix(modelID, "text-embedding-") +} + func nativeProviderError(message string, err error) *core.GatewayError { return core.NewProviderError("gemini", http.StatusBadGateway, message, err) } diff --git a/internal/providers/googleauth/googleauth.go b/internal/providers/googleauth/googleauth.go new file mode 100644 index 00000000..f6abeaed --- /dev/null +++ b/internal/providers/googleauth/googleauth.go @@ -0,0 +1,112 @@ +package googleauth + +import ( + "context" + "encoding/base64" + "fmt" + "net/http" + "os" + "strings" + + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" +) + +const DefaultScope = "https://www.googleapis.com/auth/cloud-platform" + +type Config struct { + AuthType string + ServiceAccountFile string + ServiceAccountJSON string + ServiceAccountJSONBase64 string + Scope string +} + +func NormalizeAuthType(authType string, hasServiceAccount bool) string { + switch strings.ToLower(strings.TrimSpace(authType)) { + case "gcp_service_account", "service_account": + return "gcp_service_account" + case "gcp_adc", "adc", "google_adc": + return "gcp_adc" + default: + if hasServiceAccount { + return "gcp_service_account" + } + return "gcp_adc" + } +} + +func HasServiceAccount(cfg Config) bool { + return strings.TrimSpace(cfg.ServiceAccountJSONBase64) != "" || + strings.TrimSpace(cfg.ServiceAccountJSON) != "" || + strings.TrimSpace(cfg.ServiceAccountFile) != "" +} + +func TokenSource(ctx context.Context, cfg Config) (oauth2.TokenSource, error) { + if ctx == nil { + ctx = context.Background() + } + scope := strings.TrimSpace(cfg.Scope) + if scope == "" { + scope = DefaultScope + } + + authType := NormalizeAuthType(cfg.AuthType, HasServiceAccount(cfg)) + switch authType { + case "gcp_service_account": + credentials, err := serviceAccountJSON(cfg) + if err != nil { + return nil, err + } + jwtCfg, err := google.JWTConfigFromJSON(credentials, scope) + if err != nil { + return nil, fmt.Errorf("parse service account credentials: %w", err) + } + return jwtCfg.TokenSource(ctx), nil + case "gcp_adc": + source, err := google.DefaultTokenSource(ctx, scope) + if err != nil { + return nil, fmt.Errorf("load application default credentials: %w", err) + } + return source, nil + default: + return nil, fmt.Errorf("unsupported Google auth type %q", cfg.AuthType) + } +} + +func serviceAccountJSON(cfg Config) ([]byte, error) { + if value := strings.TrimSpace(cfg.ServiceAccountJSONBase64); value != "" { + decoded, err := base64.StdEncoding.DecodeString(value) + if err != nil { + return nil, fmt.Errorf("decode service account JSON: %w", err) + } + return decoded, nil + } + if value := strings.TrimSpace(cfg.ServiceAccountJSON); value != "" { + return []byte(value), nil + } + if path := strings.TrimSpace(cfg.ServiceAccountFile); path != "" { + contents, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read service account file: %w", err) + } + return contents, nil + } + return nil, fmt.Errorf("service account credentials are required") +} + +func HTTPClient(base *http.Client, source oauth2.TokenSource) *http.Client { + if base == nil { + base = http.DefaultClient + } + transport := base.Transport + if transport == nil { + transport = http.DefaultTransport + } + clone := *base + clone.Transport = &oauth2.Transport{ + Source: oauth2.ReuseTokenSource(nil, source), + Base: transport, + } + return &clone +} diff --git a/internal/providers/vertex/vertex.go b/internal/providers/vertex/vertex.go new file mode 100644 index 00000000..3c03a7ad --- /dev/null +++ b/internal/providers/vertex/vertex.go @@ -0,0 +1,377 @@ +// Package vertex provides Google Vertex AI Gemini integration. +package vertex + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + + "gomodel/internal/core" + "gomodel/internal/httpclient" + "gomodel/internal/llmclient" + "gomodel/internal/providers" + "gomodel/internal/providers/gemini" + "gomodel/internal/providers/googleauth" +) + +// Registration provides factory registration for the Vertex AI provider. +var Registration = providers.Registration{ + Type: "vertex", + New: New, + Discovery: providers.DiscoveryConfig{ + NameSeparator: "_", + }, +} + +const ( + authTypeGCPADC = "gcp_adc" + authTypeServiceAccount = "gcp_service_account" +) + +// Provider implements Vertex AI through Gemini chat helpers plus Vertex-native +// embedding prediction. +type Provider struct { + gemini *gemini.Provider + nativeClient *llmclient.Client + authType string + configErr error +} + +// New creates a new Vertex AI provider. +func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { + return newProvider(providerCfg, opts, nil) +} + +func newProvider(providerCfg providers.ProviderConfig, opts providers.ProviderOptions, baseHTTPClient *http.Client) *Provider { + providerCfg.Backend = "vertex" + p := &Provider{ + authType: normalizeAuthType(providerCfg), + } + p.validateConfig(providerCfg) + + authClient := baseHTTPClient + if authClient == nil { + authClient = p.authHTTPClient(providerCfg, nil) + } + p.gemini = gemini.NewVertexWithHTTPClient(providerCfg, opts, authClient) + nativeBaseURL := vertexNativeBaseURL(providerCfg) + nativeCfg := llmclient.Config{ + ProviderName: "vertex", + BaseURL: nativeBaseURL, + Retry: opts.Resilience.Retry, + Hooks: opts.Hooks, + CircuitBreaker: opts.Resilience.CircuitBreaker, + } + if authClient != nil { + p.nativeClient = llmclient.NewWithHTTPClient(authClient, nativeCfg, p.setHeaders) + } else { + p.nativeClient = llmclient.New(nativeCfg, p.setHeaders) + } + return p +} + +func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { + if strings.TrimSpace(providerCfg.VertexProject) == "" || strings.TrimSpace(providerCfg.VertexLocation) == "" { + p.configErr = fmt.Errorf("vertex AI requires vertex_project and vertex_location") + return + } + switch p.authType { + case authTypeGCPADC: + return + case authTypeServiceAccount: + if googleauth.HasServiceAccount(googleauth.Config{ + ServiceAccountFile: providerCfg.ServiceAccountFile, + ServiceAccountJSON: providerCfg.ServiceAccountJSON, + ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, + }) { + return + } + p.configErr = fmt.Errorf("vertex AI service account auth requires service_account_file, service_account_json, or service_account_json_base64") + default: + p.configErr = fmt.Errorf("unsupported Vertex AI auth type %q", providerCfg.AuthType) + } +} + +func normalizeAuthType(providerCfg providers.ProviderConfig) string { + return googleauth.NormalizeAuthType(providerCfg.AuthType, googleauth.HasServiceAccount(googleauth.Config{ + ServiceAccountFile: providerCfg.ServiceAccountFile, + ServiceAccountJSON: providerCfg.ServiceAccountJSON, + ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, + })) +} + +func (p *Provider) authHTTPClient(providerCfg providers.ProviderConfig, base *http.Client) *http.Client { + if p.configErr != nil { + return base + } + source, err := googleauth.TokenSource(context.Background(), googleauth.Config{ + AuthType: p.authType, + ServiceAccountFile: providerCfg.ServiceAccountFile, + ServiceAccountJSON: providerCfg.ServiceAccountJSON, + ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, + Scope: providerCfg.GCPScope, + }) + if err != nil { + p.configErr = err + return base + } + if base == nil { + base = httpclient.NewDefaultHTTPClient() + } + return googleauth.HTTPClient(base, source) +} + +func (p *Provider) ready() error { + if p.configErr == nil { + return nil + } + return core.NewProviderError("vertex", http.StatusBadGateway, "invalid Vertex AI provider configuration: "+p.configErr.Error(), p.configErr) +} + +func (p *Provider) setHeaders(req *http.Request) { + if requestID := core.GetRequestID(req.Context()); requestID != "" { + req.Header.Set("X-Request-Id", requestID) + } +} + +// ChatCompletion sends a chat completion request to Vertex AI Gemini. +func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*core.ChatResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } + resp, err := p.gemini.ChatCompletion(ctx, req) + if resp != nil { + resp.Provider = "vertex" + } + return resp, err +} + +// StreamChatCompletion returns a raw response body for streaming. +func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { + if err := p.ready(); err != nil { + return nil, err + } + return p.gemini.StreamChatCompletion(ctx, req) +} + +// ListModels retrieves Vertex AI publisher models. +func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } + return p.gemini.ListModels(ctx) +} + +// Responses sends a Responses API request to Vertex AI Gemini. +func (p *Provider) Responses(ctx context.Context, req *core.ResponsesRequest) (*core.ResponsesResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } + return providers.ResponsesViaChat(ctx, p, req) +} + +// StreamResponses returns a raw response body for streaming Responses API. +func (p *Provider) StreamResponses(ctx context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) { + if err := p.ready(); err != nil { + return nil, err + } + return providers.StreamResponsesViaChat(ctx, p, req, "vertex") +} + +// Embeddings sends an embedding request through Vertex AI native prediction. +func (p *Provider) Embeddings(ctx context.Context, req *core.EmbeddingRequest) (*core.EmbeddingResponse, error) { + if err := p.ready(); err != nil { + return nil, err + } + if req == nil { + return nil, core.NewInvalidRequestError("embedding request is required", nil) + } + inputs, err := embeddingInputs(req.Input) + if err != nil { + return nil, err + } + + body := vertexEmbeddingPredictRequest{ + Instances: make([]vertexEmbeddingInstance, 0, len(inputs)), + Parameters: map[string]any{ + "autoTruncate": true, + }, + } + if req.Dimensions != nil && *req.Dimensions > 0 { + body.Parameters["outputDimensionality"] = *req.Dimensions + } + for _, input := range inputs { + body.Instances = append(body.Instances, vertexEmbeddingInstance{Content: input}) + } + + var resp vertexEmbeddingPredictResponse + err = p.nativeClient.Do(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: vertexPredictEndpoint(req.Model), + Body: body, + }, &resp) + if err != nil { + return nil, err + } + return openAIEmbeddingResponse(req, &resp) +} + +type vertexEmbeddingPredictRequest struct { + Instances []vertexEmbeddingInstance `json:"instances"` + Parameters map[string]any `json:"parameters,omitempty"` +} + +type vertexEmbeddingInstance struct { + Content string `json:"content"` +} + +type vertexEmbeddingPredictResponse struct { + Predictions []vertexEmbeddingPrediction `json:"predictions"` +} + +type vertexEmbeddingPrediction struct { + Embeddings vertexEmbeddingValues `json:"embeddings"` + Values []float64 `json:"values,omitempty"` +} + +type vertexEmbeddingValues struct { + Values []float64 `json:"values"` + Statistics vertexEmbeddingStatistics `json:"statistics"` +} + +type vertexEmbeddingStatistics struct { + TokenCount int `json:"token_count"` + Truncated bool `json:"truncated"` +} + +func embeddingInputs(input any) ([]string, error) { + switch v := input.(type) { + case string: + if strings.TrimSpace(v) == "" { + return nil, core.NewInvalidRequestError("embedding input is required", nil) + } + return []string{v}, nil + case []string: + return nonEmptyEmbeddingInputs(v) + case []any: + out := make([]string, 0, len(v)) + for _, item := range v { + text, ok := item.(string) + if !ok { + return nil, core.NewInvalidRequestError("vertex AI embeddings support string inputs", nil) + } + out = append(out, text) + } + return nonEmptyEmbeddingInputs(out) + default: + return nil, core.NewInvalidRequestError("vertex AI embeddings support string inputs", nil) + } +} + +func nonEmptyEmbeddingInputs(inputs []string) ([]string, error) { + out := make([]string, 0, len(inputs)) + for _, input := range inputs { + if strings.TrimSpace(input) == "" { + continue + } + out = append(out, input) + } + if len(out) == 0 { + return nil, core.NewInvalidRequestError("embedding input is required", nil) + } + return out, nil +} + +func openAIEmbeddingResponse(req *core.EmbeddingRequest, resp *vertexEmbeddingPredictResponse) (*core.EmbeddingResponse, error) { + out := &core.EmbeddingResponse{ + Object: "list", + Data: make([]core.EmbeddingData, 0, len(resp.Predictions)), + Model: req.Model, + Provider: "vertex", + } + for i, prediction := range resp.Predictions { + values := prediction.Embeddings.Values + if len(values) == 0 { + values = prediction.Values + } + embedding, err := json.Marshal(values) + if err != nil { + return nil, core.NewProviderError("vertex", http.StatusBadGateway, "failed to encode Vertex AI embedding response", err) + } + out.Data = append(out.Data, core.EmbeddingData{ + Object: "embedding", + Embedding: embedding, + Index: i, + }) + out.Usage.PromptTokens += prediction.Embeddings.Statistics.TokenCount + } + out.Usage.TotalTokens = out.Usage.PromptTokens + return out, nil +} + +func vertexNativeBaseURL(providerCfg providers.ProviderConfig) string { + _, nativeBaseURL := vertexBaseURLs(providerCfg) + return nativeBaseURL +} + +func vertexBaseURLs(providerCfg providers.ProviderConfig) (openAICompatibleBaseURL, nativeBaseURL string) { + baseURL := strings.TrimRight(strings.TrimSpace(providerCfg.BaseURL), "/") + if baseURL == "" { + project := strings.TrimSpace(providerCfg.VertexProject) + location := strings.TrimSpace(providerCfg.VertexLocation) + root := "https://aiplatform.googleapis.com/v1/projects/" + url.PathEscape(project) + "/locations/" + url.PathEscape(location) + return root + "/endpoints/openapi", root + "/publishers/google" + } + if nativeBaseURL, ok := vertexNativeBaseURLFromOpenAICompatibleBaseURL(baseURL); ok { + return baseURL, nativeBaseURL + } + if openAIBaseURL, ok := vertexOpenAICompatibleBaseURLFromNativeBaseURL(baseURL); ok { + return openAIBaseURL, baseURL + } + return baseURL, baseURL +} + +func vertexNativeBaseURLFromOpenAICompatibleBaseURL(baseURL string) (string, bool) { + const suffix = "/endpoints/openapi" + if !strings.HasSuffix(baseURL, suffix) { + return "", false + } + root := strings.TrimRight(strings.TrimSuffix(baseURL, suffix), "/") + if root == "" { + return "", false + } + return root + "/publishers/google", true +} + +func vertexOpenAICompatibleBaseURLFromNativeBaseURL(baseURL string) (string, bool) { + const suffix = "/publishers/google" + if !strings.HasSuffix(baseURL, suffix) { + return "", false + } + root := strings.TrimRight(strings.TrimSuffix(baseURL, suffix), "/") + if root == "" { + return "", false + } + return root + "/endpoints/openapi", true +} + +func vertexPredictEndpoint(model string) string { + model = normalizeVertexModelID(model) + return "/models/" + url.PathEscape(model) + ":predict" +} + +func normalizeVertexModelID(model string) string { + model = strings.TrimSpace(model) + if idx := strings.LastIndex(model, "/models/"); idx >= 0 { + model = model[idx+len("/models/"):] + } + model = strings.TrimPrefix(model, "models/") + model = strings.TrimPrefix(model, "google/") + return model +} + +var _ core.Provider = (*Provider)(nil) diff --git a/internal/providers/vertex/vertex_test.go b/internal/providers/vertex/vertex_test.go new file mode 100644 index 00000000..49bb2c03 --- /dev/null +++ b/internal/providers/vertex/vertex_test.go @@ -0,0 +1,152 @@ +package vertex + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "gomodel/internal/core" + "gomodel/internal/providers" + "gomodel/internal/providers/googleauth" + + "golang.org/x/oauth2" +) + +func TestProviderDoesNotExposeFilesOrBatches(t *testing.T) { + provider := newProvider(testConfig(), providers.ProviderOptions{}, authedTestClient(http.DefaultClient)) + + if _, ok := any(provider).(core.NativeFileProvider); ok { + t.Fatal("Vertex provider must not expose native files") + } + if _, ok := any(provider).(core.NativeBatchProvider); ok { + t.Fatal("Vertex provider must not expose native batches") + } +} + +func TestEmbeddingsUsesNativePrediction(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models/text-embedding-005:predict" { + t.Errorf("Path = %q, want Vertex native predict endpoint", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer vertex-token" { + t.Errorf("Authorization = %q, want Bearer vertex-token", got) + } + if got := r.Header.Get("x-goog-api-key"); got != "" { + t.Errorf("x-goog-api-key = %q, want empty for Vertex OAuth", got) + } + + var payload vertexEmbeddingPredictRequest + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode request: %v", err) + } + if len(payload.Instances) != 2 { + t.Fatalf("instances = %+v, want 2", payload.Instances) + } + if payload.Instances[0].Content != "hello" || payload.Instances[1].Content != "world" { + t.Fatalf("instances = %+v, want hello/world", payload.Instances) + } + if got := payload.Parameters["outputDimensionality"]; got != float64(3) { + t.Fatalf("outputDimensionality = %#v, want 3", got) + } + + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "predictions": [ + {"embeddings": {"values": [0.1, 0.2, 0.3], "statistics": {"token_count": 4}}}, + {"embeddings": {"values": [0.4, 0.5, 0.6], "statistics": {"token_count": 5}}} + ] + }`)) + })) + defer server.Close() + + dimensions := 3 + cfg := testConfig() + cfg.BaseURL = server.URL + "/v1/projects/prod-ai/locations/us-central1/publishers/google" + provider := newProvider(cfg, providers.ProviderOptions{}, authedTestClient(server.Client())) + + resp, err := provider.Embeddings(context.Background(), &core.EmbeddingRequest{ + Model: "google/text-embedding-005", + Input: []string{"hello", "world"}, + Dimensions: &dimensions, + }) + if err != nil { + t.Fatalf("Embeddings() error = %v", err) + } + if resp.Provider != "vertex" { + t.Fatalf("Provider = %q, want vertex", resp.Provider) + } + if len(resp.Data) != 2 { + t.Fatalf("data = %+v, want 2 embeddings", resp.Data) + } + if got := string(resp.Data[0].Embedding); got != `[0.1,0.2,0.3]` { + t.Fatalf("embedding = %s, want [0.1,0.2,0.3]", got) + } + if resp.Usage.PromptTokens != 9 || resp.Usage.TotalTokens != 9 { + t.Fatalf("usage = %+v, want 9 prompt/total tokens", resp.Usage) + } +} + +func TestVertexBaseURLs(t *testing.T) { + tests := []struct { + name string + cfg providers.ProviderConfig + wantCompat string + wantNative string + }{ + { + name: "derives official vertex bases from project and location", + cfg: providers.ProviderConfig{ + VertexProject: "prod-ai", + VertexLocation: "us-central1", + }, + wantCompat: "https://aiplatform.googleapis.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi", + wantNative: "https://aiplatform.googleapis.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, + { + name: "custom OpenAI-compatible vertex URL derives native sibling", + cfg: providers.ProviderConfig{ + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi/", + }, + wantCompat: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi", + wantNative: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, + { + name: "custom native vertex URL derives OpenAI-compatible sibling", + cfg: providers.ProviderConfig{ + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google/", + }, + wantCompat: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/endpoints/openapi", + wantNative: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotCompat, gotNative := vertexBaseURLs(tt.cfg) + if gotCompat != tt.wantCompat { + t.Fatalf("OpenAI-compatible base = %q, want %q", gotCompat, tt.wantCompat) + } + if gotNative != tt.wantNative { + t.Fatalf("native base = %q, want %q", gotNative, tt.wantNative) + } + }) + } +} + +func testConfig() providers.ProviderConfig { + return providers.ProviderConfig{ + Type: "vertex", + AuthType: "gcp_adc", + VertexProject: "prod-ai", + VertexLocation: "us-central1", + } +} + +func authedTestClient(base *http.Client) *http.Client { + return googleauth.HTTPClient(base, oauth2.StaticTokenSource(&oauth2.Token{ + AccessToken: "vertex-token", + TokenType: "Bearer", + })) +} From 66435b6e4e372af7d484d91167a8bdde9baae594 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 18:01:26 +0200 Subject: [PATCH 2/9] fix(providers): address vertex review feedback --- internal/providers/config.go | 26 ++++++---- internal/providers/config_test.go | 51 +++++++++++++++++++ internal/providers/gemini/gemini.go | 10 +++- internal/providers/gemini/gemini_test.go | 11 ++++ internal/providers/gemini/native.go | 3 ++ internal/providers/googleauth/googleauth.go | 43 ++++++++++++++-- .../providers/googleauth/googleauth_test.go | 43 ++++++++++++++++ internal/providers/vertex/vertex.go | 18 ++++--- internal/providers/vertex/vertex_test.go | 28 ++++++++++ 9 files changed, 212 insertions(+), 21 deletions(-) create mode 100644 internal/providers/googleauth/googleauth_test.go diff --git a/internal/providers/config.go b/internal/providers/config.go index 94f05f5a..c2d468da 100644 --- a/internal/providers/config.go +++ b/internal/providers/config.go @@ -214,22 +214,29 @@ func parseProviderEnvKey(prefix, key string, spec DiscoveryConfig) (string, prov name string field providerEnvField }{ - {name: "SERVICE_ACCOUNT_JSON_BASE64", field: providerEnvFieldServiceAccountJSONBase64}, - {name: "SERVICE_ACCOUNT_JSON", field: providerEnvFieldServiceAccountJSON}, - {name: "SERVICE_ACCOUNT_FILE", field: providerEnvFieldServiceAccountFile}, - {name: "VERTEX_PROJECT", field: providerEnvFieldVertexProject}, - {name: "VERTEX_LOCATION", field: providerEnvFieldVertexLocation}, {name: "API_VERSION", field: providerEnvFieldAPIVersion}, {name: "BASE_URL", field: providerEnvFieldBaseURL}, {name: "AUTH_TYPE", field: providerEnvFieldAuthType}, {name: "API_MODE", field: providerEnvFieldAPIMode}, - {name: "PROJECT", field: providerEnvFieldVertexProject}, - {name: "LOCATION", field: providerEnvFieldVertexLocation}, - {name: "GCP_SCOPE", field: providerEnvFieldGCPScope}, {name: "BACKEND", field: providerEnvFieldBackend}, {name: "API_KEY", field: providerEnvFieldAPIKey}, {name: "MODELS", field: providerEnvFieldModels}, } + if strings.EqualFold(prefix, "VERTEX") { + fields = append([]struct { + name string + field providerEnvField + }{ + {name: "SERVICE_ACCOUNT_JSON_BASE64", field: providerEnvFieldServiceAccountJSONBase64}, + {name: "SERVICE_ACCOUNT_JSON", field: providerEnvFieldServiceAccountJSON}, + {name: "SERVICE_ACCOUNT_FILE", field: providerEnvFieldServiceAccountFile}, + {name: "VERTEX_PROJECT", field: providerEnvFieldVertexProject}, + {name: "VERTEX_LOCATION", field: providerEnvFieldVertexLocation}, + {name: "PROJECT", field: providerEnvFieldVertexProject}, + {name: "LOCATION", field: providerEnvFieldVertexLocation}, + {name: "GCP_SCOPE", field: providerEnvFieldGCPScope}, + }, fields...) + } for _, candidate := range fields { if candidate.field == providerEnvFieldAPIVersion && !spec.SupportsAPIVersion { @@ -573,7 +580,8 @@ func isVertexProviderConfig(p config.RawProviderConfig) bool { } func validVertexProviderConfig(p config.RawProviderConfig) bool { - if !hasResolvedProviderValue(p.VertexProject) || !hasResolvedProviderValue(p.VertexLocation) { + if !hasResolvedProviderValue(p.BaseURL) && + (!hasResolvedProviderValue(p.VertexProject) || !hasResolvedProviderValue(p.VertexLocation)) { return false } authType := strings.ToLower(strings.TrimSpace(p.AuthType)) diff --git a/internal/providers/config_test.go b/internal/providers/config_test.go index ce594a3f..81767503 100644 --- a/internal/providers/config_test.go +++ b/internal/providers/config_test.go @@ -602,6 +602,27 @@ func TestResolveProviders_KeepsVertexWithProjectAndLocationUsingDefaultADC(t *te } } +func TestResolveProviders_KeepsVertexWithBaseURLWithoutProjectLocation(t *testing.T) { + t.Setenv("VERTEX_BASE_URL", "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google") + t.Setenv("VERTEX_AUTH_TYPE", "gcp_adc") + + got, filteredRaw := resolveProviders(map[string]config.RawProviderConfig{}, globalResilience, testDiscoveryConfigs) + + p, exists := got["vertex"] + if !exists { + t.Fatal("expected vertex with resolved base URL to be resolved") + } + if p.Type != "vertex" { + t.Fatalf("Type = %q, want vertex", p.Type) + } + if p.BaseURL != "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google" { + t.Fatalf("BaseURL = %q, want custom Vertex base URL", p.BaseURL) + } + if _, exists := filteredRaw["vertex"]; !exists { + t.Fatal("expected raw vertex with resolved base URL to be retained") + } +} + func TestResolveProviders_FiltersVertexServiceAccountWithoutCredentials(t *testing.T) { t.Setenv("VERTEX_PROJECT", "prod-ai") t.Setenv("VERTEX_LOCATION", "us-central1") @@ -658,6 +679,36 @@ func TestResolveProviders_FiltersVertexWithUnresolvedServiceAccountPlaceholder(t } } +func TestApplyProviderEnvVars_GeminiIgnoresVertexSpecificEnv(t *testing.T) { + t.Setenv("GEMINI_API_KEY", "gemini-key") + t.Setenv("GEMINI_PROJECT", "prod-ai") + t.Setenv("GEMINI_LOCATION", "us-central1") + t.Setenv("GEMINI_GCP_SCOPE", "scope-a") + t.Setenv("GEMINI_SERVICE_ACCOUNT_FILE", "/secrets/gemini.json") + + got := applyProviderEnvVars(map[string]config.RawProviderConfig{}, testDiscoveryConfigs) + + p, exists := got["gemini"] + if !exists { + t.Fatal("expected gemini to be discovered from GEMINI_API_KEY") + } + if p.Type != "gemini" { + t.Fatalf("Type = %q, want gemini", p.Type) + } + if p.VertexProject != "" { + t.Fatalf("VertexProject = %q, want empty", p.VertexProject) + } + if p.VertexLocation != "" { + t.Fatalf("VertexLocation = %q, want empty", p.VertexLocation) + } + if p.GCPScope != "" { + t.Fatalf("GCPScope = %q, want empty", p.GCPScope) + } + if p.ServiceAccountFile != "" { + t.Fatalf("ServiceAccountFile = %q, want empty", p.ServiceAccountFile) + } +} + func TestApplyProviderEnvVars_DiscoversVLLMFromBaseURLWithoutAPIKey(t *testing.T) { t.Setenv("VLLM_BASE_URL", "http://localhost:8000/v1") diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index c89aa55f..c2dbdf34 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -155,8 +155,9 @@ func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { p.configErr = fmt.Errorf("vertex Gemini requires gcp_adc or gcp_service_account auth") return } - if p.backend == geminiBackendVertex && (strings.TrimSpace(providerCfg.VertexProject) == "" || strings.TrimSpace(providerCfg.VertexLocation) == "") { - p.configErr = fmt.Errorf("vertex Gemini requires vertex_project and vertex_location") + if p.backend == geminiBackendVertex && !hasResolvedProviderValue(providerCfg.BaseURL) && + (!hasResolvedProviderValue(providerCfg.VertexProject) || !hasResolvedProviderValue(providerCfg.VertexLocation)) { + p.configErr = fmt.Errorf("vertex Gemini requires base_url or vertex_project and vertex_location") return } if p.backend == geminiBackendAIStudio && p.authType == geminiAuthTypeAPIKey && strings.TrimSpace(providerCfg.APIKey) == "" { @@ -231,6 +232,11 @@ func normalizeGeminiBackend(cfg providers.ProviderConfig) string { return geminiBackendAIStudio } +func hasResolvedProviderValue(value string) bool { + value = strings.TrimSpace(value) + return value != "" && !strings.Contains(value, "${") +} + func normalizeGeminiAuthType(backend string, cfg providers.ProviderConfig) string { authType := strings.ToLower(strings.TrimSpace(cfg.AuthType)) switch authType { diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 2c7f8d31..2bbca941 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -140,6 +140,17 @@ func TestVertexBaseURLs(t *testing.T) { } } +func TestNewVertexWithHTTPClientAcceptsBaseURLWithoutProjectLocation(t *testing.T) { + p := NewVertexWithHTTPClient(providers.ProviderConfig{ + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + AuthType: "gcp_adc", + }, providers.ProviderOptions{}, http.DefaultClient) + + if err := p.ready(); err != nil { + t.Fatalf("ready() error = %v, want nil for Vertex custom base URL", err) + } +} + func TestVertexModelNormalization(t *testing.T) { tests := []struct { in string diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go index daa9d563..22728f1b 100644 --- a/internal/providers/gemini/native.go +++ b/internal/providers/gemini/native.go @@ -835,6 +835,9 @@ func displayModelIDFromGemini(model, backend string) string { return model } +// isGeminiExposedModel normalizes provider model names and exposes only the +// model families reachable through Gemini/OpenAI-compatible text endpoints. +// Families such as imagen-* use different upstream endpoints. func isGeminiExposedModel(modelID string) bool { modelID = normalizeGeminiModelID(modelID) return strings.HasPrefix(modelID, "gemini-") || strings.HasPrefix(modelID, "text-embedding-") diff --git a/internal/providers/googleauth/googleauth.go b/internal/providers/googleauth/googleauth.go index f6abeaed..3eafef0d 100644 --- a/internal/providers/googleauth/googleauth.go +++ b/internal/providers/googleauth/googleauth.go @@ -69,14 +69,13 @@ func TokenSource(ctx context.Context, cfg Config) (oauth2.TokenSource, error) { return nil, fmt.Errorf("load application default credentials: %w", err) } return source, nil - default: - return nil, fmt.Errorf("unsupported Google auth type %q", cfg.AuthType) } + return nil, fmt.Errorf("unsupported Google auth type %q", authType) } func serviceAccountJSON(cfg Config) ([]byte, error) { if value := strings.TrimSpace(cfg.ServiceAccountJSONBase64); value != "" { - decoded, err := base64.StdEncoding.DecodeString(value) + decoded, err := decodeServiceAccountBase64(value) if err != nil { return nil, fmt.Errorf("decode service account JSON: %w", err) } @@ -95,6 +94,44 @@ func serviceAccountJSON(cfg Config) ([]byte, error) { return nil, fmt.Errorf("service account credentials are required") } +func decodeServiceAccountBase64(value string) ([]byte, error) { + decoded, stdErr := base64.StdEncoding.DecodeString(value) + if stdErr == nil { + return decoded, nil + } + + for _, encoding := range []*base64.Encoding{ + base64.RawStdEncoding, + base64.URLEncoding, + base64.RawURLEncoding, + } { + if decoded, err := encoding.DecodeString(value); err == nil { + return decoded, nil + } + } + if padded := paddedBase64(value); padded != value { + for _, encoding := range []*base64.Encoding{base64.StdEncoding, base64.URLEncoding} { + if decoded, err := encoding.DecodeString(padded); err == nil { + return decoded, nil + } + } + } + return nil, fmt.Errorf("standard base64 decode failed: %w; also tried raw standard, URL-safe, raw URL-safe, and padded variants", stdErr) +} + +func paddedBase64(value string) string { + switch remainder := len(value) % 4; remainder { + case 0: + return value + case 2: + return value + "==" + case 3: + return value + "=" + default: + return value + } +} + func HTTPClient(base *http.Client, source oauth2.TokenSource) *http.Client { if base == nil { base = http.DefaultClient diff --git a/internal/providers/googleauth/googleauth_test.go b/internal/providers/googleauth/googleauth_test.go new file mode 100644 index 00000000..cd329400 --- /dev/null +++ b/internal/providers/googleauth/googleauth_test.go @@ -0,0 +1,43 @@ +package googleauth + +import ( + "encoding/base64" + "strings" + "testing" +) + +func TestServiceAccountJSONDecodesURLSafeBase64(t *testing.T) { + want := []byte{0xfb, 0xff, 0xfe} + encoded := base64.RawURLEncoding.EncodeToString(want) + + got, err := serviceAccountJSON(Config{ServiceAccountJSONBase64: encoded}) + if err != nil { + t.Fatalf("serviceAccountJSON() error = %v", err) + } + if string(got) != string(want) { + t.Fatalf("decoded bytes = %q, want %q", string(got), string(want)) + } +} + +func TestServiceAccountJSONDecodesPaddedURLSafeBase64(t *testing.T) { + want := []byte{0xfb, 0xff, 0xfe} + encoded := base64.URLEncoding.EncodeToString(want) + + got, err := serviceAccountJSON(Config{ServiceAccountJSONBase64: encoded}) + if err != nil { + t.Fatalf("serviceAccountJSON() error = %v", err) + } + if string(got) != string(want) { + t.Fatalf("decoded bytes = %q, want %q", string(got), string(want)) + } +} + +func TestServiceAccountJSONReportsOriginalBase64DecodeError(t *testing.T) { + _, err := serviceAccountJSON(Config{ServiceAccountJSONBase64: "not valid base64!"}) + if err == nil { + t.Fatal("expected invalid base64 error") + } + if !strings.Contains(err.Error(), "standard base64 decode failed") { + t.Fatalf("error = %v, want standard decode context", err) + } +} diff --git a/internal/providers/vertex/vertex.go b/internal/providers/vertex/vertex.go index 3c03a7ad..46824db3 100644 --- a/internal/providers/vertex/vertex.go +++ b/internal/providers/vertex/vertex.go @@ -75,8 +75,9 @@ func newProvider(providerCfg providers.ProviderConfig, opts providers.ProviderOp } func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { - if strings.TrimSpace(providerCfg.VertexProject) == "" || strings.TrimSpace(providerCfg.VertexLocation) == "" { - p.configErr = fmt.Errorf("vertex AI requires vertex_project and vertex_location") + if !hasResolvedProviderValue(providerCfg.BaseURL) && + (!hasResolvedProviderValue(providerCfg.VertexProject) || !hasResolvedProviderValue(providerCfg.VertexLocation)) { + p.configErr = fmt.Errorf("vertex AI requires base_url or vertex_project and vertex_location") return } switch p.authType { @@ -96,6 +97,11 @@ func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { } } +func hasResolvedProviderValue(value string) bool { + value = strings.TrimSpace(value) + return value != "" && !strings.Contains(value, "${") +} + func normalizeAuthType(providerCfg providers.ProviderConfig) string { return googleauth.NormalizeAuthType(providerCfg.AuthType, googleauth.HasServiceAccount(googleauth.Config{ ServiceAccountFile: providerCfg.ServiceAccountFile, @@ -273,17 +279,15 @@ func embeddingInputs(input any) ([]string, error) { } func nonEmptyEmbeddingInputs(inputs []string) ([]string, error) { - out := make([]string, 0, len(inputs)) for _, input := range inputs { if strings.TrimSpace(input) == "" { - continue + return nil, core.NewInvalidRequestError("embedding input must not be empty", nil) } - out = append(out, input) } - if len(out) == 0 { + if len(inputs) == 0 { return nil, core.NewInvalidRequestError("embedding input is required", nil) } - return out, nil + return inputs, nil } func openAIEmbeddingResponse(req *core.EmbeddingRequest, resp *vertexEmbeddingPredictResponse) (*core.EmbeddingResponse, error) { diff --git a/internal/providers/vertex/vertex_test.go b/internal/providers/vertex/vertex_test.go index 49bb2c03..9e342fb6 100644 --- a/internal/providers/vertex/vertex_test.go +++ b/internal/providers/vertex/vertex_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "strings" "testing" "gomodel/internal/core" @@ -88,6 +89,33 @@ func TestEmbeddingsUsesNativePrediction(t *testing.T) { } } +func TestEmbeddingsRejectsEmptyStringInBatch(t *testing.T) { + provider := newProvider(testConfig(), providers.ProviderOptions{}, authedTestClient(http.DefaultClient)) + + _, err := provider.Embeddings(context.Background(), &core.EmbeddingRequest{ + Model: "google/text-embedding-005", + Input: []string{"hello", "", "world"}, + }) + if err == nil { + t.Fatal("expected empty batch embedding input to be rejected") + } + if !strings.Contains(err.Error(), "embedding input must not be empty") { + t.Fatalf("error = %v, want empty input error", err) + } +} + +func TestNewAcceptsBaseURLWithoutProjectLocation(t *testing.T) { + provider := newProvider(providers.ProviderConfig{ + Type: "vertex", + AuthType: "gcp_adc", + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, providers.ProviderOptions{}, authedTestClient(http.DefaultClient)) + + if err := provider.ready(); err != nil { + t.Fatalf("ready() error = %v, want nil for custom Vertex base URL", err) + } +} + func TestVertexBaseURLs(t *testing.T) { tests := []struct { name string From ea609464e6aeb3f155ad61cd57694c74f602c266 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 18:02:50 +0200 Subject: [PATCH 3/9] fix(providers): validate vertex auth before normalization --- internal/providers/vertex/vertex.go | 13 ++++++++++++- internal/providers/vertex/vertex_test.go | 16 ++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/internal/providers/vertex/vertex.go b/internal/providers/vertex/vertex.go index 46824db3..0d9c952b 100644 --- a/internal/providers/vertex/vertex.go +++ b/internal/providers/vertex/vertex.go @@ -80,6 +80,10 @@ func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { p.configErr = fmt.Errorf("vertex AI requires base_url or vertex_project and vertex_location") return } + if !validAuthType(providerCfg.AuthType) { + p.configErr = fmt.Errorf("unsupported vertex AI auth type %q", providerCfg.AuthType) + return + } switch p.authType { case authTypeGCPADC: return @@ -92,8 +96,15 @@ func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { return } p.configErr = fmt.Errorf("vertex AI service account auth requires service_account_file, service_account_json, or service_account_json_base64") + } +} + +func validAuthType(authType string) bool { + switch strings.ToLower(strings.TrimSpace(authType)) { + case "", "gcp_adc", "adc", "google_adc", "gcp_service_account", "service_account": + return true default: - p.configErr = fmt.Errorf("unsupported Vertex AI auth type %q", providerCfg.AuthType) + return false } } diff --git a/internal/providers/vertex/vertex_test.go b/internal/providers/vertex/vertex_test.go index 9e342fb6..44ecdd59 100644 --- a/internal/providers/vertex/vertex_test.go +++ b/internal/providers/vertex/vertex_test.go @@ -116,6 +116,22 @@ func TestNewAcceptsBaseURLWithoutProjectLocation(t *testing.T) { } } +func TestNewRejectsUnsupportedAuthType(t *testing.T) { + provider := newProvider(providers.ProviderConfig{ + Type: "vertex", + AuthType: "api_key", + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, providers.ProviderOptions{}, authedTestClient(http.DefaultClient)) + + err := provider.ready() + if err == nil { + t.Fatal("expected unsupported auth type error") + } + if !strings.Contains(err.Error(), `unsupported vertex AI auth type "api_key"`) { + t.Fatalf("error = %v, want unsupported auth type", err) + } +} + func TestVertexBaseURLs(t *testing.T) { tests := []struct { name string From 6f2a543d484badb744ca84bbfdea8c38f499f17e Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 18:31:16 +0200 Subject: [PATCH 4/9] fix(providers): address vertex auth review feedback --- docs/guides/gemini.mdx | 5 + internal/providers/gemini/gemini.go | 7 +- internal/providers/gemini/gemini_test.go | 39 ++++ .../providers/googleauth/googleauth_test.go | 170 ++++++++++++++++++ internal/providers/vertex/vertex.go | 53 ++++-- internal/providers/vertex/vertex_test.go | 45 +++++ 6 files changed, 299 insertions(+), 20 deletions(-) diff --git a/docs/guides/gemini.mdx b/docs/guides/gemini.mdx index 8789f901..8992685c 100644 --- a/docs/guides/gemini.mdx +++ b/docs/guides/gemini.mdx @@ -72,6 +72,11 @@ Vertex requires both project and location. `VERTEX_AUTH_TYPE` defaults to Application Default Credentials (`gcp_adc`) when service-account fields are not set. Use service-account JSON directly when ADC is not appropriate: + + Vertex authentication is enterprise-oriented and may become paid or licensed + in a future release. + + ```bash export VERTEX_AUTH_TYPE="gcp_service_account" export VERTEX_SERVICE_ACCOUNT_FILE="/secrets/service-account.json" diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index c2dbdf34..306d6d95 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -160,6 +160,10 @@ func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { p.configErr = fmt.Errorf("vertex Gemini requires base_url or vertex_project and vertex_location") return } + if p.backend == geminiBackendAIStudio && p.authType != geminiAuthTypeAPIKey { + p.configErr = fmt.Errorf("ai studio backend does not support GCP auth; use Vertex backend or provide an API key") + return + } if p.backend == geminiBackendAIStudio && p.authType == geminiAuthTypeAPIKey && strings.TrimSpace(providerCfg.APIKey) == "" { p.configErr = fmt.Errorf("gemini API key is required") } @@ -225,8 +229,7 @@ func normalizeGeminiBackend(cfg providers.ProviderConfig) string { return geminiBackendAIStudio } if strings.TrimSpace(cfg.VertexProject) != "" || - strings.TrimSpace(cfg.VertexLocation) != "" || - strings.HasPrefix(strings.ToLower(strings.TrimSpace(cfg.AuthType)), "gcp") { + strings.TrimSpace(cfg.VertexLocation) != "" { return geminiBackendVertex } return geminiBackendAIStudio diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 2bbca941..c6f57c40 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -42,6 +42,45 @@ func TestNew_ReturnsProvider(t *testing.T) { } } +func TestNew_AIStudioRejectsGCPAuthAliases(t *testing.T) { + for _, authType := range []string{"adc", "service_account"} { + t.Run(authType, func(t *testing.T) { + provider := New(providers.ProviderConfig{ + BaseURL: defaultOpenAICompatibleBaseURL, + AuthType: authType, + }, providers.ProviderOptions{}) + + geminiProvider, ok := provider.(*Provider) + if !ok { + t.Fatalf("provider type = %T, want *Provider", provider) + } + err := geminiProvider.ready() + if err == nil { + t.Fatal("expected GCP auth alias to be rejected for AI Studio") + } + if !strings.Contains(err.Error(), "ai studio backend does not support GCP auth") { + t.Fatalf("error = %v, want AI Studio GCP auth rejection", err) + } + }) + } +} + +func TestNew_VertexAcceptsGCPAuthAliases(t *testing.T) { + for _, authType := range []string{"adc", "service_account"} { + t.Run(authType, func(t *testing.T) { + p := NewVertexWithHTTPClient(providers.ProviderConfig{ + AuthType: authType, + VertexProject: "prod-ai", + VertexLocation: "us-central1", + }, providers.ProviderOptions{}, http.DefaultClient) + + if err := p.ready(); err != nil { + t.Fatalf("ready() error = %v, want nil for Vertex auth alias", err) + } + }) + } +} + func TestGeminiBaseURLs(t *testing.T) { tests := []struct { name string diff --git a/internal/providers/googleauth/googleauth_test.go b/internal/providers/googleauth/googleauth_test.go index cd329400..a7df8886 100644 --- a/internal/providers/googleauth/googleauth_test.go +++ b/internal/providers/googleauth/googleauth_test.go @@ -1,7 +1,18 @@ package googleauth import ( + "context" + "crypto/rand" + "crypto/rsa" + "crypto/x509" "encoding/base64" + "encoding/json" + "encoding/pem" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" "strings" "testing" ) @@ -41,3 +52,162 @@ func TestServiceAccountJSONReportsOriginalBase64DecodeError(t *testing.T) { t.Fatalf("error = %v, want standard decode context", err) } } + +func TestTokenSourceAndHTTPClientAuthSelection(t *testing.T) { + tests := []struct { + name string + cfg func(t *testing.T, tokenURL string) Config + wantToken string + wantScope string + wantADCFile bool + }{ + { + name: "service account base64 uses service account path and default scope", + cfg: func(t *testing.T, tokenURL string) Config { + credentials := serviceAccountCredentials(t, tokenURL) + encoded := base64.StdEncoding.EncodeToString([]byte(credentials)) + if _, err := serviceAccountJSON(Config{ServiceAccountJSONBase64: encoded}); err != nil { + t.Fatalf("serviceAccountJSON() error = %v", err) + } + return Config{ServiceAccountJSONBase64: encoded} + }, + wantToken: "service-account-token", + wantScope: DefaultScope, + }, + { + name: "empty service account config uses ADC path", + cfg: func(t *testing.T, tokenURL string) Config { + t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", adcCredentialsFile(t, tokenURL)) + return Config{Scope: "https://www.googleapis.com/auth/custom"} + }, + wantToken: "adc-token", + wantADCFile: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var gotForm url.Values + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Fatalf("ParseForm() error = %v", err) + } + gotForm = r.PostForm + token := tt.wantToken + if gotForm.Get("grant_type") == "urn:ietf:params:oauth:grant-type:jwt-bearer" { + token = "service-account-token" + } else if gotForm.Get("grant_type") == "refresh_token" { + token = "adc-token" + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": token, + "token_type": "Bearer", + "expires_in": 3600, + }) + })) + defer tokenServer.Close() + + source, err := TokenSource(context.Background(), tt.cfg(t, tokenServer.URL)) + if err != nil { + t.Fatalf("TokenSource() error = %v", err) + } + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer "+tt.wantToken { + t.Fatalf("Authorization = %q, want Bearer %s", got, tt.wantToken) + } + w.WriteHeader(http.StatusNoContent) + })) + defer upstream.Close() + + client := HTTPClient(upstream.Client(), source) + resp, err := client.Get(upstream.URL) + if err != nil { + t.Fatalf("HTTPClient request error = %v", err) + } + _ = resp.Body.Close() + + if tt.wantScope != "" { + scope := tokenRequestScope(t, gotForm) + if scope != tt.wantScope { + t.Fatalf("scope = %q, want %q", scope, tt.wantScope) + } + } + if tt.wantADCFile && gotForm.Get("refresh_token") != "adc-refresh-token" { + t.Fatalf("refresh_token = %q, want adc-refresh-token", gotForm.Get("refresh_token")) + } + }) + } +} + +func adcCredentialsFile(t *testing.T, tokenURL string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "adc.json") + contents := map[string]string{ + "type": "authorized_user", + "client_id": "adc-client-id", + "client_secret": "adc-client-secret", + "refresh_token": "adc-refresh-token", + "token_uri": tokenURL, + } + encoded, err := json.Marshal(contents) + if err != nil { + t.Fatalf("failed to marshal ADC credentials: %v", err) + } + if err := os.WriteFile(path, encoded, 0o600); err != nil { + t.Fatalf("failed to write ADC credentials: %v", err) + } + return path +} + +func serviceAccountCredentials(t *testing.T, tokenURL string) string { + t.Helper() + key, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("failed to generate test RSA key: %v", err) + } + keyBytes, err := x509.MarshalPKCS8PrivateKey(key) + if err != nil { + t.Fatalf("failed to marshal test RSA key: %v", err) + } + keyPEM := pem.EncodeToMemory(&pem.Block{ + Type: "PRIVATE KEY", + Bytes: keyBytes, + }) + contents := map[string]string{ + "type": "service_account", + "client_email": "service@example.com", + "private_key_id": "test-key-id", + "private_key": string(keyPEM), + "token_uri": tokenURL, + } + encoded, err := json.Marshal(contents) + if err != nil { + t.Fatalf("failed to marshal service account credentials: %v", err) + } + return string(encoded) +} + +func tokenRequestScope(t *testing.T, form url.Values) string { + t.Helper() + if scope := form.Get("scope"); scope != "" { + return scope + } + assertion := form.Get("assertion") + parts := strings.Split(assertion, ".") + if len(parts) < 2 { + t.Fatalf("JWT assertion = %q, want header.payload.signature", assertion) + } + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + t.Fatalf("failed to decode JWT payload: %v", err) + } + var claims struct { + Scope string `json:"scope"` + } + if err := json.Unmarshal(payload, &claims); err != nil { + t.Fatalf("failed to decode JWT claims: %v", err) + } + return claims.Scope +} diff --git a/internal/providers/vertex/vertex.go b/internal/providers/vertex/vertex.go index 0d9c952b..f994084a 100644 --- a/internal/providers/vertex/vertex.go +++ b/internal/providers/vertex/vertex.go @@ -3,9 +3,12 @@ package vertex import ( "context" + "encoding/base64" + "encoding/binary" "encoding/json" "fmt" "io" + "math" "net/http" "net/url" "strings" @@ -88,14 +91,12 @@ func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { case authTypeGCPADC: return case authTypeServiceAccount: - if googleauth.HasServiceAccount(googleauth.Config{ - ServiceAccountFile: providerCfg.ServiceAccountFile, - ServiceAccountJSON: providerCfg.ServiceAccountJSON, - ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, - }) { + if googleauth.HasServiceAccount(buildGoogleAuthConfig(providerCfg)) { return } p.configErr = fmt.Errorf("vertex AI service account auth requires service_account_file, service_account_json, or service_account_json_base64") + default: + p.configErr = fmt.Errorf("unsupported normalized vertex AI auth type %q", p.authType) } } @@ -114,24 +115,16 @@ func hasResolvedProviderValue(value string) bool { } func normalizeAuthType(providerCfg providers.ProviderConfig) string { - return googleauth.NormalizeAuthType(providerCfg.AuthType, googleauth.HasServiceAccount(googleauth.Config{ - ServiceAccountFile: providerCfg.ServiceAccountFile, - ServiceAccountJSON: providerCfg.ServiceAccountJSON, - ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, - })) + return googleauth.NormalizeAuthType(providerCfg.AuthType, googleauth.HasServiceAccount(buildGoogleAuthConfig(providerCfg))) } func (p *Provider) authHTTPClient(providerCfg providers.ProviderConfig, base *http.Client) *http.Client { if p.configErr != nil { return base } - source, err := googleauth.TokenSource(context.Background(), googleauth.Config{ - AuthType: p.authType, - ServiceAccountFile: providerCfg.ServiceAccountFile, - ServiceAccountJSON: providerCfg.ServiceAccountJSON, - ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, - Scope: providerCfg.GCPScope, - }) + authCfg := buildGoogleAuthConfig(providerCfg) + authCfg.AuthType = p.authType + source, err := googleauth.TokenSource(context.Background(), authCfg) if err != nil { p.configErr = err return base @@ -142,6 +135,16 @@ func (p *Provider) authHTTPClient(providerCfg providers.ProviderConfig, base *ht return googleauth.HTTPClient(base, source) } +func buildGoogleAuthConfig(providerCfg providers.ProviderConfig) googleauth.Config { + return googleauth.Config{ + AuthType: providerCfg.AuthType, + ServiceAccountFile: providerCfg.ServiceAccountFile, + ServiceAccountJSON: providerCfg.ServiceAccountJSON, + ServiceAccountJSONBase64: providerCfg.ServiceAccountJSONBase64, + Scope: providerCfg.GCPScope, + } +} + func (p *Provider) ready() error { if p.configErr == nil { return nil @@ -313,7 +316,7 @@ func openAIEmbeddingResponse(req *core.EmbeddingRequest, resp *vertexEmbeddingPr if len(values) == 0 { values = prediction.Values } - embedding, err := json.Marshal(values) + embedding, err := encodeEmbedding(values, req.EncodingFormat) if err != nil { return nil, core.NewProviderError("vertex", http.StatusBadGateway, "failed to encode Vertex AI embedding response", err) } @@ -328,11 +331,25 @@ func openAIEmbeddingResponse(req *core.EmbeddingRequest, resp *vertexEmbeddingPr return out, nil } +func encodeEmbedding(values []float64, encodingFormat string) (json.RawMessage, error) { + if strings.EqualFold(strings.TrimSpace(encodingFormat), "base64") { + buf := make([]byte, len(values)*4) + for i, value := range values { + binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(float32(value))) + } + return json.Marshal(base64.StdEncoding.EncodeToString(buf)) + } + return json.Marshal(values) +} + func vertexNativeBaseURL(providerCfg providers.ProviderConfig) string { _, nativeBaseURL := vertexBaseURLs(providerCfg) return nativeBaseURL } +// TODO: Share Vertex URL derivation with the Gemini Vertex path if this logic +// changes again. It is intentionally duplicated today to keep provider package +// boundaries simple. func vertexBaseURLs(providerCfg providers.ProviderConfig) (openAICompatibleBaseURL, nativeBaseURL string) { baseURL := strings.TrimRight(strings.TrimSpace(providerCfg.BaseURL), "/") if baseURL == "" { diff --git a/internal/providers/vertex/vertex_test.go b/internal/providers/vertex/vertex_test.go index 44ecdd59..d8741e93 100644 --- a/internal/providers/vertex/vertex_test.go +++ b/internal/providers/vertex/vertex_test.go @@ -2,7 +2,10 @@ package vertex import ( "context" + "encoding/base64" + "encoding/binary" "encoding/json" + "math" "net/http" "net/http/httptest" "strings" @@ -104,6 +107,48 @@ func TestEmbeddingsRejectsEmptyStringInBatch(t *testing.T) { } } +func TestOpenAIEmbeddingResponseSupportsBase64Encoding(t *testing.T) { + resp, err := openAIEmbeddingResponse(&core.EmbeddingRequest{ + Model: "google/text-embedding-005", + EncodingFormat: "base64", + }, &vertexEmbeddingPredictResponse{ + Predictions: []vertexEmbeddingPrediction{{ + Embeddings: vertexEmbeddingValues{ + Values: []float64{0.5, -1.25}, + Statistics: vertexEmbeddingStatistics{TokenCount: 3}, + }, + }}, + }) + if err != nil { + t.Fatalf("openAIEmbeddingResponse() error = %v", err) + } + if len(resp.Data) != 1 { + t.Fatalf("data = %+v, want one embedding", resp.Data) + } + + var encoded string + if err := json.Unmarshal(resp.Data[0].Embedding, &encoded); err != nil { + t.Fatalf("embedding is not JSON string: %v", err) + } + decoded, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + t.Fatalf("embedding is not valid base64: %v", err) + } + if len(decoded) != 8 { + t.Fatalf("decoded length = %d, want 8", len(decoded)) + } + values := []float32{ + math.Float32frombits(binary.LittleEndian.Uint32(decoded[0:4])), + math.Float32frombits(binary.LittleEndian.Uint32(decoded[4:8])), + } + if values[0] != 0.5 || values[1] != -1.25 { + t.Fatalf("decoded values = %v, want [0.5 -1.25]", values) + } + if resp.Usage.PromptTokens != 3 || resp.Usage.TotalTokens != 3 { + t.Fatalf("usage = %+v, want 3 prompt/total tokens", resp.Usage) + } +} + func TestNewAcceptsBaseURLWithoutProjectLocation(t *testing.T) { provider := newProvider(providers.ProviderConfig{ Type: "vertex", From f5f855aed286a9fed3d5ba6eed0700387291003d Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 21:23:43 +0200 Subject: [PATCH 5/9] fix(vertex): correct model discovery and stream metadata --- internal/providers/gemini/gemini.go | 73 ++++++++++++++--- internal/providers/gemini/gemini_test.go | 93 +++++++++++++++++++++- internal/providers/gemini/native.go | 7 +- internal/providers/gemini/native_stream.go | 15 ++-- 4 files changed, 167 insertions(+), 21 deletions(-) diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 306d6d95..c6bf46cc 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -47,6 +47,7 @@ const ( type Provider struct { client *llmclient.Client nativeClient *llmclient.Client + modelsClient *llmclient.Client apiKey string backend string authType string @@ -71,7 +72,8 @@ func NewVertexWithHTTPClient(providerCfg providers.ProviderConfig, opts provider func newProvider(providerCfg providers.ProviderConfig, opts providers.ProviderOptions, httpClient *http.Client, preauthenticated bool) *Provider { backend := normalizeGeminiBackend(providerCfg) authType := normalizeGeminiAuthType(backend, providerCfg) - baseURL, modelsURL := geminiBaseURLs(providerCfg, backend) + baseURL, nativeBaseURL := geminiBaseURLs(providerCfg, backend) + modelsURL := geminiModelsBaseURL(backend, nativeBaseURL) p := &Provider{ apiKey: providerCfg.APIKey, backend: backend, @@ -95,15 +97,19 @@ func newProvider(providerCfg providers.ProviderConfig, opts providers.ProviderOp Hooks: opts.Hooks, CircuitBreaker: opts.Resilience.CircuitBreaker, } + nativeCfg := clientCfg + nativeCfg.BaseURL = nativeBaseURL modelsCfg := clientCfg modelsCfg.BaseURL = modelsURL if httpClient != nil { p.client = llmclient.NewWithHTTPClient(httpClient, clientCfg, p.setHeaders) - p.nativeClient = llmclient.NewWithHTTPClient(httpClient, modelsCfg, p.setNativeHeaders) + p.nativeClient = llmclient.NewWithHTTPClient(httpClient, nativeCfg, p.setNativeHeaders) + p.modelsClient = llmclient.NewWithHTTPClient(httpClient, modelsCfg, p.setNativeHeaders) return p } p.client = llmclient.New(clientCfg, p.setHeaders) - p.nativeClient = llmclient.New(modelsCfg, p.setNativeHeaders) + p.nativeClient = llmclient.New(nativeCfg, p.setNativeHeaders) + p.modelsClient = llmclient.New(modelsCfg, p.setNativeHeaders) return p } @@ -114,7 +120,8 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H httpClient = http.DefaultClient } providerCfg := providers.ProviderConfig{APIKey: apiKey} - baseURL, modelsURL := geminiBaseURLs(providerCfg, geminiBackendAIStudio) + baseURL, nativeBaseURL := geminiBaseURLs(providerCfg, geminiBackendAIStudio) + modelsURL := geminiModelsBaseURL(geminiBackendAIStudio, nativeBaseURL) p := &Provider{ apiKey: apiKey, backend: geminiBackendAIStudio, @@ -126,18 +133,25 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H modelsCfg.Hooks = hooks cfg := llmclient.DefaultConfig("gemini", baseURL) cfg.Hooks = hooks + nativeCfg := llmclient.DefaultConfig("gemini", nativeBaseURL) + nativeCfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) - p.nativeClient = llmclient.NewWithHTTPClient(httpClient, modelsCfg, p.setNativeHeaders) + p.nativeClient = llmclient.NewWithHTTPClient(httpClient, nativeCfg, p.setNativeHeaders) + p.modelsClient = llmclient.NewWithHTTPClient(httpClient, modelsCfg, p.setNativeHeaders) return p } // SetBaseURL allows configuring a custom base URL for the provider func (p *Provider) SetBaseURL(url string) { - baseURL, modelsURL := geminiBaseURLs(providers.ProviderConfig{BaseURL: url}, p.backend) + baseURL, nativeBaseURL := geminiBaseURLs(providers.ProviderConfig{BaseURL: url}, p.backend) + modelsURL := geminiModelsBaseURL(p.backend, nativeBaseURL) p.client.SetBaseURL(baseURL) p.modelsURL = modelsURL if p.nativeClient != nil { - p.nativeClient.SetBaseURL(modelsURL) + p.nativeClient.SetBaseURL(nativeBaseURL) + } + if p.modelsClient != nil { + p.modelsClient.SetBaseURL(modelsURL) } } @@ -148,6 +162,9 @@ func (p *Provider) SetModelsURL(url string) { if p.nativeClient != nil { p.nativeClient.SetBaseURL(url) } + if p.modelsClient != nil { + p.modelsClient.SetBaseURL(url) + } } func (p *Provider) validateConfig(providerCfg providers.ProviderConfig) { @@ -197,6 +214,13 @@ func (p *Provider) ready() error { return core.NewProviderError("gemini", http.StatusBadGateway, "invalid Gemini provider configuration: "+p.configErr.Error(), p.configErr) } +func (p *Provider) responseProviderName() string { + if p.backend == geminiBackendVertex { + return "vertex" + } + return "gemini" +} + // setHeaders sets the required headers for Gemini API requests func (p *Provider) setHeaders(req *http.Request) { if p.authType == geminiAuthTypeAPIKey { @@ -310,6 +334,17 @@ func geminiBaseURLs(providerCfg providers.ProviderConfig, backend string) (openA return baseURL, baseURL } +func geminiModelsBaseURL(backend, nativeBaseURL string) string { + nativeBaseURL = strings.TrimRight(strings.TrimSpace(nativeBaseURL), "/") + if backend != geminiBackendVertex { + return nativeBaseURL + } + if modelsURL, ok := vertexPublisherModelsBaseURL(nativeBaseURL); ok { + return modelsURL + } + return nativeBaseURL +} + func vertexBaseURLs(providerCfg providers.ProviderConfig) (openAICompatibleBaseURL, nativeBaseURL string) { baseURL := strings.TrimRight(strings.TrimSpace(providerCfg.BaseURL), "/") if baseURL == "" { @@ -351,6 +386,20 @@ func vertexOpenAICompatibleBaseURLFromNativeBaseURL(baseURL string) (string, boo return root + "/endpoints/openapi", true } +func vertexPublisherModelsBaseURL(nativeBaseURL string) (string, bool) { + const projectsPath = "/v1/projects/" + nativeBaseURL = strings.TrimRight(strings.TrimSpace(nativeBaseURL), "/") + idx := strings.Index(nativeBaseURL, projectsPath) + if idx < 0 { + return "", false + } + root := strings.TrimRight(nativeBaseURL[:idx], "/") + if root == "" { + return "", false + } + return root + "/v1beta1/publishers/google", true +} + func nativeBaseURLFromOpenAICompatibleBaseURL(baseURL string) (string, bool) { const suffix = "/openai" if !strings.HasSuffix(baseURL, suffix) { @@ -466,7 +515,7 @@ func (p *Provider) nativeChatCompletion(ctx context.Context, req *core.ChatReque if err != nil { return nil, err } - return nativeChatResponse(req, &geminiResp) + return nativeChatResponse(req, &geminiResp, p.responseProviderName()) } // StreamChatCompletion returns a raw response body for streaming (caller must close) @@ -513,7 +562,7 @@ func (p *Provider) nativeStreamChatCompletion(ctx context.Context, req *core.Cha return nil, err } includeUsage := streamReq.StreamOptions != nil && streamReq.StreamOptions.IncludeUsage - return newGeminiNativeStream(stream, req.Model, includeUsage), nil + return newGeminiNativeStream(stream, req.Model, includeUsage, p.responseProviderName()), nil } // geminiModel represents a model in Gemini's native API response @@ -549,7 +598,11 @@ func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) if err := p.ready(); err != nil { return nil, err } - rawResp, err := p.nativeClient.DoRaw(ctx, llmclient.Request{ + modelsClient := p.modelsClient + if modelsClient == nil { + modelsClient = p.nativeClient + } + rawResp, err := modelsClient.DoRaw(ctx, llmclient.Request{ Method: http.MethodGet, Endpoint: "/models", }) diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index c6f57c40..b29daa43 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -179,6 +179,38 @@ func TestVertexBaseURLs(t *testing.T) { } } +func TestVertexModelsBaseURL(t *testing.T) { + tests := []struct { + name string + nativeBase string + want string + }{ + { + name: "official vertex native base", + nativeBase: "https://aiplatform.googleapis.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + want: "https://aiplatform.googleapis.com/v1beta1/publishers/google", + }, + { + name: "custom proxy vertex native base", + nativeBase: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + want: "https://proxy.example.com/v1beta1/publishers/google", + }, + { + name: "unknown custom base keeps native base", + nativeBase: "https://proxy.example.com/custom/gemini", + want: "https://proxy.example.com/custom/gemini", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := geminiModelsBaseURL(geminiBackendVertex, tt.nativeBase); got != tt.want { + t.Fatalf("geminiModelsBaseURL() = %q, want %q", got, tt.want) + } + }) + } +} + func TestNewVertexWithHTTPClientAcceptsBaseURLWithoutProjectLocation(t *testing.T) { p := NewVertexWithHTTPClient(providers.ProviderConfig{ BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", @@ -383,6 +415,55 @@ func TestVertexNativeChatUsesOAuthAuthorization(t *testing.T) { if resp == nil || resp.ID != "vertex-native" { t.Fatalf("response = %+v, want vertex-native", resp) } + if resp.Provider != "vertex" { + t.Fatalf("provider = %q, want vertex", resp.Provider) + } +} + +func TestVertexNativeStreamUsesVertexProviderName(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash:streamGenerateContent" { + t.Errorf("Path = %q, want Vertex native streamGenerateContent endpoint", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer vertex-token" { + t.Errorf("Authorization = %q, want Bearer vertex-token", got) + } + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`data: {"responseId":"vertex-stream","candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}} + +`)) + })) + defer server.Close() + + p := newVertexTestProvider(server, true) + body, err := p.StreamChatCompletion(context.Background(), &core.ChatRequest{ + Model: "google/gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + StreamOptions: &core.StreamOptions{IncludeUsage: true}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer func() { _ = body.Close() }() + + raw, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read stream: %v", err) + } + chunks := parseOpenAIStreamChunks(t, string(raw)) + if len(chunks) == 0 { + t.Fatalf("stream = %q, want chunks", string(raw)) + } + for _, chunk := range chunks { + if got := chunk["provider"]; got != "vertex" { + t.Fatalf("provider = %#v, want vertex in stream %q", got, string(raw)) + } + } } func TestVertexOpenAICompatibleChatUsesOAuthAndGoogleModelPrefix(t *testing.T) { @@ -437,7 +518,7 @@ func TestVertexOpenAICompatibleChatUsesOAuthAndGoogleModelPrefix(t *testing.T) { func TestVertexListModelsAcceptsPublisherModelsResponse(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models" { + if r.URL.Path != "/v1beta1/publishers/google/models" { t.Errorf("Path = %q, want Vertex publisher models endpoint", r.URL.Path) } if got := r.Header.Get("Authorization"); got != "Bearer vertex-token" { @@ -446,9 +527,9 @@ func TestVertexListModelsAcceptsPublisherModelsResponse(t *testing.T) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{ "publisherModels": [ - {"name": "projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash"}, - {"name": "projects/prod-ai/locations/us-central1/publishers/google/models/text-embedding-005"}, - {"name": "projects/prod-ai/locations/us-central1/publishers/google/models/imagen-4.0"} + {"name": "publishers/google/models/gemini-2.5-flash"}, + {"name": "publishers/google/models/text-embedding-005"}, + {"name": "publishers/google/models/imagen-4.0"} ] }`)) })) @@ -479,10 +560,14 @@ func newVertexTestProvider(server *httptest.Server, native bool) *Provider { } openAIBaseURL := server.URL + "/v1/projects/prod-ai/locations/us-central1/endpoints/openapi" nativeBaseURL := server.URL + "/v1/projects/prod-ai/locations/us-central1/publishers/google" + modelsBaseURL := server.URL + "/v1beta1/publishers/google" openAICfg := llmclient.DefaultConfig("gemini", openAIBaseURL) nativeCfg := llmclient.DefaultConfig("gemini", nativeBaseURL) + modelsCfg := llmclient.DefaultConfig("gemini", modelsBaseURL) p.client = llmclient.NewWithHTTPClient(tokenClient, openAICfg, p.setHeaders) p.nativeClient = llmclient.NewWithHTTPClient(tokenClient, nativeCfg, p.setNativeHeaders) + p.modelsClient = llmclient.NewWithHTTPClient(tokenClient, modelsCfg, p.setNativeHeaders) + p.modelsURL = modelsBaseURL return p } diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go index 22728f1b..e38e8e9c 100644 --- a/internal/providers/gemini/native.go +++ b/internal/providers/gemini/native.go @@ -583,10 +583,13 @@ func geminiCachedContent(req *core.ChatRequest) string { return cached } -func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContentResponse) (*core.ChatResponse, error) { +func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContentResponse, providerName string) (*core.ChatResponse, error) { if err := geminiBlockedPromptError(geminiResp); err != nil { return nil, err } + if providerName == "" { + providerName = "gemini" + } created := time.Now().Unix() respID := geminiResp.ResponseID @@ -598,7 +601,7 @@ func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContent Object: "chat.completion", Created: created, Model: req.Model, - Provider: "gemini", + Provider: providerName, Choices: make([]core.Choice, 0, len(geminiResp.Candidates)), Usage: usageFromGemini(geminiResp.UsageMetadata), } diff --git a/internal/providers/gemini/native_stream.go b/internal/providers/gemini/native_stream.go index f0d88ca7..6ead8740 100644 --- a/internal/providers/gemini/native_stream.go +++ b/internal/providers/gemini/native_stream.go @@ -16,10 +16,13 @@ type geminiNativeStream struct { body io.ReadCloser } -func newGeminiNativeStream(body io.ReadCloser, model string, includeUsage bool) io.ReadCloser { +func newGeminiNativeStream(body io.ReadCloser, model string, includeUsage bool, providerName string) io.ReadCloser { + if providerName == "" { + providerName = "gemini" + } pr, pw := io.Pipe() stream := &geminiNativeStream{reader: pr, body: body} - go convertGeminiNativeStream(body, pw, model, includeUsage) + go convertGeminiNativeStream(body, pw, model, includeUsage, providerName) return stream } @@ -32,7 +35,7 @@ func (s *geminiNativeStream) Close() error { return s.body.Close() } -func convertGeminiNativeStream(body io.ReadCloser, out *io.PipeWriter, model string, includeUsage bool) { +func convertGeminiNativeStream(body io.ReadCloser, out *io.PipeWriter, model string, includeUsage bool, providerName string) { defer func() { _ = body.Close() }() scanner := bufio.NewScanner(body) @@ -40,6 +43,7 @@ func convertGeminiNativeStream(body io.ReadCloser, out *io.PipeWriter, model str state := geminiStreamState{ model: model, + providerName: providerName, includeUsage: includeUsage, created: time.Now().Unix(), } @@ -81,6 +85,7 @@ func convertGeminiNativeStream(body io.ReadCloser, out *io.PipeWriter, model str type geminiStreamState struct { model string + providerName string includeUsage bool created int64 responseID string @@ -126,7 +131,7 @@ func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { "object": "chat.completion.chunk", "created": s.created, "model": s.model, - "provider": "gemini", + "provider": s.providerName, "choices": []map[string]any{choice}, } if err := writeOpenAIStreamChunk(out, chunk); err != nil { @@ -149,7 +154,7 @@ func (s *geminiStreamState) writeFinalUsage(out io.Writer) error { "object": "chat.completion.chunk", "created": s.created, "model": s.model, - "provider": "gemini", + "provider": s.providerName, "choices": []map[string]any{}, "usage": s.latestUsage, } From 2315eb59830ac4ec47a1e03fb8595ff6ea299a8c Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 21:55:04 +0200 Subject: [PATCH 6/9] fix(gemini): preserve vertex provider metadata --- internal/providers/gemini/gemini.go | 2 +- internal/providers/gemini/gemini_test.go | 75 ++++++++++++++++++++++ internal/providers/gemini/native.go | 17 +++-- internal/providers/gemini/native_stream.go | 4 +- 4 files changed, 88 insertions(+), 10 deletions(-) diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index c6bf46cc..2b248fc0 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -727,7 +727,7 @@ func (p *Provider) StreamResponses(ctx context.Context, req *core.ResponsesReque if err := p.ready(); err != nil { return nil, err } - return providers.StreamResponsesViaChat(ctx, p, req, "gemini") + return providers.StreamResponsesViaChat(ctx, p, req, p.responseProviderName()) } // CreateBatch creates a native Gemini batch job through its OpenAI-compatible endpoint. diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index b29daa43..2fd69268 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -420,6 +420,46 @@ func TestVertexNativeChatUsesOAuthAuthorization(t *testing.T) { } } +func TestVertexNativeBlockedPromptUsesVertexProviderName(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash:generateContent" { + t.Errorf("Path = %q, want Vertex native generateContent endpoint", r.URL.Path) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "vertex-blocked", + "promptFeedback": { + "blockReason": "SAFETY", + "blockReasonMessage": "unsafe prompt" + } + }`)) + })) + defer server.Close() + + p := newVertexTestProvider(server, true) + _, err := p.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "google/gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "blocked"}, + }, + }) + if err == nil { + t.Fatal("expected blocked prompt error, got nil") + } + var gatewayErr *core.GatewayError + if !errors.As(err, &gatewayErr) { + t.Fatalf("error = %T %[1]v, want *core.GatewayError", err) + } + if gatewayErr.Provider != "vertex" { + t.Fatalf("provider = %q, want vertex", gatewayErr.Provider) + } + if !strings.Contains(gatewayErr.Message, "SAFETY: unsafe prompt") { + t.Fatalf("message = %q, want block reason", gatewayErr.Message) + } +} + func TestVertexNativeStreamUsesVertexProviderName(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "true") @@ -466,6 +506,41 @@ func TestVertexNativeStreamUsesVertexProviderName(t *testing.T) { } } +func TestVertexNativeStreamResponsesUsesVertexProviderName(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "true") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash:streamGenerateContent" { + t.Errorf("Path = %q, want Vertex native streamGenerateContent endpoint", r.URL.Path) + } + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`data: {"responseId":"vertex-responses-stream","candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]} + +`)) + })) + defer server.Close() + + p := newVertexTestProvider(server, true) + body, err := p.StreamResponses(context.Background(), &core.ResponsesRequest{ + Model: "google/gemini-2.5-flash", + Input: "Hello", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer func() { _ = body.Close() }() + + raw, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read stream: %v", err) + } + stream := string(raw) + if !strings.Contains(stream, `"provider":"vertex"`) { + t.Fatalf("stream = %q, want vertex provider metadata", stream) + } +} + func TestVertexOpenAICompatibleChatUsesOAuthAndGoogleModelPrefix(t *testing.T) { t.Setenv(useNativeAPIEnvVar, "false") diff --git a/internal/providers/gemini/native.go b/internal/providers/gemini/native.go index e38e8e9c..dbac2faa 100644 --- a/internal/providers/gemini/native.go +++ b/internal/providers/gemini/native.go @@ -584,12 +584,12 @@ func geminiCachedContent(req *core.ChatRequest) string { } func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContentResponse, providerName string) (*core.ChatResponse, error) { - if err := geminiBlockedPromptError(geminiResp); err != nil { - return nil, err - } if providerName == "" { providerName = "gemini" } + if err := geminiBlockedPromptError(geminiResp, providerName); err != nil { + return nil, err + } created := time.Now().Unix() respID := geminiResp.ResponseID @@ -624,7 +624,7 @@ func nativeChatResponse(req *core.ChatRequest, geminiResp *geminiGenerateContent return resp, nil } -func geminiBlockedPromptError(resp *geminiGenerateContentResponse) *core.GatewayError { +func geminiBlockedPromptError(resp *geminiGenerateContentResponse, providerName string) *core.GatewayError { if resp == nil || len(resp.Candidates) > 0 { return nil } @@ -632,7 +632,7 @@ func geminiBlockedPromptError(resp *geminiGenerateContentResponse) *core.Gateway if reason == "" { return nil } - return nativeProviderError("Gemini blocked prompt: "+reason, nil) + return nativeProviderError(providerName, "Gemini blocked prompt: "+reason, nil) } func geminiPromptBlockReason(raw json.RawMessage) string { @@ -846,6 +846,9 @@ func isGeminiExposedModel(modelID string) bool { return strings.HasPrefix(modelID, "gemini-") || strings.HasPrefix(modelID, "text-embedding-") } -func nativeProviderError(message string, err error) *core.GatewayError { - return core.NewProviderError("gemini", http.StatusBadGateway, message, err) +func nativeProviderError(providerName, message string, err error) *core.GatewayError { + if providerName == "" { + providerName = "gemini" + } + return core.NewProviderError(providerName, http.StatusBadGateway, message, err) } diff --git a/internal/providers/gemini/native_stream.go b/internal/providers/gemini/native_stream.go index 6ead8740..04e6a80d 100644 --- a/internal/providers/gemini/native_stream.go +++ b/internal/providers/gemini/native_stream.go @@ -107,7 +107,7 @@ func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { var event geminiGenerateContentResponse if err := json.Unmarshal([]byte(raw), &event); err != nil { - return nativeProviderError("failed to parse native Gemini stream event", err) + return nativeProviderError(s.providerName, "failed to parse native Gemini stream event", err) } if s.responseID == "" { s.responseID = event.ResponseID @@ -116,7 +116,7 @@ func (s *geminiStreamState) consumeEvent(out io.Writer, raw string) error { } } - if err := geminiBlockedPromptError(&event); err != nil { + if err := geminiBlockedPromptError(&event, s.providerName); err != nil { s.stopped = true return writeOpenAIStreamError(out, err) } From 18acc1d6a4dd746669e46d368610f1f003697f3a Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 22:15:22 +0200 Subject: [PATCH 7/9] test(vertex): cover google auth forms --- internal/providers/vertex/vertex_test.go | 146 +++++++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/internal/providers/vertex/vertex_test.go b/internal/providers/vertex/vertex_test.go index d8741e93..90401044 100644 --- a/internal/providers/vertex/vertex_test.go +++ b/internal/providers/vertex/vertex_test.go @@ -2,12 +2,18 @@ package vertex import ( "context" + "crypto/rand" + "crypto/rsa" + "crypto/x509" "encoding/base64" "encoding/binary" "encoding/json" + "encoding/pem" "math" "net/http" "net/http/httptest" + "os" + "path/filepath" "strings" "testing" @@ -177,6 +183,98 @@ func TestNewRejectsUnsupportedAuthType(t *testing.T) { } } +func TestNewAuthFormsInjectBearerToken(t *testing.T) { + tests := []struct { + name string + authType string + token string + grantType string + configure func(t *testing.T, cfg *providers.ProviderConfig, tokenURL string) + }{ + { + name: "ADC", + authType: "gcp_adc", + token: "adc-token", + grantType: "refresh_token", + configure: func(t *testing.T, cfg *providers.ProviderConfig, tokenURL string) { + t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", vertexADCCredentialsFile(t, tokenURL)) + }, + }, + { + name: "service account JSON", + authType: "gcp_service_account", + token: "service-account-token", + grantType: "urn:ietf:params:oauth:grant-type:jwt-bearer", + configure: func(t *testing.T, cfg *providers.ProviderConfig, tokenURL string) { + cfg.ServiceAccountJSON = vertexServiceAccountCredentials(t, tokenURL) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var gotGrantType string + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Fatalf("ParseForm() error = %v", err) + } + gotGrantType = r.PostForm.Get("grant_type") + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": tt.token, + "token_type": "Bearer", + "expires_in": 3600, + }) + })) + defer tokenServer.Close() + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/publishers/google/models/gemini-2.5-flash:generateContent" { + t.Errorf("Path = %q, want Vertex native generateContent endpoint", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer "+tt.token { + t.Errorf("Authorization = %q, want Bearer %s", got, tt.token) + } + if got := r.Header.Get("x-goog-api-key"); got != "" { + t.Errorf("x-goog-api-key = %q, want empty for Vertex OAuth", got) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "responseId": "vertex-auth", + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP" + }] + }`)) + })) + defer upstream.Close() + + cfg := testConfig() + cfg.AuthType = tt.authType + cfg.APIMode = "native" + cfg.BaseURL = upstream.URL + "/v1/projects/prod-ai/locations/us-central1/publishers/google" + tt.configure(t, &cfg, tokenServer.URL) + + provider := New(cfg, providers.ProviderOptions{}) + resp, err := provider.ChatCompletion(context.Background(), &core.ChatRequest{ + Model: "google/gemini-2.5-flash", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + }) + if err != nil { + t.Fatalf("ChatCompletion() error = %v", err) + } + if resp == nil || resp.Provider != "vertex" { + t.Fatalf("response = %+v, want vertex response", resp) + } + if gotGrantType != tt.grantType { + t.Fatalf("grant_type = %q, want %q", gotGrantType, tt.grantType) + } + }) + } +} + func TestVertexBaseURLs(t *testing.T) { tests := []struct { name string @@ -239,3 +337,51 @@ func authedTestClient(base *http.Client) *http.Client { TokenType: "Bearer", })) } + +func vertexADCCredentialsFile(t *testing.T, tokenURL string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "adc.json") + contents := map[string]string{ + "type": "authorized_user", + "client_id": "adc-client-id", + "client_secret": "adc-client-secret", + "refresh_token": "adc-refresh-token", + "token_uri": tokenURL, + } + encoded, err := json.Marshal(contents) + if err != nil { + t.Fatalf("failed to marshal ADC credentials: %v", err) + } + if err := os.WriteFile(path, encoded, 0o600); err != nil { + t.Fatalf("failed to write ADC credentials: %v", err) + } + return path +} + +func vertexServiceAccountCredentials(t *testing.T, tokenURL string) string { + t.Helper() + key, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("failed to generate test RSA key: %v", err) + } + keyBytes, err := x509.MarshalPKCS8PrivateKey(key) + if err != nil { + t.Fatalf("failed to marshal test RSA key: %v", err) + } + keyPEM := pem.EncodeToMemory(&pem.Block{ + Type: "PRIVATE KEY", + Bytes: keyBytes, + }) + contents := map[string]string{ + "type": "service_account", + "client_email": "service@example.com", + "private_key_id": "test-key-id", + "private_key": string(keyPEM), + "token_uri": tokenURL, + } + encoded, err := json.Marshal(contents) + if err != nil { + t.Fatalf("failed to marshal service account credentials: %v", err) + } + return string(encoded) +} From 49e143b82230082ac180bd7d2f869bb369811702 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 22:15:53 +0200 Subject: [PATCH 8/9] fix(gemini): use vertex labels for reused paths --- internal/providers/gemini/gemini.go | 6 +- internal/providers/gemini/gemini_test.go | 124 ++++++++++++++++++++++- 2 files changed, 124 insertions(+), 6 deletions(-) diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 2b248fc0..77770942 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -211,7 +211,7 @@ func (p *Provider) ready() error { if p.configErr == nil { return nil } - return core.NewProviderError("gemini", http.StatusBadGateway, "invalid Gemini provider configuration: "+p.configErr.Error(), p.configErr) + return core.NewProviderError(p.responseProviderName(), http.StatusBadGateway, "invalid Gemini provider configuration: "+p.configErr.Error(), p.configErr) } func (p *Provider) responseProviderName() string { @@ -622,7 +622,7 @@ func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) if err := json.Unmarshal(rawResp.Body, &nativeProbe); err == nil && (nativeProbe.Models != nil || nativeProbe.PublisherModels != nil) { var geminiResp geminiModelsResponse if err := json.Unmarshal(rawResp.Body, &geminiResp); err != nil { - return nil, core.NewProviderError("gemini", http.StatusBadGateway, "failed to parse native Gemini models response", err) + return nil, core.NewProviderError(p.responseProviderName(), http.StatusBadGateway, "failed to parse native Gemini models response", err) } modelEntries := append(geminiResp.Models, geminiResp.PublisherModels...) if len(modelEntries) == 0 { @@ -684,7 +684,7 @@ func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) if len(responsePreview) > 512 { responsePreview = responsePreview[:512] + "...(truncated)" } - return nil, core.NewProviderError("gemini", http.StatusBadGateway, "unexpected Gemini models response format", fmt.Errorf("models response body: %s", responsePreview)) + return nil, core.NewProviderError(p.responseProviderName(), http.StatusBadGateway, "unexpected Gemini models response format", fmt.Errorf("models response body: %s", responsePreview)) } // Responses sends a Responses API request to Gemini (converted to chat format) diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 2fd69268..073c69b0 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -81,6 +81,25 @@ func TestNew_VertexAcceptsGCPAuthAliases(t *testing.T) { } } +func TestNew_VertexConfigErrorUsesVertexProviderName(t *testing.T) { + p := NewVertexWithHTTPClient(providers.ProviderConfig{ + AuthType: "api_key", + BaseURL: "https://proxy.example.com/v1/projects/prod-ai/locations/us-central1/publishers/google", + }, providers.ProviderOptions{}, http.DefaultClient) + + err := p.ready() + if err == nil { + t.Fatal("expected config error, got nil") + } + var gatewayErr *core.GatewayError + if !errors.As(err, &gatewayErr) { + t.Fatalf("error = %T %[1]v, want *core.GatewayError", err) + } + if gatewayErr.Provider != "vertex" { + t.Fatalf("provider = %q, want vertex", gatewayErr.Provider) + } +} + func TestGeminiBaseURLs(t *testing.T) { tests := []struct { name string @@ -591,6 +610,66 @@ func TestVertexOpenAICompatibleChatUsesOAuthAndGoogleModelPrefix(t *testing.T) { } } +func TestVertexOpenAICompatibleEmbeddingsUsesOAuthAndGoogleModelPrefix(t *testing.T) { + t.Setenv(useNativeAPIEnvVar, "false") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/projects/prod-ai/locations/us-central1/endpoints/openapi/embeddings" { + t.Errorf("Path = %q, want Vertex OpenAI-compatible embeddings endpoint", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer vertex-token" { + t.Errorf("Authorization = %q, want Bearer vertex-token", got) + } + if got := r.Header.Get("x-goog-api-key"); got != "" { + t.Errorf("x-goog-api-key = %q, want empty for Vertex OAuth", got) + } + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode request: %v", err) + } + if got := payload["model"]; got != "google/text-embedding-005" { + t.Fatalf("model = %#v, want google/text-embedding-005", got) + } + if got := payload["input"]; got != "text" { + t.Fatalf("input = %#v, want text", got) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "object": "list", + "model": "google/text-embedding-005", + "data": [{ + "object": "embedding", + "embedding": [0.1, 0.2], + "index": 0 + }], + "usage": {"prompt_tokens": 1, "total_tokens": 1} + }`)) + })) + defer server.Close() + + p := newVertexTestProvider(server, false) + resp, err := p.Embeddings(context.Background(), &core.EmbeddingRequest{ + Model: "text-embedding-005", + Input: "text", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp == nil || resp.Model != "google/text-embedding-005" { + t.Fatalf("response = %+v, want google/text-embedding-005", resp) + } + if len(resp.Data) != 1 { + t.Fatalf("data = %+v, want one embedding", resp.Data) + } + var embedding []float64 + if err := json.Unmarshal(resp.Data[0].Embedding, &embedding); err != nil { + t.Fatalf("embedding = %s, want float array: %v", string(resp.Data[0].Embedding), err) + } + if len(embedding) != 2 || embedding[0] != 0.1 || embedding[1] != 0.2 { + t.Fatalf("embedding = %v, want [0.1 0.2]", embedding) + } +} + func TestVertexListModelsAcceptsPublisherModelsResponse(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1beta1/publishers/google/models" { @@ -623,6 +702,45 @@ func TestVertexListModelsAcceptsPublisherModelsResponse(t *testing.T) { } } +func TestVertexListModelsErrorsUseVertexProviderName(t *testing.T) { + tests := []struct { + name string + body string + }{ + { + name: "native parse error", + body: `{"publisherModels":"bad"}`, + }, + { + name: "unexpected format", + body: `{"unexpected":true}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(tt.body)) + })) + defer server.Close() + + p := newVertexTestProvider(server, true) + _, err := p.ListModels(context.Background()) + if err == nil { + t.Fatal("expected models error, got nil") + } + var gatewayErr *core.GatewayError + if !errors.As(err, &gatewayErr) { + t.Fatalf("error = %T %[1]v, want *core.GatewayError", err) + } + if gatewayErr.Provider != "vertex" { + t.Fatalf("provider = %q, want vertex", gatewayErr.Provider) + } + }) + } +} + func newVertexTestProvider(server *httptest.Server, native bool) *Provider { tokenClient := googleauth.HTTPClient(server.Client(), oauth2.StaticTokenSource(&oauth2.Token{ AccessToken: "vertex-token", @@ -636,9 +754,9 @@ func newVertexTestProvider(server *httptest.Server, native bool) *Provider { openAIBaseURL := server.URL + "/v1/projects/prod-ai/locations/us-central1/endpoints/openapi" nativeBaseURL := server.URL + "/v1/projects/prod-ai/locations/us-central1/publishers/google" modelsBaseURL := server.URL + "/v1beta1/publishers/google" - openAICfg := llmclient.DefaultConfig("gemini", openAIBaseURL) - nativeCfg := llmclient.DefaultConfig("gemini", nativeBaseURL) - modelsCfg := llmclient.DefaultConfig("gemini", modelsBaseURL) + openAICfg := llmclient.DefaultConfig("vertex", openAIBaseURL) + nativeCfg := llmclient.DefaultConfig("vertex", nativeBaseURL) + modelsCfg := llmclient.DefaultConfig("vertex", modelsBaseURL) p.client = llmclient.NewWithHTTPClient(tokenClient, openAICfg, p.setHeaders) p.nativeClient = llmclient.NewWithHTTPClient(tokenClient, nativeCfg, p.setNativeHeaders) p.modelsClient = llmclient.NewWithHTTPClient(tokenClient, modelsCfg, p.setNativeHeaders) From 236859a71ecfb45e8dea7cbef6f6774587dd24ee Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Tue, 5 May 2026 23:22:09 +0200 Subject: [PATCH 9/9] fix(gemini): preserve provider labels in compatible responses --- internal/providers/gemini/gemini.go | 14 +++++++--- internal/providers/gemini/gemini_test.go | 9 ++++++ internal/providers/vertex/vertex_test.go | 28 +++++++++++++++++++ .../golden/gemini/chat_completion.golden.json | 2 +- .../gemini/chat_with_params.golden.json | 2 +- .../golden/gemini/chat_with_tools.golden.json | 2 +- .../golden/gemini/responses.golden.json | 2 +- 7 files changed, 51 insertions(+), 8 deletions(-) diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 77770942..5584dbd6 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -498,6 +498,9 @@ func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (* if resp.Model == "" { resp.Model = req.Model } + if resp.Provider == "" { + resp.Provider = p.responseProviderName() + } return &resp, nil } @@ -719,6 +722,9 @@ func (p *Provider) Embeddings(ctx context.Context, req *core.EmbeddingRequest) ( if resp.Model == "" { resp.Model = req.Model } + if resp.Provider == "" { + resp.Provider = p.responseProviderName() + } return &resp, nil } @@ -826,7 +832,7 @@ func (p *Provider) GetBatchResults(ctx context.Context, id string) (*core.BatchR if err := p.ready(); err != nil { return nil, err } - return providers.FetchBatchResultsFromOutputFile(ctx, p.client, "gemini", id) + return providers.FetchBatchResultsFromOutputFile(ctx, p.client, p.responseProviderName(), id) } // CreateFile uploads a file through Gemini's OpenAI-compatible /files API. @@ -838,7 +844,7 @@ func (p *Provider) CreateFile(ctx context.Context, req *core.FileCreateRequest) if err != nil { return nil, err } - resp.Provider = "gemini" + resp.Provider = p.responseProviderName() return resp, nil } @@ -852,7 +858,7 @@ func (p *Provider) ListFiles(ctx context.Context, purpose string, limit int, aft return nil, err } for i := range resp.Data { - resp.Data[i].Provider = "gemini" + resp.Data[i].Provider = p.responseProviderName() } return resp, nil } @@ -866,7 +872,7 @@ func (p *Provider) GetFile(ctx context.Context, id string) (*core.FileObject, er if err != nil { return nil, err } - resp.Provider = "gemini" + resp.Provider = p.responseProviderName() return resp, nil } diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 073c69b0..478b08f4 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -608,6 +608,9 @@ func TestVertexOpenAICompatibleChatUsesOAuthAndGoogleModelPrefix(t *testing.T) { if resp == nil || resp.ID != "vertex-openai" { t.Fatalf("response = %+v, want vertex-openai", resp) } + if resp.Provider != "vertex" { + t.Fatalf("provider = %q, want vertex", resp.Provider) + } } func TestVertexOpenAICompatibleEmbeddingsUsesOAuthAndGoogleModelPrefix(t *testing.T) { @@ -658,6 +661,9 @@ func TestVertexOpenAICompatibleEmbeddingsUsesOAuthAndGoogleModelPrefix(t *testin if resp == nil || resp.Model != "google/text-embedding-005" { t.Fatalf("response = %+v, want google/text-embedding-005", resp) } + if resp.Provider != "vertex" { + t.Fatalf("provider = %q, want vertex", resp.Provider) + } if len(resp.Data) != 1 { t.Fatalf("data = %+v, want one embedding", resp.Data) } @@ -804,6 +810,9 @@ func TestChatCompletion(t *testing.T) { if resp.Model != "gemini-2.0-flash" { t.Errorf("Model = %q, want %q", resp.Model, "gemini-2.0-flash") } + if resp.Provider != "gemini" { + t.Errorf("Provider = %q, want gemini", resp.Provider) + } if len(resp.Choices) != 1 { t.Fatalf("len(Choices) = %d, want 1", len(resp.Choices)) } diff --git a/internal/providers/vertex/vertex_test.go b/internal/providers/vertex/vertex_test.go index 90401044..1e048d6b 100644 --- a/internal/providers/vertex/vertex_test.go +++ b/internal/providers/vertex/vertex_test.go @@ -209,6 +209,25 @@ func TestNewAuthFormsInjectBearerToken(t *testing.T) { cfg.ServiceAccountJSON = vertexServiceAccountCredentials(t, tokenURL) }, }, + { + name: "service account file", + authType: "gcp_service_account", + token: "service-account-token", + grantType: "urn:ietf:params:oauth:grant-type:jwt-bearer", + configure: func(t *testing.T, cfg *providers.ProviderConfig, tokenURL string) { + cfg.ServiceAccountFile = vertexServiceAccountCredentialsFile(t, tokenURL) + }, + }, + { + name: "service account JSON base64", + authType: "gcp_service_account", + token: "service-account-token", + grantType: "urn:ietf:params:oauth:grant-type:jwt-bearer", + configure: func(t *testing.T, cfg *providers.ProviderConfig, tokenURL string) { + credentials := vertexServiceAccountCredentials(t, tokenURL) + cfg.ServiceAccountJSONBase64 = base64.StdEncoding.EncodeToString([]byte(credentials)) + }, + }, } for _, tt := range tests { @@ -358,6 +377,15 @@ func vertexADCCredentialsFile(t *testing.T, tokenURL string) string { return path } +func vertexServiceAccountCredentialsFile(t *testing.T, tokenURL string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "service-account.json") + if err := os.WriteFile(path, []byte(vertexServiceAccountCredentials(t, tokenURL)), 0o600); err != nil { + t.Fatalf("failed to write service account credentials: %v", err) + } + return path +} + func vertexServiceAccountCredentials(t *testing.T, tokenURL string) string { t.Helper() key, err := rsa.GenerateKey(rand.Reader, 2048) diff --git a/tests/contract/testdata/golden/gemini/chat_completion.golden.json b/tests/contract/testdata/golden/gemini/chat_completion.golden.json index 8d86e97e..e6d266dc 100644 --- a/tests/contract/testdata/golden/gemini/chat_completion.golden.json +++ b/tests/contract/testdata/golden/gemini/chat_completion.golden.json @@ -13,7 +13,7 @@ "id": "QTpyabuaIfaevdIPzqPKkQU", "model": "gemini-2.5-flash-preview-09-2025", "object": "chat.completion", - "provider": "", + "provider": "gemini", "usage": { "completion_tokens": 2, "prompt_tokens": 4, diff --git a/tests/contract/testdata/golden/gemini/chat_with_params.golden.json b/tests/contract/testdata/golden/gemini/chat_with_params.golden.json index fedaf073..3a49f9c6 100644 --- a/tests/contract/testdata/golden/gemini/chat_with_params.golden.json +++ b/tests/contract/testdata/golden/gemini/chat_with_params.golden.json @@ -13,7 +13,7 @@ "id": "4TpyaZG4A5fJvdIPi6WJ-Ac", "model": "gemini-2.5-flash-preview-09-2025", "object": "chat.completion", - "provider": "", + "provider": "gemini", "usage": { "completion_tokens": 5, "prompt_tokens": 10, diff --git a/tests/contract/testdata/golden/gemini/chat_with_tools.golden.json b/tests/contract/testdata/golden/gemini/chat_with_tools.golden.json index 193962e7..20ae0e23 100644 --- a/tests/contract/testdata/golden/gemini/chat_with_tools.golden.json +++ b/tests/contract/testdata/golden/gemini/chat_with_tools.golden.json @@ -23,7 +23,7 @@ "id": "4zpyaaPLGYXlxN8Pq7jUkQs", "model": "gemini-2.5-flash-preview-09-2025", "object": "chat.completion", - "provider": "", + "provider": "gemini", "usage": { "completion_tokens": 19, "prompt_tokens": 44, diff --git a/tests/contract/testdata/golden/gemini/responses.golden.json b/tests/contract/testdata/golden/gemini/responses.golden.json index b94fc193..4c7dfd92 100644 --- a/tests/contract/testdata/golden/gemini/responses.golden.json +++ b/tests/contract/testdata/golden/gemini/responses.golden.json @@ -17,7 +17,7 @@ "type": "message" } ], - "provider": "", + "provider": "gemini", "status": "completed", "usage": { "input_tokens": 4,