lsp logs formatting and small fixes

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -133,7 +133,7 @@ def optimize_python_code(  # noqa: D417
             "repo_name": git_repo_name,
         }
 
-        logger.info("!lsp|tags|Generating optimized candidates…")
+        logger.info("!lsp|Generating optimized candidates…")
         console.rule()
         try:
             response = self.make_ai_service_request("/optimize", payload=payload, timeout=600)
@@ -144,10 +144,10 @@ def optimize_python_code(  # noqa: D417
 
         if response.status_code == 200:
             optimizations_json = response.json()["optimizations"]
-            logger.info(f"!lsp|tags|Generated {len(optimizations_json)} candidate optimizations.")
+            logger.info(f"!lsp|Generated {len(optimizations_json)} candidate optimizations.")
             console.rule()
             end_time = time.perf_counter()
-            logger.debug(f"!lsp|tags|Generating optimizations took {end_time - start_time:.2f} seconds.")
+            logger.debug(f"!lsp|Generating possible optimizations took {end_time - start_time:.2f} seconds.")
             return self._get_valid_candidates(optimizations_json)
         try:
             error = response.json()["error"]
@@ -209,7 +209,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
         if response.status_code == 200:
             optimizations_json = response.json()["optimizations"]
             logger.info(
-                f"!lsp|tags|Generated {len(optimizations_json)} candidate optimizations using line profiler information."
+                f"!lsp|Generated {len(optimizations_json)} candidate optimizations using line profiler information."
             )
             console.rule()
             return self._get_valid_candidates(optimizations_json)
@@ -332,7 +332,7 @@ def get_new_explanation(  # noqa: D417
             "original_explanation": original_explanation,
             "dependency_code": dependency_code,
         }
-        logger.info("loading|tags|Generating explanation")
+        logger.info("loading|Generating explanation")
         console.rule()
         try:
             response = self.make_ai_service_request("/explain", payload=payload, timeout=60)
@@ -377,7 +377,7 @@ def generate_ranking(  # noqa: D417
             "optimization_ids": optimization_ids,
             "python_version": platform.python_version(),
         }
-        logger.info("Generating ranking")
+        logger.info("loading|Generating ranking")
         console.rule()
         try:
             response = self.make_ai_service_request("/rank", payload=payload, timeout=60)

codeflash/code_utils/git_worktree_utils.py

@@ -62,7 +62,7 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]:
     )
 
     if not uni_diff_text.strip():
-        logger.info("!lsp|tags|No uncommitted changes to copy to worktree.")
+        logger.info("!lsp|No uncommitted changes to copy to worktree.")
         return worktree_dir
 
     # Write the diff to a temporary file

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|tags|Finding all functions in the module '%s'…", optimize_all)
+            logger.info("!lsp|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|tags|Finding all functions in the file '%s'…", file)
+            logger.info("!lsp|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|tags|Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
+        logger.info(f"!lsp|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/helpers.py

@@ -39,7 +39,9 @@ def report_to_markdown_table(report: dict[TestType, dict[str, int]], title: str)
             continue
         lines.append(f"| {test_type.to_name()} | {passed} | {failed} |")
     table = "\n".join(lines)
-    return f"### {title}\n{table}"
+    if title:
+        return f"### {title}\n{table}"
+    return table
 
 
 def simplify_worktree_paths(msg: str, highlight: bool = True) -> str:  # noqa: FBT001, FBT002

codeflash/lsp/lsp_logger.py

@@ -14,13 +14,19 @@ class LspMessageTags:
     lsp: bool = False  # lsp                (lsp only)
     force_lsp: bool = False  # force_lsp    (you can use this to force a message to be sent to the LSP even if the level is not supported)
     loading: bool = False  # loading        (you can use this to indicate that the message is a loading message)
-
+    highlight: bool = False  # highlight    (you can use this to highlight the message by wrapping it in ``)
     h1: bool = False  # h1
     h2: bool = False  # h2
     h3: bool = False  # h3
     h4: bool = False  # h4
 
 
+def add_highlight_tags(msg: str, tags: LspMessageTags) -> str:
+    if tags.highlight:
+        return "`" + msg + "`"
+    return msg
+
+
 def add_heading_tags(msg: str, tags: LspMessageTags) -> str:
     if tags.h1:
         return "# " + msg
@@ -33,9 +39,9 @@ def add_heading_tags(msg: str, tags: LspMessageTags) -> str:
     return msg
 
 
-# TODO: Make this work when optimizing extract_tags from lsp code message that has |tags|
 def extract_tags(msg: str) -> tuple[Optional[LspMessageTags], str]:
-    parts = msg.split("|tags|")
+    delimiter = "|"
+    parts = msg.split(delimiter)
     if len(parts) == 2:
         message_tags = LspMessageTags()
         tags = {tag.strip() for tag in parts[0].split(",")}
@@ -47,6 +53,8 @@ def extract_tags(msg: str) -> tuple[Optional[LspMessageTags], str]:
             message_tags.force_lsp = True
         if "loading" in tags:
             message_tags.loading = True
+        if "highlight" in tags:
+            message_tags.highlight = True
         if "h1" in tags:
             message_tags.h1 = True
         if "h2" in tags:
@@ -55,7 +63,7 @@ def extract_tags(msg: str) -> tuple[Optional[LspMessageTags], str]:
             message_tags.h3 = True
         if "h4" in tags:
             message_tags.h4 = True
-        return message_tags, parts[1]
+        return message_tags, delimiter.join(parts[1:])
 
     return None, msg
 
@@ -74,17 +82,21 @@ def enhanced_log(
         actual_log_fn(msg, *args, **kwargs)
         return
 
-    tags, clean_msg = extract_tags(msg)
+    is_lsp_json_message = msg.startswith('{"type"')
+    is_normal_text_message = not is_lsp_json_message
+
+    # extract tags only from the text messages (not the json ones)
+    tags, clean_msg = extract_tags(msg) if is_normal_text_message else (None, msg)
+
     lsp_enabled = is_LSP_enabled()
     lsp_only = tags and tags.lsp
 
     if not lsp_enabled and not lsp_only:
+        # normal logging
         actual_log_fn(clean_msg, *args, **kwargs)
         return
 
     #### LSP mode ####
-    is_lsp_json_message = clean_msg.startswith('{"type"')
-    is_normal_text_message = not is_lsp_json_message
     final_tags = tags if tags else LspMessageTags()
 
     unsupported_level = level not in supported_lsp_log_levels
@@ -93,6 +105,7 @@ def enhanced_log(
 
     if is_normal_text_message:
         clean_msg = add_heading_tags(clean_msg, final_tags)
+        clean_msg = add_highlight_tags(clean_msg, final_tags)
         clean_msg = LspTextMessage(text=clean_msg, takes_time=final_tags.loading).serialize()
 
     actual_log_fn(clean_msg, *args, **kwargs)