xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-02 11:50:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

protections

Browse Source
This commit is contained in:
mrT23
2024-08-09 21:44:00 +03:00
parent e531245f4a
commit 84b80f792d
1 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
pr_agent/algo/token_handler.py
View File

@ -3,6 +3,8 @@ from tiktoken import encoding_for_model, get_encoding
from pr_agent.config_loader import get_settings from pr_agent.config_loader import get_settings
from threading import Lock from threading import Lock
from pr_agent.log import get_logger
class TokenEncoder: class TokenEncoder:
_encoder_instance = None _encoder_instance = None
@ -62,12 +64,16 @@ class TokenHandler:
Returns: Returns:
The sum of the number of tokens in the system and user strings. The sum of the number of tokens in the system and user strings.
""" """
environment = Environment(undefined=StrictUndefined) try:
system_prompt = environment.from_string(system).render(vars) environment = Environment(undefined=StrictUndefined)
user_prompt = environment.from_string(user).render(vars) system_prompt = environment.from_string(system).render(vars)
system_prompt_tokens = len(encoder.encode(system_prompt)) user_prompt = environment.from_string(user).render(vars)
user_prompt_tokens = len(encoder.encode(user_prompt)) system_prompt_tokens = len(encoder.encode(system_prompt))
return system_prompt_tokens + user_prompt_tokens 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: def count_tokens(self, patch: str) -> int:
""" """