docstring

pr_agent/algo/token_handler.py

@@ -6,12 +6,43 @@ from pr_agent.config_loader import settings
 
 
 class TokenHandler:
+    """
+    A class for handling tokens in the context of a pull request.
+
+    Attributes:
+    - encoder: An object of the encoding_for_model class from the tiktoken module. Used to encode strings and count the number of tokens in them.
+    - limit: The maximum number of tokens allowed for the given model, as defined in the MAX_TOKENS dictionary in the pr_agent.algo module.
+    - prompt_tokens: The number of tokens in the system and user strings, as calculated by the _get_system_user_tokens method.
+    """
+
     def __init__(self, pr, vars: dict, system, user):
+        """
+        Initializes the TokenHandler object.
+
+        Args:
+        - pr: The pull request object.
+        - vars: A dictionary of variables.
+        - system: The system string.
+        - user: The user string.
+        """
         self.encoder = encoding_for_model(settings.config.model)
         self.limit = MAX_TOKENS[settings.config.model]
         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):
+        """
+        Calculates the number of tokens in the system and user strings.
+
+        Args:
+        - pr: The pull request object.
+        - encoder: An object of the encoding_for_model class from the tiktoken module.
+        - vars: A dictionary of variables.
+        - system: The system string.
+        - user: The user string.
+
+        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)
@@ -21,4 +52,13 @@ class TokenHandler:
         return system_prompt_tokens + user_prompt_tokens
 
     def count_tokens(self, patch: str) -> int:
+        """
+        Counts the number of tokens in a given patch string.
+
+        Args:
+        - patch: The patch string.
+
+        Returns:
+        The number of tokens in the patch string.
+        """
         return len(self.encoder.encode(patch, disallowed_special=()))