unit tests fixing

Author: mohammedahmed18

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(f"Payload: {json_payload}")
-            print("======================")
+            logger.debug(f"========JSON PAYLOAD FOR {url}==============")
+            logger.debug(json_payload)
+            logger.debug("======================")
             headers = {**self.headers, "Content-Type": "application/json"}
             response = requests.post(url, data=json_payload, headers=headers, timeout=timeout)
         else:

codeflash/code_utils/code_replacer.py

@@ -19,7 +19,7 @@
     from pathlib import Path
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-    from codeflash.models.models import CodeOptimizationContext, OptimizedCandidate, ValidCode
+    from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, OptimizedCandidate, ValidCode
 
 ASTNodeT = TypeVar("ASTNodeT", bound=ast.AST)
 
@@ -408,16 +408,17 @@ def replace_functions_and_add_imports(
 
 def replace_function_definitions_in_module(
     function_names: list[str],
-    optimized_code: str,
+    optimized_code: CodeStringsMarkdown,
     module_abspath: Path,
     preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")
+    code_to_apply = get_optimized_code_for_module(module_abspath.relative_to(project_root_path), optimized_code)
     new_code: str = replace_functions_and_add_imports(
-        add_global_assignments(optimized_code, source_code),
+        add_global_assignments(code_to_apply, source_code),
         function_names,
-        optimized_code,
+        code_to_apply,
         module_abspath,
         preexisting_objects,
         project_root_path,
@@ -428,6 +429,19 @@ def replace_function_definitions_in_module(
     return True
 
 
+def get_optimized_code_for_module(relative_path: Path, optimized_code: CodeStringsMarkdown) -> str:
+    file_to_code_context = optimized_code.file_to_path()
+    module_optimized_code = file_to_code_context.get(str(relative_path))
+    if module_optimized_code is None:
+        logger.warning(
+            f"Optimized code not found for {relative_path} In the context\n-------\n{optimized_code}\n-------\n"
+            "re-check your 'markdown code structure'"
+            f"existing files are {file_to_code_context.keys()}"
+        )
+        module_optimized_code = ""
+    return module_optimized_code
+
+
 def is_zero_diff(original_code: str, new_code: str) -> bool:
     return normalize_code(original_code) == normalize_code(new_code)
 

codeflash/context/unused_definition_remover.py

@@ -3,16 +3,14 @@
 import ast
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from pathlib import Path
+from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
 import libcst as cst
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_replacer import replace_function_definitions_in_module
+from codeflash.models.models import CodeString, CodeStringsMarkdown
 
 if TYPE_CHECKING:
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
@@ -530,7 +528,11 @@ def revert_unused_helper_functions(
                 helper_names = [helper.qualified_name for helper in helpers_in_file]
                 reverted_code = replace_function_definitions_in_module(
                     function_names=helper_names,
-                    optimized_code=original_code,  # Use original code as the "optimized" code to revert
+                    optimized_code=CodeStringsMarkdown(
+                        code_strings=[
+                            CodeString(code=original_code, file_path=Path(file_path).relative_to(project_root))
+                        ]
+                    ),  # Use original code as the "optimized" code to revert
                     module_abspath=file_path,
                     preexisting_objects=set(),  # Empty set since we're reverting
                     project_root_path=project_root,

codeflash/models/models.py

@@ -174,6 +174,15 @@ class CodeStringsMarkdown(BaseModel):
 
     @property
     def flat(self) -> str:
+        """Returns the combined Python module from all code blocks.
+
+        Each block is prefixed by a file path comment to indicate its origin.
+        This representation is syntactically valid Python code.
+
+        Returns:
+            str: The concatenated code of all blocks with file path annotations.
+
+        """
         if self._cache.get("flat") is not None:
             return self._cache["flat"]
         self._cache["flat"] = "\n".join(
@@ -183,7 +192,15 @@ def flat(self) -> str:
 
     @property
     def markdown(self) -> str:
-        """Returns the markdown representation of the code, including the file path where possible."""
+        """Returns a Markdown-formatted string containing all code blocks.
+
+        Each block is enclosed in a triple-backtick code block with an optional
+        file path suffix (e.g., ```python:filename.py).
+
+        Returns:
+            str: Markdown representation of the code blocks.
+
+        """
         return "\n".join(
             [
                 f"```python{':' + str(code_string.file_path) if code_string.file_path else ''}\n{code_string.code.strip()}\n```"
@@ -192,6 +209,12 @@ def markdown(self) -> str:
         )
 
     def file_to_path(self) -> dict[str, str]:
+        """Return a dictionary mapping file paths to their corresponding code blocks.
+
+        Returns:
+            dict[str, str]: Mapping from file path (as string) to code.
+
+        """
         if self._cache.get("file_to_path") is not None:
             return self._cache["file_to_path"]
         self._cache["file_to_path"] = {
@@ -201,6 +224,17 @@ def file_to_path(self) -> dict[str, str]:
 
     @staticmethod
     def parse_markdown_code(markdown_code: str) -> CodeStringsMarkdown:
+        """Parse a Markdown string into a CodeStringsMarkdown object.
+
+        Extracts code blocks and their associated file paths and constructs a new CodeStringsMarkdown instance.
+
+        Args:
+            markdown_code (str): The Markdown-formatted string to parse.
+
+        Returns:
+            CodeStringsMarkdown: Parsed object containing code blocks.
+
+        """
         matches = markdown_pattern.findall(markdown_code)
         results = CodeStringsMarkdown()
         for file_path, code in matches:

codeflash/optimization/function_optimizer.py

@@ -720,29 +720,13 @@ def replace_function_and_helpers_with_optimized_code(
         read_writable_functions_by_file_path[self.function_to_optimize.file_path].add(
             self.function_to_optimize.qualified_name
         )
-
-        file_to_code_context = optimized_code.file_to_path()
-
         for helper_function in code_context.helper_functions:
             if helper_function.jedi_definition.type != "class":
                 read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name)
-
         for module_abspath, qualified_names in read_writable_functions_by_file_path.items():
-            relative_module_path = str(module_abspath.relative_to(self.project_root))
-            logger.debug(f"applying optimized code to: {relative_module_path}")
-
-            scoped_optimized_code = file_to_code_context.get(relative_module_path)
-            if scoped_optimized_code is None:
-                logger.warning(
-                    f"Optimized code not found for {relative_module_path} In the context\n-------\n{optimized_code}\n-------\n"
-                    "re-check your 'split markers'"
-                    f"existing files are {file_to_code_context.keys()}"
-                )
-                scoped_optimized_code = ""
-
             did_update |= replace_function_definitions_in_module(
                 function_names=list(qualified_names),
-                optimized_code=scoped_optimized_code,
+                optimized_code=optimized_code,
                 module_abspath=module_abspath,
                 preexisting_objects=code_context.preexisting_objects,
                 project_root_path=self.project_root,