Skip to content

Commit d776f4c

Browse files
maskarbclaude
andauthored
feat: add Gemini model discovery to model-discovery script (#811)
## Summary - Refactor the model-discovery script from a hardcoded Anthropic-only model list to API-driven multi-publisher discovery supporting both Anthropic and Google Gemini models - Auto-discover new models via the Model Garden list API — no code changes needed when new models are released - Version info extracted from the list API response (eliminates per-model GET calls and the 403 errors they caused) - Model flags created enabled with 100% rollout by default (previously disabled at 0%) - Smart version limiting with Gemini semver parsing, temporal qualifier stripping, and configurable version cutoff ## Changes ### Multi-publisher probing (`model-discovery.py`) - `_build_probe_request()` builds publisher-specific probe requests: `rawPredict` for Anthropic, `generateContent` for Gemini; raises `ValueError` for unknown publishers - `probe_model()` guards against unknown publishers (returns `"unknown"` instead of crashing) - New model entries use the correct `provider` field instead of hardcoded `"anthropic"` ### API-driven auto-discovery - `list_publisher_models()` queries `GET /v1beta1/publishers/{publisher}/models` with pagination and retry/backoff (3 attempts with exponential backoff for transient errors), returns `(model_id, version_id)` tuples - `discover_models()` filters API results by prefix patterns, applies exclude regexes, deduplicates, and falls back to `SEED_MODELS` when the API is unavailable - `PublisherConfig` TypedDict for type-safe publisher configuration - Exclude patterns filter out: pinned versions (`-001`), experimental models, embedding/non-text models, base aliases without minor version (`claude-opus-4`) - `version_cutoff` config per publisher — Google models with version <= 2.0 are excluded - Seed models also respect `version_cutoff` for consistency ### Gemini semver parsing - `parse_model_family()` recognizes Gemini's mid-string semver (`gemini-2.5-flash` → family `gemini-flash`, version `(2, 5)`) - Temporal qualifiers (`preview`, `exp`, date stamps like `04-17`) are stripped before determining family, so preview variants group with their stable counterpart - Claude's trailing-digit convention still works (`claude-opus-4-6` → family `claude-opus`, version `(4, 6)`) ### Version limiting - `MAX_VERSIONS_PER_FAMILY = 2` — keeps only the N most recent versions per model family - Provider default models (`defaultModel` + `providerDefaults` values from `models.json`) are exempt from version limiting - Models without a parseable version are always kept ### Flags created enabled by default (`sync_flags.go`) - `FlagSpec.EnabledByDefault` — model flags are created with 100% rollout and enabled in the environment - `enableFlagInEnv()` — calls `POST .../environments/{env}/on` after flag creation; reads response body only on error - `addRolloutStrategy()` now accepts a rollout percentage parameter - Non-model flags (from `flags.json`) retain the existing behavior (disabled, 0% rollout) ### Removed: `resolve_version()` - The Model Garden list API already returns `versionId` for each model - Eliminates 403 errors on models not exposing version info via the individual GET endpoint - Reduces API calls from N+1 to 1 per publisher - `roles/serviceusage.serviceUsageConsumer` IAM role is no longer required ### Observability - Summary table logged for every publisher showing each model's final disposition: `KEEP`, `KEEP (default)`, `EXCLUDE`, `EXCLUDE (version)`, `SKIP (prefix)`, `SKIP (version limit)` ### Example output ``` Discovering models from Vertex AI Model Garden... claude-opus: keeping 2 latest, dropping ['claude-opus-4-1'] anthropic: 8 model(s) from API claude-haiku-4-5 KEEP claude-opus-4 EXCLUDE claude-opus-4-1 SKIP (version limit) claude-opus-4-5 KEEP claude-opus-4-6 KEEP claude-sonnet-4 EXCLUDE claude-sonnet-4-5 KEEP (default) claude-sonnet-4-6 KEEP google: 10 model(s) from API gemini-1.5-pro-002 EXCLUDE gemini-2.0-flash-001 EXCLUDE gemini-2.0-flash-lite-001 EXCLUDE gemini-2.0-flash-preview-image-generation EXCLUDE (version) gemini-2.5-flash KEEP (default) gemini-2.5-flash-lite KEEP gemini-2.5-flash-preview-04-17 KEEP gemini-2.5-pro KEEP gemini-2.5-pro-exp-03-25 EXCLUDE gemini-3.1-flash-image-preview KEEP Processing 10 model(s) in us-east5/ambient-code-platform... claude-haiku-4-5: unchanged (available=True) claude-opus-4-5: unchanged (available=True) claude-opus-4-6: unchanged (available=True) claude-sonnet-4-5: unchanged (available=True) claude-sonnet-4-6: unchanged (available=True) gemini-2.5-flash: vertexId updated -> gemini-2.5-flash@default gemini-2.5-flash: unchanged (available=True) gemini-2.5-flash-lite: NEW model added (available=True) gemini-2.5-flash-preview-04-17: NEW model added (available=False) gemini-2.5-pro: vertexId updated -> gemini-2.5-pro@default gemini-2.5-pro: unchanged (available=True) gemini-3.1-flash-image-preview: NEW model added (available=False) 5 change(s) written to models.json: gemini-2.5-flash: vertexId updated gemini-2.5-flash -> gemini-2.5-flash@default gemini-2.5-flash-lite: added (available=True) gemini-2.5-flash-preview-04-17: added (available=False) gemini-2.5-pro: vertexId updated gemini-2.5-pro -> gemini-2.5-pro@default gemini-3.1-flash-image-preview: added (available=False) ``` ### Example `models.json` diff Running the script against the current manifest produces the following changes — existing Anthropic models are unchanged, Gemini vertexIds get `@default` suffix, and three new Gemini models are auto-discovered: ```diff { "id": "gemini-2.5-flash", "label": "Gemini 2.5 Flash", - "vertexId": "gemini-2.5-flash", + "vertexId": "gemini-2.5-flash@default", "provider": "google", "available": true, "featureGated": false }, { "id": "gemini-2.5-pro", "label": "Gemini 2.5 Pro", - "vertexId": "gemini-2.5-pro", + "vertexId": "gemini-2.5-pro@default", "provider": "google", "available": true, "featureGated": true + }, + { + "id": "gemini-2.5-flash-lite", + "label": "Gemini 2.5 Flash Lite", + "vertexId": "gemini-2.5-flash-lite@default", + "provider": "google", + "available": true, + "featureGated": true + }, + { + "id": "gemini-2.5-flash-preview-04-17", + "label": "Gemini 2.5 Flash Preview 04.17", + "vertexId": "gemini-2.5-flash-preview-04-17@default", + "provider": "google", + "available": false, + "featureGated": true + }, + { + "id": "gemini-3.1-flash-image-preview", + "label": "Gemini 3.1 Flash Image Preview", + "vertexId": "gemini-3.1-flash-image-preview@default", + "provider": "google", + "available": false, + "featureGated": true } ``` New models are added with `featureGated: true` (requires explicit opt-in). Unavailable models (`available: false`) are still tracked so they appear automatically when Google enables them. ### GCP IAM requirements Only one role is needed: | Role | Permission | Used by | |------|-----------|---------| | `roles/aiplatform.user` | `aiplatform.endpoints.predict` | `probe_model` (rawPredict / generateContent) | The list API (`/v1beta1/publishers/{publisher}/models`) requires authentication but no specific IAM role beyond basic project access. ### Test coverage **Python** (`scripts/test_model_discovery.py`) — 19 unit tests: - `parse_model_family`: Claude trailing digits, Gemini semver, qualifier stripping, preview/exp/date removal, versionless models - `model_id_to_label`: both naming conventions - `keep_latest_versions`: version limiting, protected models, semver grouping, empty input **Go** (`cmd/sync_flags_test.go`) — 2 new tests: - `enableFlagInEnv` happy path (POST .../on returns 200) - `enableFlagInEnv` non-2xx error path (returns error with status code) ## Test plan - [x] Run script locally against `us-east5/ambient-code-platform` — verified filtering, version limiting, and probing for both publishers - [x] Verified v1 API returns 404 for list endpoint — seed list fallback works correctly - [x] Verified Gemini semver parsing groups families correctly (preview variants collapse with stable) - [x] Verified `version_cutoff: (2, 0)` excludes gemini-2.0 models - [x] Verified seed models respect `version_cutoff` - [x] 19 Python unit tests pass (`python3 scripts/test_model_discovery.py`) - [x] Go `go vet` and `go test ./cmd/` pass for sync_flags changes - [x] Verified `models.json` diff: 3 new Gemini models discovered, vertexIds updated with `@default` suffix - [ ] Verify GHA workflow runs successfully on `workflow_dispatch` - [ ] Verify new model flags are created enabled with 100% rollout in Unleash 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 52b5efb commit d776f4c

5 files changed

Lines changed: 918 additions & 325 deletions

File tree

0 commit comments

Comments
 (0)