From de0a4e1e89c0c72912cd7b3442dea2f5c86d4480 Mon Sep 17 00:00:00 2001 From: Shmulik Cohen Date: Mon, 6 Jul 2026 23:03:56 +0300 Subject: [PATCH 1/4] refactor(algo): key model capability tables by base model name Collapse the provider-prefixed capability lists in pr_agent/algo/__init__.py into base-model-keyed sets, matched via a new base_model_name() normalizer that strips the provider prefix, Bedrock region/vendor infixes, and cloud-version / vertex date suffixes. Adding a model is now one line per capability instead of one per provider x region, removing the silent failure where a missing prefix (e.g. bedrock/eu.) sent an unsupported arg to a model. Add a litellm fallback to get_max_tokens() for models absent from the MAX_TOKENS table (the table is still checked first, so every model already in it keeps its exact value), route the direct MAX_TOKENS[...] readers through get_max_tokens(), and bump litellm 1.84.0 -> 1.91.0 so the fallback covers more models. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01SqHrpoPezxDkD5dg4VuQL5 --- pr_agent/algo/__init__.py | 124 +++++++----------- .../algo/ai_handlers/litellm_ai_handler.py | 22 ++-- pr_agent/algo/token_handler.py | 6 +- pr_agent/algo/utils.py | 26 +++- pr_agent/tools/pr_code_suggestions.py | 8 +- pr_agent/tools/pr_help_docs.py | 8 +- pr_agent/tools/pr_help_message.py | 7 +- pr_agent/tools/pr_similar_issue.py | 7 +- requirements.txt | 2 +- tests/unittest/test_base_model_name.py | 45 +++++++ .../test_litellm_claude_extended_thinking.py | 7 +- 11 files changed, 137 insertions(+), 125 deletions(-) create mode 100644 tests/unittest/test_base_model_name.py diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py index 8611c47cd4..dae7ccef5d 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -1,3 +1,27 @@ +import re + +# Deployment-decoration tokens that appear between the provider prefix and the model +# identity: Bedrock cross-region prefixes and cloud vendor names. +_PROVIDER_INFIX = {"us", "eu", "au", "jp", "apac", "global", + "anthropic", "meta", "amazon", "ai21", "cohere", "mistral"} + + +def base_model_name(model: str) -> str: + """Strip deployment decoration so every provider spelling of a model maps to one key. + + e.g. ``anthropic/claude-opus-4-8``, ``bedrock/us.anthropic.claude-opus-4-8-v1:0`` and + ``vertex_ai/claude-opus-4-8`` all collapse to ``claude-opus-4-8``. Used for the capability + sets below, which are keyed by base model rather than by every provider-prefixed variant. + """ + m = model.split("/", 1)[-1] # drop provider prefix (anthropic/, bedrock/, vertex_ai/, ...) + m = m.replace("@", "-") # vertex date separator -> '-' (matches anthropic spelling) + m = re.sub(r"-v\d+:\d+$", "", m) # bedrock cloud-version suffix (-v1:0) + parts = m.split(".") + while len(parts) > 1 and parts[0] in _PROVIDER_INFIX: # bedrock region + vendor infixes + parts.pop(0) + return ".".join(parts) # dotted model versions (gpt-5.2) are preserved + + MAX_TOKENS = { 'text-embedding-ada-002': 8000, 'gpt-3.5-turbo': 16000, @@ -298,15 +322,19 @@ "codestral/codestral-2405": 8191, } -USER_MESSAGE_ONLY_MODELS = [ - "deepseek/deepseek-reasoner", +# The capability sets below are keyed by base model name (see base_model_name); membership is +# checked with `base_model_name(model) in `, so each model is listed once regardless of +# provider prefix, Bedrock region, or cloud-version/date suffix. + +USER_MESSAGE_ONLY_MODELS = { + "deepseek-reasoner", "o1-mini", "o1-mini-2024-09-12", - "o1-preview" -] + "o1-preview", +} -NO_SUPPORT_TEMPERATURE_MODELS = [ - "deepseek/deepseek-reasoner", +NO_SUPPORT_TEMPERATURE_MODELS = { + "deepseek-reasoner", "o1-mini", "o1-mini-2024-09-12", "o1", @@ -323,104 +351,40 @@ "gpt-5.2-codex", "gpt-5.3-codex", "gpt-5-mini", - # Anthropic Claude Opus 4-7 — temperature is deprecated (Issue #2400), (Issue #2449) + # Anthropic Claude — temperature is deprecated (Issue #2400), (Issue #2449) "claude-opus-4-7", - "anthropic/claude-opus-4-7", "claude-opus-4-8", - "anthropic/claude-opus-4-8", - "vertex_ai/claude-opus-4-8", - "bedrock/anthropic.claude-opus-4-8", - "bedrock/global.anthropic.claude-opus-4-8", - "bedrock/us.anthropic.claude-opus-4-8", - "bedrock/eu.anthropic.claude-opus-4-8", - "bedrock/au.anthropic.claude-opus-4-8", - "bedrock/jp.anthropic.claude-opus-4-8", "claude-fable-5", - "anthropic/claude-fable-5", "claude-sonnet-5", - "anthropic/claude-sonnet-5", - "vertex_ai/claude-sonnet-5", - "bedrock/anthropic.claude-sonnet-5", - "bedrock/global.anthropic.claude-sonnet-5", - "bedrock/us.anthropic.claude-sonnet-5", - "bedrock/au.anthropic.claude-sonnet-5", - "bedrock/eu.anthropic.claude-sonnet-5", - "bedrock/jp.anthropic.claude-sonnet-5", - "vertex_ai/claude-opus-4-7", - "bedrock/anthropic.claude-opus-4-7", - "bedrock/anthropic.claude-opus-4-7-v1:0", - "bedrock/us.anthropic.claude-opus-4-7", - "bedrock/global.anthropic.claude-opus-4-7", -] +} -SUPPORT_REASONING_EFFORT_MODELS = [ +SUPPORT_REASONING_EFFORT_MODELS = { "o3-mini", "o3-mini-2025-01-31", "o3", "o3-2025-04-16", "o4-mini", "o4-mini-2025-04-16", -] +} # Claude models that support "extended thinking" through the manual # thinking={"type": "enabled", "budget_tokens": ...} request built by # LiteLLMAIHandler._configure_claude_extended_thinking(). Only models that # accept budget_tokens belong here. Adaptive-only models (Claude Opus 4.7/4.8, # Sonnet 5, Fable 5) reject budget_tokens with an HTTP 400 and must not be added -# without also adding an adaptive-thinking code path. This list is the built-in +# without also adding an adaptive-thinking code path. This set is the built-in # default; it can be replaced via the `claude_extended_thinking_models_override` # configuration option. -CLAUDE_EXTENDED_THINKING_MODELS = [ - "anthropic/claude-3-7-sonnet-20250219", +CLAUDE_EXTENDED_THINKING_MODELS = { "claude-3-7-sonnet-20250219", - "anthropic/claude-sonnet-4-6", "claude-sonnet-4-6", - "vertex_ai/claude-sonnet-4-6", - "bedrock/anthropic.claude-sonnet-4-6", - "bedrock/us.anthropic.claude-sonnet-4-6", - "bedrock/au.anthropic.claude-sonnet-4-6", - "bedrock/eu.anthropic.claude-sonnet-4-6", - "bedrock/jp.anthropic.claude-sonnet-4-6", - "bedrock/global.anthropic.claude-sonnet-4-6", - "anthropic/claude-sonnet-4-5-20250929", "claude-sonnet-4-5-20250929", - "vertex_ai/claude-sonnet-4-5@20250929", - "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/au.anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/jp.anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/global.anthropic.claude-sonnet-4-5-20250929-v1:0", - "anthropic/claude-opus-4-5-20251101", "claude-opus-4-5-20251101", - "vertex_ai/claude-opus-4-5@20251101", - "bedrock/anthropic.claude-opus-4-5-20251101-v1:0", - "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0", - "bedrock/au.anthropic.claude-opus-4-5-20251101-v1:0", - "bedrock/eu.anthropic.claude-opus-4-5-20251101-v1:0", - "bedrock/jp.anthropic.claude-opus-4-5-20251101-v1:0", - "bedrock/global.anthropic.claude-opus-4-5-20251101-v1:0", - "anthropic/claude-opus-4-6", "claude-opus-4-6", - "vertex_ai/claude-opus-4-6", - "bedrock/anthropic.claude-opus-4-6-v1:0", - "bedrock/us.anthropic.claude-opus-4-6-v1:0", - "bedrock/au.anthropic.claude-opus-4-6-v1:0", - "bedrock/eu.anthropic.claude-opus-4-6-v1:0", - "bedrock/jp.anthropic.claude-opus-4-6-v1:0", - "bedrock/global.anthropic.claude-opus-4-6-v1:0", - "anthropic/claude-haiku-4-5-20251001", "claude-haiku-4-5-20251001", - "vertex_ai/claude-haiku-4-5@20251001", - "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/au.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/eu.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/jp.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/global.anthropic.claude-haiku-4-5-20251001-v1:0", -] +} # Models that require streaming mode -STREAMING_REQUIRED_MODELS = [ - "openai/qwq-plus" -] +STREAMING_REQUIRED_MODELS = { + "qwq-plus", +} diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 6a306bc90a..888dbb62d8 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -14,7 +14,7 @@ NO_SUPPORT_TEMPERATURE_MODELS, STREAMING_REQUIRED_MODELS, SUPPORT_REASONING_EFFORT_MODELS, - USER_MESSAGE_ONLY_MODELS) + USER_MESSAGE_ONLY_MODELS, base_model_name) from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_helpers import ( MockResponse, _get_azure_ad_token, _handle_streaming_response, @@ -257,11 +257,11 @@ def __init__(self): "Falling back to the built-in Claude extended-thinking model list." ) override = [] - # Store stripped names so exact-match checks against the model succeed even when the config - # entries contain surrounding whitespace (validation above already used model.strip()). - self.claude_extended_thinking_models = ( - [model.strip() for model in override] if override else CLAUDE_EXTENDED_THINKING_MODELS - ) + # Normalize to base model names so membership checks match regardless of provider prefix, + # whether the entries come from config override or the built-in default. + self.claude_extended_thinking_models = { + base_model_name(model.strip()) for model in override + } if override else set(CLAUDE_EXTENDED_THINKING_MODELS) # Models that require streaming self.streaming_required_models = STREAMING_REQUIRED_MODELS @@ -508,7 +508,7 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: # Currently, some models do not support a separate system and user prompts - if model in self.user_message_only_models or get_settings().config.custom_reasoning_model: + if base_model_name(model) in self.user_message_only_models or get_settings().config.custom_reasoning_model: user = f"{system}\n\n\n{user}" system = "" get_logger().info(f"Using model {model}, combining system and user prompts") @@ -528,7 +528,7 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: } # Add temperature only if model supports it - if model not in self.no_support_temperature_models and not get_settings().config.custom_reasoning_model: + if base_model_name(model) not in self.no_support_temperature_models and not get_settings().config.custom_reasoning_model: # get_logger().info(f"Adding temperature with value {temperature} to model {model}.") kwargs["temperature"] = temperature @@ -538,7 +538,7 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: del kwargs['temperature'] # Add reasoning_effort if model supports it - if model in self.support_reasoning_models: + if base_model_name(model) in self.support_reasoning_models: config_effort = get_settings().config.reasoning_effort try: ReasoningEffort(config_effort) @@ -555,7 +555,7 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: kwargs["reasoning_effort"] = reasoning_effort # https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking - if (model in self.claude_extended_thinking_models) and get_settings().config.get("enable_claude_extended_thinking", False): + if (base_model_name(model) in self.claude_extended_thinking_models) and get_settings().config.get("enable_claude_extended_thinking", False): kwargs = self._configure_claude_extended_thinking(model, kwargs) if get_settings().litellm.get("enable_callbacks", False): @@ -642,7 +642,7 @@ async def _get_completion(self, **kwargs): Wrapper that automatically handles streaming for required models. """ model = kwargs["model"] - if model in self.streaming_required_models: + if base_model_name(model) in self.streaming_required_models: kwargs["stream"] = True get_logger().info(f"Using streaming mode for model {model}") response = await acompletion(**kwargs) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index cb313f023f..cf4d767c82 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -99,10 +99,10 @@ def _get_system_user_tokens(self, pr, encoder, vars: dict, system, user): def _calc_claude_tokens(self, patch: str) -> int: try: import anthropic - from pr_agent.algo import MAX_TOKENS - + from pr_agent.algo.utils import get_max_tokens + client = anthropic.Anthropic(api_key=get_settings(use_context=False).get('anthropic.key')) - max_tokens = MAX_TOKENS[get_settings().config.model] + max_tokens = get_max_tokens(get_settings().config.model) if len(patch.encode('utf-8')) > self.CLAUDE_MAX_CONTENT_SIZE: get_logger().warning( diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index 35bf9e0a36..dfc929ced1 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -989,15 +989,27 @@ def get_user_labels(current_labels: List[str] = None): return user_labels -def get_max_tokens(model): +def _litellm_max_tokens(model): + """Context window from litellm's built-in registry; None if litellm doesn't know the model.""" + try: + import litellm + info = litellm.get_model_info(model) + return info.get("max_input_tokens") or info.get("max_tokens") + except Exception: + return None + + +def get_max_tokens(model, cap=True): """ Get the maximum number of tokens allowed for a model. logic: (1) If the model is in './pr_agent/algo/__init__.py', use the value from there. - (2) else, the user needs to define explicitly 'config.custom_model_max_tokens' + (2) else, if 'config.custom_model_max_tokens' is set explicitly, use it. + (3) else, fall back to litellm's built-in context-window registry. - For both cases, we further limit the number of tokens to 'config.max_model_tokens' if it is set. + When cap=True, further limit the number of tokens to 'config.max_model_tokens' if it is set. This aims to improve the algorithmic quality, as the AI model degrades in performance when the input is too long. + Pass cap=False when the uncapped context window is wanted (e.g. fitting full documentation into the prompt). """ settings = get_settings() if model in MAX_TOKENS: @@ -1005,10 +1017,12 @@ def get_max_tokens(model): elif settings.config.custom_model_max_tokens > 0: max_tokens_model = settings.config.custom_model_max_tokens else: - get_logger().error(f"Model {model} is not defined in MAX_TOKENS in ./pr_agent/algo/__init__.py and no custom_model_max_tokens is set") - raise Exception(f"Ensure {model} is defined in MAX_TOKENS in ./pr_agent/algo/__init__.py or set a positive value for it in config.custom_model_max_tokens") + max_tokens_model = _litellm_max_tokens(model) + if not max_tokens_model: + get_logger().error(f"Model {model} is not defined in MAX_TOKENS in ./pr_agent/algo/__init__.py, not known to litellm, and no custom_model_max_tokens is set") + raise Exception(f"Ensure {model} is defined in MAX_TOKENS in ./pr_agent/algo/__init__.py or set a positive value for it in config.custom_model_max_tokens") - if settings.config.max_model_tokens and settings.config.max_model_tokens > 0: + if cap and settings.config.max_model_tokens and settings.config.max_model_tokens > 0: max_tokens_model = min(settings.config.max_model_tokens, max_tokens_model) return max_tokens_model diff --git a/pr_agent/tools/pr_code_suggestions.py b/pr_agent/tools/pr_code_suggestions.py index f950ce9e01..fc5a163563 100644 --- a/pr_agent/tools/pr_code_suggestions.py +++ b/pr_agent/tools/pr_code_suggestions.py @@ -10,7 +10,6 @@ from jinja2 import Environment, StrictUndefined -from pr_agent.algo import MAX_TOKENS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler from pr_agent.algo.git_patch_processing import decouple_and_convert_to_hunks_with_lines_numbers @@ -758,11 +757,8 @@ async def convert_to_decoupled_with_line_numbers(self, patches_diff_list_no_line file=None).strip() patches_new[i] = patches_new[i].strip() patch_final = "\n\n\n".join(patches_new) - if model in MAX_TOKENS: - max_tokens_full = MAX_TOKENS[ - model] # note - here we take the actual max tokens, without any reductions. we do aim to get the full documentation website in the prompt - else: - max_tokens_full = get_max_tokens(model) + # uncapped context window - we aim to fit the full content into the prompt + max_tokens_full = get_max_tokens(model, cap=False) delta_output = 2000 token_count = self.token_handler.count_tokens(patch_final) if token_count > max_tokens_full - delta_output: diff --git a/pr_agent/tools/pr_help_docs.py b/pr_agent/tools/pr_help_docs.py index aa7002189c..25609bdd67 100644 --- a/pr_agent/tools/pr_help_docs.py +++ b/pr_agent/tools/pr_help_docs.py @@ -7,7 +7,6 @@ import re from tempfile import TemporaryDirectory -from pr_agent.algo import MAX_TOKENS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler from pr_agent.algo.pr_processing import retry_with_fallback_models @@ -466,11 +465,8 @@ def _trim_docs_input(self, docs_input: str, max_allowed_txt_input: int, only_ret token_count = self.token_handler.count_tokens(docs_input, force_accurate=True) get_logger().debug(f"Estimated token count of documentation to send to model: {token_count}") model = get_settings().config.model - if model in MAX_TOKENS: - max_tokens_full = MAX_TOKENS[ - model] # note - here we take the actual max tokens, without any reductions. we do aim to get the full documentation website in the prompt - else: - max_tokens_full = get_max_tokens(model) + # uncapped context window - we aim to fit the full documentation website into the prompt + max_tokens_full = get_max_tokens(model, cap=False) delta_output = 5000 # Elbow room to reduce chance of exceeding token limit or model paying less attention to prompt guidelines. if token_count > max_tokens_full - delta_output: if only_return_if_trim_needed: diff --git a/pr_agent/tools/pr_help_message.py b/pr_agent/tools/pr_help_message.py index a83e03e0c7..4b0dee84d1 100644 --- a/pr_agent/tools/pr_help_message.py +++ b/pr_agent/tools/pr_help_message.py @@ -5,7 +5,6 @@ from jinja2 import Environment, StrictUndefined -from pr_agent.algo import MAX_TOKENS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler from pr_agent.algo.pr_processing import retry_with_fallback_models @@ -134,10 +133,8 @@ async def run(self): get_logger().debug(f"Token count of full documentation website: {token_count}") model = get_settings().config.model - if model in MAX_TOKENS: - max_tokens_full = MAX_TOKENS[model] # note - here we take the actual max tokens, without any reductions. we do aim to get the full documentation website in the prompt - else: - max_tokens_full = get_max_tokens(model) + # uncapped context window - we aim to fit the full documentation website into the prompt + max_tokens_full = get_max_tokens(model, cap=False) delta_output = 2000 if token_count > max_tokens_full - delta_output: get_logger().info(f"Token count {token_count} exceeds the limit {max_tokens_full - delta_output}. Skipping the PR Help message.") diff --git a/pr_agent/tools/pr_similar_issue.py b/pr_agent/tools/pr_similar_issue.py index f48756d793..82cc9953f8 100644 --- a/pr_agent/tools/pr_similar_issue.py +++ b/pr_agent/tools/pr_similar_issue.py @@ -5,7 +5,6 @@ import openai from pydantic import BaseModel, Field -from pr_agent.algo import MAX_TOKENS from pr_agent.algo.token_handler import TokenHandler from pr_agent.algo.utils import get_max_tokens from pr_agent.config_loader import get_settings @@ -443,7 +442,7 @@ def _update_index_with_issues(self, issues_list, repo_name_for_index, upsert=Fal continue if len(comment_body) < 8000 or \ - self.token_handler.count_tokens(comment_body) < MAX_TOKENS[MODEL]: + self.token_handler.count_tokens(comment_body) < get_max_tokens(MODEL): comment_record = Record( id=issue_key + ".comment_" + str(j + 1), text=comment_body, @@ -539,7 +538,7 @@ def _update_table_with_issues(self, issues_list, repo_name_for_index, ingest=Fal continue if len(comment_body) < 8000 or \ - self.token_handler.count_tokens(comment_body) < MAX_TOKENS[MODEL]: + self.token_handler.count_tokens(comment_body) < get_max_tokens(MODEL): comment_record = Record( id=issue_key + ".comment_" + str(j + 1), text=comment_body, @@ -637,7 +636,7 @@ def _update_qdrant_with_issues(self, issues_list, repo_name_for_index, ingest=Fa continue if len(comment_body) < 8000 or \ - self.token_handler.count_tokens(comment_body) < MAX_TOKENS[MODEL]: + self.token_handler.count_tokens(comment_body) < get_max_tokens(MODEL): comment_record = Record( id=issue_key + ".comment_" + str(j + 1), text=comment_body, diff --git a/requirements.txt b/requirements.txt index a9c58f0b69..67edefda84 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,7 +16,7 @@ google-cloud-storage==2.10.0; python_version < "3.13" google-cloud-storage==3.12.0; python_version >= "3.13" protobuf==6.33.6 # grpcio-status 1.80.0 forces >=6.31.1,<7; pinned to lock the resolved transitive version Jinja2==3.1.6 -litellm==1.84.0 +litellm==1.91.0 loguru==0.7.2 msrest==0.7.1 openai>=1.55.3 diff --git a/tests/unittest/test_base_model_name.py b/tests/unittest/test_base_model_name.py new file mode 100644 index 0000000000..e9b44e3619 --- /dev/null +++ b/tests/unittest/test_base_model_name.py @@ -0,0 +1,45 @@ +from pr_agent.algo import ( + CLAUDE_EXTENDED_THINKING_MODELS, + NO_SUPPORT_TEMPERATURE_MODELS, + STREAMING_REQUIRED_MODELS, + base_model_name, +) + + +class TestBaseModelName: + def test_provider_prefix_stripped(self): + """anthropic/, bedrock/ (with region+vendor+version), and vertex_ai/ collapse to one base""" + for variant in [ + "claude-opus-4-8", + "anthropic/claude-opus-4-8", + "vertex_ai/claude-opus-4-8", + "bedrock/anthropic.claude-opus-4-8", + "bedrock/us.anthropic.claude-opus-4-8", + "bedrock/global.anthropic.claude-opus-4-8", + ]: + assert base_model_name(variant) == "claude-opus-4-8" + + def test_date_and_cloud_version_suffixes(self): + """vertex '@date' and bedrock '-v1:0' normalize to the anthropic spelling""" + expected = "claude-sonnet-4-5-20250929" + for variant in [ + "claude-sonnet-4-5-20250929", + "anthropic/claude-sonnet-4-5-20250929", + "vertex_ai/claude-sonnet-4-5@20250929", + "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + ]: + assert base_model_name(variant) == expected + + def test_dotted_version_preserved(self): + """A model version containing a dot must not be truncated""" + assert base_model_name("gpt-5.2-codex") == "gpt-5.2-codex" + assert base_model_name("openai/qwq-plus") == "qwq-plus" + + def test_capability_set_membership(self): + """Provider-prefixed models resolve into the base-name capability sets""" + assert base_model_name("bedrock/eu.anthropic.claude-opus-4-8") in NO_SUPPORT_TEMPERATURE_MODELS + assert base_model_name("vertex_ai/claude-opus-4-6") in CLAUDE_EXTENDED_THINKING_MODELS + assert base_model_name("openai/qwq-plus") in STREAMING_REQUIRED_MODELS + # a model with no special capability resolves out of every set + assert base_model_name("gpt-4o") not in NO_SUPPORT_TEMPERATURE_MODELS diff --git a/tests/unittest/test_litellm_claude_extended_thinking.py b/tests/unittest/test_litellm_claude_extended_thinking.py index 3925b562bd..8137ffa55e 100644 --- a/tests/unittest/test_litellm_claude_extended_thinking.py +++ b/tests/unittest/test_litellm_claude_extended_thinking.py @@ -55,17 +55,18 @@ def test_valid_claude_extended_thinking_override_replaces_built_in_models(monkey handler = litellm_handler.LiteLLMAIHandler() - assert handler.claude_extended_thinking_models == ["custom-claude-model"] + assert handler.claude_extended_thinking_models == {"custom-claude-model"} assert handler.claude_extended_thinking_models is not override logger.warning.assert_not_called() def test_claude_extended_thinking_override_entries_are_stripped(monkeypatch, logger): - # Entries with surrounding whitespace must be stored stripped so exact model matches succeed. + # Entries with surrounding whitespace are stripped and normalized to base model names so + # membership checks succeed regardless of provider prefix. override = [" custom-claude-model ", "another-model\n"] monkeypatch.setattr(litellm_handler, "get_settings", lambda: settings_with_claude_override(override)) handler = litellm_handler.LiteLLMAIHandler() - assert handler.claude_extended_thinking_models == ["custom-claude-model", "another-model"] + assert handler.claude_extended_thinking_models == {"custom-claude-model", "another-model"} logger.warning.assert_not_called() From 684d529e1266fbdc7eed026cb9b31c5533f764ff Mon Sep 17 00:00:00 2001 From: Shmulik Cohen Date: Tue, 7 Jul 2026 00:01:20 +0300 Subject: [PATCH 2/4] feat(algo): never fail on unknown models; resolve context window via litellm get_max_tokens now resolves the input context window as: MAX_TOKENS override (exact or base model name) -> litellm's built-in, self-updating registry (exact or base model name) -> config.custom_model_max_tokens -> a conservative default (DEFAULT_MODEL_MAX_TOKENS) with a warning. It no longer raises on a model that is absent from MAX_TOKENS, so any model runs without a code edit. The MAX_TOKENS table stays as the authoritative, deterministic source for the models it lists (checked first); litellm and the default only cover models that previously errored out. base_model_name lets the litellm lookup resolve provider-prefixed / region / cloud-version spellings (e.g. bedrock/us.anthropic.claude-sonnet-4-6-v1:0 -> claude-sonnet-4-6). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01SqHrpoPezxDkD5dg4VuQL5 --- pr_agent/algo/utils.py | 60 +++++++++++++++++---------- tests/unittest/test_get_max_tokens.py | 9 ++-- 2 files changed, 43 insertions(+), 26 deletions(-) diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index dfc929ced1..610efbfd0f 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -23,7 +23,7 @@ from pydantic import BaseModel from starlette_context import context -from pr_agent.algo import MAX_TOKENS +from pr_agent.algo import MAX_TOKENS, base_model_name from pr_agent.algo.git_patch_processing import extract_hunk_lines_from_patch from pr_agent.algo.token_handler import TokenEncoder from pr_agent.algo.types import FilePatchInfo @@ -989,38 +989,54 @@ def get_user_labels(current_labels: List[str] = None): return user_labels +# Context window used when a model is unknown to both the MAX_TOKENS overrides and litellm. +# Conservative but non-zero so any model still runs (token clipping stays sane) rather than failing. +DEFAULT_MODEL_MAX_TOKENS = 32000 + + def _litellm_max_tokens(model): - """Context window from litellm's built-in registry; None if litellm doesn't know the model.""" - try: - import litellm - info = litellm.get_model_info(model) - return info.get("max_input_tokens") or info.get("max_tokens") - except Exception: - return None + """Context window from litellm's built-in, self-updating registry; None if litellm doesn't know it. + + Tries the exact model string first, then the base model name (see base_model_name) since litellm + keys its registry by canonical names and often doesn't recognize provider-prefixed / region / + cloud-version spellings (e.g. 'bedrock/us.anthropic.claude-sonnet-4-6-v1:0' -> 'claude-sonnet-4-6'). + """ + import litellm + for name in (model, base_model_name(model)): + try: + info = litellm.get_model_info(name) + ctx = info.get("max_input_tokens") or info.get("max_tokens") + if ctx: + return ctx + except Exception: + continue + return None def get_max_tokens(model, cap=True): """ - Get the maximum number of tokens allowed for a model. - logic: - (1) If the model is in './pr_agent/algo/__init__.py', use the value from there. - (2) else, if 'config.custom_model_max_tokens' is set explicitly, use it. - (3) else, fall back to litellm's built-in context-window registry. + Get the maximum number of input tokens allowed for a model. Resolution order: + (1) MAX_TOKENS in './pr_agent/algo/__init__.py' — a small override layer for values litellm + gets wrong or lacks (matched by exact then base model name). + (2) litellm's built-in, self-updating model registry (exact then base model name). + (3) 'config.custom_model_max_tokens' if set (>0) — user-declared budget for private/unknown models. + (4) DEFAULT_MODEL_MAX_TOKENS with a warning — never raises, so any model runs. When cap=True, further limit the number of tokens to 'config.max_model_tokens' if it is set. This aims to improve the algorithmic quality, as the AI model degrades in performance when the input is too long. Pass cap=False when the uncapped context window is wanted (e.g. fitting full documentation into the prompt). """ settings = get_settings() - if model in MAX_TOKENS: - max_tokens_model = MAX_TOKENS[model] - elif settings.config.custom_model_max_tokens > 0: - max_tokens_model = settings.config.custom_model_max_tokens - else: - max_tokens_model = _litellm_max_tokens(model) - if not max_tokens_model: - get_logger().error(f"Model {model} is not defined in MAX_TOKENS in ./pr_agent/algo/__init__.py, not known to litellm, and no custom_model_max_tokens is set") - raise Exception(f"Ensure {model} is defined in MAX_TOKENS in ./pr_agent/algo/__init__.py or set a positive value for it in config.custom_model_max_tokens") + max_tokens_model = (MAX_TOKENS.get(model) or MAX_TOKENS.get(base_model_name(model)) + or _litellm_max_tokens(model)) + if not max_tokens_model: + if settings.config.custom_model_max_tokens > 0: + max_tokens_model = settings.config.custom_model_max_tokens + else: + max_tokens_model = DEFAULT_MODEL_MAX_TOKENS + get_logger().warning( + f"Model {model} is unknown to MAX_TOKENS and litellm; using default {DEFAULT_MODEL_MAX_TOKENS}. " + f"Set config.custom_model_max_tokens to specify its context window.") if cap and settings.config.max_model_tokens and settings.config.max_model_tokens > 0: max_tokens_model = min(settings.config.max_model_tokens, max_tokens_model) diff --git a/tests/unittest/test_get_max_tokens.py b/tests/unittest/test_get_max_tokens.py index 7474b226e8..bd18f9fe22 100644 --- a/tests/unittest/test_get_max_tokens.py +++ b/tests/unittest/test_get_max_tokens.py @@ -1,7 +1,7 @@ import pytest import pr_agent.algo.utils as utils -from pr_agent.algo.utils import MAX_TOKENS, get_max_tokens +from pr_agent.algo.utils import DEFAULT_MODEL_MAX_TOKENS, MAX_TOKENS, get_max_tokens class TestGetMaxTokens: @@ -110,6 +110,8 @@ def test_gpt_codex_models_max_tokens(self, monkeypatch, model): assert get_max_tokens(model) == expected def test_model_not_max_tokens_and_not_has_custom(self, monkeypatch): + # A model unknown to both MAX_TOKENS and litellm, with no custom_model_max_tokens set, + # falls back to the conservative default instead of raising, so any model still runs. fake_settings = type('', (), { 'config': type('', (), { 'custom_model_max_tokens': 0, @@ -119,10 +121,9 @@ def test_model_not_max_tokens_and_not_has_custom(self, monkeypatch): monkeypatch.setattr(utils, "get_settings", lambda: fake_settings) - model = "custom-model" + model = "some-model-unknown-to-litellm-and-max-tokens" - with pytest.raises(Exception): - get_max_tokens(model) + assert get_max_tokens(model) == DEFAULT_MODEL_MAX_TOKENS def test_model_max_tokens_with__limit(self, monkeypatch): fake_settings = type('', (), { From d73a9d1956da8ea5870f75b2460977882942002f Mon Sep 17 00:00:00 2001 From: Shmulik Cohen Date: Tue, 7 Jul 2026 00:08:16 +0300 Subject: [PATCH 3/4] fix(algo): correct stale context windows; drop retired/Claude-3 models Correct MAX_TOKENS values verified wrong against OpenAI/Anthropic docs and litellm's registry: - gpt-5 / gpt-5.1 200K -> 272K, gpt-5.2 / gpt-5.2-codex / gpt-5.3-codex 400K -> 272K (272K is the enforced input cap; 400K over-budgeted and risked 400 errors), gpt-5.1-chat-latest 200K -> 128K (200K exceeded the model). - o1 / o1-2024-12-17 / o3-mini 204800 -> 200000 (204800 exceeded the real limit). - claude-sonnet-4-6 / claude-opus-4-6 (all provider spellings) 200K -> 1M. Drop retired / no-longer-used models (litellm + the never-fail fallback still serve any stragglers): the whole Claude 3.x line (incl. 3.7), Claude 2 / instant, the retired Opus 4 / Sonnet 4 20250514 snapshots, and dead gpt-3.5 snapshots. Also remove claude-3-7-sonnet from CLAUDE_EXTENDED_THINKING_MODELS. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01SqHrpoPezxDkD5dg4VuQL5 --- pr_agent/algo/__init__.py | 142 +++++++++----------------- tests/unittest/test_get_max_tokens.py | 4 +- 2 files changed, 51 insertions(+), 95 deletions(-) diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py index dae7ccef5d..6af24b9c41 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -26,10 +26,8 @@ def base_model_name(model: str) -> str: 'text-embedding-ada-002': 8000, 'gpt-3.5-turbo': 16000, 'gpt-3.5-turbo-0125': 16000, - 'gpt-3.5-turbo-0613': 4000, 'gpt-3.5-turbo-1106': 16000, 'gpt-3.5-turbo-16k': 16000, - 'gpt-3.5-turbo-16k-0613': 16000, 'gpt-4': 8000, 'gpt-4-0613': 8000, 'gpt-4-32k': 32000, @@ -52,20 +50,20 @@ def base_model_name(model: str) -> str: 'gpt-4.1-mini-2025-04-14': 1047576, 'gpt-4.1-nano': 1047576, 'gpt-4.1-nano-2025-04-14': 1047576, - 'gpt-5-nano': 200000, # 200K, but may be limited by config.max_model_tokens - 'gpt-5-mini': 200000, # 200K, but may be limited by config.max_model_tokens - 'gpt-5': 200000, - 'gpt-5-2025-08-07': 200000, - 'gpt-5.1': 200000, - 'gpt-5.1-2025-11-13': 200000, - 'gpt-5.1-chat-latest': 200000, - 'gpt-5.1-codex': 200000, - 'gpt-5.1-codex-mini': 200000, - 'gpt-5.2': 400000, # 400K, but may be limited by config.max_model_tokens - 'gpt-5.2-2025-12-11': 400000, # 400K, but may be limited by config.max_model_tokens + 'gpt-5-nano': 272000, # 272K, but may be limited by config.max_model_tokens + 'gpt-5-mini': 272000, # 272K, but may be limited by config.max_model_tokens + 'gpt-5': 272000, + 'gpt-5-2025-08-07': 272000, + 'gpt-5.1': 272000, + 'gpt-5.1-2025-11-13': 272000, + 'gpt-5.1-chat-latest': 128000, + 'gpt-5.1-codex': 272000, + 'gpt-5.1-codex-mini': 272000, + 'gpt-5.2': 272000, # 272K, but may be limited by config.max_model_tokens + 'gpt-5.2-2025-12-11': 272000, # 272K, but may be limited by config.max_model_tokens 'gpt-5.2-chat-latest': 128000, # 128K, but may be limited by config.max_model_tokens - 'gpt-5.2-codex': 400000, # 400K, but may be limited by config.max_model_tokens - 'gpt-5.3-codex': 400000, # 400K, but may be limited by config.max_model_tokens + 'gpt-5.2-codex': 272000, # 272K, but may be limited by config.max_model_tokens + 'gpt-5.3-codex': 272000, # 272K, but may be limited by config.max_model_tokens 'gpt-5.3-chat': 128000, # 128K, but may be limited by config.max_model_tokens 'gpt-5.4': 272000, # 272K safe default without opt-in 1M context parameters 'gpt-5.4-2026-03-05': 272000, # 272K safe default without opt-in 1M context parameters @@ -79,16 +77,14 @@ def base_model_name(model: str) -> str: 'o1-mini-2024-09-12': 128000, # 128K, but may be limited by config.max_model_tokens 'o1-preview': 128000, # 128K, but may be limited by config.max_model_tokens 'o1-preview-2024-09-12': 128000, # 128K, but may be limited by config.max_model_tokens - 'o1-2024-12-17': 204800, # 200K, but may be limited by config.max_model_tokens - 'o1': 204800, # 200K, but may be limited by config.max_model_tokens - 'o3-mini': 204800, # 200K, but may be limited by config.max_model_tokens - 'o3-mini-2025-01-31': 204800, # 200K, but may be limited by config.max_model_tokens + 'o1-2024-12-17': 200000, # 200K, but may be limited by config.max_model_tokens + 'o1': 200000, # 200K, but may be limited by config.max_model_tokens + 'o3-mini': 200000, # 200K, but may be limited by config.max_model_tokens + 'o3-mini-2025-01-31': 200000, # 200K, but may be limited by config.max_model_tokens 'o3': 200000, # 200K, but may be limited by config.max_model_tokens 'o3-2025-04-16': 200000, # 200K, but may be limited by config.max_model_tokens 'o4-mini': 200000, # 200K, but may be limited by config.max_model_tokens 'o4-mini-2025-04-16': 200000, # 200K, but may be limited by config.max_model_tokens - 'claude-instant-1': 100000, - 'claude-2': 100000, 'command-nightly': 4096, 'deepseek/deepseek-chat': 128000, # 128K, but may be limited by config.max_model_tokens 'deepseek/deepseek-reasoner': 64000, # 64K, but may be limited by config.max_model_tokens @@ -99,24 +95,15 @@ def base_model_name(model: str) -> str: 'meta-llama/Llama-2-7b-chat-hf': 4096, 'vertex_ai/codechat-bison': 6144, 'vertex_ai/codechat-bison-32k': 32000, - 'vertex_ai/claude-3-haiku@20240307': 100000, - 'vertex_ai/claude-3-5-haiku@20241022': 100000, 'vertex_ai/claude-haiku-4-5@20251001': 200000, - 'vertex_ai/claude-3-sonnet@20240229': 100000, - 'vertex_ai/claude-3-opus@20240229': 100000, - 'vertex_ai/claude-opus-4@20250514': 200000, 'vertex_ai/claude-opus-4-1@20250805': 200000, 'vertex_ai/claude-opus-4-5@20251101': 200000, - 'vertex_ai/claude-opus-4-6@20260120': 200000, - 'vertex_ai/claude-opus-4-6': 200000, + 'vertex_ai/claude-opus-4-6@20260120': 1000000, + 'vertex_ai/claude-opus-4-6': 1000000, 'vertex_ai/claude-opus-4-7': 1000000, 'vertex_ai/claude-opus-4-8': 1000000, - 'vertex_ai/claude-3-5-sonnet@20240620': 100000, - 'vertex_ai/claude-3-5-sonnet-v2@20241022': 100000, - 'vertex_ai/claude-3-7-sonnet@20250219': 200000, - 'vertex_ai/claude-sonnet-4@20250514': 200000, 'vertex_ai/claude-sonnet-4-5@20250929': 200000, - 'vertex_ai/claude-sonnet-4-6': 200000, + 'vertex_ai/claude-sonnet-4-6': 1000000, 'vertex_ai/claude-sonnet-5': 1000000, 'vertex_ai/gemini-1.5-pro': 1048576, 'vertex_ai/gemini-2.5-pro-preview-03-25': 1048576, @@ -157,74 +144,52 @@ def base_model_name(model: str) -> str: 'gemini/gemini-3.5-pro': 1048576, 'codechat-bison': 6144, 'codechat-bison-32k': 32000, - 'anthropic.claude-instant-v1': 100000, - 'anthropic.claude-v1': 100000, - 'anthropic.claude-v2': 100000, - 'anthropic/claude-3-opus-20240229': 100000, - 'anthropic/claude-opus-4-20250514': 200000, 'anthropic/claude-opus-4-1-20250805': 200000, 'anthropic/claude-opus-4-5-20251101': 200000, - 'anthropic/claude-opus-4-6': 200000, - 'anthropic/claude-opus-4-6-20260120': 200000, + 'anthropic/claude-opus-4-6': 1000000, + 'anthropic/claude-opus-4-6-20260120': 1000000, 'anthropic/claude-opus-4-7': 1000000, 'anthropic/claude-opus-4-8': 1000000, - 'anthropic/claude-3-5-sonnet-20240620': 100000, - 'anthropic/claude-3-5-sonnet-20241022': 100000, - 'anthropic/claude-3-7-sonnet-20250219': 200000, - 'anthropic/claude-sonnet-4-20250514': 200000, 'anthropic/claude-sonnet-4-5-20250929': 200000, - 'anthropic/claude-sonnet-4-6': 200000, + 'anthropic/claude-sonnet-4-6': 1000000, 'anthropic/claude-sonnet-5': 1000000, 'claude-opus-4-1-20250805': 200000, 'claude-opus-4-5-20251101': 200000, - 'claude-opus-4-6': 200000, - 'claude-opus-4-6-20260120': 200000, + 'claude-opus-4-6': 1000000, + 'claude-opus-4-6-20260120': 1000000, 'claude-opus-4-7': 1000000, 'claude-opus-4-8': 1000000, - 'claude-3-7-sonnet-20250219': 200000, - 'claude-sonnet-4-6': 200000, + 'claude-sonnet-4-6': 1000000, 'claude-sonnet-5': 1000000, - 'anthropic/claude-3-5-haiku-20241022': 100000, 'anthropic/claude-haiku-4-5-20251001': 200000, 'claude-haiku-4-5-20251001': 200000, - 'bedrock/anthropic.claude-instant-v1': 100000, - 'bedrock/anthropic.claude-v2': 100000, 'bedrock/anthropic.claude-v2:1': 100000, - 'bedrock/anthropic.claude-3-sonnet-20240229-v1:0': 100000, - 'bedrock/anthropic.claude-opus-4-20250514-v1:0': 200000, 'bedrock/anthropic.claude-opus-4-1-20250805-v1:0': 200000, - 'bedrock/anthropic.claude-opus-4-6-20260120-v1:0': 200000, - 'bedrock/anthropic.claude-opus-4-6-v1:0': 200000, + 'bedrock/anthropic.claude-opus-4-6-20260120-v1:0': 1000000, + 'bedrock/anthropic.claude-opus-4-6-v1:0': 1000000, 'bedrock/anthropic.claude-opus-4-7': 1000000, 'bedrock/anthropic.claude-opus-4-7-v1:0': 1000000, 'bedrock/anthropic.claude-opus-4-8': 1000000, - 'bedrock/anthropic.claude-3-haiku-20240307-v1:0': 100000, - 'bedrock/anthropic.claude-3-5-haiku-20241022-v1:0': 100000, 'bedrock/anthropic.claude-haiku-4-5-20251001-v1:0': 200000, - 'bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0': 100000, - 'bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0': 100000, - 'bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0': 200000, - 'bedrock/anthropic.claude-sonnet-4-20250514-v1:0': 200000, 'bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0': 200000, - 'bedrock/anthropic.claude-sonnet-4-6': 200000, - 'bedrock/anthropic.claude-sonnet-4-6-v1:0': 200000, + 'bedrock/anthropic.claude-sonnet-4-6': 1000000, + 'bedrock/anthropic.claude-sonnet-4-6-v1:0': 1000000, 'bedrock/anthropic.claude-sonnet-5': 1000000, 'bedrock/anthropic.claude-opus-4-5-20251101-v1:0': 200000, - "bedrock/us.anthropic.claude-opus-4-20250514-v1:0": 200000, "bedrock/us.anthropic.claude-opus-4-1-20250805-v1:0": 200000, - "bedrock/us.anthropic.claude-opus-4-6-20260120-v1:0": 200000, + "bedrock/us.anthropic.claude-opus-4-6-20260120-v1:0": 1000000, "bedrock/global.anthropic.claude-opus-4-5-20251101-v1:0": 200000, "bedrock/eu.anthropic.claude-opus-4-5-20251101-v1:0": 200000, "bedrock/au.anthropic.claude-opus-4-5-20251101-v1:0": 200000, "bedrock/jp.anthropic.claude-opus-4-5-20251101-v1:0": 200000, "bedrock/apac.anthropic.claude-opus-4-5-20251101-v1:0": 200000, "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0": 200000, - "bedrock/global.anthropic.claude-opus-4-6-v1:0": 200000, - "bedrock/eu.anthropic.claude-opus-4-6-v1:0": 200000, - "bedrock/au.anthropic.claude-opus-4-6-v1:0": 200000, - "bedrock/jp.anthropic.claude-opus-4-6-v1:0": 200000, - "bedrock/apac.anthropic.claude-opus-4-6-v1:0": 200000, - "bedrock/us.anthropic.claude-opus-4-6-v1:0": 200000, + "bedrock/global.anthropic.claude-opus-4-6-v1:0": 1000000, + "bedrock/eu.anthropic.claude-opus-4-6-v1:0": 1000000, + "bedrock/au.anthropic.claude-opus-4-6-v1:0": 1000000, + "bedrock/jp.anthropic.claude-opus-4-6-v1:0": 1000000, + "bedrock/apac.anthropic.claude-opus-4-6-v1:0": 1000000, + "bedrock/us.anthropic.claude-opus-4-6-v1:0": 1000000, "bedrock/global.anthropic.claude-opus-4-7": 1000000, "bedrock/us.anthropic.claude-opus-4-7": 1000000, "bedrock/global.anthropic.claude-opus-4-8": 1000000, @@ -232,42 +197,34 @@ def base_model_name(model: str) -> str: "bedrock/eu.anthropic.claude-opus-4-8": 1000000, "bedrock/au.anthropic.claude-opus-4-8": 1000000, "bedrock/jp.anthropic.claude-opus-4-8": 1000000, - "bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0": 100000, "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, "bedrock/eu.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, "bedrock/au.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, "bedrock/jp.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, "bedrock/apac.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, "bedrock/global.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, - "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0": 200000, - "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0": 200000, - "bedrock/global.anthropic.claude-sonnet-4-20250514-v1:0": 200000, "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, "bedrock/au.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/us.anthropic.claude-sonnet-4-6": 200000, - "bedrock/us.anthropic.claude-sonnet-4-6-v1:0": 200000, - "bedrock/au.anthropic.claude-sonnet-4-6": 200000, - "bedrock/au.anthropic.claude-sonnet-4-6-v1:0": 200000, - "bedrock/apac.anthropic.claude-3-5-sonnet-20241022-v2:0": 100000, - "bedrock/apac.anthropic.claude-3-7-sonnet-20250219-v1:0": 200000, - "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0": 200000, + "bedrock/us.anthropic.claude-sonnet-4-6": 1000000, + "bedrock/us.anthropic.claude-sonnet-4-6-v1:0": 1000000, + "bedrock/au.anthropic.claude-sonnet-4-6": 1000000, + "bedrock/au.anthropic.claude-sonnet-4-6-v1:0": 1000000, "bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/eu.anthropic.claude-sonnet-4-6": 200000, - "bedrock/eu.anthropic.claude-sonnet-4-6-v1:0": 200000, + "bedrock/eu.anthropic.claude-sonnet-4-6": 1000000, + "bedrock/eu.anthropic.claude-sonnet-4-6-v1:0": 1000000, "bedrock/jp.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/jp.anthropic.claude-sonnet-4-6": 200000, - "bedrock/jp.anthropic.claude-sonnet-4-6-v1:0": 200000, - "bedrock/apac.anthropic.claude-sonnet-4-6": 200000, - "bedrock/apac.anthropic.claude-sonnet-4-6-v1:0": 200000, + "bedrock/jp.anthropic.claude-sonnet-4-6": 1000000, + "bedrock/jp.anthropic.claude-sonnet-4-6-v1:0": 1000000, + "bedrock/apac.anthropic.claude-sonnet-4-6": 1000000, + "bedrock/apac.anthropic.claude-sonnet-4-6-v1:0": 1000000, "bedrock/global.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/global.anthropic.claude-sonnet-4-6": 200000, - "bedrock/global.anthropic.claude-sonnet-4-6-v1:0": 200000, + "bedrock/global.anthropic.claude-sonnet-4-6": 1000000, + "bedrock/global.anthropic.claude-sonnet-4-6-v1:0": 1000000, "bedrock/us.anthropic.claude-sonnet-5": 1000000, "bedrock/au.anthropic.claude-sonnet-5": 1000000, "bedrock/eu.anthropic.claude-sonnet-5": 1000000, "bedrock/jp.anthropic.claude-sonnet-5": 1000000, "bedrock/global.anthropic.claude-sonnet-5": 1000000, - 'claude-3-5-sonnet': 100000, 'bedrock/us.meta.llama4-scout-17b-instruct-v1:0': 128000, 'bedrock/us.meta.llama4-maverick-17b-instruct-v1:0': 128000, 'groq/openai/gpt-oss-120b': 131072, @@ -376,7 +333,6 @@ def base_model_name(model: str) -> str: # default; it can be replaced via the `claude_extended_thinking_models_override` # configuration option. CLAUDE_EXTENDED_THINKING_MODELS = { - "claude-3-7-sonnet-20250219", "claude-sonnet-4-6", "claude-sonnet-4-5-20250929", "claude-opus-4-5-20251101", diff --git a/tests/unittest/test_get_max_tokens.py b/tests/unittest/test_get_max_tokens.py index bd18f9fe22..d3bf5eb31d 100644 --- a/tests/unittest/test_get_max_tokens.py +++ b/tests/unittest/test_get_max_tokens.py @@ -238,7 +238,7 @@ def test_claude_opus_4_6_model_max_tokens(self, monkeypatch, model): monkeypatch.setattr(utils, "get_settings", lambda: fake_settings) - assert get_max_tokens(model) == 200000 + assert get_max_tokens(model) == 1000000 @pytest.mark.parametrize( "model", @@ -290,4 +290,4 @@ def test_claude_sonnet_4_6_model_max_tokens(self, monkeypatch, model): monkeypatch.setattr(utils, "get_settings", lambda: fake_settings) - assert get_max_tokens(model) == 200000 + assert get_max_tokens(model) == 1000000 From 3943da6cd5675db96b350cff99d3552203036d5d Mon Sep 17 00:00:00 2001 From: Shmulik Cohen Date: Tue, 7 Jul 2026 00:17:20 +0300 Subject: [PATCH 4/4] refactor: collapse MAX_TOKENS Claude/Gemini rows to base-name keys get_max_tokens() resolves via MAX_TOKENS.get(model) or MAX_TOKENS.get(base_model_name(model)), so every provider prefix, Bedrock region+vendor infix, cloud-version (-v1:0) and vertex @date spelling of a model already collapses to one base key. The Claude and Gemini blocks encoded that single fact ~130 times (90 Claude rows -> 10 keys, 36 Gemini rows -> 18 keys). Dedup to base-name keys. Behavior is identical: the parametrized bedrock/vertex/anthropic tests in test_get_max_tokens.py still pass against the bare keys. Also drops the lone retired Claude 2.1 bedrock row, whose base name is degenerate and which litellm still covers. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01SqHrpoPezxDkD5dg4VuQL5 --- pr_agent/algo/__init__.py | 146 +++++++------------------------------- 1 file changed, 24 insertions(+), 122 deletions(-) diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py index 6af24b9c41..ebe44d1db5 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -93,138 +93,40 @@ def base_model_name(model: str) -> str: 'openai/qwq-plus': 131072, # 131K context length, but may be limited by config.max_model_tokens 'replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1': 4096, 'meta-llama/Llama-2-7b-chat-hf': 4096, - 'vertex_ai/codechat-bison': 6144, - 'vertex_ai/codechat-bison-32k': 32000, - 'vertex_ai/claude-haiku-4-5@20251001': 200000, - 'vertex_ai/claude-opus-4-1@20250805': 200000, - 'vertex_ai/claude-opus-4-5@20251101': 200000, - 'vertex_ai/claude-opus-4-6@20260120': 1000000, - 'vertex_ai/claude-opus-4-6': 1000000, - 'vertex_ai/claude-opus-4-7': 1000000, - 'vertex_ai/claude-opus-4-8': 1000000, - 'vertex_ai/claude-sonnet-4-5@20250929': 200000, - 'vertex_ai/claude-sonnet-4-6': 1000000, - 'vertex_ai/claude-sonnet-5': 1000000, - 'vertex_ai/gemini-1.5-pro': 1048576, - 'vertex_ai/gemini-2.5-pro-preview-03-25': 1048576, - 'vertex_ai/gemini-2.5-pro-preview-05-06': 1048576, - 'vertex_ai/gemini-2.5-pro-preview-06-05': 1048576, - 'vertex_ai/gemini-2.5-pro': 1048576, - 'vertex_ai/gemini-1.5-flash': 1048576, - 'vertex_ai/gemini-2.0-flash': 1048576, - 'vertex_ai/gemini-2.5-flash-preview-04-17': 1048576, - 'vertex_ai/gemini-2.5-flash-preview-05-20': 1048576, - 'vertex_ai/gemini-2.5-flash': 1048576, - 'vertex_ai/gemini-3-flash-preview': 1048576, - 'vertex_ai/gemini-3-pro-preview': 1048576, - 'vertex_ai/gemini-3.1-flash': 1048576, - 'vertex_ai/gemini-3.1-pro': 1048576, - 'vertex_ai/gemini-3.1-flash-lite-preview': 1048576, - 'vertex_ai/gemini-3.1-pro-preview': 1048576, - 'vertex_ai/gemini-3.5-flash': 1048576, - 'vertex_ai/gemini-3.5-pro': 1048576, - 'vertex_ai/gemma2': 8200, - 'gemini/gemini-1.5-pro': 1048576, - 'gemini/gemini-1.5-flash': 1048576, - 'gemini/gemini-2.0-flash': 1048576, - 'gemini/gemini-2.5-flash-preview-04-17': 1048576, - 'gemini/gemini-2.5-flash-preview-05-20': 1048576, - 'gemini/gemini-2.5-flash': 1048576, - 'gemini/gemini-2.5-pro-preview-03-25': 1048576, - 'gemini/gemini-2.5-pro-preview-05-06': 1048576, - 'gemini/gemini-2.5-pro-preview-06-05': 1048576, - 'gemini/gemini-2.5-pro': 1048576, - 'gemini/gemini-3-flash-preview': 1048576, - 'gemini/gemini-3-pro-preview': 1048576, - 'gemini/gemini-3.1-flash': 1048576, - 'gemini/gemini-3.1-pro': 1048576, - 'gemini/gemini-3.1-flash-lite-preview': 1048576, - 'gemini/gemini-3.1-pro-preview': 1048576, - 'gemini/gemini-3.5-flash': 1048576, - 'gemini/gemini-3.5-pro': 1048576, + # codechat / gemini / claude below are keyed by base model name: get_max_tokens() normalizes any + # provider prefix (anthropic/, vertex_ai/, gemini/, bedrock/...), Bedrock region + vendor infix, + # cloud-version (-v1:0) and vertex @date via base_model_name(), so one key covers every spelling. 'codechat-bison': 6144, 'codechat-bison-32k': 32000, - 'anthropic/claude-opus-4-1-20250805': 200000, - 'anthropic/claude-opus-4-5-20251101': 200000, - 'anthropic/claude-opus-4-6': 1000000, - 'anthropic/claude-opus-4-6-20260120': 1000000, - 'anthropic/claude-opus-4-7': 1000000, - 'anthropic/claude-opus-4-8': 1000000, - 'anthropic/claude-sonnet-4-5-20250929': 200000, - 'anthropic/claude-sonnet-4-6': 1000000, - 'anthropic/claude-sonnet-5': 1000000, + 'gemini-1.5-pro': 1048576, + 'gemini-1.5-flash': 1048576, + 'gemini-2.0-flash': 1048576, + 'gemini-2.5-flash-preview-04-17': 1048576, + 'gemini-2.5-flash-preview-05-20': 1048576, + 'gemini-2.5-flash': 1048576, + 'gemini-2.5-pro-preview-03-25': 1048576, + 'gemini-2.5-pro-preview-05-06': 1048576, + 'gemini-2.5-pro-preview-06-05': 1048576, + 'gemini-2.5-pro': 1048576, + 'gemini-3-flash-preview': 1048576, + 'gemini-3-pro-preview': 1048576, + 'gemini-3.1-flash': 1048576, + 'gemini-3.1-flash-lite-preview': 1048576, + 'gemini-3.1-pro': 1048576, + 'gemini-3.1-pro-preview': 1048576, + 'gemini-3.5-flash': 1048576, + 'gemini-3.5-pro': 1048576, + 'vertex_ai/gemma2': 8200, + 'claude-haiku-4-5-20251001': 200000, 'claude-opus-4-1-20250805': 200000, 'claude-opus-4-5-20251101': 200000, 'claude-opus-4-6': 1000000, 'claude-opus-4-6-20260120': 1000000, 'claude-opus-4-7': 1000000, 'claude-opus-4-8': 1000000, + 'claude-sonnet-4-5-20250929': 200000, 'claude-sonnet-4-6': 1000000, 'claude-sonnet-5': 1000000, - 'anthropic/claude-haiku-4-5-20251001': 200000, - 'claude-haiku-4-5-20251001': 200000, - 'bedrock/anthropic.claude-v2:1': 100000, - 'bedrock/anthropic.claude-opus-4-1-20250805-v1:0': 200000, - 'bedrock/anthropic.claude-opus-4-6-20260120-v1:0': 1000000, - 'bedrock/anthropic.claude-opus-4-6-v1:0': 1000000, - 'bedrock/anthropic.claude-opus-4-7': 1000000, - 'bedrock/anthropic.claude-opus-4-7-v1:0': 1000000, - 'bedrock/anthropic.claude-opus-4-8': 1000000, - 'bedrock/anthropic.claude-haiku-4-5-20251001-v1:0': 200000, - 'bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0': 200000, - 'bedrock/anthropic.claude-sonnet-4-6': 1000000, - 'bedrock/anthropic.claude-sonnet-4-6-v1:0': 1000000, - 'bedrock/anthropic.claude-sonnet-5': 1000000, - 'bedrock/anthropic.claude-opus-4-5-20251101-v1:0': 200000, - "bedrock/us.anthropic.claude-opus-4-1-20250805-v1:0": 200000, - "bedrock/us.anthropic.claude-opus-4-6-20260120-v1:0": 1000000, - "bedrock/global.anthropic.claude-opus-4-5-20251101-v1:0": 200000, - "bedrock/eu.anthropic.claude-opus-4-5-20251101-v1:0": 200000, - "bedrock/au.anthropic.claude-opus-4-5-20251101-v1:0": 200000, - "bedrock/jp.anthropic.claude-opus-4-5-20251101-v1:0": 200000, - "bedrock/apac.anthropic.claude-opus-4-5-20251101-v1:0": 200000, - "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0": 200000, - "bedrock/global.anthropic.claude-opus-4-6-v1:0": 1000000, - "bedrock/eu.anthropic.claude-opus-4-6-v1:0": 1000000, - "bedrock/au.anthropic.claude-opus-4-6-v1:0": 1000000, - "bedrock/jp.anthropic.claude-opus-4-6-v1:0": 1000000, - "bedrock/apac.anthropic.claude-opus-4-6-v1:0": 1000000, - "bedrock/us.anthropic.claude-opus-4-6-v1:0": 1000000, - "bedrock/global.anthropic.claude-opus-4-7": 1000000, - "bedrock/us.anthropic.claude-opus-4-7": 1000000, - "bedrock/global.anthropic.claude-opus-4-8": 1000000, - "bedrock/us.anthropic.claude-opus-4-8": 1000000, - "bedrock/eu.anthropic.claude-opus-4-8": 1000000, - "bedrock/au.anthropic.claude-opus-4-8": 1000000, - "bedrock/jp.anthropic.claude-opus-4-8": 1000000, - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, - "bedrock/eu.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, - "bedrock/au.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, - "bedrock/jp.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, - "bedrock/apac.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, - "bedrock/global.anthropic.claude-haiku-4-5-20251001-v1:0": 200000, - "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/au.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/us.anthropic.claude-sonnet-4-6": 1000000, - "bedrock/us.anthropic.claude-sonnet-4-6-v1:0": 1000000, - "bedrock/au.anthropic.claude-sonnet-4-6": 1000000, - "bedrock/au.anthropic.claude-sonnet-4-6-v1:0": 1000000, - "bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/eu.anthropic.claude-sonnet-4-6": 1000000, - "bedrock/eu.anthropic.claude-sonnet-4-6-v1:0": 1000000, - "bedrock/jp.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/jp.anthropic.claude-sonnet-4-6": 1000000, - "bedrock/jp.anthropic.claude-sonnet-4-6-v1:0": 1000000, - "bedrock/apac.anthropic.claude-sonnet-4-6": 1000000, - "bedrock/apac.anthropic.claude-sonnet-4-6-v1:0": 1000000, - "bedrock/global.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000, - "bedrock/global.anthropic.claude-sonnet-4-6": 1000000, - "bedrock/global.anthropic.claude-sonnet-4-6-v1:0": 1000000, - "bedrock/us.anthropic.claude-sonnet-5": 1000000, - "bedrock/au.anthropic.claude-sonnet-5": 1000000, - "bedrock/eu.anthropic.claude-sonnet-5": 1000000, - "bedrock/jp.anthropic.claude-sonnet-5": 1000000, - "bedrock/global.anthropic.claude-sonnet-5": 1000000, 'bedrock/us.meta.llama4-scout-17b-instruct-v1:0': 128000, 'bedrock/us.meta.llama4-maverick-17b-instruct-v1:0': 128000, 'groq/openai/gpt-oss-120b': 131072,