Merge pull request #680 from codeflash-ai/fix/validate-generated-optimized-code

[FIX] skip invalid generated canidates

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,14 +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.")
-            return [
-                OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
-                    explanation=opt["explanation"],
-                    optimization_id=opt["optimization_id"],
-                )
-                for opt in optimizations_json
-            ]
+            return self._get_valid_candidates(optimizations_json)
         try:
             error = response.json()["error"]
         except Exception:
@@ -205,14 +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()
-            return [
-                OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
-                    explanation=opt["explanation"],
-                    optimization_id=opt["optimization_id"],
-                )
-                for opt in optimizations_json
-            ]
+            return self._get_valid_candidates(optimizations_json)
         try:
             error = response.json()["error"]
         except Exception:
@@ -262,14 +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()
+
+            refinements = self._get_valid_candidates(refined_optimizations)
             return [
                 OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
-                    explanation=opt["explanation"],
-                    optimization_id=opt["optimization_id"][:-4] + "refi",
+                    source_code=c.source_code,
+                    explanation=c.explanation,
+                    optimization_id=c.optimization_id[:-4] + "refi",
                 )
-                for opt in refined_optimizations
+                for c in refinements
             ]
+
         try:
             error = response.json()["error"]
         except Exception:

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, PrivateAttr
+from pydantic import AfterValidator, BaseModel, ConfigDict, PrivateAttr, ValidationError
 from pydantic.dataclasses import dataclass
 
 from codeflash.cli_cmds.console import console, logger
@@ -239,10 +239,14 @@ def parse_markdown_code(markdown_code: str) -> CodeStringsMarkdown:
         """
         matches = markdown_pattern.findall(markdown_code)
         results = CodeStringsMarkdown()
-        for file_path, code in matches:
-            path = file_path.strip()
-            results.code_strings.append(CodeString(code=code, file_path=Path(path)))
-        return results
+        try:
+            for file_path, code in matches:
+                path = file_path.strip()
+                results.code_strings.append(CodeString(code=code, file_path=Path(path)))
+            return results  # noqa: TRY300
+        except ValidationError:
+            # if any file is invalid, return an empty CodeStringsMarkdown for the entire context
+            return CodeStringsMarkdown()
 
 
 class CodeOptimizationContext(BaseModel):

codeflash/optimization/function_optimizer.py

@@ -1351,6 +1351,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