divertsity

Author: KRRT7

codeflash/code_utils/config_consts.py

@@ -32,6 +32,32 @@
 MAX_N_CANDIDATES = 5
 MAX_N_CANDIDATES_LP = 6
 
+# Multi-model diversity configuration
+# Each tuple is (model_name, num_calls) where each call returns 1 candidate
+# Standard mode: 3 GPT-4.1 + 2 Claude Sonnet = 5 candidates
+MODEL_DISTRIBUTION: list[tuple[str, int]] = [
+    ("gpt-4.1", 3),
+    ("claude-sonnet-4-5", 2),
+]
+
+# LSP mode: fewer candidates for faster response
+MODEL_DISTRIBUTION_LSP: list[tuple[str, int]] = [
+    ("gpt-4.1", 2),
+    ("claude-sonnet-4-5", 1),
+]
+
+# Line profiler mode: 6 candidates total
+MODEL_DISTRIBUTION_LP: list[tuple[str, int]] = [
+    ("gpt-4.1", 4),
+    ("claude-sonnet-4-5", 2),
+]
+
+# Line profiler LSP mode
+MODEL_DISTRIBUTION_LP_LSP: list[tuple[str, int]] = [
+    ("gpt-4.1", 2),
+    ("claude-sonnet-4-5", 1),
+]
+
 try:
     from codeflash.lsp.helpers import is_LSP_enabled
 
@@ -43,5 +69,7 @@
 N_CANDIDATES_LP_EFFECTIVE = min(N_CANDIDATES_LP_LSP if _IS_LSP_ENABLED else N_CANDIDATES_LP, MAX_N_CANDIDATES_LP)
 N_TESTS_TO_GENERATE_EFFECTIVE = N_TESTS_TO_GENERATE_LSP if _IS_LSP_ENABLED else N_TESTS_TO_GENERATE
 TOTAL_LOOPING_TIME_EFFECTIVE = TOTAL_LOOPING_TIME_LSP if _IS_LSP_ENABLED else TOTAL_LOOPING_TIME
+MODEL_DISTRIBUTION_EFFECTIVE = MODEL_DISTRIBUTION_LSP if _IS_LSP_ENABLED else MODEL_DISTRIBUTION
+MODEL_DISTRIBUTION_LP_EFFECTIVE = MODEL_DISTRIBUTION_LP_LSP if _IS_LSP_ENABLED else MODEL_DISTRIBUTION_LP
 
 MAX_CONTEXT_LEN_REVIEW = 1000

codeflash/models/models.py

@@ -464,6 +464,7 @@ class OptimizedCandidate:
     optimization_id: str
     source: OptimizedCandidateSource
     parent_id: str | None = None
+    model: str | None = None  # Which LLM model generated this candidate
 
 
 @dataclass(frozen=True)

codeflash/optimization/function_optimizer.py

@@ -46,6 +46,8 @@
     COVERAGE_THRESHOLD,
     INDIVIDUAL_TESTCASE_TIMEOUT,
     MAX_REPAIRS_PER_TRACE,
+    MODEL_DISTRIBUTION_EFFECTIVE,
+    MODEL_DISTRIBUTION_LP_EFFECTIVE,
     N_CANDIDATES_EFFECTIVE,
     N_CANDIDATES_LP_EFFECTIVE,
     N_TESTS_TO_GENERATE_EFFECTIVE,
@@ -921,18 +923,20 @@ def determine_best_candidate(
         ai_service_client = self.aiservice_client if exp_type == "EXP0" else self.local_aiservice_client
         assert ai_service_client is not None, "AI service client must be set for optimization"
 
+        # Use multi-model approach for line profiler optimization
         future_line_profile_results = self.executor.submit(
-            ai_service_client.optimize_python_code_line_profiler,
+            ai_service_client.optimize_python_code_line_profiler_multi_model,
             source_code=code_context.read_writable_code.markdown,
             dependency_code=code_context.read_only_context_code,
-            trace_id=self.get_trace_id(exp_type),
+            base_trace_id=self.get_trace_id(exp_type),
             line_profiler_results=original_code_baseline.line_profile_results["str_out"],
-            num_candidates=N_CANDIDATES_LP_EFFECTIVE,
+            model_distribution=MODEL_DISTRIBUTION_LP_EFFECTIVE,
             experiment_metadata=ExperimentMetadata(
                 id=self.experiment_id, group="control" if exp_type == "EXP0" else "experiment"
             )
             if self.experiment_id
             else None,
+            executor=self.executor,
         )
 
         processor = CandidateProcessor(
@@ -1353,17 +1357,17 @@ def generate_optimizations(
         read_only_context_code: str,
         run_experiment: bool = False,  # noqa: FBT001, FBT002
     ) -> Result[tuple[OptimizationSet, str], str]:
-        """Generate optimization candidates for the function."""
-        n_candidates = N_CANDIDATES_EFFECTIVE
-
+        """Generate optimization candidates for the function using multiple models in parallel."""
+        # Use multi-model approach for diversity
         future_optimization_candidates = self.executor.submit(
-            self.aiservice_client.optimize_python_code,
+            self.aiservice_client.optimize_python_code_multi_model,
             read_writable_code.markdown,
             read_only_context_code,
             self.function_trace_id[:-4] + "EXP0" if run_experiment else self.function_trace_id,
-            n_candidates,
+            MODEL_DISTRIBUTION_EFFECTIVE,
             ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None,
             is_async=self.function_to_optimize.is_async,
+            executor=self.executor,
         )
 
         future_references = self.executor.submit(
@@ -1380,13 +1384,14 @@ def generate_optimizations(
 
         if run_experiment:
             future_candidates_exp = self.executor.submit(
-                self.local_aiservice_client.optimize_python_code,
+                self.local_aiservice_client.optimize_python_code_multi_model,
                 read_writable_code.markdown,
                 read_only_context_code,
                 self.function_trace_id[:-4] + "EXP1",
-                n_candidates,
+                MODEL_DISTRIBUTION_EFFECTIVE,
                 ExperimentMetadata(id=self.experiment_id, group="experiment"),
                 is_async=self.function_to_optimize.is_async,
+                executor=self.executor,
             )
             futures.append(future_candidates_exp)
 
@@ -1395,7 +1400,7 @@ def generate_optimizations(
 
         # Retrieve results
         candidates: list[OptimizedCandidate] = future_optimization_candidates.result()
-        logger.info(f"!lsp|Generated '{len(candidates)}' candidate optimizations.")
+        logger.info(f"!lsp|Generated '{len(candidates)}' candidate optimizations from multiple models.")
 
         if not candidates:
             return Failure(f"/!\\ NO OPTIMIZATIONS GENERATED for {self.function_to_optimize.function_name}")