send&recieve markdown code

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -73,6 +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(f"Payload: {json_payload}")
+            print("======================")
             headers = {**self.headers, "Content-Type": "application/json"}
             response = requests.post(url, data=json_payload, headers=headers, timeout=timeout)
         else:
@@ -136,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_flattened_code(opt["source_code"]),
+                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"],
                 )
@@ -206,7 +209,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             console.rule()
             return [
                 OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_flattened_code(opt["source_code"]),
+                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"],
                 )
@@ -263,7 +266,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
             console.rule()
             return [
                 OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_flattened_code(opt["source_code"]),
+                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"][:-4] + "refi",
                 )

codeflash/models/models.py

@@ -165,7 +165,7 @@ def get_code_block_splitter(file_path: Path) -> str:
     return f"{LINE_SPLITTER_MARKER_PREFIX}{file_path}"
 
 
-splitter_pattern = re.compile(f"^{LINE_SPLITTER_MARKER_PREFIX}([^\n]+)\n", re.MULTILINE | re.DOTALL)
+markdown_pattern = re.compile(r"```python:([^\n]+)\n(.*?)\n```", re.DOTALL)
 
 
 class CodeStringsMarkdown(BaseModel):
@@ -200,15 +200,12 @@ def file_to_path(self) -> dict[str, str]:
         return self._cache["file_to_path"]
 
     @staticmethod
-    def parse_flattened_code(flat_code: str) -> CodeStringsMarkdown:
-        matches = list(splitter_pattern.finditer(flat_code))
+    def parse_markdown_code(markdown_code: str) -> CodeStringsMarkdown:
+        matches = markdown_pattern.findall(markdown_code)
         results = CodeStringsMarkdown()
-        for i, match in enumerate(matches):
-            start = match.end()
-            end = matches[i + 1].start() if i + 1 < len(matches) else len(flat_code)
-            file_path = match.group(1).strip()
-            code = flat_code[start:end].lstrip("\n")
-            results.code_strings.append(CodeString(code=code, file_path=Path(file_path)))
+        for file_path, code in matches:
+            path = file_path.strip()
+            results.code_strings.append(CodeString(code=code, file_path=Path(path)))
         return results
 
 

codeflash/optimization/function_optimizer.py

@@ -384,7 +384,7 @@ def determine_best_candidate(
             ai_service_client = self.aiservice_client if exp_type == "EXP0" else self.local_aiservice_client
             future_line_profile_results = executor.submit(
                 ai_service_client.optimize_python_code_line_profiler,
-                source_code=code_context.read_writable_code.flat,
+                source_code=code_context.read_writable_code.markdown,
                 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,
                 line_profiler_results=original_code_baseline.line_profile_results["str_out"],
@@ -611,10 +611,10 @@ def refine_optimizations(
         request = [
             AIServiceRefinerRequest(
                 optimization_id=opt.candidate.optimization_id,
-                original_source_code=code_context.read_writable_code.flat,
+                original_source_code=code_context.read_writable_code.markdown,
                 read_only_dependency_code=code_context.read_only_context_code,
                 original_code_runtime=humanize_runtime(original_code_baseline.runtime),
-                optimized_source_code=opt.candidate.source_code.flat,
+                optimized_source_code=opt.candidate.source_code.markdown,
                 optimized_explanation=opt.candidate.explanation,
                 optimized_code_runtime=humanize_runtime(opt.runtime),
                 speedup=f"{int(performance_gain(original_runtime_ns=original_code_baseline.runtime, optimized_runtime_ns=opt.runtime) * 100)}%",
@@ -894,7 +894,7 @@ def generate_tests_and_optimizations(
             )
             future_optimization_candidates = executor.submit(
                 self.aiservice_client.optimize_python_code,
-                read_writable_code.flat,
+                read_writable_code.markdown,
                 read_only_context_code,
                 self.function_trace_id[:-4] + "EXP0" if run_experiment else self.function_trace_id,
                 N_CANDIDATES,
@@ -913,7 +913,7 @@ def generate_tests_and_optimizations(
             if run_experiment:
                 future_candidates_exp = executor.submit(
                     self.local_aiservice_client.optimize_python_code,
-                    read_writable_code.flat,
+                    read_writable_code.markdown,
                     read_only_context_code,
                     self.function_trace_id[:-4] + "EXP1",
                     N_CANDIDATES,

tests/test_code_replacement.py

@@ -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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code=original_helper_code
     )
     new_code = code_path.read_text(encoding="utf-8")
     code_path.unlink(missing_ok=True)

tests/test_multi_file_code_replacement.py

@@ -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=CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code=original_helper_code
+        code_context=code_context, optimized_code=CodeStringsMarkdown.parse_markdown_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")

tests/test_unused_helper_revert.py

@@ -121,7 +121,7 @@ def helper_function_2(x):
 
     # Apply optimization and test reversion
     optimizer.replace_function_and_helpers_with_optimized_code(
-        code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code_with_modified_helper), original_helper_code
+        code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper), original_helper_code
     )
     # Check final file content
     final_content = main_file.read_text()
@@ -140,7 +140,7 @@ def helper_function_2(x):
     original_helper_code = {main_file: main_file.read_text()}
 
     # Apply optimization and test reversion
-    optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+    optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
     # Check final file content
     final_content = main_file.read_text()
@@ -203,7 +203,7 @@ def helper_function_2(x):
     # 1. Apply the optimization
     # 2. Detect unused helpers
     # 3. Revert unused helpers to original definitions
-    optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+    optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
     # Check final file content
     final_content = main_file.read_text()
@@ -267,7 +267,7 @@ def helper_function_2(x):
     assert len(unused_helpers) == 0, "No helpers should be detected as unused"
 
     # Apply optimization
-    optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+    optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
     # Check final file content - should contain the optimized versions
     final_content = main_file.read_text()
@@ -388,7 +388,7 @@ def helper_function_2(x):
         }
 
         # Apply optimization and test reversion
-        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
         # Check main file content
         main_content = main_file.read_text()
         assert "result1 + n * 3" in main_content, "Entrypoint function should be optimized"
@@ -436,7 +436,7 @@ def helper_function_2(x):
         }
 
         # Apply optimization and test reversion
-        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
         # Check main file content
         main_content = main_file.read_text()
@@ -563,7 +563,7 @@ def helper_method_2(self, x):
 
         # Apply optimization and test reversion
         optimizer.replace_function_and_helpers_with_optimized_code(
-            code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code_with_modified_helper), original_helper_code
+            code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper), original_helper_code
         )
 
         # Check final file content
@@ -582,7 +582,7 @@ def helper_method_2(self, x):
         # Test reversion
         original_helper_code = {main_file: main_file.read_text()}
 
-        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
         # Check final file content
         final_content = main_file.read_text()
@@ -706,7 +706,7 @@ def process_data(self, n):
 
         # Apply optimization and test reversion
         optimizer.replace_function_and_helpers_with_optimized_code(
-            code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code_with_modified_helper), original_helper_code
+            code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper), original_helper_code
         )
 
         # Check final file content
@@ -745,7 +745,7 @@ def process_data(self, n):
 
         # Apply optimization and test reversion
         optimizer.replace_function_and_helpers_with_optimized_code(
-            code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code_with_modified_helper), original_helper_code
+            code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper), original_helper_code
         )
 
         # Check final file content
@@ -1066,7 +1066,7 @@ def subtract(x, y):
         }
 
         # Apply optimization and test reversion
-        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
         # Check main file content
         main_content = main_file.read_text()
@@ -1216,7 +1216,7 @@ def divide_numbers(x, y):
         }
 
         # Apply optimization and test reversion
-        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
         # Check main file content
         main_content = main_file.read_text()
@@ -1275,7 +1275,7 @@ def divide_numbers(x, y):
         }
 
         # Apply optimization and test reversion
-        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_flattened_code(optimized_code), original_helper_code)
+        optimizer.replace_function_and_helpers_with_optimized_code(code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code)
 
         # Check main file content
         main_content = main_file.read_text()
@@ -1424,7 +1424,7 @@ def calculate_class(cls, n):
 
         # Apply optimization and test reversion
         optimizer.replace_function_and_helpers_with_optimized_code(
-            code_context, CodeStringsMarkdown.parse_flattened_code(optimized_static_code_with_modified_helper), original_helper_code
+            code_context, CodeStringsMarkdown.parse_markdown_code(optimized_static_code_with_modified_helper), original_helper_code
         )
 
         # Check final file content