fix markdown context for formatting and more refactoring

mohammedahmed18 · mohammedahmed18 · commit 3eee162d4116 · 2025-08-04T15:35:28.000+03:00
diff --git a/codeflash/api/aiservice.py b/codeflash/api/aiservice.py
@@ -73,9 +73,9 @@ def make_ai_service_request(
         url = f"{self.base_url}/ai{endpoint}"
         if method.upper() == "POST":
             json_payload = json.dumps(payload, indent=None, default=pydantic_encoder)
-            print(f"------------------------JSON PAYLOAD for {url}--------------------")
-            print(json_payload)
-            print("-------------------END OF JSON PAYLOAD--------------------")
+            # print(f"------------------------JSON PAYLOAD for {url}--------------------")
+            # print(json_payload)
+            # print("-------------------END OF JSON PAYLOAD--------------------")
             headers = {**self.headers, "Content-Type": "application/json"}
             response = requests.post(url, data=json_payload, headers=headers, timeout=timeout)
         else:
@@ -139,7 +139,7 @@ def optimize_python_code(  # noqa: D417
             logger.debug(f"Generating optimizations took {end_time - start_time:.2f} seconds.")
             return [
                 OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_splitter_markers(opt["source_code"]),
+                    source_code=CodeStringsMarkdown.parse_flattened_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"],
                 )
@@ -209,7 +209,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             console.rule()
             return [
                 OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_splitter_markers(opt["source_code"]),
+                    source_code=CodeStringsMarkdown.parse_flattened_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"],
                 )
@@ -266,7 +266,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
             console.rule()
             return [
                 OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_splitter_markers(opt["source_code"]),
+                    source_code=CodeStringsMarkdown.parse_flattened_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"][:-4] + "refi",
                 )
diff --git a/codeflash/code_utils/formatter.py b/codeflash/code_utils/formatter.py
@@ -121,7 +121,6 @@ def format_code(
 
         if check_diff and original_code_lines > 50:
             # we dont' count the formatting diff for the optimized function as it should be well-formatted
-            # TODO: This is not correct, optimized_code is not continuous, Think of a better way for doing this.
             original_code_without_opfunc = original_code.replace(optimized_code, "")
 
             original_temp = Path(test_dir_str) / "original_temp.py"
diff --git a/codeflash/models/models.py b/codeflash/models/models.py
@@ -19,7 +19,7 @@
 from typing import Annotated, Optional, cast
 
 from jedi.api.classes import Name
-from pydantic import AfterValidator, BaseModel, ConfigDict
+from pydantic import AfterValidator, BaseModel, ConfigDict, PrivateAttr
 from pydantic.dataclasses import dataclass
 
 from codeflash.cli_cmds.console import console, logger
@@ -167,16 +167,16 @@ def get_code_block_splitter(file_path: Path) -> str:
 
 class CodeStringsMarkdown(BaseModel):
     code_strings: list[CodeString] = []
-    cached_code: Optional[str] = None
+    _cache: dict = PrivateAttr(default_factory=dict)
 
     @property
     def flat(self) -> str:
-        if self.cached_code is not None:
-            return self.cached_code
-        self.cached_code = "\n".join(
+        if self._cache.get("flat") is not None:
+            return self._cache["flat"]
+        self._cache["flat"] = "\n".join(
             get_code_block_splitter(block.file_path) + "\n" + block.code for block in self.code_strings
         )
-        return self.cached_code
+        return self._cache["flat"]
 
     @property
     def markdown(self) -> str:
@@ -188,17 +188,25 @@ def markdown(self) -> str:
             ]
         )
 
+    def file_to_path(self) -> dict[str, str]:
+        if self._cache.get("file_to_path") is not None:
+            return self._cache["file_to_path"]
+        self._cache["file_to_path"] = {
+            str(code_string.file_path): code_string.code for code_string in self.code_strings
+        }
+        return self._cache["file_to_path"]
+
     @staticmethod
-    def parse_splitter_markers(code_with_markers: str) -> CodeStringsMarkdown:
+    def parse_flattened_code(flat_code: str) -> CodeStringsMarkdown:
         pattern = rf"{LINE_SPLITTER_MARKER_PREFIX}([^\n]+)\n"
-        matches = list(re.finditer(pattern, code_with_markers))
+        matches = list(re.finditer(pattern, flat_code))
 
         results = CodeStringsMarkdown()
         for i, match in enumerate(matches):
             start = match.end()
-            end = matches[i + 1].start() if i + 1 < len(matches) else len(code_with_markers)
+            end = matches[i + 1].start() if i + 1 < len(matches) else len(flat_code)
             file_path = match.group(1).strip()
-            code = code_with_markers[start:end].lstrip("\n")
+            code = flat_code[start:end].lstrip("\n")
             results.code_strings.append(CodeString(code=code, file_path=Path(file_path)))
         return results
 
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
@@ -65,7 +65,6 @@
     LINE_SPLITTER_MARKER_PREFIX,
     BestOptimization,
     CodeOptimizationContext,
-    CodeStringsMarkdown,
     GeneratedTests,
     GeneratedTestsList,
     OptimizationSet,
@@ -97,6 +96,7 @@
     from codeflash.either import Result
     from codeflash.models.models import (
         BenchmarkKey,
+        CodeStringsMarkdown,
         CoverageData,
         FunctionCalledInTest,
         FunctionSource,
@@ -419,7 +419,7 @@ def determine_best_candidate(
                     try:
                         did_update = self.replace_function_and_helpers_with_optimized_code(
                             code_context=code_context,
-                            optimized_code=candidate.source_code.flat,
+                            optimized_code=candidate.source_code,
                             original_helper_code=original_helper_code,
                         )
                         if not did_update:
@@ -680,12 +680,21 @@ def write_code_and_helpers(original_code: str, original_helper_code: dict[Path,
                 f.write(helper_code)
 
     def reformat_code_and_helpers(
-        self, helper_functions: list[FunctionSource], path: Path, original_code: str, optimized_code: str
+        self,
+        helper_functions: list[FunctionSource],
+        path: Path,
+        original_code: str,
+        optimized_context: CodeStringsMarkdown,
     ) -> tuple[str, dict[Path, str]]:
         should_sort_imports = not self.args.disable_imports_sorting
         if should_sort_imports and isort.code(original_code) != original_code:
             should_sort_imports = False
 
+        optimized_code = ""
+        if optimized_context is not None:
+            file_to_code_context = optimized_context.file_to_path()
+            optimized_code = file_to_code_context.get(str(path.relative_to(self.project_root)), "")
+
         new_code = format_code(self.args.formatter_cmds, path, optimized_code=optimized_code, check_diff=True)
         if should_sort_imports:
             new_code = sort_imports(new_code)
@@ -704,16 +713,15 @@ def reformat_code_and_helpers(
         return new_code, new_helper_code
 
     def replace_function_and_helpers_with_optimized_code(
-        self, code_context: CodeOptimizationContext, optimized_code: str, original_helper_code: str
+        self, code_context: CodeOptimizationContext, optimized_code: CodeStringsMarkdown, original_helper_code: str
     ) -> bool:
         did_update = False
         read_writable_functions_by_file_path = defaultdict(set)
         read_writable_functions_by_file_path[self.function_to_optimize.file_path].add(
             self.function_to_optimize.qualified_name
         )
 
-        code_strings = CodeStringsMarkdown.parse_splitter_markers(optimized_code).code_strings
-        file_to_code_context = {str(code_string.file_path): code_string.code for code_string in code_strings}
+        file_to_code_context = optimized_code.file_to_path()
 
         for helper_function in code_context.helper_functions:
             if helper_function.jedi_definition.type != "class":
@@ -739,7 +747,7 @@ def replace_function_and_helpers_with_optimized_code(
                 preexisting_objects=code_context.preexisting_objects,
                 project_root_path=self.project_root,
             )
-        unused_helpers = detect_unused_helper_functions(self.function_to_optimize, code_context, optimized_code)
+        unused_helpers = detect_unused_helper_functions(self.function_to_optimize, code_context, optimized_code.flat)
 
         # Revert unused helper functions to their original definitions
         if unused_helpers:
@@ -1092,15 +1100,15 @@ def find_and_process_best_optimization(
 
                 self.replace_function_and_helpers_with_optimized_code(
                     code_context=code_context,
-                    optimized_code=best_optimization.candidate.source_code.flat,
+                    optimized_code=best_optimization.candidate.source_code,
                     original_helper_code=original_helper_code,
                 )
 
                 new_code, new_helper_code = self.reformat_code_and_helpers(
                     code_context.helper_functions,
                     explanation.file_path,
                     self.function_to_optimize_source_code,
-                    optimized_code=best_optimization.candidate.source_code.flat,
+                    optimized_context=best_optimization.candidate.source_code,
                 )
 
                 original_code_combined = original_helper_code.copy()
@@ -1173,8 +1181,9 @@ def process_review(
         )
         new_explanation_raw_str = self.aiservice_client.get_new_explanation(
             source_code=code_context.read_writable_code.flat.replace(
-                LINE_SPLITTER_MARKER_PREFIX, "# file: "
-            ),  # for better readability to the LLM
+                LINE_SPLITTER_MARKER_PREFIX,
+                "# file: ",  # for better readability
+            ),
             dependency_code=code_context.read_only_context_code,
             trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
             optimized_code=best_optimization.candidate.source_code.flat.replace(
diff --git a/tests/test_code_replacement.py b/tests/test_code_replacement.py
@@ -13,7 +13,7 @@
     replace_functions_in_file,
 )
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.models import CodeOptimizationContext, FunctionParent, get_code_block_splitter
+from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, FunctionParent, get_code_block_splitter
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.verification_utils import TestConfig
 
@@ -73,7 +73,7 @@ def sorter(arr):
             original_helper_code[helper_function_path] = helper_code
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     final_output = code_path.read_text(encoding="utf-8")
     assert "inconsequential_var = '123'" in final_output
@@ -1742,7 +1742,7 @@ def new_function2(value):
             original_helper_code[helper_function_path] = helper_code
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = code_path.read_text(encoding="utf-8")
     code_path.unlink(missing_ok=True)
@@ -1819,7 +1819,7 @@ def new_function2(value):
             original_helper_code[helper_function_path] = helper_code
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = code_path.read_text(encoding="utf-8")
     code_path.unlink(missing_ok=True)
@@ -1897,7 +1897,7 @@ def new_function2(value):
             original_helper_code[helper_function_path] = helper_code
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = code_path.read_text(encoding="utf-8")
     code_path.unlink(missing_ok=True)
@@ -1974,7 +1974,7 @@ def new_function2(value):
             original_helper_code[helper_function_path] = helper_code
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = code_path.read_text(encoding="utf-8")
     code_path.unlink(missing_ok=True)
@@ -2052,7 +2052,7 @@ def new_function2(value):
             original_helper_code[helper_function_path] = helper_code
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = code_path.read_text(encoding="utf-8")
     code_path.unlink(missing_ok=True)
@@ -2141,7 +2141,7 @@ def new_function2(value):
             original_helper_code[helper_function_path] = helper_code
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = code_path.read_text(encoding="utf-8")
     code_path.unlink(missing_ok=True)
diff --git a/tests/test_formatter.py b/tests/test_formatter.py
@@ -10,6 +10,7 @@
 from codeflash.code_utils.formatter import format_code, sort_imports
 
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import CodeString, CodeStringsMarkdown
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.verification_utils import TestConfig
 
@@ -263,7 +264,12 @@ def _run_formatting_test(source_code: str, should_content_change: bool, expected
             helper_functions=[],
             path=target_path,
             original_code=optimizer.function_to_optimize_source_code,
-            optimized_code=optimized_function,
+            optimized_context=CodeStringsMarkdown(code_strings=[
+                CodeString(
+                    code=optimized_function,
+                    file_path=target_path.relative_to(test_dir)
+                )
+            ]),
         )
 
         content = target_path.read_text(encoding="utf8")
@@ -796,7 +802,7 @@ def _is_valid(self, item):
         return isinstance(item, dict) and "id" in item
 '''
 
-    optimization_function = """    def process(self,data):
+    optimization_function = """def process(self,data):
         '''Single quote docstring with formatting issues.'''
         return{'result':[item for item in data if self._is_valid(item)]}"""
     _run_formatting_test(source_code, True, optimized_function=optimization_function, expected=expected)
diff --git a/tests/test_multi_file_code_replacement.py b/tests/test_multi_file_code_replacement.py
@@ -1,6 +1,6 @@
 from pathlib import Path
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.models import CodeOptimizationContext, get_code_block_splitter
+from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, get_code_block_splitter
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.verification_utils import TestConfig
 
@@ -117,7 +117,7 @@ def _get_string_usage(text: str) -> Usage:
 
     func_optimizer.args = Args()
     func_optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context=code_context, optimized_code=optimized_code, original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = main_file.read_text(encoding="utf-8")
     new_helper_code = helper_file.read_text(encoding="utf-8")
diff --git a/tests/test_unused_helper_revert.py b/tests/test_unused_helper_revert.py
