Added context file storage and retrieval features.

Author: codeastra2

src/main.py

@@ -4,6 +4,7 @@
 import requests
 import logging
 import sys
+import json
 import tiktoken
 
 import matplotlib.pyplot as plt
@@ -179,15 +180,56 @@ def plot_rsp_times(self):
         ax.set_ylabel("Response Time (s)")
         ax.set_title("ChatGPT API Response Time")
         ax.tick_params(axis='x', rotation=45)
-        ax.xaxis.labelpad = 35
+        ax.xaxis.labelpad = 85
         ax.yaxis.labelpad = 35
 
         # Increase the bottom and left margins of the plot
-        plt.subplots_adjust(bottom=0.2, left=0.15, fontsize=8)
+        plt.subplots_adjust(bottom=0.2, left=0.15)
+
+        plt.xticks(rotation=45, fontsize=6)
 
         # Save the figure as a PNG file
         plt.savefig("response_times.png")
 
         # Show the plot
         plt.show()
 
+    def dump_context_to_a_file(self, filename="context"):
+        self.dicts_to_jsonl(data_list=self.prev_msgs, filename=filename)
+
+        
+    def dicts_to_jsonl(self, data_list: list, filename: str, compress: bool = True) -> None:
+        """
+        Method saves list of dicts into jsonl file.
+        :param data_list: (list) list of dicts to be stored,
+        :param filename: (str) path to the output file. If suffix .jsonl is not given then methods appends
+            .jsonl suffix into the file.
+        :param compress: (bool) should file be compressed into a gzip archive?
+        """
+        sjsonl = '.jsonl'
+
+        # Check filename
+        if not filename.endswith(sjsonl):
+            filename = filename + sjsonl
+
+        # Save data
+        with open(filename, 'w') as out:
+            for ddict in data_list:
+                jout = json.dumps(ddict) + '\n'
+                out.write(jout)
+    
+    def load_context_from_a_file(self, file):
+
+        sjsonl = '.jsonl'
+
+        # Check filename
+        if not filename.endswith(sjsonl):
+            filename = filename + sjsonl
+        
+        self.prev_msgs = []
+        
+        with open(file, encoding='utf-8') as json_file:
+            for line in json_file.readlines():
+                self.prev_msgs.append(line)
+
+        return self.prev_msgs

tests/main.py

@@ -4,6 +4,7 @@
 import requests
 import logging
 import sys
+import json
 import tiktoken
 
 import matplotlib.pyplot as plt
@@ -187,9 +188,49 @@ def plot_rsp_times(self):
 
         plt.xticks(rotation=45, fontsize=6)
 
+
         # Save the figure as a PNG file
         plt.savefig("response_times.png")
 
         # Show the plot
         plt.show()
 
+    def dump_context_to_a_file(self, filename="context"):
+        self.dicts_to_jsonl(data_list=self.prev_msgs, filename=filename)
+
+        
+    def dicts_to_jsonl(self, data_list: list, filename: str, compress: bool = True) -> None:
+        """
+        Method saves list of dicts into jsonl file.
+        :param data_list: (list) list of dicts to be stored,
+        :param filename: (str) path to the output file. If suffix .jsonl is not given then methods appends
+            .jsonl suffix into the file.
+        :param compress: (bool) should file be compressed into a gzip archive?
+        """
+        sjsonl = '.jsonl'
+
+        # Check filename
+        if not filename.endswith(sjsonl):
+            filename = filename + sjsonl
+
+        # Save data
+        with open(filename, 'w') as out:
+            for ddict in data_list:
+                jout = json.dumps(ddict) + '\n'
+                out.write(jout)
+    
+    def load_context_from_a_file(self, filename):
+
+        sjsonl = '.jsonl'
+
+        # Check filename
+        if not filename.endswith(sjsonl):
+            filename = filename + sjsonl
+        
+        self.prev_msgs = []
+        
+        with open(filename, encoding='utf-8') as json_file:
+            for line in json_file.readlines():
+                self.prev_msgs.append(line)
+
+        return self.prev_msgs

tests/test_main.py

@@ -10,4 +10,7 @@
     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")
 
+chatgptsmtclient.dump_context_to_a_file("context")
+chatgptsmtclient.load_context_from_a_file("context")
+
 chatgptsmtclient.print_metrics()