refactoring and test

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -81,6 +81,19 @@ def make_ai_service_request(
         # response.raise_for_status()  # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
         return response
 
+    def _get_valid_candidates(self, optimizations_json: list[dict[str, Any]]) -> list[OptimizedCandidate]:
+        candidates: list[OptimizedCandidate] = []
+        for opt in optimizations_json:
+            code = CodeStringsMarkdown.parse_markdown_code(opt["source_code"])
+            if not code.code_strings:
+                continue
+            candidates.append(
+                OptimizedCandidate(
+                    source_code=code, explanation=opt["explanation"], optimization_id=opt["optimization_id"]
+                )
+            )
+        return candidates
+
     def optimize_python_code(  # noqa: D417
         self,
         source_code: str,
@@ -135,18 +148,7 @@ def optimize_python_code(  # noqa: D417
             console.rule()
             end_time = time.perf_counter()
             logger.debug(f"Generating optimizations took {end_time - start_time:.2f} seconds.")
-            candidates = []
-
-            for opt in optimizations_json:
-                code = CodeStringsMarkdown.parse_markdown_code(opt["source_code"])
-                if not code.code_strings:
-                    continue
-                candidates.append(
-                    OptimizedCandidate(
-                        source_code=code, explanation=opt["explanation"], optimization_id=opt["optimization_id"]
-                    )
-                )
-            return candidates
+            return self._get_valid_candidates(optimizations_json)
         try:
             error = response.json()["error"]
         except Exception:
@@ -209,17 +211,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             optimizations_json = response.json()["optimizations"]
             logger.info(f"Generated {len(optimizations_json)} candidate optimizations using line profiler information.")
             console.rule()
-            candidates = []
-            for opt in optimizations_json:
-                code = CodeStringsMarkdown.parse_markdown_code(opt["source_code"])
-                if not code.code_strings:
-                    continue
-                candidates.append(
-                    OptimizedCandidate(
-                        source_code=code, explanation=opt["explanation"], optimization_id=opt["optimization_id"]
-                    )
-                )
-            return candidates
+            return self._get_valid_candidates(optimizations_json)
         try:
             error = response.json()["error"]
         except Exception:
@@ -269,19 +261,17 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
             refined_optimizations = response.json()["refinements"]
             logger.debug(f"Generated {len(refined_optimizations)} candidate refinements.")
             console.rule()
-            candidates = []
-            for opt in refined_optimizations:
-                code = CodeStringsMarkdown.parse_markdown_code(opt["source_code"])
-                if not code.code_strings:
-                    continue
-                candidates.append(
-                    OptimizedCandidate(
-                        source_code=code,
-                        explanation=opt["explanation"],
-                        optimization_id=opt["optimization_id"][:-4] + "refi",
-                    )
+
+            refinements = self._get_valid_candidates(refined_optimizations)
+            return [
+                OptimizedCandidate(
+                    source_code=c.source_code,
+                    explanation=c.explanation,
+                    optimization_id=c.optimization_id[:-4] + "refi",
                 )
-            return candidates
+                for c in refinements
+            ]
+
         try:
             error = response.json()["error"]
         except Exception:

codeflash/optimization/function_optimizer.py

@@ -1300,6 +1300,7 @@ def process_review(
             return
 
     def revert_code_and_helpers(self, original_helper_code: dict[Path, str]) -> None:
+        logger.info("Reverting code and helpers...")
         self.write_code_and_helpers(
             self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
         )

tests/test_validate_python_code.py

@@ -1,6 +1,7 @@
 import pytest
 from pydantic import ValidationError
 
+from codeflash.api.aiservice import AiServiceClient
 from codeflash.models.models import CodeString
 
 
@@ -41,3 +42,30 @@ def test_whitespace_only():
     whitespace_code = "    "
     cs = CodeString(code=whitespace_code)
     assert cs.code == whitespace_code
+
+def test_generated_candidates_validation():
+    ai_service = AiServiceClient()
+    code = """```python:file.py
+print name
+```"""
+    mock_generate_candidates = [
+        {
+            "source_code": code,
+            "explanation": "",
+            "optimization_id": ""
+        }
+    ]
+    candidates = ai_service._get_valid_candidates(mock_generate_candidates)
+    assert len(candidates) == 0
+    code = """```python:file.py
+print('Hello, World!')
+```"""
+    mock_generate_candidates = [
+        {
+            "source_code": code,
+            "explanation": "",
+            "optimization_id": ""
+        }
+    ]
+    candidates = ai_service._get_valid_candidates(mock_generate_candidates)
+    assert len(candidates) == 1