From 05ab5f699faf2c211602dcf925bd68bb778bbf16 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Fri, 16 May 2025 17:51:22 +0900 Subject: [PATCH 01/12] Improve token calculation logic based on model type - Rename calc_tokens to get_token_count_by_model_type for clearer intent - Separate model type detection logic to improve maintainability --- pr_agent/algo/token_handler.py | 44 +++++++++++++++++++--------------- 1 file changed, 25 insertions(+), 19 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 60cf2c84..0c8851e8 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -107,25 +107,37 @@ class TokenHandler: get_logger().error( f"Error in Anthropic token counting: {e}") return MaxTokens - def estimate_token_count_for_non_anth_claude_models(self, model, default_encoder_estimate): + def is_openai_model(self, model_name): + from re import match + + return 'gpt' in model_name or match(r"^o[1-9](-mini|-preview)?$", model_name) + + def apply_estimation_factor(self, model_name, default_estimate): from math import ceil - import re - model_is_from_o_series = re.match(r"^o[1-9](-mini|-preview)?$", model) - if ('gpt' in get_settings().config.model.lower() or model_is_from_o_series) and get_settings(use_context=False).get('openai.key'): - return default_encoder_estimate - #else: Model is not an OpenAI one - therefore, cannot provide an accurate token count and instead, return a higher number as best effort. + factor = 1 + get_settings().get('config.model_token_count_estimate_factor', 0) + get_logger().warning(f"{model_name}'s token count cannot be accurately estimated. Using factor of {factor}") + + return ceil(factor * default_estimate) - elbow_factor = 1 + get_settings().get('config.model_token_count_estimate_factor', 0) - get_logger().warning(f"{model}'s expected token count cannot be accurately estimated. Using {elbow_factor} of encoder output as best effort estimate") - return ceil(elbow_factor * default_encoder_estimate) - - def count_tokens(self, patch: str, force_accurate=False) -> int: + def get_token_count_by_model_type(self, patch: str, default_estimate: int) -> int: + model_name = get_settings().config.model.lower() + + if 'claude' in model_name and get_settings(use_context=False).get('anthropic.key'): + return self.calc_claude_tokens(patch) + + if self.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): + return default_estimate + + return self.apply_estimation_factor(model_name, default_estimate) + + def count_tokens(self, patch: str, force_accurate: bool = False) -> int: """ Counts the number of tokens in a given patch string. Args: - patch: The patch string. + - force_accurate: If True, uses a more precise calculation method. Returns: The number of tokens in the patch string. @@ -135,11 +147,5 @@ class TokenHandler: #If an estimate is enough (for example, in cases where the maximal allowed tokens is way below the known limits), return it. if not force_accurate: return encoder_estimate - - #else, force_accurate==True: User requested providing an accurate estimation: - model = get_settings().config.model.lower() - if 'claude' in model and get_settings(use_context=False).get('anthropic.key'): - return self.calc_claude_tokens(patch) # API call to Anthropic for accurate token counting for Claude models - - #else: Non Anthropic provided model: - return self.estimate_token_count_for_non_anth_claude_models(model, encoder_estimate) + else: + return self.get_token_count_by_model_type(patch, encoder_estimate=encoder_estimate) From 81fa22e4df6c8fd4ce3da0466bd181019bb6d108 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Tue, 20 May 2025 13:47:15 +0900 Subject: [PATCH 02/12] Add model name validation --- pr_agent/algo/token_handler.py | 44 ++++++++++++++++++++-------------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 0c8851e8..2781fc5c 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -1,4 +1,6 @@ from threading import Lock +from math import ceil +import re from jinja2 import Environment, StrictUndefined from tiktoken import encoding_for_model, get_encoding @@ -7,6 +9,16 @@ from pr_agent.config_loader import get_settings from pr_agent.log import get_logger +class ModelTypeValidator: + @staticmethod + def is_openai_model(model_name: str) -> bool: + return 'gpt' in model_name or re.match(r"^o[1-9](-mini|-preview)?$", model_name) + + @staticmethod + def is_claude_model(model_name: str) -> bool: + return 'claude' in model_name + + class TokenEncoder: _encoder_instance = None _model = None @@ -51,6 +63,9 @@ class TokenHandler: - user: The user string. """ self.encoder = TokenEncoder.get_token_encoder() + self.settings = get_settings() + self.model_validator = ModelTypeValidator() + if pr is not None: self.prompt_tokens = self._get_system_user_tokens(pr, self.encoder, vars, system, user) @@ -79,19 +94,20 @@ class TokenHandler: get_logger().error(f"Error in _get_system_user_tokens: {e}") return 0 - def calc_claude_tokens(self, patch): + def calc_claude_tokens(self, patch: str) -> int: try: import anthropic from pr_agent.algo import MAX_TOKENS - client = anthropic.Anthropic(api_key=get_settings(use_context=False).get('anthropic.key')) - MaxTokens = MAX_TOKENS[get_settings().config.model] + + client = anthropic.Anthropic(api_key=self.settings.get('anthropic.key')) + max_tokens = MAX_TOKENS[self.settings.config.model] # Check if the content size is too large (9MB limit) if len(patch.encode('utf-8')) > 9_000_000: get_logger().warning( "Content too large for Anthropic token counting API, falling back to local tokenizer" ) - return MaxTokens + return max_tokens response = client.messages.count_tokens( model="claude-3-7-sonnet-20250219", @@ -104,29 +120,21 @@ class TokenHandler: return response.input_tokens except Exception as e: - get_logger().error( f"Error in Anthropic token counting: {e}") - return MaxTokens + get_logger().error(f"Error in Anthropic token counting: {e}") + return max_tokens - def is_openai_model(self, model_name): - from re import match - - return 'gpt' in model_name or match(r"^o[1-9](-mini|-preview)?$", model_name) - - def apply_estimation_factor(self, model_name, default_estimate): - from math import ceil - - factor = 1 + get_settings().get('config.model_token_count_estimate_factor', 0) + def apply_estimation_factor(self, model_name: str, default_estimate: int) -> int: + factor = 1 + self.settings.get('config.model_token_count_estimate_factor', 0) get_logger().warning(f"{model_name}'s token count cannot be accurately estimated. Using factor of {factor}") - return ceil(factor * default_estimate) def get_token_count_by_model_type(self, patch: str, default_estimate: int) -> int: model_name = get_settings().config.model.lower() - if 'claude' in model_name and get_settings(use_context=False).get('anthropic.key'): + if self.model_validator.is_claude_model(model_name) and get_settings(use_context=False).get('anthropic.key'): return self.calc_claude_tokens(patch) - if self.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): + if self.model_validator.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): return default_estimate return self.apply_estimation_factor(model_name, default_estimate) From e72bb28c4e9cad740f9240b25a5adb7339943044 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Tue, 20 May 2025 13:50:30 +0900 Subject: [PATCH 03/12] Replace get_settings() with self.settings --- pr_agent/algo/token_handler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 2781fc5c..99dbc635 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -129,12 +129,12 @@ class TokenHandler: return ceil(factor * default_estimate) def get_token_count_by_model_type(self, patch: str, default_estimate: int) -> int: - model_name = get_settings().config.model.lower() + model_name = self.settings.config.model.lower() - if self.model_validator.is_claude_model(model_name) and get_settings(use_context=False).get('anthropic.key'): + if self.model_validator.is_claude_model(model_name) and self.settings.get('anthropic.key'): return self.calc_claude_tokens(patch) - if self.model_validator.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): + if self.model_validator.is_openai_model(model_name) and self.settings.get('openai.key'): return default_estimate return self.apply_estimation_factor(model_name, default_estimate) From f198e6fa097a9ad27a08da7b8f9e14dbe7be5a9b Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Tue, 20 May 2025 14:12:24 +0900 Subject: [PATCH 04/12] Add constants and improve token calculation logic --- pr_agent/algo/token_handler.py | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 99dbc635..8bd3f115 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -52,6 +52,10 @@ class TokenHandler: method. """ + # Constants + CLAUDE_MODEL = "claude-3-7-sonnet-20250219" + CLAUDE_MAX_CONTENT_SIZE = 9_000_000 # Maximum allowed content size (9MB) for Claude API + def __init__(self, pr=None, vars: dict = {}, system="", user=""): """ Initializes the TokenHandler object. @@ -102,15 +106,14 @@ class TokenHandler: client = anthropic.Anthropic(api_key=self.settings.get('anthropic.key')) max_tokens = MAX_TOKENS[self.settings.config.model] - # Check if the content size is too large (9MB limit) - if len(patch.encode('utf-8')) > 9_000_000: + if len(patch.encode('utf-8')) > self.CLAUDE_MAX_CONTENT_SIZE: get_logger().warning( "Content too large for Anthropic token counting API, falling back to local tokenizer" ) return max_tokens response = client.messages.count_tokens( - model="claude-3-7-sonnet-20250219", + model=self.CLAUDE_MODEL, system="system", messages=[{ "role": "user", @@ -126,9 +129,20 @@ class TokenHandler: def apply_estimation_factor(self, model_name: str, default_estimate: int) -> int: factor = 1 + self.settings.get('config.model_token_count_estimate_factor', 0) get_logger().warning(f"{model_name}'s token count cannot be accurately estimated. Using factor of {factor}") + return ceil(factor * default_estimate) def get_token_count_by_model_type(self, patch: str, default_estimate: int) -> int: + """ + Get token count based on model type. + + Args: + patch: The text to count tokens for. + default_estimate: The default token count estimate. + + Returns: + int: The calculated token count. + """ model_name = self.settings.config.model.lower() if self.model_validator.is_claude_model(model_name) and self.settings.get('anthropic.key'): @@ -152,8 +166,8 @@ class TokenHandler: """ encoder_estimate = len(self.encoder.encode(patch, disallowed_special=())) - #If an estimate is enough (for example, in cases where the maximal allowed tokens is way below the known limits), return it. - if not force_accurate: - return encoder_estimate - else: + if force_accurate: return self.get_token_count_by_model_type(patch, encoder_estimate=encoder_estimate) + + # If an estimate is enough (for example, in cases where the maximal allowed tokens is way below the known limits), return it. + return encoder_estimate From 97f2b6f7360117355b46da6436e7f88ecfbd16c8 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Tue, 20 May 2025 15:29:27 +0900 Subject: [PATCH 05/12] Fix TypeError --- pr_agent/algo/token_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 8bd3f115..0a1aeb89 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -167,7 +167,7 @@ class TokenHandler: encoder_estimate = len(self.encoder.encode(patch, disallowed_special=())) if force_accurate: - return self.get_token_count_by_model_type(patch, encoder_estimate=encoder_estimate) + return self.get_token_count_by_model_type(patch, encoder_estimate) # If an estimate is enough (for example, in cases where the maximal allowed tokens is way below the known limits), return it. return encoder_estimate From 648829b770773c4db8b7dd671acb2b6c57310d6e Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Wed, 21 May 2025 17:51:03 +0900 Subject: [PATCH 06/12] Rename method --- pr_agent/algo/token_handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 0a1aeb89..b8eaac89 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -15,7 +15,7 @@ class ModelTypeValidator: return 'gpt' in model_name or re.match(r"^o[1-9](-mini|-preview)?$", model_name) @staticmethod - def is_claude_model(model_name: str) -> bool: + def is_anthropic_model(model_name: str) -> bool: return 'claude' in model_name @@ -145,7 +145,7 @@ class TokenHandler: """ model_name = self.settings.config.model.lower() - if self.model_validator.is_claude_model(model_name) and self.settings.get('anthropic.key'): + if self.model_validator.is_anthropic_model(model_name) and self.settings.get('anthropic.key'): return self.calc_claude_tokens(patch) if self.model_validator.is_openai_model(model_name) and self.settings.get('openai.key'): From c3ea048b718510f04637027f0f3ef6a504209a94 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Wed, 21 May 2025 17:52:51 +0900 Subject: [PATCH 07/12] Restore original return logic for force_accurate condition --- pr_agent/algo/token_handler.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index b8eaac89..5a29da36 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -166,8 +166,8 @@ class TokenHandler: """ encoder_estimate = len(self.encoder.encode(patch, disallowed_special=())) - if force_accurate: - return self.get_token_count_by_model_type(patch, encoder_estimate) - # If an estimate is enough (for example, in cases where the maximal allowed tokens is way below the known limits), return it. - return encoder_estimate + if not force_accurate: + return encoder_estimate + + return self.get_token_count_by_model_type(patch, encoder_estimate) From df0355d827b0f83f50031bb4133e8acc4929e1d8 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Wed, 21 May 2025 18:07:47 +0900 Subject: [PATCH 08/12] Remove member variable for restroring get_settings() --- pr_agent/algo/token_handler.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 5a29da36..da239cfb 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -67,7 +67,6 @@ class TokenHandler: - user: The user string. """ self.encoder = TokenEncoder.get_token_encoder() - self.settings = get_settings() self.model_validator = ModelTypeValidator() if pr is not None: @@ -103,8 +102,8 @@ class TokenHandler: import anthropic from pr_agent.algo import MAX_TOKENS - client = anthropic.Anthropic(api_key=self.settings.get('anthropic.key')) - max_tokens = MAX_TOKENS[self.settings.config.model] + client = anthropic.Anthropic(api_key=get_settings().get('anthropic.key')) + max_tokens = MAX_TOKENS[get_settings().config.model] if len(patch.encode('utf-8')) > self.CLAUDE_MAX_CONTENT_SIZE: get_logger().warning( @@ -127,7 +126,7 @@ class TokenHandler: return max_tokens def apply_estimation_factor(self, model_name: str, default_estimate: int) -> int: - factor = 1 + self.settings.get('config.model_token_count_estimate_factor', 0) + factor = 1 + get_settings().get('config.model_token_count_estimate_factor', 0) get_logger().warning(f"{model_name}'s token count cannot be accurately estimated. Using factor of {factor}") return ceil(factor * default_estimate) @@ -143,12 +142,12 @@ class TokenHandler: Returns: int: The calculated token count. """ - model_name = self.settings.config.model.lower() + model_name = get_settings().config.model.lower() - if self.model_validator.is_anthropic_model(model_name) and self.settings.get('anthropic.key'): return self.calc_claude_tokens(patch) + if self.model_validator.is_anthropic_model(model_name) and get_settings(use_context=False).get('anthropic.key'): - if self.model_validator.is_openai_model(model_name) and self.settings.get('openai.key'): + if self.model_validator.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): return default_estimate return self.apply_estimation_factor(model_name, default_estimate) From ead7491ca91e24ab185b33c2867dd3f4a8baad0a Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Wed, 21 May 2025 18:08:48 +0900 Subject: [PATCH 09/12] Apply convention for marking private --- pr_agent/algo/token_handler.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index da239cfb..e45d611f 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -97,7 +97,7 @@ class TokenHandler: get_logger().error(f"Error in _get_system_user_tokens: {e}") return 0 - def calc_claude_tokens(self, patch: str) -> int: + def _calc_claude_tokens(self, patch: str) -> int: try: import anthropic from pr_agent.algo import MAX_TOKENS @@ -125,13 +125,13 @@ class TokenHandler: get_logger().error(f"Error in Anthropic token counting: {e}") return max_tokens - def apply_estimation_factor(self, model_name: str, default_estimate: int) -> int: + def _apply_estimation_factor(self, model_name: str, default_estimate: int) -> int: factor = 1 + get_settings().get('config.model_token_count_estimate_factor', 0) get_logger().warning(f"{model_name}'s token count cannot be accurately estimated. Using factor of {factor}") return ceil(factor * default_estimate) - def get_token_count_by_model_type(self, patch: str, default_estimate: int) -> int: + def _get_token_count_by_model_type(self, patch: str, default_estimate: int) -> int: """ Get token count based on model type. @@ -144,13 +144,13 @@ class TokenHandler: """ model_name = get_settings().config.model.lower() - return self.calc_claude_tokens(patch) if self.model_validator.is_anthropic_model(model_name) and get_settings(use_context=False).get('anthropic.key'): + return self._calc_claude_tokens(patch) if self.model_validator.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): return default_estimate - return self.apply_estimation_factor(model_name, default_estimate) + return self._apply_estimation_factor(model_name, default_estimate) def count_tokens(self, patch: str, force_accurate: bool = False) -> int: """ @@ -169,4 +169,4 @@ class TokenHandler: if not force_accurate: return encoder_estimate - return self.get_token_count_by_model_type(patch, encoder_estimate) + return self._get_token_count_by_model_type(patch, encoder_estimate) From cc686ef26d961bd515b731c447f98cd23a240633 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Thu, 22 May 2025 13:12:04 +0900 Subject: [PATCH 10/12] Reorder model check: OpenAI before Anthropic OpenAI is the default in most cases, so checking it first skips unnecessary Anthropic logic. --- pr_agent/algo/token_handler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index e45d611f..67a6eed9 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -144,11 +144,11 @@ class TokenHandler: """ model_name = get_settings().config.model.lower() - if self.model_validator.is_anthropic_model(model_name) and get_settings(use_context=False).get('anthropic.key'): - return self._calc_claude_tokens(patch) - if self.model_validator.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): return default_estimate + + if self.model_validator.is_anthropic_model(model_name) and get_settings(use_context=False).get('anthropic.key'): + return self._calc_claude_tokens(patch) return self._apply_estimation_factor(model_name, default_estimate) From facfb5f46b1be200b7dd5cfe87d3c745fa0cd85e Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Thu, 22 May 2025 13:32:20 +0900 Subject: [PATCH 11/12] Add missing code: use_context=False --- pr_agent/algo/token_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 67a6eed9..998391b1 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -102,7 +102,7 @@ class TokenHandler: import anthropic from pr_agent.algo import MAX_TOKENS - client = anthropic.Anthropic(api_key=get_settings().get('anthropic.key')) + client = anthropic.Anthropic(api_key=get_settings(use_context=False).get('anthropic.key')) max_tokens = MAX_TOKENS[get_settings().config.model] if len(patch.encode('utf-8')) > self.CLAUDE_MAX_CONTENT_SIZE: From 84f2f4fe3d8b601f55c39de1e5e03988703990d9 Mon Sep 17 00:00:00 2001 From: kkan9ma Date: Sun, 25 May 2025 18:00:38 +0900 Subject: [PATCH 12/12] Fix: use ModelTypeValidator static methods directly --- pr_agent/algo/token_handler.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 998391b1..cb313f02 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -67,7 +67,6 @@ class TokenHandler: - user: The user string. """ self.encoder = TokenEncoder.get_token_encoder() - self.model_validator = ModelTypeValidator() if pr is not None: self.prompt_tokens = self._get_system_user_tokens(pr, self.encoder, vars, system, user) @@ -144,10 +143,10 @@ class TokenHandler: """ model_name = get_settings().config.model.lower() - if self.model_validator.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): + if ModelTypeValidator.is_openai_model(model_name) and get_settings(use_context=False).get('openai.key'): return default_estimate - - if self.model_validator.is_anthropic_model(model_name) and get_settings(use_context=False).get('anthropic.key'): + + if ModelTypeValidator.is_anthropic_model(model_name) and get_settings(use_context=False).get('anthropic.key'): return self._calc_claude_tokens(patch) return self._apply_estimation_factor(model_name, default_estimate)