Print model response time for each step

merico/unit_tests/assistants/directory_structure/relevant_file_finder.py

@@ -15,6 +15,7 @@ from llm_conf import (
 from openai_util import create_chat_completion_content
 from tools.directory_viewer import ListViewer
 from tools.tiktoken_util import get_encoding
+from tools.time_util import print_exec_time
 
 MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-4-turbo-preview"  # "gpt-3.5-turbo"
 ENCODING = (
@@ -89,6 +90,7 @@ class RelevantFileFinder(DirectoryStructureBase):
 
         return message
 
+    @print_exec_time("Model response time")
     def _find_relevant_files(self, objective: str, dir_structure_pages: List[str]) -> List[str]:
         files: List[str] = []
         for dir_structure in dir_structure_pages:

merico/unit_tests/assistants/recommend_test_context.py

@@ -12,6 +12,7 @@ from llm_conf import (
 from model import FuncToTest
 from openai_util import create_chat_completion_content
 from tools.tiktoken_util import get_encoding
+from tools.time_util import print_exec_time
 
 MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-4-turbo-preview"
 ENCODING = (
@@ -82,6 +83,7 @@ def _mk_user_msg(func_to_test: FuncToTest, contexts: List) -> str:
     return msg
 
 
+@print_exec_time("Model response time")
 def get_recommended_symbols(
     func_to_test: FuncToTest, known_context: Optional[List] = None
 ) -> List[str]:

merico/unit_tests/propose_test.py

@@ -15,6 +15,7 @@ from model import FuncToTest, TokenBudgetExceededException
 from openai_util import create_chat_completion_content
 from prompts import PROPOSE_TEST_PROMPT
 from tools.tiktoken_util import get_encoding
+from tools.time_util import print_exec_time
 
 MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-4-turbo-preview"  # "gpt-3.5-turbo"
 ENCODING = (
@@ -82,6 +83,7 @@ def _mk_user_msg(
     )
 
 
+@print_exec_time("Model response time")
 def propose_test(
     user_prompt: str,
     func_to_test: FuncToTest,

merico/unit_tests/tools/time_util.py

@@ -0,0 +1,18 @@
+import time
+from functools import wraps
+
+
+def print_exec_time(message: str):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            result = func(*args, **kwargs)
+            end_time = time.time()
+            duration = end_time - start_time
+            print(f"{message} ({duration:.3f} s)", flush=True)
+            return result
+
+        return wrapper
+
+    return decorator