From 52f8d3a3a5a0cef8a48c47e3148e570dc1ec255f Mon Sep 17 00:00:00 2001 From: Xi Zhang <106144707+X-iZhang@users.noreply.github.com> Date: Mon, 27 Apr 2026 19:28:57 +0200 Subject: [PATCH] Feat/llm context window patch table (#191) * feat(context-window): add model context window patch table and apply function * fix(tests): clean up formatting in context window tests * fix(tests): update context window tests for Claude model exceptions --- EvoScientist/llm/context_window.py | 86 ++++++++++++++++- EvoScientist/llm/models.py | 3 + tests/test_context_window.py | 143 +++++++++++++++++++++++++++++ 3 files changed, 231 insertions(+), 1 deletion(-) diff --git a/EvoScientist/llm/context_window.py b/EvoScientist/llm/context_window.py index f1ce23c..d080493 100644 --- a/EvoScientist/llm/context_window.py +++ b/EvoScientist/llm/context_window.py @@ -7,6 +7,41 @@ from typing import Any DEFAULT_CONTEXT_WINDOW_FALLBACK = 200_000 +# Patch table for new models that providers haven't registered profile data +# for yet. Remove an entry once langchain/provider exposes max_input_tokens +# via ``model.profile`` — the attribute-reading layer always wins. +# Keys are matched against ``model.model_name`` (or ``model.model`` / +# ``model.name``); lookup tries exact match first, then ``split('/')[-1]`` +# to also accept OpenRouter-style ``vendor/model`` IDs. +_KNOWN_MODEL_CONTEXT_WINDOWS: dict[str, int] = { + # Qwen 3.6 open-source variants — exceptions to the ``qwen3.6`` family. + "qwen3.6-27b": 262_000, + "qwen3.6-35b-a3b": 262_000, + # Claude Haiku 4.5 — exception to the ``claude-`` family (200K, not 1M). + "claude-haiku-4-5": 200_000, +} + +# Family-level fallbacks: tried only after exact-name lookup misses. +# Each entry is (substring_to_match_in_lowercased_model_id, window). +# Order matters — first match wins; put more specific patterns first. +_KNOWN_MODEL_FAMILIES: list[tuple[str, int]] = [ + # All Claude — 1M via the ``context-1m-2025-08-07`` beta header. + ("claude-", 1_000_000), + # OpenAI GPT-5.5 family — base, pro, future variants + ("gpt-5.5", 1_050_000), + # Moonshot Kimi K2 family — k2.5, k2.6, k2-thinking, k2-thinking-turbo + ("kimi-k2", 262_000), + # Zhipu GLM-5 family — base, 5.1, 5-turbo, 5v-turbo, etc. + ("glm-5", 203_000), + # DeepSeek V4 family — pro, flash, future variants + ("deepseek-v4", 1_050_000), + # Xiaomi MiMo v2.5 family — base, pro, future variants + ("mimo-v2.5", 1_050_000), + # Qwen 3.6 closed-source family — flash, plus, max-preview, etc. + # Open-source ``-b`` variants are 262K — listed in the dict above. + ("qwen3.6", 1_000_000), +] + _DIRECT_WINDOW_ATTRS = ( "context_window", "context_length", @@ -19,6 +54,7 @@ _CONTAINER_ATTRS = ( "model_kwargs", "metadata", ) +_NAME_ATTRS = ("model_name", "model", "name") def _coerce_positive_int(value: Any) -> int | None: @@ -49,6 +85,32 @@ def _resolve_from_mapping(mapping: Mapping[str, Any]) -> int | None: return None +def _lookup_by_model_name(model_obj: Any) -> int | None: + """Return a patched context window by model name, if registered. + + Lookup is case-insensitive. Matching order: + 1. exact lowercased ``model_name`` value; + 2. last ``/``-segment of the lowercased value (handles + OpenRouter-style ``vendor/model`` and SiliconFlow-style + ``Pro/vendor/Model`` IDs); + 3. family-level substring patterns (e.g. all ``claude-*``). + """ + for attr in _NAME_ATTRS: + value = getattr(model_obj, attr, None) + if not isinstance(value, str) or not value: + continue + lowered = value.lower() + if lowered in _KNOWN_MODEL_CONTEXT_WINDOWS: + return _KNOWN_MODEL_CONTEXT_WINDOWS[lowered] + short = lowered.split("/")[-1] + if short != lowered and short in _KNOWN_MODEL_CONTEXT_WINDOWS: + return _KNOWN_MODEL_CONTEXT_WINDOWS[short] + for pattern, window in _KNOWN_MODEL_FAMILIES: + if pattern in lowered: + return window + return None + + def get_context_window(model_obj: Any | None) -> int | None: """Return the best available context-window value from a model object.""" if model_obj is None: @@ -66,7 +128,7 @@ def get_context_window(model_obj: Any | None) -> int | None: if resolved is not None: return resolved - return None + return _lookup_by_model_name(model_obj) def resolve_context_window( @@ -79,3 +141,25 @@ def resolve_context_window( if resolved is not None: return resolved return fallback + + +def apply_known_context_window(model: Any) -> None: + """Inject patched ``max_input_tokens`` into ``model.profile`` for new + models that providers haven't registered profile data for yet. + + Without this, deepagents' ``SummarizationMiddleware`` falls back to a + hardcoded 170K trigger, ignoring the true (e.g. 1M) context window. + Real provider-supplied profile data always wins — we only fill gaps. + """ + patched = _lookup_by_model_name(model) + if patched is None: + return + profile = getattr(model, "profile", None) + if isinstance(profile, dict) and "max_input_tokens" in profile: + return + base = profile if isinstance(profile, dict) else {} + new_profile = {**base, "max_input_tokens": patched} + try: + model.profile = new_profile + except Exception: + pass diff --git a/EvoScientist/llm/models.py b/EvoScientist/llm/models.py index 3603530..3329677 100644 --- a/EvoScientist/llm/models.py +++ b/EvoScientist/llm/models.py @@ -14,6 +14,7 @@ from typing import Any from langchain.chat_models import init_chat_model +from .context_window import apply_known_context_window from .patches import ( _is_ccproxy_codex, _patch_ccproxy_system_to_developer, @@ -466,6 +467,8 @@ def get_chat_model( if _is_openai_proxy: _patch_ccproxy_system_to_developer(chat_model) + apply_known_context_window(chat_model) + return chat_model diff --git a/tests/test_context_window.py b/tests/test_context_window.py index 70a676a..aed67d0 100644 --- a/tests/test_context_window.py +++ b/tests/test_context_window.py @@ -4,6 +4,7 @@ from types import SimpleNamespace from EvoScientist.llm.context_window import ( DEFAULT_CONTEXT_WINDOW_FALLBACK, + apply_known_context_window, get_context_window, resolve_context_window, ) @@ -59,3 +60,145 @@ def test_resolve_context_window_falls_back_when_missing(): assert get_context_window(model) is None assert resolve_context_window(model) == DEFAULT_CONTEXT_WINDOW_FALLBACK assert resolve_context_window(model, fallback=42_000) == 42_000 + + +def test_patch_table_resolves_known_model_name(): + model = SimpleNamespace(model_name="claude-opus-4.7", profile=None) + + assert get_context_window(model) == 1_000_000 + + +def test_patch_table_strips_provider_prefix(): + full = SimpleNamespace(model_name="openai/gpt-5.5", profile=None) + short = SimpleNamespace(model_name="gpt-5.5", profile=None) + + assert get_context_window(full) == 1_050_000 + assert get_context_window(short) == 1_050_000 + + +def test_real_attribute_beats_patch_table(): + model = SimpleNamespace( + model_name="claude-opus-4.7", + context_window=500_000, + profile=None, + ) + + assert get_context_window(model) == 500_000 + + +def test_unknown_model_falls_through_to_fallback(): + model = SimpleNamespace(model_name="some-unreleased-model", profile=None) + + assert get_context_window(model) is None + assert resolve_context_window(model) == DEFAULT_CONTEXT_WINDOW_FALLBACK + + +def test_claude_family_pattern_covers_all_variants(): + # Native Anthropic (dash form) + native = SimpleNamespace(model_name="claude-opus-4-7", profile=None) + # OpenRouter (dot form, with vendor prefix) + openrouter = SimpleNamespace(model_name="anthropic/claude-sonnet-4.6", profile=None) + sonnet = SimpleNamespace(model_name="claude-sonnet-4-5", profile=None) + # Haiku 4.5 is a dict-level exception (200K, not 1M). + haiku = SimpleNamespace(model_name="claude-haiku-4-5", profile=None) + + assert get_context_window(native) == 1_000_000 + assert get_context_window(openrouter) == 1_000_000 + assert get_context_window(sonnet) == 1_000_000 + assert get_context_window(haiku) == 200_000 + + +def test_dict_entry_overrides_family_pattern(): + # Qwen 3.6 closed-source default is 1M (family), but open-source -27b + # variant is 262K (dict override). Dict wins because it's checked first. + closed = SimpleNamespace(model_name="qwen/qwen3.6-flash", profile=None) + open_27b = SimpleNamespace(model_name="qwen/qwen3.6-27b", profile=None) + open_plus = SimpleNamespace(model_name="qwen3.6-plus", profile=None) + + assert get_context_window(closed) == 1_000_000 + assert get_context_window(open_27b) == 262_000 + assert get_context_window(open_plus) == 1_000_000 + + +def test_claude_3x_matches_family_but_upstream_profile_wins_in_practice(): + # The broad ``claude-`` family pattern intentionally also matches older + # Claude 3.x. This is safe because langchain-anthropic ships upstream + # profiles for those (200K), and profile lookup runs before the patch + # table — so the family value never actually reaches downstream callers + # for normal Anthropic usage. + old_sonnet_no_profile = SimpleNamespace( + model_name="claude-3-5-sonnet-20241022", profile=None + ) + old_sonnet_with_profile = SimpleNamespace( + model_name="claude-3-5-sonnet-20241022", + profile={"max_input_tokens": 200_000}, + ) + + # No profile (rare custom-anthropic edge case): family fires. + assert get_context_window(old_sonnet_no_profile) == 1_000_000 + # Normal langchain-anthropic path: upstream profile wins. + assert get_context_window(old_sonnet_with_profile) == 200_000 + + +def test_apply_injects_into_empty_profile(): + model = SimpleNamespace(model_name="claude-opus-4-7", profile=None) + + apply_known_context_window(model) + + assert model.profile == {"max_input_tokens": 1_000_000} + + +def test_apply_preserves_existing_profile_keys(): + model = SimpleNamespace( + model_name="qwen3.6-flash", + profile={"name": "Qwen 3.6 Flash"}, + ) + + apply_known_context_window(model) + + assert model.profile == { + "name": "Qwen 3.6 Flash", + "max_input_tokens": 1_000_000, + } + + +def test_apply_does_not_overwrite_existing_max_input_tokens(): + model = SimpleNamespace( + model_name="claude-opus-4-7", + profile={"max_input_tokens": 500_000}, + ) + + apply_known_context_window(model) + + assert model.profile == {"max_input_tokens": 500_000} + + +def test_apply_skips_unknown_models(): + model = SimpleNamespace(model_name="some-unknown-model", profile=None) + + apply_known_context_window(model) + + assert model.profile is None + + +def test_apply_handles_non_dict_profile_safely(): + # Non-dict profile (e.g. namespace from a future langchain version) must + # not raise — fall back to a fresh dict so deepagents can still read it. + model = SimpleNamespace( + model_name="claude-opus-4-7", + profile=SimpleNamespace(name="opaque object"), + ) + + apply_known_context_window(model) + + assert model.profile == {"max_input_tokens": 1_000_000} + + +def test_lookup_is_case_insensitive(): + # SiliconFlow uses capitalized IDs like "Pro/moonshotai/Kimi-K2.5" + siliconflow = SimpleNamespace(model_name="Pro/moonshotai/Kimi-K2.5", profile=None) + assert get_context_window(siliconflow) == 262_000 + + # GLM-5 capitalized variant + glm = SimpleNamespace(model_name="Pro/zai-org/GLM-5", profile=None) + assert get_context_window(glm) == 203_000