logging enhancement

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -133,6 +133,7 @@ def optimize_python_code(  # noqa: D417
             "repo_name": git_repo_name,
         }
 
+        logger.info("!lsp|tags|Generating optimized candidates…")
         console.rule()
         try:
             response = self.make_ai_service_request("/optimize", payload=payload, timeout=600)
@@ -143,8 +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.")
+            console.rule()
             end_time = time.perf_counter()
-            logger.debug(f"Generating optimizations took {end_time - start_time:.2f} seconds.")
+            logger.debug(f"!lsp|tags|Generating optimizations took {end_time - start_time:.2f} seconds.")
             return self._get_valid_candidates(optimizations_json)
         try:
             error = response.json()["error"]
@@ -191,7 +194,6 @@ def optimize_python_code_line_profiler(  # noqa: D417
             "lsp_mode": is_LSP_enabled(),
         }
 
-        logger.info("loading|tags|Generating optimized candidates using line profiler information…")
         console.rule()
         if line_profiler_results == "":
             logger.info("No LineProfiler results were provided, Skipping optimization.")
@@ -206,6 +208,9 @@ 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."
+            )
             console.rule()
             return self._get_valid_candidates(optimizations_json)
         try:

codeflash/cli_cmds/console.py

@@ -87,10 +87,8 @@ def code_print(code_str: str, file_name: Optional[str] = None, function_name: Op
     """Print code with syntax highlighting."""
     from rich.syntax import Syntax
 
-    formatted_code = Syntax(code_str, "python", line_numbers=True, theme="github-dark")
-
     console.rule()
-    console.print(formatted_code)
+    console.print(Syntax(code_str, "python", line_numbers=True, theme="github-dark"))
     console.rule()
 
 

codeflash/lsp/lsp_logger.py

@@ -11,7 +11,8 @@
 class LspMessageTags:
     # always set default values for message tags
     not_lsp: bool = False  # !lsp           (prevent the message from being sent to the LSP)
-    force_lsp: bool = False  # lsp            (you can use this to force a message to be sent to the LSP even if the level is not supported)
+    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)
 
     h1: bool = False  # h1
@@ -33,16 +34,15 @@ def add_heading_tags(msg: str, tags: LspMessageTags) -> str:
 
 
 def extract_tags(msg: str) -> tuple[Optional[LspMessageTags], str]:
-    if not isinstance(msg, str):
-        return None, 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 "lsp" in tags:
+            message_tags.lsp = True
+        if "force_lsp" in tags:
             message_tags.force_lsp = True
         if "loading" in tags:
             message_tags.loading = True
@@ -63,16 +63,21 @@ def extract_tags(msg: str) -> tuple[Optional[LspMessageTags], str]:
 
 
 def enhanced_log(
-    msg: str,
+    msg: str | Any,  # noqa: ANN401
     actual_log_fn: Callable[[str, Any, Any], None],
     level: str,
     *args: Any,  # noqa: ANN401
     **kwargs: Any,  # noqa: ANN401
 ) -> None:
-    lsp_enabled = is_LSP_enabled()
+    if not isinstance(msg, str):
+        actual_log_fn(msg, *args, **kwargs)
+        return
+
     tags, clean_msg = extract_tags(msg)
+    lsp_enabled = is_LSP_enabled()
+    lsp_only = tags and tags.lsp
 
-    if not lsp_enabled or not isinstance(clean_msg, str):
+    if not lsp_enabled and not lsp_only:
         actual_log_fn(clean_msg, *args, **kwargs)
         return
 

codeflash/lsp/server_entry.py

@@ -16,7 +16,7 @@
 # Configure logging to stderr for VS Code output channel
 def setup_logging() -> logging.Logger:
     # Clear any existing handlers to prevent conflicts
-    root_logger = logging.getLogger("rich")
+    root_logger = logging.getLogger()
     root_logger.handlers.clear()
 
     # Set up stderr handler for VS Code output channel with [LSP-Server] prefix

codeflash/optimization/function_optimizer.py

@@ -152,8 +152,6 @@ def _process_line_profiler_results(self) -> OptimizedCandidate | None:
         concurrent.futures.wait([self.future_line_profile_results])
         line_profile_results = self.future_line_profile_results.result()
 
-        logger.info(f"Generated {len(line_profile_results)} candidate optimizations using line profiler information.")
-
         for candidate in line_profile_results:
             self.candidate_queue.put(candidate)
 
@@ -517,7 +515,7 @@ def determine_best_candidate(
                     )
                     if not did_update:
                         logger.warning(
-                            "lsp|tags|No functions were replaced in the optimized code. Skipping optimization candidate."
+                            "force_lsp|tags|No functions were replaced in the optimized code. Skipping optimization candidate."
                         )
                         console.rule()
                         continue
@@ -1085,7 +1083,7 @@ def generate_tests_and_optimizations(
 
         # Retrieve results
         candidates: list[OptimizedCandidate] = future_optimization_candidates.result()
-        logger.info(f"Generated '{len(candidates)}' candidate optimizations.")
+        logger.info(f"lsp|tags|Generated '{len(candidates)}' candidate optimizations.")
         console.rule()
 
         if not candidates: