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

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

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,

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

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

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('', (), {
@@ -237,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",
@@ -289,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