feat(algo): let any model run — litellm-backed context window + base-name capability tables

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 <set>`, 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",
+}

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)

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(

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,26 +989,56 @@ def get_user_labels(current_labels: List[str] = None):
     return user_labels
 
 
-def get_max_tokens(model):
+# 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, 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').
     """
-    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'
+    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
 
-    For both cases, we further limit the number of tokens to 'config.max_model_tokens' if it is set.
+
+def get_max_tokens(model, cap=True):
+    """
+    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:
-        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 = (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 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
 

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:

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:

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.")