Skip to content

Commit d7d60b8

Browse files
lukecameronclaude
authored andcommitted
fix(discovery): fall back to serving-endpoints for OSS models
OSS model discovery only queried UC model-services for `system.ai.*` ids and, unlike the claude/gemini/codex families, had no AI-Gateway fallback. On workspaces that expose OSS foundation models as regular `databricks-*` serving endpoints (rather than `system.ai.*` UC model-services), `oss_models` came back empty and pi/opencode never offered GLM, Qwen, Gemma, inkling, etc. Add `discover_oss_models`, which lists foundation-model endpoints advertising `mlflow/v1/chat/completions` and keeps the OSS chat families via the existing `_is_oss_chat_model` filter (the family filter is required because Claude/Gemini endpoints can advertise the same dialect on some workspaces). Wire it as the fallback in cli.py, mirroring the other three families. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent c5daf76 commit d7d60b8

3 files changed

Lines changed: 107 additions & 0 deletions

File tree

src/ucode/cli.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
discover_codex_models,
3838
discover_gemini_models,
3939
discover_model_services,
40+
discover_oss_models,
4041
ensure_ai_gateway_v2,
4142
ensure_databricks_auth,
4243
ensure_pat_bearer,
@@ -391,6 +392,8 @@ def configure_shared_state(
391392
codex_models, codex_reason = discover_codex_models(workspace, token)
392393
if want_oss:
393394
oss_models, oss_reason = ms_oss, ms_reason
395+
if not oss_models:
396+
oss_models, oss_reason = discover_oss_models(workspace, token)
394397
if claude_models:
395398
opencode_models["anthropic"] = list(claude_models.values())
396399
if gemini_models:

src/ucode/databricks.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2056,6 +2056,34 @@ def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str |
20562056
return discover_endpoints_with_api_type(workspace, token, "openai/v1/responses")
20572057

20582058

2059+
def discover_oss_models(workspace: str, token: str) -> tuple[list[str], str | None]:
2060+
"""Discover OSS chat models served as AI Gateway foundation-model endpoints.
2061+
2062+
Fallback for workspaces that don't register OSS foundation models as
2063+
`system.ai.*` UC model-services (see `discover_model_services`): those
2064+
workspaces expose the same models as regular `databricks-*` serving
2065+
endpoints instead. Lists every endpoint advertising the
2066+
`mlflow/v1/chat/completions` dialect, then keeps only the OSS chat families
2067+
(`_is_oss_chat_model`) — on some workspaces the Claude/Gemini endpoints also
2068+
advertise that dialect, so the family filter is what separates the OSS
2069+
cohort from them. Mirrors the AI-Gateway fallback the other families use
2070+
when the UC model-services listing is empty.
2071+
"""
2072+
endpoints, reason = discover_endpoints_with_api_type(
2073+
workspace, token, "mlflow/v1/chat/completions"
2074+
)
2075+
if not endpoints:
2076+
return [], reason
2077+
oss = [e for e in endpoints if _is_oss_chat_model(e)]
2078+
if oss:
2079+
return oss, None
2080+
sample = ", ".join(endpoints[:5])
2081+
return [], (
2082+
"foundation-models exposing `mlflow/v1/chat/completions` matched no OSS "
2083+
f"chat family (got: {sample})"
2084+
)
2085+
2086+
20592087
def fetch_gemini_models(workspace: str, token: str) -> list[str]:
20602088
models, _ = discover_gemini_models(workspace, token)
20612089
return models

tests/test_databricks.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -923,6 +923,82 @@ def test_codex_discovery_keeps_alphabetical_order(self, monkeypatch):
923923
assert models == ["databricks-gpt-4-1", "databricks-gpt-5-2-codex"]
924924

925925

926+
def _mlflow_chat_payload(names, *, api_type="mlflow/v1/chat/completions", v2=True):
927+
return {
928+
"endpoints": [
929+
{
930+
"name": name,
931+
"config": {
932+
"served_entities": [
933+
{
934+
"foundation_model": {
935+
"ai_gateway_v2_supported": v2,
936+
"api_types": [api_type],
937+
}
938+
}
939+
]
940+
},
941+
}
942+
for name in names
943+
]
944+
}
945+
946+
947+
class TestDiscoverOssModels:
948+
def test_finds_oss_endpoints_via_foundation_models(self, monkeypatch):
949+
# Mirrors a workspace with no system.ai UC model-services: OSS models are
950+
# plain databricks-* serving endpoints under the mlflow chat dialect.
951+
payload = _mlflow_chat_payload(
952+
[
953+
"databricks-glm-5-2",
954+
"databricks-inkling",
955+
"databricks-qwen35-122b-a10b",
956+
"databricks-gemma-3-12b",
957+
]
958+
)
959+
monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None))
960+
961+
models, reason = db_mod.discover_oss_models(WS, "token")
962+
963+
assert reason is None
964+
assert "databricks-glm-5-2" in models
965+
assert set(models) == {
966+
"databricks-glm-5-2",
967+
"databricks-inkling",
968+
"databricks-qwen35-122b-a10b",
969+
"databricks-gemma-3-12b",
970+
}
971+
972+
def test_excludes_claude_and_gemini_sharing_the_mlflow_dialect(self, monkeypatch):
973+
# On some workspaces every foundation model advertises the mlflow chat
974+
# dialect, so the api_type filter alone is too broad — the OSS family
975+
# filter must drop Claude/Gemini and keep only the OSS cohort.
976+
payload = _mlflow_chat_payload(
977+
[
978+
"databricks-claude-opus-4-8",
979+
"databricks-gemini-2-5-pro",
980+
"databricks-glm-5-2",
981+
"databricks-qwen3-embedding-0-6b",
982+
]
983+
)
984+
monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None))
985+
986+
models, reason = db_mod.discover_oss_models(WS, "token")
987+
988+
assert reason is None
989+
assert models == ["databricks-glm-5-2"]
990+
991+
def test_reports_reason_when_no_oss_family_matches(self, monkeypatch):
992+
payload = _mlflow_chat_payload(["databricks-claude-opus-4-8"])
993+
monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None))
994+
995+
models, reason = db_mod.discover_oss_models(WS, "token")
996+
997+
assert models == []
998+
assert reason is not None
999+
assert "no OSS" in reason
1000+
1001+
9261002
class TestResolvePatToken:
9271003
def test_reads_pat_profile_token_from_cfg(self, monkeypatch, tmp_path):
9281004
cfg = tmp_path / "databrickscfg"

0 commit comments

Comments
 (0)