need to test now

Author: Codeflash Bot

codeflash/api/aiservice.py

@@ -27,7 +27,7 @@
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.models.ExperimentMetadata import ExperimentMetadata
-    from codeflash.models.models import AIServiceRefinerRequest
+    from codeflash.models.models import AIServiceCodeRepairRequest, AIServiceRefinerRequest
     from codeflash.result.explanation import Explanation
 
 
@@ -294,6 +294,59 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
         console.rule()
         return []
 
+    def optimize_python_code_repair(self, request: list[AIServiceCodeRepairRequest]) -> list[OptimizedCandidate]:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Args:
+        request: A list of optimization candidate details for refinement
+
+        Returns:
+        -------
+        - List[OptimizationCandidate]: A list of Optimization Candidates.
+
+        """
+        payload = [
+            {
+                "optimization_id": opt.optimization_id,
+                "original_source_code": opt.original_source_code,
+                "modified_source_code": opt.modified_source_code,
+                "trace_id": opt.trace_id,
+            }
+            for opt in request
+        ]
+        # logger.debug(f"Repair {len(request)} optimizations…")
+        console.rule()
+        try:
+            response = self.make_ai_service_request("/code_repair", payload=payload, timeout=120)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating optimization repair: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return []
+
+        if response.status_code == 200:
+            refined_optimizations = response.json()["code_repairs"]
+            logger.debug(f"Generated {len(refined_optimizations)} candidate refinements.")
+            console.rule()
+
+            refinements = self._get_valid_candidates(refined_optimizations)
+            return [
+                OptimizedCandidate(
+                    source_code=c.source_code,
+                    explanation=c.explanation,
+                    optimization_id=c.optimization_id[:-4] + "cdrp",
+                )
+                for c in refinements
+            ]
+
+        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 get_new_explanation(  # noqa: D417
         self,
         source_code: str,

codeflash/models/models.py

@@ -48,6 +48,15 @@ class AIServiceRefinerRequest:
     function_references: str | None = None
 
 
+@dataclass(frozen=True)
+class AIServiceCodeRepairRequest:
+    optimization_id: str
+    original_source_code: str
+    modified_source_code: str
+    test_details: str
+    trace_id: str
+
+
 # If the method spam is in the class Ham, which is at the top level of the module eggs in the package foo, the fully
 # qualified name of the method is foo.eggs.Ham.spam, its qualified name is Ham.spam, and its name is spam. The full name
 # of the module is foo.eggs.

codeflash/optimization/function_optimizer.py

@@ -70,6 +70,7 @@
 from codeflash.lsp.lsp_message import LspCodeMessage, LspMarkdownMessage, LSPMessageId
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import (
+    AIServiceCodeRepairRequest,
     BestOptimization,
     CodeOptimizationContext,
     GeneratedTests,
@@ -862,6 +863,26 @@ def refine_optimizations(
         ]
         return executor.submit(ai_service_client.optimize_python_code_refinement, request=request)
 
+    def code_repair_optimizations(
+        self,
+        original_source_code: str,
+        modified_source_code: str,
+        test_details: str,
+        trace_id: str,
+        ai_service_client: AiServiceClient,
+        executor: concurrent.futures.ThreadPoolExecutor,
+    ) -> concurrent.futures.Future:
+        request = [
+            AIServiceCodeRepairRequest(
+                optimization_id="",
+                original_source_code=original_source_code,
+                modified_source_code=modified_source_code,
+                test_details=test_details,
+                trace_id=trace_id,
+            )
+        ]
+        return executor.submit(ai_service_client.optimize_python_code_repair, request=request)
+
     def log_successful_optimization(
         self, explanation: Explanation, generated_tests: GeneratedTestsList, exp_type: str
     ) -> None: