log tags for lsp

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -331,7 +331,7 @@ def get_new_explanation(  # noqa: D417
             "original_explanation": original_explanation,
             "dependency_code": dependency_code,
         }
-        logger.info("Generating explanation")
+        logger.info("loading|tags|Generating explanation")
         console.rule()
         try:
             response = self.make_ai_service_request("/explain", payload=payload, timeout=60)

codeflash/discovery/functions_to_optimize.py

@@ -173,15 +173,15 @@ def get_functions_to_optimize(
     with warnings.catch_warnings():
         warnings.simplefilter(action="ignore", category=SyntaxWarning)
         if optimize_all:
-            logger.info("!lsp:Finding all functions in the module '%s'…", optimize_all)
+            logger.info("!lsp|tags|Finding all functions in the module '%s'…", optimize_all)
             console.rule()
             functions = get_all_files_and_functions(Path(optimize_all))
         elif replay_test:
             functions, trace_file_path = get_all_replay_test_functions(
                 replay_test=replay_test, test_cfg=test_cfg, project_root_path=project_root
             )
         elif file is not None:
-            logger.info("!lsp:Finding all functions in the file '%s'…", file)
+            logger.info("!lsp|tags|Finding all functions in the file '%s'…", file)
             console.rule()
             functions = find_all_functions_in_file(file)
             if only_get_this_function is not None:
@@ -219,7 +219,7 @@ def get_functions_to_optimize(
             functions, test_cfg.tests_root, ignore_paths, project_root, module_root, previous_checkpoint_functions
         )
 
-        logger.info(f"!lsp:Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
+        logger.info(f"!lsp|tags|Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
         if optimize_all:
             three_min_in_ns = int(1.8e11)
             console.rule()

codeflash/lsp/lsp_logger.py

@@ -1,32 +1,61 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
 from typing import Any, Callable
 
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.lsp.lsp_message import LspTextMessage
 
-skip_lsp_log_prefix = "!lsp:"
+
+@dataclass
+class LspMessageTags:
+    # always set default values for message tags
+    not_lsp: bool = False
+    loading: bool = False
+
+
+def extract_tags(msg: str) -> tuple[LspMessageTags, str]:
+    if not isinstance(msg, str):
+        return LspMessageTags(), msg
+
+    parts = msg.split("|tags|")
+    if len(parts) == 2:
+        message_tags = LspMessageTags()
+        tags = [tag.strip() for tag in parts[0].split(",")]
+        if "!lsp" in tags:
+            message_tags.not_lsp = True
+        if "loading" in tags:
+            message_tags.loading = True
+        return message_tags, parts[1]
+
+    return LspMessageTags(), msg
 
 
 def enhanced_log(msg: str, actual_log_fn: Callable[[str, Any, Any], None], *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
     lsp_enabled = is_LSP_enabled()
-    str_msg = isinstance(msg, str)
-    # if the message starts with !lsp:, it won't be sent to the client
-    skip_lsp_log = str_msg and msg.strip().startswith(skip_lsp_log_prefix)
+    if not lsp_enabled or not isinstance(msg, str):
+        actual_log_fn(msg, *args, **kwargs)
+        return
+
+    is_lsp_json_message = msg.startswith('{"type"')
+    is_normal_text_message = not is_lsp_json_message
 
-    if skip_lsp_log:
-        # get the message without the prefix
-        msg = msg[len(skip_lsp_log_prefix) :]
+    tags = LspMessageTags()
+    clean_msg = msg
+
+    if is_normal_text_message:
+        tags, clean_msg = extract_tags(msg)
 
     # normal cli mode
     if not lsp_enabled:
-        actual_log_fn(msg, *args, **kwargs)
+        actual_log_fn(clean_msg, *args, **kwargs)
         return
 
     #### LSP mode ####
-    if skip_lsp_log or not str_msg:
+    if tags.not_lsp:
         return
 
-    if not msg.startswith("{"):
-        # it is not a json message, use a text message
-        msg = LspTextMessage(text=msg).serialize()
+    if is_normal_text_message:
+        clean_msg = LspTextMessage(text=clean_msg, takes_time=tags.loading).serialize()
 
-    actual_log_fn(msg, *args, **kwargs)
+    actual_log_fn(clean_msg, *args, **kwargs)

codeflash/optimization/function_optimizer.py

@@ -1552,6 +1552,8 @@ def run_optimized_candidate(
                 console.rule()
                 return Failure("Test results did not match the test results of the original code.")
 
+            # TODO: Add here a log message for benchmarking with takes_time tag.
+
             if test_framework == "pytest":
                 candidate_benchmarking_results, _ = self.run_and_parse_tests(
                     testing_type=TestingMode.PERFORMANCE,
@@ -1681,7 +1683,7 @@ def run_and_parse_tests(
             return TestResults(), None
         if run_result.returncode != 0 and testing_type == TestingMode.BEHAVIOR:
             logger.debug(
-                f"!lsp:Nonzero return code {run_result.returncode} when running tests in "
+                f"!lsp|tags|Nonzero return code {run_result.returncode} when running tests in "
                 f"{', '.join([str(f.instrumented_behavior_file_path) for f in test_files.test_files])}.\n"
                 f"stdout: {run_result.stdout}\n"
                 f"stderr: {run_result.stderr}\n"