Add UCODE_MODEL_PREFERRED_PREFIX to prefer custom gateway endpoints#164
Open
tplass-ias wants to merge 1 commit into
Open
Add UCODE_MODEL_PREFERRED_PREFIX to prefer custom gateway endpoints#164tplass-ias wants to merge 1 commit into
tplass-ias wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR updates Databricks model discovery to support a user-configured preferred model prefix and adjusts ordering so the “best”/newest model is selected first more consistently (including Codex discovery).
Changes:
- Add
UCODE_MODEL_PREFERRED_PREFIXsupport to promote matching model IDs/endpoints to the top of candidate lists. - Centralize ordering behavior via
_sort_candidates_by_prefix()and apply it across Claude/Codex/Gemini discovery paths. - Update Codex discovery test expectations to reflect newest-first ordering.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| tests/test_databricks.py | Updates Codex discovery test to expect newest-first ordering. |
| src/ucode/databricks.py | Adds preferred-prefix support and consolidates sorting logic across model discovery functions. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
tplass-ias
force-pushed
the
tplass/ucode-model-preferred-prefix
branch
from
3dba5cc to
d2164cd
Compare
Set `UCODE_MODEL_PREFERRED_PREFIX` (e.g. `ias-`) to promote any discovered model endpoint whose name contains that substring to the front of each family's candidate list during `ucode configure`. Preference is applied across all model families (Claude, Gemini, Codex) and both discovery paths (UC model-services and AI Gateway v2 fallback). - Add `_preferred_model_prefix()` to read the env var (empty string treated as unset). - Add `_sort_candidates_by_prefix(candidates, preferred_prefix, *, sort_key)` that groups preferred items first while preserving the existing within-group ordering (reverse-alpha for Claude/Codex, `model_version_sort_key` for Gemini). - Apply preference in `discover_model_services` for all three families. - Thread `preferred_prefix` through `discover_endpoints_with_api_type` so `discover_gemini_models` and `discover_codex_models` (AI Gateway fallback) also respect the setting. - Update the one test that asserted alphabetical Codex order, which was the unintentional pre-existing default; Codex now sorts newest-first consistently with the other families. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
tplass-ias
force-pushed
the
tplass/ucode-model-preferred-prefix
branch
from
d2164cd to
44abf96
Compare
Comment on lines
+1327
to
+1337
| # Accept both the standard databricks-claude-* prefix and any custom | ||
| # preferred prefix so that e.g. companyname-claude-haiku-4-5 is discoverable. | ||
| family_ids = [ | ||
| m | ||
| for m in raw_ids | ||
| if f"claude-{family}-" in m | ||
| and ( | ||
| f"databricks-claude-{family}-" in m | ||
| or (preferred_prefix and _model_bare_name(m).startswith(preferred_prefix)) | ||
| ) | ||
| ] |
Comment on lines
1396
to
+1398
| describes why the list is empty. `sort_key` overrides the default | ||
| alphabetical ordering of the returned names. | ||
| newest-first (reverse-alphabetical) ordering of the returned names. `preferred_prefix` promotes | ||
| matching endpoints to the front of the list (see ``_preferred_model_prefix``). |
| out.append(name) | ||
| if out: | ||
| return sorted(out, key=sort_key), None | ||
| return _sort_candidates_by_prefix(out, preferred_prefix, sort_key=sort_key), None |
Comment on lines
+89
to
+97
| preferred = [m for m in candidates if _model_bare_name(m).startswith(preferred_prefix)] | ||
| others = [m for m in candidates if not _model_bare_name(m).startswith(preferred_prefix)] | ||
| if sort_key: | ||
| return sorted(preferred, key=sort_key) + sorted(others, key=sort_key) | ||
| return sorted( | ||
| candidates, | ||
| key=lambda m: (_model_bare_name(m).startswith(preferred_prefix), m), | ||
| reverse=True, | ||
| ) |
Comment on lines
1396
to
+1398
| describes why the list is empty. `sort_key` overrides the default | ||
| alphabetical ordering of the returned names. | ||
| newest-first (reverse-alphabetical) ordering of the returned names. `preferred_prefix` promotes | ||
| matching endpoints to the front of the list (see ``_preferred_model_prefix``). |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Some deployments route through a custom AI Gateway prefix (e.g.
companyname-claude-haiku-4-5) rather than the standarddatabricks-claude-*endpoints. There was no way to tellucode configureto prefer those variants when both are available on the same workspace.This PR adds
UCODE_MODEL_PREFERRED_PREFIX— an env var that, when set, promotes any discovered model endpoint whose bare name starts with that prefix to the front of each family's candidate list. The normal newest-first ordering is preserved within the preferred and non-preferred groups._preferred_model_prefix()to read the env var (empty string treated as unset)_model_bare_name()to strip anysystem.ai.catalog prefix before matching, so that both UC model-service IDs (system.ai.companyname-claude-haiku-4-5) and AI Gateway IDs (companyname-claude-haiku-4-5) match correctly_sort_candidates_by_prefix(candidates, preferred_prefix, *, sort_key)— groups preferred items first using an anchored prefix check (startswith), while preserving existing within-group ordering (reverse-alpha for Claude/Codex,model_version_sort_keyfor Gemini)discover_endpoints_with_api_type)Usage
export UCODE_MODEL_PREFERRED_PREFIX=companyname- ucode configureWith
companyname-set,companyname-claude-haiku-4-5will be selected overdatabricks-claude-haiku-4-5when both are available, across Claude, Gemini, and Codex families.Test plan
uv run ruff check .— cleanuv run ruff format --check .— cleanuv run pytest— 759 passed, 30 skippedtest_preferred_prefix_promotes_matching_modeltoTestDiscoverClaudeModelstest_preferred_prefix_promotes_matching_claude_modeltoTestDiscoverModelServicestest_preferred_prefix_promotes_matching_gemini_modeltoTestDiscoverGeminiModelstest_preferred_prefix_promotes_matching_codex_modeltoTestDiscoverGeminiModelstest_codex_discovery_orders_newest_firstto reflect corrected Codex sort order