further streamlining

Author: aseembits93

codeflash/api/aiservice.py

@@ -7,38 +7,22 @@
 from typing import TYPE_CHECKING, Any
 
 import requests
-from pydantic.dataclasses import dataclass
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.env_utils import get_codeflash_api_key, is_LSP_enabled
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
-from codeflash.models.models import OptimizedCandidate
+from codeflash.models.models import AIServiceRefinerRequest, OptimizedCandidate
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.version import __version__ as codeflash_version
 
 if TYPE_CHECKING:
     from pathlib import Path
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-
-
-@dataclass(frozen=True)
-class AIServiceRefinerRequest:
-    optimization_id: str
-    original_source_code: str
-    read_only_dependency_code: str
-    original_code_runtime: str
-    optimized_source_code: str
-    optimized_explanation: str
-    optimized_code_runtime: str
-    speedup: str
-    trace_id: str
-    fto_name: str
-    original_line_profiler_results: str
-    optimized_line_profiler_results: str
-    experiment_metadata: ExperimentMetadata | None
+    from codeflash.models.ExperimentMetadata import ExperimentMetadata
+    from codeflash.models.models import AIServiceRefinerRequest
 
 
 class AiServiceClient:
@@ -238,7 +222,16 @@ def optimize_python_code_line_profiler(  # noqa: D417
         return []
 
     def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]) -> list[OptimizedCandidate]:
-        git_repo_owner, git_repo_name = safe_get_repo_owner_and_name()
+        """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,
@@ -251,34 +244,10 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
                 "optimized_line_profiler_results": opt.optimized_line_profiler_results,
                 "optimized_code_runtime": opt.optimized_code_runtime,
                 "speedup": opt.speedup,
-                "python_version": platform.python_version(),
-                "experiment_metadata": opt.experiment_metadata,
-                "codeflash_version": codeflash_version,
-                "lsp_mode": is_LSP_enabled(),
-                # needed for tracking the refinement behavior
                 "trace_id": opt.trace_id,
-                "function_to_optimize": opt.fto_name,
-                "repo_owner": git_repo_owner,
-                "repo_name": git_repo_name,
             }
             for opt in request
         ]
-        """Optimize the given python code for performance by making a request to the Django endpoint.
-
-        Parameters
-        ----------
-        - source_code (str): The python code to optimize.
-        - dependency_code (str): The dependency code used as read-only context for the optimization
-        - 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
-
-        Returns
-        -------
-        - List[OptimizationCandidate]: A list of Optimization Candidates.
-
-        """
-
         logger.info(f"Refining {len(request)} optimizations…")
         console.rule()
         try:

codeflash/models/models.py

@@ -27,6 +27,22 @@
 from codeflash.code_utils.env_utils import is_end_to_end
 from codeflash.verification.comparator import comparator
 
+
+@dataclass(frozen=True)
+class AIServiceRefinerRequest:
+    optimization_id: str
+    original_source_code: str
+    read_only_dependency_code: str
+    original_code_runtime: str
+    optimized_source_code: str
+    optimized_explanation: str
+    optimized_code_runtime: str
+    speedup: str
+    trace_id: str
+    original_line_profiler_results: str
+    optimized_line_profiler_results: 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.