add diversity

Author: KRRT7

codeflash/api/aiservice.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import concurrent.futures
 import json
 import os
 import platform
@@ -12,7 +13,6 @@
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.code_replacer import is_zero_diff
 from codeflash.code_utils.code_utils import unified_diff_strings
-from codeflash.code_utils.config_consts import N_CANDIDATES_EFFECTIVE, N_CANDIDATES_LP_EFFECTIVE
 from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
 from codeflash.code_utils.time_utils import humanize_runtime
@@ -35,6 +35,8 @@
     from codeflash.models.models import AIServiceCodeRepairRequest, AIServiceRefinerRequest
     from codeflash.result.explanation import Explanation
 
+multi_model_executor = concurrent.futures.ThreadPoolExecutor(max_workers=10, thread_name_prefix="multi_model")
+
 
 class AiServiceClient:
     def __init__(self) -> None:
@@ -92,7 +94,7 @@ def make_ai_service_request(
         return response
 
     def _get_valid_candidates(
-        self, optimizations_json: list[dict[str, Any]], source: OptimizedCandidateSource
+        self, optimizations_json: list[dict[str, Any]], source: OptimizedCandidateSource, model: str | None = None
     ) -> list[OptimizedCandidate]:
         candidates: list[OptimizedCandidate] = []
         for opt in optimizations_json:
@@ -106,6 +108,7 @@ def _get_valid_candidates(
                     optimization_id=opt["optimization_id"],
                     source=source,
                     parent_id=opt.get("parent_id", None),
+                    model=model,
                 )
             )
         return candidates
@@ -119,6 +122,7 @@ def optimize_python_code(  # noqa: D417
         experiment_metadata: ExperimentMetadata | None = None,
         *,
         is_async: bool = False,
+        model: str | None = None,
     ) -> list[OptimizedCandidate]:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -129,6 +133,7 @@ def optimize_python_code(  # noqa: D417
         - trace_id (str): Trace id of optimization run
         - num_candidates (int): Number of optimization variants to generate. Default is 10.
         - experiment_metadata (Optional[ExperimentalMetadata, None]): Any available experiment metadata for this optimization
+        - model (str | None): Model name to use ("gpt-4.1" or "claude-sonnet-4-5"). Default is None (server default).
 
         Returns
         -------
@@ -149,8 +154,9 @@ def optimize_python_code(  # noqa: D417
             "current_username": get_last_commit_author_if_pr_exists(None),
             "repo_owner": git_repo_owner,
             "repo_name": git_repo_name,
-            "n_candidates": N_CANDIDATES_EFFECTIVE,
+            "n_candidates": num_candidates,
             "is_async": is_async,
+            "model": model,
         }
 
         logger.info("!lsp|Generating optimized candidates…")
@@ -167,7 +173,7 @@ def optimize_python_code(  # noqa: D417
             console.rule()
             end_time = time.perf_counter()
             logger.debug(f"!lsp|Generating possible optimizations took {end_time - start_time:.2f} seconds.")
-            return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE)
+            return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE, model=model)
         try:
             error = response.json()["error"]
         except Exception:
@@ -185,6 +191,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
         line_profiler_results: str,
         num_candidates: int = 10,
         experiment_metadata: ExperimentMetadata | None = None,
+        model: str | None = None,
     ) -> list[OptimizedCandidate]:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -195,6 +202,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
         - trace_id (str): Trace id of optimization run
         - num_candidates (int): Number of optimization variants to generate. Default is 10.
         - experiment_metadata (Optional[ExperimentalMetadata, None]): Any available experiment metadata for this optimization
+        - model (str | None): Model name to use ("gpt-4.1" or "claude-sonnet-4-5"). Default is None (server default).
 
         Returns
         -------
@@ -211,7 +219,8 @@ def optimize_python_code_line_profiler(  # noqa: D417
             "experiment_metadata": experiment_metadata,
             "codeflash_version": codeflash_version,
             "lsp_mode": is_LSP_enabled(),
-            "n_candidates_lp": N_CANDIDATES_LP_EFFECTIVE,
+            "n_candidates_lp": num_candidates,
+            "model": model,
         }
 
         console.rule()
@@ -232,7 +241,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
                 f"!lsp|Generated {len(optimizations_json)} candidate optimizations using line profiler information."
             )
             console.rule()
-            return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE_LP)
+            return self._get_valid_candidates(optimizations_json, OptimizedCandidateSource.OPTIMIZE_LP, model=model)
         try:
             error = response.json()["error"]
         except Exception:
@@ -242,6 +251,95 @@ def optimize_python_code_line_profiler(  # noqa: D417
         console.rule()
         return []
 
+    def optimize_python_code_multi_model(
+        self,
+        source_code: str,
+        dependency_code: str,
+        base_trace_id: str,
+        model_distribution: list[tuple[str, int]],
+        experiment_metadata: ExperimentMetadata | None = None,
+        *,
+        is_async: bool = False,
+    ) -> list[OptimizedCandidate]:
+        """Generate optimizations using multiple models in parallel."""
+        futures: list[tuple[concurrent.futures.Future[list[OptimizedCandidate]], str]] = []
+        call_index = 0
+
+        for model_name, num_calls in model_distribution:
+            for _ in range(num_calls):
+                call_trace_id = f"{base_trace_id[:-4]}M{call_index:02d}"
+                call_index += 1
+
+                future = multi_model_executor.submit(
+                    self.optimize_python_code,
+                    source_code,
+                    dependency_code,
+                    call_trace_id,
+                    num_candidates=1,  # Each call returns 1 candidate
+                    experiment_metadata=experiment_metadata,
+                    is_async=is_async,
+                    model=model_name,
+                )
+                futures.append((future, model_name))
+
+        # Wait for all calls to complete
+        concurrent.futures.wait([f for f, _ in futures])
+
+        # Collect results
+        all_candidates: list[OptimizedCandidate] = []
+        for future, model_name in futures:
+            try:
+                candidates = future.result()
+                all_candidates.extend(candidates)
+            except Exception as e:
+                logger.warning(f"Model {model_name} call failed: {e}")
+                continue
+
+        return all_candidates
+
+    def optimize_python_code_line_profiler_multi_model(
+        self,
+        source_code: str,
+        dependency_code: str,
+        base_trace_id: str,
+        line_profiler_results: str,
+        model_distribution: list[tuple[str, int]],
+        experiment_metadata: ExperimentMetadata | None = None,
+    ) -> list[OptimizedCandidate]:
+        """Generate line profiler optimizations using multiple models in parallel."""
+        futures: list[tuple[concurrent.futures.Future[list[OptimizedCandidate]], str]] = []
+        call_index = 0
+
+        for model_name, num_calls in model_distribution:
+            for _ in range(num_calls):
+                call_trace_id = f"{base_trace_id[:-4]}L{call_index:02d}"
+                call_index += 1
+
+                future = multi_model_executor.submit(
+                    self.optimize_python_code_line_profiler,
+                    source_code,
+                    dependency_code,
+                    call_trace_id,
+                    line_profiler_results,
+                    num_candidates=1,
+                    experiment_metadata=experiment_metadata,
+                    model=model_name,
+                )
+                futures.append((future, model_name))
+
+        concurrent.futures.wait([f for f, _ in futures])
+
+        all_candidates: list[OptimizedCandidate] = []
+        for future, model_name in futures:
+            try:
+                candidates = future.result()
+                all_candidates.extend(candidates)
+            except Exception as e:
+                logger.warning(f"Line profiler model {model_name} call failed: {e}")
+                continue
+
+        return all_candidates
+
     def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]) -> list[OptimizedCandidate]:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 

codeflash/code_utils/config_consts.py

@@ -35,28 +35,16 @@
 # 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),
-]
+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),
-]
+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),
-]
+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),
-]
+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

codeflash/optimization/function_optimizer.py

@@ -48,8 +48,6 @@
     MAX_REPAIRS_PER_TRACE,
     MODEL_DISTRIBUTION_EFFECTIVE,
     MODEL_DISTRIBUTION_LP_EFFECTIVE,
-    N_CANDIDATES_EFFECTIVE,
-    N_CANDIDATES_LP_EFFECTIVE,
     N_TESTS_TO_GENERATE_EFFECTIVE,
     REFINE_ALL_THRESHOLD,
     REFINED_CANDIDATE_RANKING_WEIGHTS,
@@ -936,7 +934,6 @@ def determine_best_candidate(
             )
             if self.experiment_id
             else None,
-            executor=self.executor,
         )
 
         processor = CandidateProcessor(
@@ -1367,7 +1364,6 @@ def generate_optimizations(
             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(
@@ -1391,7 +1387,6 @@ def generate_optimizations(
                 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)