more enhancements

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -144,7 +144,7 @@ def optimize_python_code(  # noqa: D417
 
         if response.status_code == 200:
             optimizations_json = response.json()["optimizations"]
-            logger.info(f"Generated {len(optimizations_json)} candidate optimizations.")
+            logger.info(f"Generated `{len(optimizations_json)}` candidate optimizations.")
             console.rule()
             end_time = time.perf_counter()
             logger.debug(f"Generating optimizations took {end_time - start_time:.2f} seconds.")

codeflash/cli_cmds/console.py

@@ -21,7 +21,7 @@
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.lsp.lsp_logger import enhanced_log
-from codeflash.lsp.lsp_message import LspCodeMessage
+from codeflash.lsp.lsp_message import LspCodeMessage, LspTextMessage
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -100,6 +100,11 @@ def progress_bar(
     If revert_to_print is True, falls back to printing a single logger.info message
     instead of showing a progress bar.
     """
+    if is_LSP_enabled():
+        lsp_log(LspTextMessage(text=message, takes_time=True))
+        yield
+        return
+
     if revert_to_print:
         logger.info(message)
 

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("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("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"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

@@ -1,7 +1,28 @@
 import os
+import re
 from functools import lru_cache
 
 
 @lru_cache(maxsize=1)
 def is_LSP_enabled() -> bool:
     return os.getenv("CODEFLASH_LSP", default="false").lower() == "true"
+
+
+worktree_path_regex = re.compile(r'\/[^"]*worktrees\/[^"]\S*')
+
+
+def simplify_worktree_paths(msg: str, highlight: bool = True) -> str:  # noqa: FBT001, FBT002
+    path_in_msg = worktree_path_regex.search(msg)
+    if path_in_msg:
+        last_part_of_path = path_in_msg.group(0).split("/")[-1]
+        if highlight:
+            last_part_of_path = f"`{last_part_of_path}`"
+        return msg.replace(path_in_msg.group(0), last_part_of_path)
+    return msg
+
+
+def replace_quotes_with_backticks(text: str) -> str:
+    # double-quoted strings
+    text = re.sub(r'"(.*?)"', r"`\1`", text)
+    # single-quoted strings
+    return re.sub(r"'(.*?)'", r"`\1`", text)

codeflash/lsp/lsp_message.py

@@ -5,7 +5,10 @@
 from pathlib import Path
 from typing import Any, Optional
 
+from codeflash.lsp.helpers import replace_quotes_with_backticks, simplify_worktree_paths
+
 json_primitive_types = (str, float, int, bool)
+max_code_lines_before_collapse = 45
 
 
 @dataclass
@@ -28,6 +31,15 @@ def type(self) -> str:
         raise NotImplementedError
 
     def serialize(self) -> str:
+        if isinstance(self, LspTextMessage):
+            self.text = simplify_worktree_paths(self.text)
+            self.text = replace_quotes_with_backticks(self.text)
+        if isinstance(self, LspCodeMessage):
+            self.file_name = simplify_worktree_paths(str(self.file_name), highlight=False)
+        if isinstance(self, LspMarkdownMessage):
+            self.markdown = simplify_worktree_paths(self.markdown)
+            self.markdown = replace_quotes_with_backticks(self.markdown)
+
         data = self._loop_through(asdict(self))
         # Important: keep type as the first key, for making it easy and fast for the client to know if this is a lsp message before parsing it
         ordered = {"type": self.type(), **data}
@@ -47,10 +59,19 @@ class LspCodeMessage(LspMessage):
     code: str
     file_name: Optional[Path] = None
     function_name: Optional[str] = None
+    collapsed: bool = False
+    lines_count: Optional[int] = None
 
     def type(self) -> str:
         return "code"
 
+    def serialize(self) -> str:
+        code_lines_length = len(self.code.split("\n"))
+        self.lines_count = code_lines_length
+        if code_lines_length > max_code_lines_before_collapse:
+            self.collapsed = True
+        return super().serialize()
+
 
 @dataclass
 class LspMarkdownMessage(LspMessage):

codeflash/optimization/optimizer.py

@@ -192,7 +192,7 @@ def prepare_module_for_optimization(
         from codeflash.code_utils.code_replacer import normalize_code, normalize_node
         from codeflash.code_utils.static_analysis import analyze_imported_modules
 
-        logger.info(f"Examining file {original_module_path!s}…")
+        logger.info(f"Examining file {original_module_path!s} …")
         console.rule()
 
         original_module_code: str = original_module_path.read_text(encoding="utf8")