From 366637f7601a2d08e9c6b4717aef9d0d933faa6d Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 21 Jul 2026 14:41:43 +0000 Subject: [PATCH] pi: support OSS models (kimi, glm) via managed databricks-oss provider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pi previously exposed only claude/openai/gemini. OSS families (kimi, glm) were left unmanaged, so an early build's `databricks-kimi` provider — wired with the wrong `openai-responses` dialect against the MLflow chat-completions route — lingered with a frozen, never-refreshed token and 401'd on every call. Add a managed `databricks-oss` provider (api: openai-completions → /ai-gateway/mlflow/v1) with per-model contextWindow/maxTokens from model_token_limits (so GLM's 25k output cap is honored) and supportsStore/supportsStrictMode compat flags the MLflow route requires. `ucode configure pi` now discovers and refreshes OSS models, and the stale `databricks-kimi` block is stripped on every write via LEGACY_PROVIDER_NAMES. Co-authored-by: Isaac --- src/ucode/agents/pi.py | 62 +++++++++++++++++++++++++++----- src/ucode/cli.py | 4 +-- src/ucode/databricks.py | 6 ++-- tests/test_agent_pi.py | 79 ++++++++++++++++++++++++++++++++++++----- 4 files changed, 130 insertions(+), 21 deletions(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e7c1760..eb082f1 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -7,6 +7,7 @@ - `databricks-claude` (api: anthropic-messages) → /ai-gateway/anthropic - `databricks-openai` (api: openai-responses) → /ai-gateway/codex/v1 - `databricks-gemini` (api: google-generative-ai) → /ai-gateway/gemini/v1beta +- `databricks-oss` (api: openai-completions) → /ai-gateway/mlflow/v1 Per-provider `compat` flags work around fields the gateway translators reject: @@ -15,11 +16,11 @@ pi uses for every request. With this flag pi omits the per-tool field and sends the legacy `anthropic-beta: fine-grained-tool-streaming-...` header instead, which the gateway accepts. - -OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via -pi today — they live behind /ai-gateway/mlflow/v1 with per-model -`max_tokens` caps that pi has no global way to honor without per-model -config we don't currently maintain. +- oss: `supportsStore: false` + `supportsStrictMode: false` — the MLflow + chat-completions route rejects the OpenAI `store` field and + `tools[].function.strict`. OSS models also carry per-model `contextWindow` + and `maxTokens` (from `model_token_limits`) so pi clamps output to a value + the gateway accepts (e.g. GLM's 25k cap) instead of pi's 16k default. The bearer token is baked into the file and refreshed by a background thread while the session runs (same pattern as OpenCode/Copilot). @@ -45,6 +46,7 @@ TOKEN_REFRESH_INTERVAL_SECONDS, build_pi_base_urls, get_databricks_token, + model_token_limits, ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version @@ -68,13 +70,17 @@ "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-oss", ) PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES] # Old provider names earlier ucode versions wrote; cleaned up on each write so # users don't end up with stale entries pointing at routes that 400. -LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-oss") +# `databricks-kimi` was written by an early build with `api: openai-responses` +# against the MLflow route (which speaks openai-completions), so its requests +# 404/401 and its baked token was never refreshed — strip it on every write. +LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-kimi") def is_update_available() -> tuple[str, str] | None: @@ -86,6 +92,7 @@ def _resolve_model_selector( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + oss_models: list[str], ) -> str: """Return a Pi model selector in `/` form when possible.""" for name in PROVIDER_NAMES: @@ -97,9 +104,26 @@ def _resolve_model_selector( return f"databricks-openai/{model}" if model in gemini_models: return f"databricks-gemini/{model}" + if model in oss_models: + return f"databricks-oss/{model}" return model +def _oss_model_entry(model: str) -> dict: + """Per-model entry for an OSS model. + + Pins `contextWindow`/`maxTokens` from the shared limits table when known so + pi clamps `max_tokens` to a value the MLflow route accepts (e.g. GLM's 25k + output cap) rather than pi's 16k default. Both fields are supplied together + because the limits table always provides both.""" + entry: dict = {"id": model} + limits = model_token_limits(model) + if limits is not None: + entry["contextWindow"] = limits["context"] + entry["maxTokens"] = limits["output"] + return entry + + def render_overlay( model: str, token: str, @@ -107,6 +131,7 @@ def render_overlay( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + oss_models: list[str], ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json.""" providers: dict = {} @@ -150,8 +175,23 @@ def render_overlay( "models": [{"id": m} for m in gemini_models], } keys.append(["providers", "databricks-gemini"]) + if oss_models: + providers["databricks-oss"] = { + "baseUrl": pi_base_urls["oss"], + "api": "openai-completions", + "apiKey": token, + "authHeader": True, + # The MLflow chat-completions route rejects the OpenAI `store` field + # and `tools[].function.strict`; opt out of both. + "compat": {"supportsStore": False, "supportsStrictMode": False}, + "headers": ua_headers, + "models": [_oss_model_entry(m) for m in oss_models], + } + keys.append(["providers", "databricks-oss"]) overlay: dict = { - "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models), + "model": _resolve_model_selector( + model, claude_models, codex_models, gemini_models, oss_models + ), } if providers: overlay["providers"] = providers @@ -178,6 +218,7 @@ def write_tool_config( state.get("claude_models") or {}, state.get("codex_models") or [], state.get("gemini_models") or [], + state.get("oss_models") or [], ) existing = read_json_safe(PI_CONFIG_PATH) providers = existing.get("providers") @@ -206,7 +247,7 @@ def _write_settings(model_selector: str) -> None: def default_model(state: dict) -> str | None: - """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini.""" + """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini, oss.""" claude_models = state.get("claude_models") or {} for family in ("opus", "sonnet", "haiku"): if claude_models.get(family): @@ -215,7 +256,10 @@ def default_model(state: dict) -> str | None: if codex_models: return codex_models[0] gemini_models = state.get("gemini_models") or [] - return gemini_models[0] if gemini_models else None + if gemini_models: + return gemini_models[0] + oss_models = state.get("oss_models") or [] + return oss_models[0] if oss_models else None def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a0ef7c4..899b25a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -90,7 +90,7 @@ "claude": ("claude", "opencode", "copilot", "pi"), "codex": ("codex", "copilot", "pi"), "gemini": ("gemini", "opencode", "pi"), - "oss": ("opencode",), + "oss": ("opencode", "pi"), } @@ -337,7 +337,7 @@ def configure_shared_state( ) want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools - want_oss = fetch_all or "opencode" in tools + want_oss = fetch_all or "opencode" in tools or "pi" in tools claude_reason: str | None = None gemini_reason: str | None = None diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b38bca9..327d4b3 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2135,12 +2135,14 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]: # - openai-completions appends `/chat/completions` # # So the baseUrls below stop just before the suffix Pi will tack on. - # Compat flags applied per-provider in agents/pi.py; required for `oss` - # only (MLflow rejects `store` and `tools[].function.strict`). + # OSS families (kimi, glm) speak openai-completions to the MLflow route; + # compat flags applied per-provider in agents/pi.py (MLflow rejects `store` + # and `tools[].function.strict`). return { "claude": build_tool_base_url("claude", workspace), "openai": build_tool_base_url("codex", workspace), "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", + "oss": f"{workspace}/ai-gateway/mlflow/v1", } diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 0afc5fb..d68316e 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -20,12 +20,17 @@ def _base_urls() -> dict[str, str]: } +def _base_urls_with_oss() -> dict[str, str]: + return {**_base_urls(), "oss": f"{WS}/ai-gateway/mlflow/v1"} + + def _empty() -> dict: """No-models input bundle for render_overlay.""" return { "claude_models": {}, "codex_models": [], "gemini_models": [], + "oss_models": [], } @@ -35,10 +40,11 @@ def _overlay(model: str, token: str = "tok", **kwargs): return pi.render_overlay( model, token, - _base_urls(), + _base_urls_with_oss(), bundle["claude_models"], bundle["codex_models"], bundle["gemini_models"], + bundle["oss_models"], ) @@ -81,17 +87,25 @@ def test_gemini_provider_uses_google_generative_ai(self): assert provider["api"] == "google-generative-ai" assert provider["baseUrl"] == f"{WS}/ai-gateway/gemini/v1beta" - def test_all_three_providers_when_all_present(self): + def test_oss_provider_uses_openai_completions(self): + overlay, _ = _overlay("system.ai.kimi-k2-7-code", oss_models=["system.ai.kimi-k2-7-code"]) + provider = overlay["providers"]["databricks-oss"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{WS}/ai-gateway/mlflow/v1" + + def test_all_providers_when_all_present(self): overlay, _ = _overlay( "claude-sonnet", claude_models={"sonnet": "claude-sonnet"}, codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.kimi-k2-7-code"], ) assert set(overlay["providers"].keys()) == { "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-oss", } @@ -129,6 +143,14 @@ def test_openai_and_gemini_have_no_compat_flags(self): assert "compat" not in overlay["providers"]["databricks-openai"] assert "compat" not in overlay["providers"]["databricks-gemini"] + def test_oss_disables_store_and_strict_mode(self): + # The MLflow chat-completions route rejects `store` and + # `tools[].function.strict`. + overlay, _ = _overlay("system.ai.kimi-k2-7-code", oss_models=["system.ai.kimi-k2-7-code"]) + compat = overlay["providers"]["databricks-oss"]["compat"] + assert compat["supportsStore"] is False + assert compat["supportsStrictMode"] is False + class TestRenderOverlayAuthAndModels: def test_token_in_api_key(self): @@ -163,6 +185,29 @@ def test_gemini_models_listed(self): ids = {m["id"] for m in overlay["providers"]["databricks-gemini"]["models"]} assert ids == {"gemini-2", "gemini-2-pro"} + def test_oss_models_listed(self): + oss = ["system.ai.kimi-k2-7-code", "system.ai.glm-5-2"] + overlay, _ = _overlay("system.ai.kimi-k2-7-code", oss_models=oss) + ids = {m["id"] for m in overlay["providers"]["databricks-oss"]["models"]} + assert ids == {"system.ai.kimi-k2-7-code", "system.ai.glm-5-2"} + + def test_oss_glm_carries_token_limits(self): + # GLM's output is capped well below pi's 16k default on the MLflow route; + # pin contextWindow/maxTokens from the shared limits table. + overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"]) + entry = next( + m for m in overlay["providers"]["databricks-oss"]["models"] if "glm" in m["id"] + ) + assert entry["contextWindow"] == 200_000 + assert entry["maxTokens"] == 25_000 + + def test_oss_kimi_omits_limits_when_unknown(self): + # Kimi has no entry in the limits table, so pi uses its own defaults. + overlay, _ = _overlay("system.ai.kimi-k2-7-code", oss_models=["system.ai.kimi-k2-7-code"]) + entry = overlay["providers"]["databricks-oss"]["models"][0] + assert "contextWindow" not in entry + assert "maxTokens" not in entry + class TestRenderOverlayManagedKeys: def test_managed_keys_include_model(self): @@ -193,6 +238,10 @@ def test_prefixes_gemini_model(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) assert overlay["model"] == "databricks-gemini/gemini-2" + def test_prefixes_oss_model(self): + overlay, _ = _overlay("system.ai.kimi-k2-7-code", oss_models=["system.ai.kimi-k2-7-code"]) + assert overlay["model"] == "databricks-oss/system.ai.kimi-k2-7-code" + def test_preserves_already_prefixed_model(self): overlay, _ = _overlay( "databricks-claude/claude-sonnet", @@ -228,10 +277,22 @@ def test_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} assert pi.default_model(state) == "gemini-2" + def test_falls_back_to_oss(self): + state = { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + "oss_models": ["system.ai.kimi-k2-7-code"], + } + assert pi.default_model(state) == "system.ai.kimi-k2-7-code" + def test_returns_none_when_empty(self): assert pi.default_model({}) is None assert ( - pi.default_model({"claude_models": {}, "codex_models": [], "gemini_models": []}) is None + pi.default_model( + {"claude_models": {}, "codex_models": [], "gemini_models": [], "oss_models": []} + ) + is None ) @@ -279,10 +340,11 @@ def _setup(self, tmp_path, monkeypatch): def _state(self, **overrides) -> dict: state = { "workspace": WS, - "base_urls": {"pi": _base_urls()}, + "base_urls": {"pi": _base_urls_with_oss()}, "claude_models": {"sonnet": "claude-sonnet"}, "codex_models": [], "gemini_models": [], + "oss_models": [], "managed_configs": {}, } state.update(overrides) @@ -315,8 +377,9 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): """Earlier ucode versions wrote `databricks-anthropic`, `databricks-codex`, - and `databricks-oss` providers. They must be stripped on the next write - so users don't end up with stale entries pointing at routes that 400.""" + and `databricks-kimi` providers. They must be stripped on the next write + so users don't end up with stale entries pointing at routes that 400 (or + carrying a frozen, never-refreshed token).""" pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) config_file.write_text( @@ -325,7 +388,7 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): "providers": { "databricks-anthropic": {"api": "anthropic-messages"}, "databricks-codex": {"api": "openai-responses"}, - "databricks-oss": {"api": "openai-completions"}, + "databricks-kimi": {"api": "openai-responses"}, } } ), @@ -339,7 +402,7 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") written_providers = json.loads(config_file.read_text()).get("providers", {}) - for legacy in ("databricks-anthropic", "databricks-codex", "databricks-oss"): + for legacy in ("databricks-anthropic", "databricks-codex", "databricks-kimi"): assert legacy not in written_providers assert "databricks-claude" in written_providers