diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 0dae0bd1ed..a365f08738 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -25,19 +25,29 @@ class TokenEncoder: _lock = Lock() # Create a lock object @classmethod - def get_token_encoder(cls): - model = get_settings().config.model + def get_token_encoder(cls, model=None): + configured_model = get_settings().config.model + model = model or configured_model + + # Use a fresh tokenizer for explicit fallback models without replacing + # the cached tokenizer for the configured primary model. + if model != configured_model: + return cls._create_encoder(model) + if cls._encoder_instance is None or model != cls._model: # Check without acquiring the lock for performance with cls._lock: # Lock acquisition to ensure thread safety if cls._encoder_instance is None or model != cls._model: cls._model = model - try: - cls._encoder_instance = encoding_for_model(cls._model) if "gpt" in cls._model else get_encoding( - "o200k_base") - except: - cls._encoder_instance = get_encoding("o200k_base") + cls._encoder_instance = cls._create_encoder(cls._model) return cls._encoder_instance + @staticmethod + def _create_encoder(model): + try: + return encoding_for_model(model) if "gpt" in model else get_encoding("o200k_base") + except Exception: + return get_encoding("o200k_base") + class TokenHandler: """ diff --git a/pr_agent/tools/pr_line_questions.py b/pr_agent/tools/pr_line_questions.py index 5c0d6970d4..c6cfeffa07 100644 --- a/pr_agent/tools/pr_line_questions.py +++ b/pr_agent/tools/pr_line_questions.py @@ -2,15 +2,19 @@ import copy from functools import partial -from jinja2 import Environment, StrictUndefined +from jinja2 import Environment, StrictUndefined, select_autoescape +from litellm import token_counter 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, extract_hunk_lines_from_patch) -from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models -from pr_agent.algo.token_handler import TokenHandler -from pr_agent.algo.utils import ModelType + decouple_and_convert_to_hunks_with_lines_numbers, + extract_hunk_lines_from_patch) +from pr_agent.algo.pr_processing import (OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD, + get_pr_diff, + retry_with_fallback_models) +from pr_agent.algo.token_handler import TokenEncoder, TokenHandler +from pr_agent.algo.utils import ModelType, get_max_tokens from pr_agent.config_loader import get_settings from pr_agent.git_providers import get_git_provider from pr_agent.git_providers.git_provider import get_main_pr_language @@ -18,6 +22,7 @@ from pr_agent.log import get_logger from pr_agent.servers.help import HelpMessage + class PR_LineQuestions: def __init__(self, pr_url: str, args=None, ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler): self.question_str = self.parse_args(args) @@ -137,26 +142,26 @@ async def run(self): self.git_provider.publish_comment(no_hunk_message) return "" - + def _load_conversation_history(self) -> str: """Generate conversation history from the code review thread - + Returns: str: The formatted conversation history """ comment_id = get_settings().get('comment_id', '') file_path = get_settings().get('file_name', '') line_number = get_settings().get('line_end', '') - + # early return if any required parameter is missing if not all([comment_id, file_path, line_number]): get_logger().error("Missing required parameters for conversation history") return "" - + try: # retrieve thread comments thread_comments = self.git_provider.get_review_thread_comments(comment_id) - + # filter and prepare comments filtered_comments = [] for comment in thread_comments: @@ -165,23 +170,23 @@ def _load_conversation_history(self) -> str: # skip empty comments, current comment(will be added as a question at prompt) if not body or not body.strip() or comment_id == comment.id: continue - + user = comment.user author = user.login if hasattr(user, 'login') else 'Unknown' filtered_comments.append((author, body)) - + # transform conversation history to string using the same pattern as get_commit_messages if filtered_comments: comment_count = len(filtered_comments) get_logger().info(f"Loaded {comment_count} comments from the code review thread") - + # Format as numbered list, similar to get_commit_messages - conversation_history_str = "\n".join([f"{i + 1}. {author}: {body}" + conversation_history_str = "\n".join([f"{i + 1}. {author}: {body}" for i, (author, body) in enumerate(filtered_comments)]) return conversation_history_str - + return "" - + except Exception as e: get_logger().error(f"Error processing conversation history, error: {e}") return "" @@ -190,9 +195,8 @@ async def _get_prediction(self, model: str): variables = copy.deepcopy(self.vars) variables["full_hunk"] = self.patch_with_lines # update diff variables["selected_lines"] = self.selected_lines - environment = Environment(undefined=StrictUndefined) - system_prompt = environment.from_string(get_settings().pr_line_questions_prompt.system).render(variables) - user_prompt = environment.from_string(get_settings().pr_line_questions_prompt.user).render(variables) + variables["conversation_history"] = self._fit_conversation_history(variables, model) + system_prompt, user_prompt = self._render_prompts(variables) if get_settings().config.verbosity_level >= 2: # get_logger().info(f"\nSystem prompt:\n{system_prompt}") # get_logger().info(f"\nUser prompt:\n{user_prompt}") @@ -202,3 +206,77 @@ async def _get_prediction(self, model: str): response, finish_reason = await self.ai_handler.chat_completion( model=model, temperature=get_settings().config.temperature, system=system_prompt, user=user_prompt) return response + + def _render_prompts(self, variables): + environment = Environment( + autoescape=select_autoescape(default_for_string=False), + undefined=StrictUndefined, + ) + system_prompt = environment.from_string(get_settings().pr_line_questions_prompt.system).render(variables) + user_prompt = environment.from_string(get_settings().pr_line_questions_prompt.user).render(variables) + return system_prompt, user_prompt + + def _fit_conversation_history(self, variables, model): + conversation_history = variables.get("conversation_history", "") + encoder = None + try: + completion_tokens = int(get_settings().config.get("max_output_tokens", 0)) + except (TypeError, ValueError): + completion_tokens = 0 + if completion_tokens <= 0: + completion_tokens = OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD + max_tokens = max(get_max_tokens(model) - completion_tokens, 0) + + def render_with_history(history): + prompt_variables = copy.deepcopy(variables) + prompt_variables["conversation_history"] = history + return self._render_prompts(prompt_variables) + + def count_prompts(prompts): + nonlocal encoder + system_prompt, user_prompt = prompts + try: + model_token_count = token_counter( + model=model, + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + ) + if model_token_count > 0: + return model_token_count + except Exception as e: + get_logger().debug(f"Model-aware token counting failed for {model}: {e}") + if encoder is None: + encoder = TokenEncoder.get_token_encoder(model) + return len(encoder.encode(system_prompt, disallowed_special=())) + len( + encoder.encode(user_prompt, disallowed_special=())) + + if count_prompts(render_with_history(conversation_history)) <= max_tokens: + return conversation_history + + if count_prompts(render_with_history("")) > max_tokens: + raise ValueError( + f"The /ask_line prompt exceeds the token limit for {model} even without conversation history" + ) + + truncation_marker = "\n...(truncated)\n" + low, high = 0, len(conversation_history) + best_history = "" + while low <= high: + keep_chars = (low + high) // 2 + candidate = ( + truncation_marker + conversation_history[-keep_chars:] + if keep_chars + else "" + ) + if count_prompts(render_with_history(candidate)) <= max_tokens: + best_history = candidate + low = keep_chars + 1 + else: + high = keep_chars - 1 + + get_logger().warning( + f"Conversation history was clipped for /ask_line to fit the {max_tokens}-token input limit" + ) + return best_history diff --git a/tests/unittest/test_pr_line_questions_context_budget.py b/tests/unittest/test_pr_line_questions_context_budget.py new file mode 100644 index 0000000000..4e6f4a3fda --- /dev/null +++ b/tests/unittest/test_pr_line_questions_context_budget.py @@ -0,0 +1,237 @@ +"""Regression tests for the /ask_line conversation-context budget.""" + +from types import SimpleNamespace + +import pytest +from litellm import token_counter + +import pr_agent.tools.pr_line_questions as plq +from pr_agent.config_loader import get_settings +from tests.unittest._settings_helpers import (restore_settings, + snapshot_settings) + + +class _FakeGithubProvider: + def __init__(self, comments, patch): + self.comments = comments + self.patch = patch + self.replies = [] + + def get_review_thread_comments(self, comment_id): + return self.comments + + def reply_to_comment_from_comment_id(self, comment_id, body): + self.replies.append((comment_id, body)) + + +class _RecordingAIHandler: + def __init__(self, fail_models=()): + self.requests = [] + self.fail_models = set(fail_models) + + async def chat_completion(self, *, model, temperature, system, user): + self.requests.append({"model": model, "system": system, "user": user}) + if model in self.fail_models: + raise RuntimeError(f"simulated failure for {model}") + return "answer", "stop" + + +@pytest.mark.parametrize( + ("fallback_models", "fail_models", "expected_model"), + [ + ([], (), "gpt-4o"), + (["gpt-3.5-turbo"], ("gpt-4o",), "gpt-3.5-turbo"), + ], +) +@pytest.mark.asyncio +async def test_ask_line_keeps_final_prompt_within_budget_after_loading_history( + monkeypatch, fallback_models, fail_models, expected_model +): + settings = get_settings() + keys = ( + "config.model", + "config.model_weak", + "config.fallback_models", + "config.max_model_tokens", + "config.max_output_tokens", + "openai.deployment_id", + "openai.fallback_deployments", + "pr_questions.use_conversation_history", + "ask_diff_hunk", + "line_start", + "line_end", + "side", + "file_name", + "comment_id", + ) + saved = snapshot_settings(keys) + + patch = ( + "@@ -5,7 +5,8 @@ def main():\n" + " a = 1\n" + " b = 2\n" + "+ c = 3\n" + " return a\n" + ) + comments = [ + SimpleNamespace( + id=100, + body="current question", + user=SimpleNamespace(login="alice"), + ) + ] + [ + SimpleNamespace( + id=101 + index, + body=f"reply {index}: " + ("context " * 20), + user=SimpleNamespace(login="reviewer"), + ) + for index in range(200) + ] + provider = _FakeGithubProvider(comments, patch) + ai_handler = _RecordingAIHandler(fail_models) + question = plq.PR_LineQuestions.__new__(plq.PR_LineQuestions) + question.question_str = "Why is this change needed?" + question.git_provider = provider + question.ai_handler = ai_handler + question.resolve_threads = False + question.vars = { + "title": "Budget regression", + "branch": "feature/budget", + "question": question.question_str, + "full_hunk": "", + "selected_lines": "", + "conversation_history": "", + "resolve_threads": False, + "extra_instructions": "", + } + + try: + settings.set("config.model", "gpt-4o") + settings.set("config.model_weak", "") + settings.set("config.fallback_models", fallback_models) + settings.set("config.max_model_tokens", 700) + settings.set("config.max_output_tokens", 100) + settings.set("openai.deployment_id", None) + settings.set("openai.fallback_deployments", []) + settings.set("pr_questions.use_conversation_history", True) + settings.set("ask_diff_hunk", patch) + settings.set("line_start", 6) + settings.set("line_end", 8) + settings.set("side", "RIGHT") + settings.set("file_name", "src/example.py") + settings.set("comment_id", 100) + monkeypatch.setattr(plq, "GithubProvider", _FakeGithubProvider) + + await question.run() + + assert len(ai_handler.requests) == 1 + len(fail_models) + request = next(item for item in ai_handler.requests if item["model"] == expected_model) + prompt_tokens = token_counter( + model=expected_model, + messages=[ + {"role": "system", "content": request["system"]}, + {"role": "user", "content": request["user"]}, + ], + ) + assert prompt_tokens <= 600 + assert len(request["user"]) < len("\n".join(comment.body for comment in comments)) + assert "Why is this change needed?" in request["user"] + assert "+ c = 3" in request["user"] + assert "reply 199" in request["user"] + assert provider.replies == [(100, "answer")] + finally: + restore_settings(saved) + + +@pytest.mark.asyncio +async def test_ask_line_uses_attempted_model_for_non_gpt_prompt_budget(monkeypatch): + settings = get_settings() + keys = ( + "config.model", + "config.model_weak", + "config.fallback_models", + "config.max_model_tokens", + "config.max_output_tokens", + "openai.deployment_id", + "openai.fallback_deployments", + "pr_questions.use_conversation_history", + "ask_diff_hunk", + "line_start", + "line_end", + "side", + "file_name", + "comment_id", + ) + saved = snapshot_settings(keys) + patch = "@@ -5,2 +5,3 @@ def main():\n a = 1\n+ b = 2\n" + comments = [ + SimpleNamespace( + id=100, + body="current question", + user=SimpleNamespace(login="alice"), + ) + ] + [ + SimpleNamespace( + id=101 + index, + body=f"reply {index}: " + ("context " * 20), + user=SimpleNamespace(login="reviewer"), + ) + for index in range(200) + ] + provider = _FakeGithubProvider(comments, patch) + ai_handler = _RecordingAIHandler(("gpt-4o",)) + question = plq.PR_LineQuestions.__new__(plq.PR_LineQuestions) + question.question_str = "Why is this change needed?" + question.git_provider = provider + question.ai_handler = ai_handler + question.resolve_threads = False + question.vars = { + "title": "Budget regression", + "branch": "feature/budget", + "question": question.question_str, + "full_hunk": "", + "selected_lines": "", + "conversation_history": "", + "resolve_threads": False, + "extra_instructions": "", + } + + counter_calls = [] + + def model_aware_counter(*, model, messages): + counter_calls.append(model) + return sum(len(message["content"]) for message in messages) + + try: + settings.set("config.model", "gpt-4o") + settings.set("config.model_weak", "") + settings.set("config.fallback_models", ["claude-2"]) + settings.set("config.max_model_tokens", 3000) + settings.set("config.max_output_tokens", 100) + settings.set("openai.deployment_id", None) + settings.set("openai.fallback_deployments", []) + settings.set("pr_questions.use_conversation_history", True) + settings.set("ask_diff_hunk", patch) + settings.set("line_start", 6) + settings.set("line_end", 7) + settings.set("side", "RIGHT") + settings.set("file_name", "src/example.py") + settings.set("comment_id", 100) + monkeypatch.setattr(plq, "GithubProvider", _FakeGithubProvider) + monkeypatch.setattr(plq, "token_counter", model_aware_counter, raising=False) + + await question.run() + + request = next(item for item in ai_handler.requests if item["model"] == "claude-2") + assert "claude-2" in counter_calls + assert model_aware_counter( + model="claude-2", + messages=[ + {"role": "system", "content": request["system"]}, + {"role": "user", "content": request["user"]}, + ], + ) <= 2900 + assert "reply 199" in request["user"] + assert provider.replies == [(100, "answer")] + finally: + restore_settings(saved)