Skip to content

Commit 2422e9a

Browse files
log tags for lsp
1 parent f630096 commit 2422e9a

File tree

4 files changed

+49
-18
lines changed

4 files changed

+49
-18
lines changed

codeflash/api/aiservice.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ def get_new_explanation( # noqa: D417
331331
"original_explanation": original_explanation,
332332
"dependency_code": dependency_code,
333333
}
334-
logger.info("Generating explanation")
334+
logger.info("loading|tags|Generating explanation")
335335
console.rule()
336336
try:
337337
response = self.make_ai_service_request("/explain", payload=payload, timeout=60)

codeflash/discovery/functions_to_optimize.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -173,15 +173,15 @@ def get_functions_to_optimize(
173173
with warnings.catch_warnings():
174174
warnings.simplefilter(action="ignore", category=SyntaxWarning)
175175
if optimize_all:
176-
logger.info("!lsp:Finding all functions in the module '%s'…", optimize_all)
176+
logger.info("!lsp|tags|Finding all functions in the module '%s'…", optimize_all)
177177
console.rule()
178178
functions = get_all_files_and_functions(Path(optimize_all))
179179
elif replay_test:
180180
functions, trace_file_path = get_all_replay_test_functions(
181181
replay_test=replay_test, test_cfg=test_cfg, project_root_path=project_root
182182
)
183183
elif file is not None:
184-
logger.info("!lsp:Finding all functions in the file '%s'…", file)
184+
logger.info("!lsp|tags|Finding all functions in the file '%s'…", file)
185185
console.rule()
186186
functions = find_all_functions_in_file(file)
187187
if only_get_this_function is not None:
@@ -219,7 +219,7 @@ def get_functions_to_optimize(
219219
functions, test_cfg.tests_root, ignore_paths, project_root, module_root, previous_checkpoint_functions
220220
)
221221

222-
logger.info(f"!lsp:Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
222+
logger.info(f"!lsp|tags|Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
223223
if optimize_all:
224224
three_min_in_ns = int(1.8e11)
225225
console.rule()

codeflash/lsp/lsp_logger.py

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,61 @@
1+
from __future__ import annotations
2+
3+
from dataclasses import dataclass
14
from typing import Any, Callable
25

36
from codeflash.lsp.helpers import is_LSP_enabled
47
from codeflash.lsp.lsp_message import LspTextMessage
58

6-
skip_lsp_log_prefix = "!lsp:"
9+
10+
@dataclass
11+
class LspMessageTags:
12+
# always set default values for message tags
13+
not_lsp: bool = False
14+
loading: bool = False
15+
16+
17+
def extract_tags(msg: str) -> tuple[LspMessageTags, str]:
18+
if not isinstance(msg, str):
19+
return LspMessageTags(), msg
20+
21+
parts = msg.split("|tags|")
22+
if len(parts) == 2:
23+
message_tags = LspMessageTags()
24+
tags = [tag.strip() for tag in parts[0].split(",")]
25+
if "!lsp" in tags:
26+
message_tags.not_lsp = True
27+
if "loading" in tags:
28+
message_tags.loading = True
29+
return message_tags, parts[1]
30+
31+
return LspMessageTags(), msg
732

833

934
def enhanced_log(msg: str, actual_log_fn: Callable[[str, Any, Any], None], *args: Any, **kwargs: Any) -> None: # noqa: ANN401
1035
lsp_enabled = is_LSP_enabled()
11-
str_msg = isinstance(msg, str)
12-
# if the message starts with !lsp:, it won't be sent to the client
13-
skip_lsp_log = str_msg and msg.strip().startswith(skip_lsp_log_prefix)
36+
if not lsp_enabled or not isinstance(msg, str):
37+
actual_log_fn(msg, *args, **kwargs)
38+
return
39+
40+
is_lsp_json_message = msg.startswith('{"type"')
41+
is_normal_text_message = not is_lsp_json_message
1442

15-
if skip_lsp_log:
16-
# get the message without the prefix
17-
msg = msg[len(skip_lsp_log_prefix) :]
43+
tags = LspMessageTags()
44+
clean_msg = msg
45+
46+
if is_normal_text_message:
47+
tags, clean_msg = extract_tags(msg)
1848

1949
# normal cli mode
2050
if not lsp_enabled:
21-
actual_log_fn(msg, *args, **kwargs)
51+
actual_log_fn(clean_msg, *args, **kwargs)
2252
return
2353

2454
#### LSP mode ####
25-
if skip_lsp_log or not str_msg:
55+
if tags.not_lsp:
2656
return
2757

28-
if not msg.startswith("{"):
29-
# it is not a json message, use a text message
30-
msg = LspTextMessage(text=msg).serialize()
58+
if is_normal_text_message:
59+
clean_msg = LspTextMessage(text=clean_msg, takes_time=tags.loading).serialize()
3160

32-
actual_log_fn(msg, *args, **kwargs)
61+
actual_log_fn(clean_msg, *args, **kwargs)

codeflash/optimization/function_optimizer.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1552,6 +1552,8 @@ def run_optimized_candidate(
15521552
console.rule()
15531553
return Failure("Test results did not match the test results of the original code.")
15541554

1555+
# TODO: Add here a log message for benchmarking with takes_time tag.
1556+
15551557
if test_framework == "pytest":
15561558
candidate_benchmarking_results, _ = self.run_and_parse_tests(
15571559
testing_type=TestingMode.PERFORMANCE,
@@ -1681,7 +1683,7 @@ def run_and_parse_tests(
16811683
return TestResults(), None
16821684
if run_result.returncode != 0 and testing_type == TestingMode.BEHAVIOR:
16831685
logger.debug(
1684-
f"!lsp:Nonzero return code {run_result.returncode} when running tests in "
1686+
f"!lsp|tags|Nonzero return code {run_result.returncode} when running tests in "
16851687
f"{', '.join([str(f.instrumented_behavior_file_path) for f in test_files.test_files])}.\n"
16861688
f"stdout: {run_result.stdout}\n"
16871689
f"stderr: {run_result.stderr}\n"

0 commit comments

Comments
 (0)