validating generated optimized code

Signed-off-by: ali <mohammed18200118@gmail.com>

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -135,14 +135,18 @@ 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"],
+            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"]
+                    )
                 )
-                for opt in optimizations_json
-            ]
+            return candidates
         try:
             error = response.json()["error"]
         except Exception:
@@ -205,14 +209,17 @@ 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"],
+            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"]
+                    )
                 )
-                for opt in optimizations_json
-            ]
+            return candidates
         try:
             error = response.json()["error"]
         except Exception:
@@ -262,14 +269,19 @@ 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()
-            return [
-                OptimizedCandidate(
-                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
-                    explanation=opt["explanation"],
-                    optimization_id=opt["optimization_id"][:-4] + "refi",
+            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",
+                    )
                 )
-                for opt in refined_optimizations
-            ]
+            return candidates
         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):