diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 422e763..82821a0 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -232,8 +232,12 @@ def _maybe_add_1m_suffix(model: str) -> str: family, major_raw, minor_raw, _ = match.groups() major = int(major_raw) minor = int(minor_raw) + # 1M-token context window support: Opus from 4.6, Sonnet from 4.5 (Sonnet + # 4.5 gained the 1M beta window before Opus did). Without the `[1m]` + # suffix the gateway serves these models with the default 200k context, so + # gate each family at its own floor rather than a shared version. should_suffix = (family == "opus" and (major, minor) >= (4, 6)) or ( - family == "sonnet" and (major, minor) >= (4, 6) + family == "sonnet" and (major, minor) >= (4, 5) ) return f"{model}[1m]" if should_suffix else model diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index 78d1ba6..3982be3 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -38,6 +38,7 @@ from __future__ import annotations import os +import re import signal import subprocess import threading @@ -55,7 +56,9 @@ from ucode.databricks import ( TOKEN_REFRESH_INTERVAL_SECONDS, build_pi_base_urls, + claude_model_token_limits, get_databricks_token, + gpt_model_token_limits, model_is_reasoning, model_token_limits, ) @@ -117,6 +120,26 @@ def _resolve_model_selector( return model +def _pi_claude_model_entry(model_id: str) -> dict: + """Build a Claude entry with explicit limits. + + Databricks model ids do not match Pi's built-in Anthropic ids, so a bare + custom entry silently gets Pi's 128k context / 4k output defaults. + """ + limits = claude_model_token_limits(model_id) + entry: dict = { + "id": model_id, + "reasoning": True, + "input": ["text", "image"], + "contextWindow": limits["context"], + "maxTokens": limits["output"], + } + normalized = model_id.lower().replace(".", "-") + if re.search(r"claude-(?:opus|sonnet)-4-(?:[6-9]|[1-9]\d)", normalized): + entry["compat"] = {"forceAdaptiveThinking": True} + return entry + + def _pi_oss_model_entry(model_id: str) -> dict: """Build a Pi mlflow model entry enriched from the shared limits/reasoning tables: `reasoning:true` for reasoning models (Pi renders their streamed @@ -135,6 +158,23 @@ def _pi_oss_model_entry(model_id: str) -> dict: return entry +def _pi_gpt_model_entry(model_id: str) -> dict: + """Build a Pi openai (codex) model entry with `contextWindow`/`maxTokens` + from `databricks.gpt_model_token_limits`. GPT ids aren't in Pi's built-in + catalog, so without an explicit window Pi falls back to a small default and + truncates long sessions.""" + limits = gpt_model_token_limits(model_id) + entry: dict = { + "id": model_id, + "contextWindow": limits["context"], + "maxTokens": limits["output"], + } + if "gpt-5" in model_id.lower().replace(".", "-"): + entry["reasoning"] = True + entry["input"] = ["text", "image"] + return entry + + def render_overlay( model: str, token: str, @@ -163,7 +203,7 @@ def render_overlay( # the legacy beta header instead when this is false. "compat": {"supportsEagerToolInputStreaming": False}, "headers": ua_headers, - "models": [{"id": m} for m in claude_ids], + "models": [_pi_claude_model_entry(m) for m in claude_ids], } keys.append(["providers", "databricks-claude"]) if codex_models: @@ -173,7 +213,7 @@ def render_overlay( "apiKey": token, "authHeader": True, "headers": ua_headers, - "models": [{"id": m} for m in codex_models], + "models": [_pi_gpt_model_entry(m) for m in codex_models], } keys.append(["providers", "databricks-openai"]) if gemini_models: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 4d89201..8aa70e4 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1184,6 +1184,76 @@ def model_token_limits(model_id: str) -> dict[str, int] | None: return None +# Pi treats every custom model without explicit metadata as 128k context / 4k +# output. Gateway ids are custom ids (not Pi's built-ins), so preserve the +# upstream windows explicitly. Entries are ordered most-specific first after +# normalizing dotted OpenAI ids and hyphenated Databricks ids to one form. +# GPT-5.6 Sol/Terra/Luna support the opt-in 1.05M window; 272k is merely Pi's +# built-in short-context pricing default, not the model's hard context limit. +_GPT_TOKEN_LIMITS: tuple[tuple[str, dict[str, int]], ...] = ( + ("gpt-5-6-sol", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-6-terra", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-6-luna", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-5-pro", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-4-pro", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-5", {"context": 272_000, "output": 128_000}), + ("gpt-5-4-mini", {"context": 400_000, "output": 128_000}), + ("gpt-5-4-nano", {"context": 400_000, "output": 128_000}), + ("gpt-5-4", {"context": 272_000, "output": 128_000}), + ("gpt-5", {"context": 400_000, "output": 128_000}), + ("gpt-4-1", {"context": 1_047_576, "output": 32_768}), + ("gpt-4o", {"context": 128_000, "output": 16_384}), + ("gpt-4-turbo", {"context": 128_000, "output": 4_096}), + ("gpt-4", {"context": 8_192, "output": 8_192}), +) +_GPT_FALLBACK_LIMITS = {"context": 128_000, "output": 16_384} + + +def _normalized_foundation_model_id(model_id: str) -> str: + """Strip route prefixes and normalize dotted versions to hyphens.""" + tail = model_id.split("/")[-1] + if tail.startswith("system.ai."): + tail = tail[len("system.ai.") :] + if tail.startswith("databricks-"): + tail = tail[len("databricks-") :] + return tail.lower().replace(".", "-") + + +def gpt_model_token_limits(model_id: str) -> dict[str, int]: + """Return Pi metadata limits for a GPT (codex/openai) gateway model.""" + tail = _normalized_foundation_model_id(model_id) + for family, limits in _GPT_TOKEN_LIMITS: + if tail.startswith(family): + return dict(limits) + return dict(_GPT_FALLBACK_LIMITS) + + +_CLAUDE_FALLBACK_LIMITS = {"context": 200_000, "output": 64_000} + + +def claude_model_token_limits(model_id: str) -> dict[str, int]: + """Return Pi metadata limits for a Claude gateway model. + + Claude gateway ids are custom to Pi, so omitting these fields silently + applies Pi's 128k/4k custom-model defaults. Current 1M families mirror + Anthropic's model metadata: Opus >=4.6 and Sonnet >=4.5. Sonnet 4.5 has a + 64k output cap; later 1M models have 128k output. + """ + tail = _normalized_foundation_model_id(model_id) + match = re.match(r"claude-(opus|sonnet|haiku)-(\d+)-(\d+)", tail) + if not match: + return dict(_CLAUDE_FALLBACK_LIMITS) + family, major_raw, minor_raw = match.groups() + version = (int(major_raw), int(minor_raw)) + if family == "opus" and version >= (4, 6): + return {"context": 1_000_000, "output": 128_000} + if family == "sonnet" and version >= (4, 6): + return {"context": 1_000_000, "output": 128_000} + if family == "sonnet" and version >= (4, 5): + return {"context": 1_000_000, "output": 64_000} + return dict(_CLAUDE_FALLBACK_LIMITS) + + def _model_service_id(service: dict) -> str | None: """Extract the `system.ai.` id from one model-service entry. diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index b2d3a27..a69fb9d 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -48,6 +48,30 @@ def test_adds_1m_suffix_for_sonnet_4_6_and_later(self): overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-7[1m]" ) + def test_adds_1m_suffix_for_sonnet_4_5(self): + # Sonnet 4.5 supports the 1M context window (its 1M beta shipped before + # Opus's), so it must get the [1m] suffix even though it predates the + # Opus 4.6 floor. + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-5"} + ) + assert ( + overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-5[1m]" + ) + + def test_does_not_add_1m_suffix_for_sonnet_4_4(self): + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-4"} + ) + assert overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-4" + + def test_does_not_add_1m_suffix_for_opus_4_5(self): + # Opus's 1M window starts at 4.6, so 4.5 stays on the default context. + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"opus": "databricks-claude-opus-4-5"} + ) + assert overlay["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-5" + def test_does_not_add_1m_suffix_for_haiku(self): overlay, _ = claude.render_overlay( WS, "s4", claude_models={"haiku": "databricks-claude-haiku-4-6"} diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index de002ba..c7242cb 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -78,6 +78,55 @@ def test_openai_provider_uses_openai_responses(self): assert provider["api"] == "openai-responses" assert provider["baseUrl"] == f"{WS}/ai-gateway/codex/v1" + def test_gpt56_sol_model_entry_pins_1m_context(self): + # Gateway ids are custom to Pi, so explicit metadata is required to + # avoid its 128k custom-model default. + overlay, _ = _overlay("gpt-5-6-sol", codex_models=["gpt-5-6-sol"]) + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert entry["id"] == "gpt-5-6-sol" + assert entry["contextWindow"] == 1_050_000 + assert entry["maxTokens"] == 128_000 + assert entry["reasoning"] is True + assert entry["input"] == ["text", "image"] + + def test_gpt_model_entries_use_model_specific_windows(self): + overlay, _ = _overlay( + "system.ai.gpt-5-2", + codex_models=[ + "system.ai.gpt-5-2", + "databricks-gpt-5-4-nano", + "databricks-gpt-5-6-sol", + ], + ) + windows = { + m["id"]: m["contextWindow"] for m in overlay["providers"]["databricks-openai"]["models"] + } + assert windows == { + "system.ai.gpt-5-2": 400_000, + "databricks-gpt-5-4-nano": 400_000, + "databricks-gpt-5-6-sol": 1_050_000, + } + + def test_claude_entries_pin_limits_and_capabilities(self): + overlay, _ = _overlay( + "databricks-claude-opus-4-8", + claude_models={ + "opus": "databricks-claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-4-5", + "haiku": "databricks-claude-haiku-4-5", + }, + ) + entries = {m["id"]: m for m in overlay["providers"]["databricks-claude"]["models"]} + opus = entries["databricks-claude-opus-4-8"] + assert opus["contextWindow"] == 1_000_000 + assert opus["maxTokens"] == 128_000 + assert opus["reasoning"] is True + assert opus["input"] == ["text", "image"] + assert opus["compat"] == {"forceAdaptiveThinking": True} + assert entries["system.ai.claude-sonnet-4-5"]["contextWindow"] == 1_000_000 + assert entries["system.ai.claude-sonnet-4-5"]["maxTokens"] == 64_000 + assert entries["databricks-claude-haiku-4-5"]["contextWindow"] == 200_000 + def test_gemini_provider_uses_google_generative_ai(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) provider = overlay["providers"]["databricks-gemini"] diff --git a/tests/test_databricks.py b/tests/test_databricks.py index a7c6e69..2775c0a 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -187,6 +187,60 @@ def test_embedding_model_returns_none_not_fallback(self): assert db_mod.model_token_limits("system.ai.qwen3-embedding-0-6b") is None +class TestGptModelTokenLimits: + def test_gpt56_sol_serves_1m_context_across_id_forms(self): + for model_id in ( + "gpt-5.6-sol", + "system.ai.gpt-5-6-sol", + "databricks-gpt-5-6-sol", + "databricks-openai/gpt-5.6-sol", + ): + assert db_mod.gpt_model_token_limits(model_id) == { + "context": 1_050_000, + "output": 128_000, + } + + def test_gpt5_windows_are_model_specific(self): + assert db_mod.gpt_model_token_limits("system.ai.gpt-5-2")["context"] == 400_000 + assert db_mod.gpt_model_token_limits("databricks-gpt-5-4")["context"] == 272_000 + assert db_mod.gpt_model_token_limits("databricks-gpt-5-4-nano")["context"] == 400_000 + assert db_mod.gpt_model_token_limits("gpt-5.5-pro")["context"] == 1_050_000 + + def test_gpt41_window(self): + assert db_mod.gpt_model_token_limits("databricks-gpt-4-1") == { + "context": 1_047_576, + "output": 32_768, + } + + def test_unknown_gpt_uses_conservative_floor(self): + assert db_mod.gpt_model_token_limits("gpt-6-turbo") == { + "context": 128_000, + "output": 16_384, + } + + +class TestClaudeModelTokenLimits: + def test_1m_models_and_output_caps(self): + assert db_mod.claude_model_token_limits("databricks-claude-opus-4-8") == { + "context": 1_000_000, + "output": 128_000, + } + assert db_mod.claude_model_token_limits("system.ai.claude-sonnet-4-5") == { + "context": 1_000_000, + "output": 64_000, + } + assert db_mod.claude_model_token_limits("claude-sonnet-4-6[1m]") == { + "context": 1_000_000, + "output": 128_000, + } + + def test_older_and_unknown_claude_use_200k_floor(self): + expected = {"context": 200_000, "output": 64_000} + assert db_mod.claude_model_token_limits("claude-opus-4-5") == expected + assert db_mod.claude_model_token_limits("claude-haiku-4-5") == expected + assert db_mod.claude_model_token_limits("claude-future") == expected + + class TestModelIsReasoning: def test_reasoning_families(self): assert db_mod.model_is_reasoning("system.ai.glm-5-2") is True