count properly

Author: KRRT7

codeflash/api/aiservice.py

@@ -4,6 +4,7 @@
 import os
 import platform
 import time
+from itertools import count
 from typing import TYPE_CHECKING, Any, cast
 
 import requests
@@ -39,6 +40,11 @@ class AiServiceClient:
     def __init__(self) -> None:
         self.base_url = self.get_aiservice_base_url()
         self.headers = {"Authorization": f"Bearer {get_codeflash_api_key()}", "Connection": "close"}
+        self.llm_call_counter = count(1)
+
+    def get_next_sequence(self) -> int:
+        """Get the next LLM call sequence number."""
+        return next(self.llm_call_counter)
 
     def get_aiservice_base_url(self) -> str:
         if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local":
@@ -151,7 +157,7 @@ def optimize_python_code(  # noqa: D417
             "repo_name": git_repo_name,
             "is_async": is_async,
             "lsp_mode": is_LSP_enabled(),
-            "call_sequence": 1,
+            "call_sequence": self.get_next_sequence(),
         }
         logger.debug(f"Sending optimize request: trace_id={trace_id}, lsp_mode={payload['lsp_mode']}")
 
@@ -218,7 +224,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             "experiment_metadata": experiment_metadata,
             "codeflash_version": codeflash_version,
             "lsp_mode": is_LSP_enabled(),
-            "call_sequence": 1,
+            "call_sequence": self.get_next_sequence(),
         }
 
         try:
@@ -269,7 +275,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
                 "trace_id": opt.trace_id,
                 "function_references": opt.function_references,
                 "python_version": platform.python_version(),
-                "call_sequence": opt.call_sequence,
+                "call_sequence": self.get_next_sequence(),
             }
             for opt in request
         ]
@@ -359,7 +365,6 @@ def get_new_explanation(  # noqa: D417
         throughput_improvement: str | None = None,
         function_references: str | None = None,
         codeflash_version: str = codeflash_version,
-        call_sequence: int | None = None,
     ) -> str:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -405,7 +410,7 @@ def get_new_explanation(  # noqa: D417
             "throughput_improvement": throughput_improvement,
             "function_references": function_references,
             "codeflash_version": codeflash_version,
-            "call_sequence": call_sequence,
+            "call_sequence": self.get_next_sequence(),
         }
         logger.info("loading|Generating explanation")
         console.rule()
@@ -533,7 +538,6 @@ def generate_regression_tests(  # noqa: D417
         test_timeout: int,
         trace_id: str,
         test_index: int,
-        call_sequence: int | None = None,
     ) -> tuple[str, str, str] | None:
         """Generate regression tests for the given function by making a request to the Django endpoint.
 
@@ -569,7 +573,7 @@ def generate_regression_tests(  # noqa: D417
             "python_version": platform.python_version(),
             "codeflash_version": codeflash_version,
             "is_async": function_to_optimize.is_async,
-            "call_sequence": call_sequence,
+            "call_sequence": self.get_next_sequence(),
         }
         try:
             response = self.make_ai_service_request("/testgen", payload=payload, timeout=90)
@@ -610,7 +614,6 @@ def get_optimization_review(
         replay_tests: str,
         concolic_tests: str,  # noqa: ARG002
         calling_fn_details: str,
-        call_sequence: int | None = None,
     ) -> str:
         """Compute the optimization review of current Pull Request.
 
@@ -626,7 +629,6 @@ def get_optimization_review(
         root_dir: Path -> path of git directory
         concolic_tests: str -> concolic_tests (not used)
         calling_fn_details: str -> filenames and definitions of functions which call the function_to_optimize
-        call_sequence: int | None -> sequence number for multi-model calls
 
         Returns:
         -------
@@ -658,7 +660,7 @@ def get_optimization_review(
             "codeflash_version": codeflash_version,
             "calling_fn_details": calling_fn_details,
             "python_version": platform.python_version(),
-            "call_sequence": call_sequence,
+            "call_sequence": self.get_next_sequence(),
         }
         console.rule()
         try:

codeflash/optimization/function_optimizer.py

@@ -327,8 +327,6 @@ def __init__(
         self.optimization_review = ""
         self.future_all_code_repair: list[concurrent.futures.Future] = []
         self.repair_counter = 0  # track how many repairs we did for each function
-        # Counter for post-optimization LLM calls (explanation, review) - optimization calls are handled by backend
-        self.post_optimization_call_count = 0
 
     def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]:
         should_run_experiment = self.experiment_id is not None
@@ -1655,10 +1653,6 @@ def process_review(
             )
             throughput_improvement_str = f"{throughput_improvement_value * 100:.1f}%"
 
-        # Explanation call sequence for tracking
-        self.post_optimization_call_count += 1
-        explanation_call_sequence = self.post_optimization_call_count
-
         new_explanation_raw_str = self.aiservice_client.get_new_explanation(
             source_code=code_context.read_writable_code.flat,
             dependency_code=code_context.read_only_context_code,
@@ -1676,7 +1670,6 @@ def process_review(
             optimized_throughput=optimized_throughput_str,
             throughput_improvement=throughput_improvement_str,
             function_references=function_references,
-            call_sequence=explanation_call_sequence,
         )
         new_explanation = Explanation(
             raw_explanation_message=new_explanation_raw_str or explanation.raw_explanation_message,
@@ -1712,13 +1705,9 @@ def process_review(
         staging_review = self.args.staging_review
         opt_review_response = ""
         # this will now run regardless of pr, staging review flags
-        # Review call sequence for tracking
-        self.post_optimization_call_count += 1
-        review_call_sequence = self.post_optimization_call_count
-
         try:
             opt_review_response = self.aiservice_client.get_optimization_review(
-                **data, calling_fn_details=function_references, call_sequence=review_call_sequence
+                **data, calling_fn_details=function_references
             )
         except Exception as e:
             logger.debug(f"optimization review response failed, investigate {e}")
@@ -2223,7 +2212,6 @@ def submit_test_generation_tasks(
                 test_index,
                 test_path,
                 test_perf_path,
-                call_sequence=test_index + 1,
             )
             for test_index, (test_path, test_perf_path) in enumerate(
                 zip(generated_test_paths, generated_perf_test_paths)

codeflash/verification/verifier.py

@@ -27,7 +27,6 @@ def generate_tests(
     test_index: int,
     test_path: Path,
     test_perf_path: Path,
-    call_sequence: int | None = None,
 ) -> tuple[str, str, Path] | None:
     # TODO: Sometimes this recreates the original Class definition. This overrides and messes up the original
     #  class import. Remove the recreation of the class definition
@@ -43,7 +42,6 @@ def generate_tests(
         test_timeout=test_timeout,
         trace_id=function_trace_id,
         test_index=test_index,
-        call_sequence=call_sequence,
     )
     if response and isinstance(response, tuple) and len(response) == 3:
         generated_test_source, instrumented_behavior_test_source, instrumented_perf_test_source = response