more enhancements

Author: mohammedahmed18

codeflash/cli_cmds/console.py

@@ -3,7 +3,7 @@
 import logging
 from contextlib import contextmanager
 from itertools import cycle
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Optional
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -19,7 +19,8 @@
 
 from codeflash.cli_cmds.console_constants import SPINNER_TYPES
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
-from codeflash.lsp.helpers import enhanced_log, is_LSP_enabled
+from codeflash.lsp.helpers import is_LSP_enabled
+from codeflash.lsp.lsp_logger import enhanced_log
 from codeflash.lsp.lsp_message import LspCodeMessage
 
 if TYPE_CHECKING:
@@ -51,11 +52,11 @@
     setattr(logger, level, lambda msg, *args, _real_fn=real_fn, **kwargs: enhanced_log(msg, _real_fn, *args, **kwargs))
 
 
-def lsp_log(message: LspMessage, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
+def lsp_log(message: LspMessage) -> None:
     if not is_LSP_enabled():
         return
     json_msg = message.serialize()
-    logger.info(json_msg, *args, **kwargs)
+    logger.info(json_msg)
 
 
 def paneled_text(
@@ -81,9 +82,6 @@ def code_print(code_str: str, file_name: Optional[str] = None, function_name: Op
     from rich.syntax import Syntax
 
     formatted_code = Syntax(code_str, "python", line_numbers=True, theme="github-dark")
-    # if is_LSP_enabled():
-    #     logger.debug(formatted_code.__str__())
-    #     return
 
     console.rule()
     console.print(formatted_code)

codeflash/lsp/beta.py

@@ -11,6 +11,7 @@
 
 from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
 from codeflash.cli_cmds.cli import process_pyproject_config
+from codeflash.cli_cmds.console import code_print
 from codeflash.code_utils.git_worktree_utils import (
     create_diff_patch_from_worktree,
     get_patches_metadata,
@@ -321,6 +322,12 @@ def perform_function_optimization(  # noqa: PLR0911
 
         should_run_experiment, code_context, original_helper_code = initialization_result.unwrap()
 
+        code_print(
+            code_context.read_writable_code.flat,
+            file_name=current_function.file_path,
+            function_name=current_function.function_name,
+        )
+
         test_setup_result = function_optimizer.generate_and_instrument_tests(
             code_context, should_run_experiment=should_run_experiment
         )

codeflash/lsp/helpers.py

@@ -1,30 +1,7 @@
 import os
 from functools import lru_cache
-from typing import Any, Callable
-
-skip_lsp_log_prefix = "!lsp:"
 
 
 @lru_cache(maxsize=1)
 def is_LSP_enabled() -> bool:
     return os.getenv("CODEFLASH_LSP", default="false").lower() == "true"
-
-
-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)
-    skip_lsp_log = str_msg and msg.strip().startswith(skip_lsp_log_prefix)
-
-    if skip_lsp_log:
-        msg = msg[len(skip_lsp_log_prefix) :]
-
-    # normal cli mode
-    if not lsp_enabled:
-        actual_log_fn(msg, *args, **kwargs)
-        return
-
-    #### LSP mode ####
-    if skip_lsp_log or not str_msg:
-        return
-
-    actual_log_fn(f"::::{msg}", *args, **kwargs)

codeflash/lsp/lsp_logger.py

@@ -0,0 +1,32 @@
+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:"
+
+
+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 skip_lsp_log:
+        # get the message without the prefix
+        msg = msg[len(skip_lsp_log_prefix) :]
+
+    # normal cli mode
+    if not lsp_enabled:
+        actual_log_fn(msg, *args, **kwargs)
+        return
+
+    #### LSP mode ####
+    if skip_lsp_log or not str_msg:
+        return
+
+    if not msg.startswith("{"):
+        # it is not a json message, use a text message
+        msg = LspTextMessage(text=msg).serialize()
+
+    actual_log_fn(msg, *args, **kwargs)

codeflash/lsp/lsp_message.py

@@ -10,9 +10,8 @@
 
 @dataclass
 class LspMessage:
-    takes_time: bool = field(
-        default=False, kw_only=True
-    )  # to show a loading indicator if the operation is taking time like generating candidates or tests
+    # to show a loading indicator if the operation is taking time like generating candidates or tests
+    takes_time: bool = field(default=False, kw_only=True)
 
     def _loop_through(self, obj: Any) -> Any:  # noqa: ANN401
         if isinstance(obj, list):
@@ -29,7 +28,7 @@ def type(self) -> str:
         raise NotImplementedError
 
     def serialize(self) -> str:
-        data = asdict(self)
+        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}
         return json.dumps(ordered)
@@ -67,8 +66,3 @@ class LspStatsMessage(LspMessage):
 
     def type(self) -> str:
         return "stats"
-
-
-if __name__ == "__main__":
-    msg = LspTextMessage(text="Hello World")
-    print(msg.serialize())

codeflash/optimization/function_optimizer.py

@@ -365,6 +365,7 @@ def generate_and_instrument_tests(
             )
         )
 
+    # note: this isn't called by the lsp, only called by cli
     def optimize_function(self) -> Result[BestOptimization, str]:
         initialization_result = self.can_be_optimized()
         if not is_successful(initialization_result):