diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e7c1760..a9e1806 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -77,6 +77,93 @@ LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-oss") +# pi's full thinking-level ladder, in order. +PI_THINKING_LEVELS = ("off", "minimal", "low", "medium", "high", "xhigh") + +# Levels each Databricks-hosted model accepts, in pi's vocabulary. +# "off" == reasoning disabled (OpenAI effort "none"). () == no thinking at all. +# Keyed by model id with any ".." prefix stripped. +MODEL_THINKING_LEVELS: dict[str, tuple[str, ...]] = { + # ---- databricks-claude (anthropic-messages) ---- + "claude-opus-4-8": ("low", "medium", "high", "xhigh"), + "claude-sonnet-5": ("low", "medium", "high", "xhigh"), + "claude-haiku-4-5": (), + # ---- databricks-openai (openai-responses) ---- + "gpt-5": ("minimal", "low", "medium", "high"), + "gpt-5-mini": ("minimal", "low", "medium", "high"), + "gpt-5-nano": ("minimal", "low", "medium", "high"), + "gpt-5-1": ("off", "low", "medium", "high"), + "gpt-5-1-codex-max": ("off", "low", "medium", "high", "xhigh"), + "gpt-5-1-codex-mini": ("off", "low", "medium", "high"), + "gpt-5-2": ("off", "low", "medium", "high", "xhigh"), + "gpt-5-2-codex": ("low", "medium", "high", "xhigh"), + "gpt-5-3-codex": ("low", "medium", "high", "xhigh"), + "gpt-5-4": ("off", "low", "medium", "high", "xhigh"), + "gpt-5-4-mini": ("off", "low", "medium", "high", "xhigh"), + "gpt-5-4-nano": ("off", "low", "medium", "high", "xhigh"), + "gpt-5-5": ("off", "low", "medium", "high", "xhigh"), + "gpt-5-5-pro": ("medium", "high", "xhigh"), + "gpt-oss-120b": ("low", "medium", "high"), + "gpt-oss-20b": ("low", "medium", "high"), + # ---- databricks-gemini (google-generative-ai) ---- + # Only gemini-3-pro and gemini-3-flash get levels. Pi's google handler picks + # thinkingLevel (Gemini 3) vs thinkingBudget (2.5) by regex-matching model.id + # (isGemini3{Pro,Flash}Model / getGoogleBudget in pi's + # packages/ai/src/api/google-generative-ai.ts), and those regexes expect + # DOTTED minor versions (gemini-3.1-pro, 2.5-flash). Databricks serves DASHED + # ids (gemini-3-1-pro, gemini-2-5-flash), so only the unversioned + # gemini-3-pro/-flash match; every other id falls through to a dynamic budget + # (-1) that ignores the requested effort. () (reasoning: False) hides the + # effort selector — these models still think, at their native default + # (Gemini 3 can't disable it, 2.5 is on by default), since pi then sends no + # thinkingConfig at all. + "gemini-3-5-flash": (), + "gemini-3-1-flash-lite": (), + "gemini-3-1-pro": (), + "gemini-3-flash": ("minimal", "low", "medium", "high"), + "gemini-3-pro": ("low", "high"), + "gemini-2-5-flash": (), + "gemini-2-5-pro": (), +} + + +def _thinking_level_map(supported: tuple[str, ...], family: str) -> dict[str, str | None]: + m: dict[str, str | None] = {} + for lvl in PI_THINKING_LEVELS: + if lvl in supported: + if lvl == "off": + m["off"] = "none" if family == "openai" else None # openai disables via "none" + elif lvl == "xhigh": + m["xhigh"] = "xhigh" # must be explicit or pi hides it + # minimal/low/medium/high: leave unset -> pi sends the level verbatim + else: + m[lvl] = None # hide a level the model rejects + return m + + +def _pi_model_entry(model_id: str, family: str) -> dict: + # The legacy AI-Gateway route doesn't support these thinking parameters and + # will be deprecated, so skip those models. + if not model_id.startswith("system.ai."): + return {"id": model_id, "reasoning": False} + + norm = model_id.removeprefix("system.ai.").lower() + + if "image" in norm: # safety net for image-gen models + return {"id": model_id, "reasoning": False} + + # Don't set thinking levels for unknown/new models. + supported = MODEL_THINKING_LEVELS.get(norm, ()) + if not supported: + return {"id": model_id, "reasoning": False} + + return { + "id": model_id, + "reasoning": True, + "thinkingLevelMap": _thinking_level_map(supported, family), + } + + def is_update_available() -> tuple[str, str] | None: return available_npm_package_update(SPEC["package"]) @@ -125,9 +212,9 @@ def render_overlay( # Gateway's Anthropic translator rejects per-tool # `eager_input_streaming` on the streaming + tools path. Pi sends # the legacy beta header instead when this is false. - "compat": {"supportsEagerToolInputStreaming": False}, + "compat": {"supportsEagerToolInputStreaming": False, "forceAdaptiveThinking": True}, "headers": ua_headers, - "models": [{"id": m} for m in claude_ids], + "models": [_pi_model_entry(m, "claude") for m in claude_ids], } keys.append(["providers", "databricks-claude"]) if codex_models: @@ -137,7 +224,7 @@ def render_overlay( "apiKey": token, "authHeader": True, "headers": ua_headers, - "models": [{"id": m} for m in codex_models], + "models": [_pi_model_entry(m, "openai") for m in codex_models], } keys.append(["providers", "databricks-openai"]) if gemini_models: @@ -147,7 +234,7 @@ def render_overlay( "apiKey": token, "authHeader": True, "headers": ua_headers, - "models": [{"id": m} for m in gemini_models], + "models": [_pi_model_entry(m, "gemini") for m in gemini_models], } keys.append(["providers", "databricks-gemini"]) overlay: dict = { diff --git a/tests/conftest.py b/tests/conftest.py index 0cc7932..9cfaa22 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,6 +8,7 @@ from ucode.databricks import ( build_shared_base_urls, + discover_model_services, fetch_ai_gateway_claude_models, fetch_codex_models, fetch_gemini_models, @@ -72,3 +73,16 @@ def e2e_state(e2e_workspace, e2e_token): "base_urls": build_shared_base_urls(e2e_workspace), "managed_configs": {}, } + + +@pytest.fixture(scope="session") +def e2e_uc_models(e2e_workspace, e2e_token): + """Models served through Unity Catalog (system.ai.*).""" + claude_models, codex_models, gemini_models, _oss, _ = discover_model_services( + e2e_workspace, e2e_token + ) + return { + "claude_models": claude_models, + "codex_models": codex_models, + "gemini_models": gemini_models, + } diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 604452a..42b3a97 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -882,23 +882,33 @@ class TestPiLaunch: test exercises each one end-to-end through the validation path. """ - def _all_models(self, e2e_state: dict) -> list[tuple[str, str]]: - out: list[tuple[str, str]] = [] + def _all_models(self, e2e_state: dict, e2e_uc_models: dict) -> list[tuple[str, str, str]]: + out: list[tuple[str, str, str]] = [] + # AI-Gateway route (databricks-*). claude_models: dict = e2e_state.get("claude_models") or {} for family, model_id in _launchable_model_items(claude_models): - out.append((f"claude-{family}", model_id)) + out.append(("aig", f"claude-{family}", model_id)) for model in e2e_state.get("codex_models") or []: - out.append(("codex", model)) + out.append(("aig", "codex", model)) for model in e2e_state.get("gemini_models") or []: - out.append(("gemini", model)) + out.append(("aig", "gemini", model)) + # UC model route (system.ai.*) + for family, model_id in _launchable_model_items(e2e_uc_models["claude_models"]): + out.append(("uc", f"claude-{family}", model_id)) + for model in e2e_uc_models.get("codex_models") or []: + out.append(("uc", "codex", model)) + for model in e2e_uc_models.get("gemini_models") or []: + out.append(("uc", "gemini", model)) return out - def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token): + def test_launch_pi_per_model( + self, tmp_path, monkeypatch, e2e_state, e2e_uc_models, e2e_workspace, e2e_token + ): import ucode.config_io as config_io_mod from ucode.agents import pi _require_binary("pi") - models = self._all_models(e2e_state) + models = self._all_models(e2e_state, e2e_uc_models) if not models: pytest.skip("No Pi-compatible models available on this workspace") @@ -913,8 +923,13 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) monkeypatch.setattr(pi, "PI_BACKUP_PATH", backup_path) + state_by_route = { + "aig": {**e2e_state, "workspace": e2e_workspace}, + "uc": {**e2e_state, "workspace": e2e_workspace, **e2e_uc_models}, + } + failures = [] - for family, model in models: + for route, family, model in models: if config_path.exists(): config_path.unlink() @@ -925,7 +940,7 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa lambda ws, profile=None, **kwargs: e2e_token, ) pi.write_tool_config( - {**e2e_state, "workspace": e2e_workspace}, + state_by_route[route], model, token=e2e_token, ) @@ -936,7 +951,7 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa combined = (result.stdout + result.stderr).strip() if result.returncode != 0 or not combined: failures.append( - f"family={family} model={model} rc={result.returncode} " + f"route={route} family={family} model={model} rc={result.returncode} " f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" )