From 84b80f792d6be24a90a473831cb57bb8040cf3f2 Mon Sep 17 00:00:00 2001 From: mrT23 Date: Fri, 9 Aug 2024 21:44:00 +0300 Subject: [PATCH] protections --- pr_agent/algo/token_handler.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/pr_agent/algo/token_handler.py b/pr_agent/algo/token_handler.py index 9cc3b41f..935e6da5 100644 --- a/pr_agent/algo/token_handler.py +++ b/pr_agent/algo/token_handler.py @@ -3,6 +3,8 @@ from tiktoken import encoding_for_model, get_encoding from pr_agent.config_loader import get_settings from threading import Lock +from pr_agent.log import get_logger + class TokenEncoder: _encoder_instance = None @@ -62,12 +64,16 @@ class TokenHandler: Returns: The sum of the number of tokens in the system and user strings. """ - environment = Environment(undefined=StrictUndefined) - system_prompt = environment.from_string(system).render(vars) - user_prompt = environment.from_string(user).render(vars) - system_prompt_tokens = len(encoder.encode(system_prompt)) - user_prompt_tokens = len(encoder.encode(user_prompt)) - return system_prompt_tokens + user_prompt_tokens + try: + environment = Environment(undefined=StrictUndefined) + system_prompt = environment.from_string(system).render(vars) + user_prompt = environment.from_string(user).render(vars) + system_prompt_tokens = len(encoder.encode(system_prompt)) + user_prompt_tokens = len(encoder.encode(user_prompt)) + return system_prompt_tokens + user_prompt_tokens + except Exception as e: + get_logger().error(f"Error in _get_system_user_tokens: {e}") + return 0 def count_tokens(self, patch: str) -> int: """