change TimeStatistic to object instead of class

Core/Agents/CustomAgentBase.py

@@ -33,11 +33,12 @@ def _response(self, data):
     async def _get_model_response(self, _check_valid_response):
         while True:
             try:
+                timer = TimeStatistic()
                 prompt = self.model.format(Msg("system", self.prompt.system(), role="system"), self.memory.get_memory())
-                TimeStatistic.start(self.name + "Model request")
+                timer.start("Model request")
                 loop = asyncio.get_event_loop()
                 model_response = await loop.run_in_executor(model_response_thread_pool, self.model, prompt)
-                TimeStatistic.end(self.name + "Model request", self.recorder, output=True)
+                self.recorder.time(timer.end("Model request"))
 
                 if _check_valid_response(model_response):
                     if model_response.text is not None:

Core/Api/ApiRetrieverBase.py

@@ -14,6 +14,7 @@ def __init__(self, config: Config, context: Context):
         self.config = config
         self.api_index: ApiIndex = index_manager.get_index(TextConstants.API, context.get_embed_model_name(config))
         self.api_manager = api_manager
+        self.context = context
         self.recorder = Recorder(config, context.session_id)
 
     @abstractmethod

Core/Api/SimpleApiRetriever.py

@@ -11,7 +11,8 @@ def __init__(self, config: Config, context: Context):
         super().__init__(config, context)
 
     async def retrieve_apis(self, user_question, model_keywords, embed_model):
-        TimeStatistic.start("RAG retrieval")
+        timer = TimeStatistic()
+        timer.start("RAG retrieval")
         nodes = []
         nodes.extend(await self.api_index.retrieval(user_question, embed_model))
         for keyword in model_keywords:
@@ -23,5 +24,5 @@ async def retrieve_apis(self, user_question, model_keywords, embed_model):
             api = self.api_manager.get_api(api_json["name"])
             if api not in apis:
                 apis.append(api)
-        TimeStatistic.end("RAG retrieval", self.recorder, output=True)
+        self.recorder.time(timer.end("RAG retrieval"))
         return apis

Core/Common/Context.py

@@ -60,7 +60,7 @@ def _parser_model(cls, x):
 
     @classmethod
     def _parser_test_data(cls, x):
-        return x["isTest"], x["testAnswer"]
+        return x.get("isTest", False), x.get("testAnswer", None)
 
     def _get_envs_prompt(self, environments):
         prompt = ""

Core/Common/Recoder.py

@@ -42,6 +42,9 @@ def rag(self, message: str, record: bool = True):
     def model(self, message: str, record: bool = True):
         self.info(message, record, label="Model")
 
+    def time(self, message: str, record: bool = True):
+        self.info(message, record, label="Time")
+
     def info(self, message: str, record: bool = True, label="Info"):
         if record:
             self.messages.append("[{}]".format(label) + message)

Core/Common/TimeStatistic.py

@@ -1,53 +1,27 @@
 import time
-from threading import current_thread, Lock
-
-from pandas import DataFrame
-
-from Common.Recoder import Recorder
 
 
 class TimeStatistic:
-    _start_time = {}
-    _count = {}
-    _total_time = {}
-    _lock = Lock()
-
-    @classmethod
-    def start(cls, name):
-        with cls._lock:
-            thread_id = current_thread().ident
-            cls._start_time[(thread_id, name)] = time.time()
-
-    @classmethod
-    def end(cls, name, recorder: Recorder, output=True):
-        with cls._lock:
-            thread_id = current_thread().ident
-            if (thread_id, name) in cls._start_time:
-                inc_time = time.time() - cls._start_time[(thread_id, name)]
-                cls._add_time(name, inc_time)
-                del cls._start_time[(thread_id, name)]
-            else:
-                raise RuntimeError("end but no start")
-
-            if output:
-                recorder.info(f"{name} cost time: {inc_time}")
-
-    @classmethod
-    def _add_time(cls, name, inc_time):
-        if name not in cls._total_time:
-            cls._total_time[name] = 0
-            cls._count[name] = 0
-        cls._total_time[name] += inc_time
-        cls._count[name] += 1
-
-    @classmethod
-    def add_time(cls, name, inc_time):
-        with cls._lock:
-            cls._add_time(name, inc_time)
-
-    @classmethod
-    def clear(cls):
-        cls._start_time = {}
-        cls._count = {}
-        cls._total_time = {}
-        cls._lock = Lock()
+    def __init__(self):
+        self._start_time = {}
+        self._count = {}
+        self._total_time = {}
+
+    def start(self, name):
+        self._start_time[name] = time.time()
+
+    def end(self, name):
+        if name in self._start_time:
+            inc_time = time.time() - self._start_time[name]
+            self._add_time(name, inc_time)
+            del self._start_time[name]
+        else:
+            raise RuntimeError(f"TimeStatistic: {name} not started")
+        return str(inc_time)
+
+    def _add_time(self, name, inc_time):
+        if name not in self._total_time:
+            self._total_time[name] = 0
+            self._count[name] = 0
+        self._total_time[name] += inc_time
+        self._count[name] += 1

Core/Index/BaseIndex.py

@@ -22,7 +22,8 @@ def __init__(self, config: Config, identifier, data_dir, index_dir):
     def build_index(self, emb_model_config_name):
         Settings.embed_model = _EmbeddingModel(load_model_by_config_name(emb_model_config_name))
         time_identifier = "build_{}_index".format(self.identifier.lower())
-        TimeStatistic.start(time_identifier)
+        timer = TimeStatistic()
+        timer.start(time_identifier)
 
         if self.exist_index():
             sys_logger.info("Loading index for {}".format(self.identifier))
@@ -31,7 +32,8 @@ def build_index(self, emb_model_config_name):
             self.index = VectorStoreIndex([])
         self._update_index()
         self._storage_index()
-        TimeStatistic.end(time_identifier, sys_logger, output=True)
+
+        sys_logger.info("The time taken to build the index is {}".format(timer.end(time_identifier)))
 
     def exist_index(self):
         return os.path.exists(self.index_dir)