Added context trimming.

Author: codeastra2

README.md

@@ -27,7 +27,15 @@ response  = chatgpt_client.query(prompt, w_context=True, add_to_context=False)
 ```
 
 ## Features
-We currently have the features of
+
+
+- [x] Save Conversations to a file
+- [x] Resume conversations by loading context from a file. 
+- [x] Retry logic in case of API failures.
+- [x] Regular Trimming of context to 4000 tokens so that limit of 4097 is not breached. 
+- [x] Total token and token vs time metric. 
+
+
 - Retries: This is incase of failures like connection based request exceptions, API errors.
 ```
     (openai) C:\Users\Srinivas\OneDrive\Desktop\StartupSearchGPT\tests>python test_main.py
@@ -39,6 +47,9 @@ We currently have the features of
     Retrying after 12 seconds...
     Error occurred: API error , please try later
 ``` 
+- Context trimming: Context is trimmed as needed when the limit breaches 4000 tokens. 
+![Trimming and printing metrics](printed_metrics.png)
+
 - Tracking metrics such as average time per response and total token usage.
 ```
     04-10-2023 10:26:44 | INFO | The time taken for this response is : 7.85 seconds

printed_metrics.png

@@ -73,6 +73,8 @@ class ChatGptSmartClient(object):
     read more on that here: https://platform.openai.com/docs/guides/chat .
     """
 
+    CONTEXT_TOKEN_LIMIT = 4000
+
     def __init__(self, api_key: str, model: str, log_info: bool=False):
         openai.api_key = api_key
 
@@ -89,6 +91,10 @@ def __init__(self, api_key: str, model: str, log_info: bool=False):
         self.rsp_tstamp_list = []
         self.total_token_cnt_list = []
 
+        self.rspid_vs_tottokens_dict = {}
+
+        self.total_token_cnt = 0 
+
         self.log_info = log_info
         self.logger = logging.getLogger("chatgptlogger")
         self.logger.setLevel(logging.INFO)
@@ -105,6 +111,9 @@ def query(self, query: str, w_context=True, add_to_context=True):
         query = {"role": "user", "content": query}
         rsp_id = None
 
+        if sum(self.total_token_cnt_list) >= self.CONTEXT_TOKEN_LIMIT:
+            self.trim_conversation()
+
         if w_context:
             msgs = self.prev_msgs[:]
             msgs.append(query)
@@ -126,13 +135,20 @@ def query(self, query: str, w_context=True, add_to_context=True):
         self.rsp_time_list.append(end_time - start_time)
         self.total_token_cnt_list.append(tot_token_cnt)
 
+        self.total_token_cnt += tot_token_cnt
+
+
         if self.log_info:
-            self.logger.info(f"The total token count currently is {sum(self.total_token_cnt_list)}")
+            self.logger.info(f"The total token count currently is {self.total_token_cnt}")
 
         if add_to_context:
+            self.prev_msgs.append(query)
             self.prev_msgs.append(f_resp)
             self.rsp_id += 1
             rsp_id = self.rsp_id
+            self.rspid_vs_tottokens_dict[self.rsp_id] = tot_token_cnt
+
+        #print(self.prev_msgs)
 
         return f_resp, rsp_id
 
@@ -163,7 +179,7 @@ def rollback_conversation(self, rsp_id):
         self.rsp_id = len(self.prev_msgs)
 
     def print_metrics(self):
-        self.logger.info(f"The total tokens used up-till now is: {sum(self.total_token_cnt_list)}")
+        self.logger.info(f"The total tokens used up-till now is: {self.total_token_cnt}")
         self.logger.info(f"The average response time is: {sum(self.rsp_time_list)/len(self.rsp_time_list)} sec")
 
         self.plot_rsp_times()
@@ -218,7 +234,7 @@ def dicts_to_jsonl(self, data_list: list, filename: str, compress: bool = True)
                 jout = json.dumps(ddict) + '\n'
                 out.write(jout)
 
-    def load_context_from_a_file(self, file):
+    def load_context_from_a_file(self, filename):
 
         sjsonl = '.jsonl'
 
@@ -228,8 +244,17 @@ def load_context_from_a_file(self, file):
 
         self.prev_msgs = []
 
-        with open(file, encoding='utf-8') as json_file:
+        with open(filename, encoding='utf-8') as json_file:
             for line in json_file.readlines():
                 self.prev_msgs.append(line)
 
-        return self.prev_msgs
+        return self.prev_msgs
+
+    def trim_conversation(self):
+        while sum(self.total_token_cnt_list) >= self.CONTEXT_TOKEN_LIMIT:
+
+            del  self.total_token_cnt_list[0]
+            del self.prev_msgs[1]
+            del self.prev_msgs[2]
+        
+        print(f"Trimmed the context list to length: {sum(self.total_token_cnt_list)}")

response_times.png

@@ -73,6 +73,8 @@ class ChatGptSmartClient(object):
     read more on that here: https://platform.openai.com/docs/guides/chat .
     """
 
+    CONTEXT_TOKEN_LIMIT = 4000
+
     def __init__(self, api_key: str, model: str, log_info: bool=False):
         openai.api_key = api_key
 
@@ -89,6 +91,10 @@ def __init__(self, api_key: str, model: str, log_info: bool=False):
         self.rsp_tstamp_list = []
         self.total_token_cnt_list = []
 
+        self.rspid_vs_tottokens_dict = {}
+
+        self.total_token_cnt = 0 
+
         self.log_info = log_info
         self.logger = logging.getLogger("chatgptlogger")
         self.logger.setLevel(logging.INFO)
@@ -105,6 +111,9 @@ def query(self, query: str, w_context=True, add_to_context=True):
         query = {"role": "user", "content": query}
         rsp_id = None
 
+        if sum(self.total_token_cnt_list) >= self.CONTEXT_TOKEN_LIMIT:
+            self.trim_conversation()
+
         if w_context:
             msgs = self.prev_msgs[:]
             msgs.append(query)
@@ -126,13 +135,20 @@ def query(self, query: str, w_context=True, add_to_context=True):
         self.rsp_time_list.append(end_time - start_time)
         self.total_token_cnt_list.append(tot_token_cnt)
 
+        self.total_token_cnt += tot_token_cnt
+
+
         if self.log_info:
-            self.logger.info(f"The total token count currently is {sum(self.total_token_cnt_list)}")
+            self.logger.info(f"The total token count currently is {self.total_token_cnt}")
 
         if add_to_context:
+            self.prev_msgs.append(query)
             self.prev_msgs.append(f_resp)
             self.rsp_id += 1
             rsp_id = self.rsp_id
+            self.rspid_vs_tottokens_dict[self.rsp_id] = tot_token_cnt
+
+        #print(self.prev_msgs)
 
         return f_resp, rsp_id
 
@@ -163,7 +179,7 @@ def rollback_conversation(self, rsp_id):
         self.rsp_id = len(self.prev_msgs)
 
     def print_metrics(self):
-        self.logger.info(f"The total tokens used up-till now is: {sum(self.total_token_cnt_list)}")
+        self.logger.info(f"The total tokens used up-till now is: {self.total_token_cnt}")
         self.logger.info(f"The average response time is: {sum(self.rsp_time_list)/len(self.rsp_time_list)} sec")
 
         self.plot_rsp_times()
@@ -188,7 +204,6 @@ def plot_rsp_times(self):
 
         plt.xticks(rotation=45, fontsize=6)
 
-
         # Save the figure as a PNG file
         plt.savefig("response_times.png")
 
@@ -233,4 +248,13 @@ def load_context_from_a_file(self, filename):
             for line in json_file.readlines():
                 self.prev_msgs.append(line)
 
-        return self.prev_msgs
+        return self.prev_msgs
+
+    def trim_conversation(self):
+        while sum(self.total_token_cnt_list) >= self.CONTEXT_TOKEN_LIMIT:
+
+            del  self.total_token_cnt_list[0]
+            del self.prev_msgs[1]
+            del self.prev_msgs[2]
+        
+        print(f"Trimmed the context list to length: {sum(self.total_token_cnt_list)}")

src/main.py

@@ -1,12 +1,12 @@
-api_key = "your_api_key"
+api_key = "api_key"
 model_name = "gpt-3.5-turbo"
 
 from main import ChatGptSmartClient
 
 
 chatgptsmtclient = ChatGptSmartClient(api_key=api_key, model=model_name, log_info=True)
 
-for _ in range(2):
+for _ in range(10):
     chatgptsmtclient.query("List the top 10 upcoming startups in India?")
     chatgptsmtclient.query("Ok thanks, can you giv me the valuation of these startups in tabuar format")