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

A less hacky way

Browse Source
This commit is contained in:
Ori Kotek
2023-08-09 12:17:54 +03:00
parent e3274af831
commit e0f295659d
8 changed files with 43 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
pr_agent/algo/token_handler.py
View File

@ -4,6 +4,10 @@ from tiktoken import encoding_for_model, get_encoding
from pr_agent.config_loader import get_settings
def get_token_encoder():
return encoding_for_model(get_settings().config.model) if "gpt" in get_settings().config.model else get_encoding(
"cl100k_base")
class TokenHandler:
"""
A class for handling tokens in the context of a pull request.
@ -27,7 +31,7 @@ class TokenHandler:
- system: The system string.
- user: The user string.
"""
self.encoder = encoding_for_model(get_settings().config.model) if "gpt" in get_settings().config.model else get_encoding("cl100k_base")
self.encoder = get_token_encoder()
self.prompt_tokens = self._get_system_user_tokens(pr, self.encoder, vars, system, user)
def _get_system_user_tokens(self, pr, encoder, vars: dict, system, user):