Integrate Line Profiler in Codeflash CF-470

.env

@@ -0,0 +1,2 @@
+CODEFLASH_AIS_SERVER=local
+CODEFLASH_CFAPI_SERVER=prod

codeflash/api/aiservice.py

@@ -135,6 +135,72 @@ def optimize_python_code(
         console.rule()
         return []
 
+    def optimize_python_code_line_profiler(
+        self,
+        source_code: str,
+        dependency_code: str,
+        trace_id: str,
+        line_profiler_results: str,
+        num_candidates: int = 10,
+        experiment_metadata: ExperimentMetadata | None = None,
+    ) -> list[OptimizedCandidate]:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Parameters
+        ----------
+        - source_code (str): The python code to optimize.
+        - dependency_code (str): The dependency code used as read-only context for the optimization
+        - trace_id (str): Trace id of optimization run
+        - num_candidates (int): Number of optimization variants to generate. Default is 10.
+        - experiment_metadata (Optional[ExperimentalMetadata, None]): Any available experiment metadata for this optimization
+
+        Returns
+        -------
+        - List[OptimizationCandidate]: A list of Optimization Candidates.
+
+        """
+        payload = {
+            "source_code": source_code,
+            "dependency_code": dependency_code,
+            "num_variants": num_candidates,
+            "line_profiler_results": line_profiler_results,
+            "trace_id": trace_id,
+            "python_version": platform.python_version(),
+            "experiment_metadata": experiment_metadata,
+            "codeflash_version": codeflash_version,
+        }
+
+        logger.info("Generating optimized candidates…")
+        console.rule()
+        try:
+            response = self.make_ai_service_request("/optimize-line-profiler", payload=payload, timeout=600)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating optimized candidates: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return []
+
+        if response.status_code == 200:
+            optimizations_json = response.json()["optimizations"]
+            logger.info(f"Generated {len(optimizations_json)} candidates.")
+            console.rule()
+            return [
+                OptimizedCandidate(
+                    source_code=opt["source_code"],
+                    explanation=opt["explanation"],
+                    optimization_id=opt["optimization_id"],
+                )
+                for opt in optimizations_json
+            ]
+        try:
+            error = response.json()["error"]
+        except Exception:
+            error = response.text
+        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
+        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        console.rule()
+        return []
+
+
     def log_results(
         self,
         function_trace_id: str,

codeflash/code_utils/code_replacer.py

@@ -10,6 +10,7 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import add_needed_imports_from_module
 from codeflash.models.models import FunctionParent
+import isort
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -336,3 +337,86 @@ def function_to_optimize_original_worktree_fqn(
         + "."
         + function_to_optimize.qualified_name
     )
+
+
+def add_decorator_cst(module_node, function_name, decorator_name):
+    """Adds a decorator to a function definition in a LibCST module node."""
+
+    class AddDecoratorTransformer(cst.CSTTransformer):
+        def leave_FunctionDef(self, original_node, updated_node):
+            if original_node.name.value == function_name:
+                new_decorator = cst.Decorator(
+                    decorator=cst.Name(value=decorator_name)
+                )
+
+                updated_decorators = list(updated_node.decorators)
+                updated_decorators.insert(0, new_decorator)
+
+                return updated_node.with_changes(decorators=updated_decorators)
+            return updated_node
+
+    transformer = AddDecoratorTransformer()
+    updated_module = module_node.visit(transformer)
+    return updated_module
+
+def add_decorator_imports(file_paths, fn_list, db_file):
+    """Adds a decorator to a function in a Python file."""
+    for file_path, fn_name in zip(file_paths, fn_list):
+        #open file
+        with open(file_path, "r", encoding="utf-8") as file:
+            file_contents = file.read()
+
+        # parse to cst
+        module_node = cst.parse_module(file_contents)
+        # add decorator
+        module_node = add_decorator_cst(module_node, fn_name, 'profile')
+        # add imports
+        # Create a transformer to add the import
+        transformer = ImportAdder("from line_profiler import profile")
+
+        # Apply the transformer to add the import
+        module_node = module_node.visit(transformer)
+        modified_code = isort.code(module_node.code, float_to_top=True)
+        # write to file
+        with open(file_path, "w", encoding="utf-8") as file:
+            file.write(modified_code)
+    #do this only for the main file and not the helper files, can use libcst but will go just with some simple string manipulation
+    with open(file_paths[0],'r') as f:
+        file_contents = f.readlines()
+    for idx, line in enumerate(file_contents):
+        if 'from line_profiler import profile' in line:
+            file_contents.insert(idx+1, f"profile.enable(output_prefix='{db_file}')\n")
+            break
+    with open(file_paths[0],'w') as f:
+        f.writelines(file_contents)
+
+
+
+
+class ImportAdder(cst.CSTTransformer):
+    def __init__(self, import_statement='from line_profiler import profile'):
+        self.import_statement = import_statement
+        self.has_import = False
+
+    def leave_Module(self, original_node, updated_node):
+        # If the import is already there, don't add it again
+        if self.has_import:
+            return updated_node
+
+        # Parse the import statement into a CST node
+        import_node = cst.parse_statement(self.import_statement)
+
+        # Add the import to the module's body
+        return updated_node.with_changes(
+            body=[import_node] + list(updated_node.body)
+        )
+
+    def visit_Import(self, node):
+        pass
+
+    def visit_ImportFrom(self, node):
+        # Check if the profile is already imported from line_profiler
+        if node.module and node.module.value == "line_profiler":
+            for import_alias in node.names:
+                if import_alias.name.value == "profile":
+                    self.has_import = True

codeflash/code_utils/coverage_utils.py

@@ -59,3 +59,8 @@ def prepare_coverage_files() -> tuple[Path, Path]:
     coveragerc_content = f"[run]\n branch = True\ndata_file={coverage_database_file}\n"
     coveragercfile.write_text(coveragerc_content)
     return coverage_database_file, coveragercfile
+
+def prepare_lprofiler_files(prefix="") -> tuple[Path]:
+    """Prepare line profiler output file."""
+    lprofiler_database_file = get_run_tmp_file(prefix)
+    return lprofiler_database_file

codeflash/models/models.py

@@ -218,6 +218,7 @@ class OriginalCodeBaseline(BaseModel):
     benchmarking_test_results: TestResults
     runtime: int
     coverage_results: Optional[CoverageData]
+    lprof_results: str
 
 
 class CoverageStatus(Enum):