formatting code

bot/chatgpt/chat_gpt_session.py

@@ -1,20 +1,23 @@
 from bot.session_manager import Session
 from common.log import logger
-'''
+
+"""
     e.g.  [
         {"role": "system", "content": "You are a helpful assistant."},
         {"role": "user", "content": "Who won the world series in 2020?"},
         {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
         {"role": "user", "content": "Where was it played?"}
     ]
-'''
+"""
+
+
 class ChatGPTSession(Session):
-    def __init__(self, session_id, system_prompt=None, model= "gpt-3.5-turbo"):
+    def __init__(self, session_id, system_prompt=None, model="gpt-3.5-turbo"):
         super().__init__(session_id, system_prompt)
         self.model = model
         self.reset()
-    
-    def discard_exceeding(self, max_tokens, cur_tokens= None):
+
+    def discard_exceeding(self, max_tokens, cur_tokens=None):
         precise = True
         try:
             cur_tokens = self.calc_tokens()
@@ -22,7 +25,9 @@ class ChatGPTSession(Session):
             precise = False
             if cur_tokens is None:
                 raise e
-            logger.debug("Exception when counting tokens precisely for query: {}".format(e))
+            logger.debug(
+                "Exception when counting tokens precisely for query: {}".format(e)
+            )
         while cur_tokens > max_tokens:
             if len(self.messages) > 2:
                 self.messages.pop(1)
@@ -34,25 +39,32 @@ class ChatGPTSession(Session):
                     cur_tokens = cur_tokens - max_tokens
                 break
             elif len(self.messages) == 2 and self.messages[1]["role"] == "user":
-                logger.warn("user message exceed max_tokens. total_tokens={}".format(cur_tokens))
+                logger.warn(
+                    "user message exceed max_tokens. total_tokens={}".format(cur_tokens)
+                )
                 break
             else:
-                logger.debug("max_tokens={}, total_tokens={}, len(messages)={}".format(max_tokens, cur_tokens, len(self.messages)))
+                logger.debug(
+                    "max_tokens={}, total_tokens={}, len(messages)={}".format(
+                        max_tokens, cur_tokens, len(self.messages)
+                    )
+                )
                 break
             if precise:
                 cur_tokens = self.calc_tokens()
             else:
                 cur_tokens = cur_tokens - max_tokens
         return cur_tokens
-    
+
     def calc_tokens(self):
         return num_tokens_from_messages(self.messages, self.model)
-    
+
 
 # refer to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
 def num_tokens_from_messages(messages, model):
     """Returns the number of tokens used by a list of messages."""
     import tiktoken
+
     try:
         encoding = tiktoken.encoding_for_model(model)
     except KeyError:
@@ -63,13 +75,17 @@ def num_tokens_from_messages(messages, model):
     elif model == "gpt-4":
         return num_tokens_from_messages(messages, model="gpt-4-0314")
     elif model == "gpt-3.5-turbo-0301":
-        tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
+        tokens_per_message = (
+            4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
+        )
         tokens_per_name = -1  # if there's a name, the role is omitted
     elif model == "gpt-4-0314":
         tokens_per_message = 3
         tokens_per_name = 1
     else:
-        logger.warn(f"num_tokens_from_messages() is not implemented for model {model}. Returning num tokens assuming gpt-3.5-turbo-0301.")
+        logger.warn(
+            f"num_tokens_from_messages() is not implemented for model {model}. Returning num tokens assuming gpt-3.5-turbo-0301."
+        )
         return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301")
     num_tokens = 0
     for message in messages: