Merge branch 'main' into chore/asyncio-optimization

Author: HeshamHM28

code_to_optimize/code_directories/nested_module_root/pyproject.toml

@@ -0,0 +1,8 @@
+[tool.codeflash]
+# All paths are relative to this pyproject.toml's directory.
+module-root = "src/app"
+tests-root = "src/tests"
+test-framework = "pytest"
+ignore-paths = []
+disable-telemetry = true
+formatter-cmds = ["disabled"]

code_to_optimize/code_directories/nested_module_root/src/app/main.py

@@ -0,0 +1,10 @@
+def sorter(arr):
+    print("codeflash stdout: Sorting list")
+    for i in range(len(arr)):
+        for j in range(len(arr) - 1):
+            if arr[j] > arr[j + 1]:
+                temp = arr[j]
+                arr[j] = arr[j + 1]
+                arr[j + 1] = temp
+    print(f"result: {arr}")
+    return arr

code_to_optimize/code_directories/nested_module_root/src/tests/.gitkeep

@@ -4,27 +4,30 @@
 import os
 import platform
 import time
-from typing import TYPE_CHECKING, Any
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
 
 import requests
 from pydantic.json import pydantic_encoder
 
 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
 from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import AIServiceRefinerRequest, CodeStringsMarkdown, 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
     from codeflash.models.ExperimentMetadata import ExperimentMetadata
     from codeflash.models.models import AIServiceRefinerRequest
+    from codeflash.result.explanation import Explanation
 
 
 class AiServiceClient:
@@ -529,6 +532,85 @@ def generate_regression_tests(  # noqa: D417
             ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": response.text})
             return None
 
+    def get_optimization_impact(
+        self,
+        original_code: dict[Path, str],
+        new_code: dict[Path, str],
+        explanation: Explanation,
+        existing_tests_source: str,
+        generated_original_test_source: str,
+        function_trace_id: str,
+        coverage_message: str,
+        replay_tests: str,
+        root_dir: Path,
+        concolic_tests: str,  # noqa: ARG002
+    ) -> str:
+        """Compute the optimization impact of current Pull Request.
+
+        Args:
+        original_code: dict -> data structure mapping file paths to function definition for original code
+        new_code: dict -> data structure mapping file paths to function definition for optimized code
+        explanation: Explanation -> data structure containing runtime information
+        existing_tests_source: str -> existing tests table
+        generated_original_test_source: str -> annotated generated tests
+        function_trace_id: str -> traceid of function
+        coverage_message: str -> coverage information
+        replay_tests: str -> replay test table
+        root_dir: Path -> path of git directory
+        concolic_tests: str -> concolic_tests (not used)
+
+        Returns:
+        -------
+        - 'high' or 'low' optimization impact
+
+        """
+        diff_str = "\n".join(
+            [
+                unified_diff_strings(
+                    code1=original_code[p],
+                    code2=new_code[p],
+                    fromfile=Path(p).relative_to(root_dir).as_posix(),
+                    tofile=Path(p).relative_to(root_dir).as_posix(),
+                )
+                for p in original_code
+                if not is_zero_diff(original_code[p], new_code[p])
+            ]
+        )
+        code_diff = f"```diff\n{diff_str}\n```"
+        logger.info("!lsp|Computing Optimization Impact…")
+        payload = {
+            "code_diff": code_diff,
+            "explanation": explanation.raw_explanation_message,
+            "existing_tests": existing_tests_source,
+            "generated_tests": generated_original_test_source,
+            "trace_id": function_trace_id,
+            "coverage_message": coverage_message,
+            "replay_tests": replay_tests,
+            "speedup": f"{(100 * float(explanation.speedup)):.2f}%",
+            "loop_count": explanation.winning_benchmarking_test_results.number_of_loops(),
+            "benchmark_details": explanation.benchmark_details if explanation.benchmark_details else None,
+            "optimized_runtime": humanize_runtime(explanation.best_runtime_ns),
+            "original_runtime": humanize_runtime(explanation.original_runtime_ns),
+        }
+        console.rule()
+        try:
+            response = self.make_ai_service_request("/optimization_impact", payload=payload, timeout=600)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating optimization refinements: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return ""
+
+        if response.status_code == 200:
+            return cast("str", response.json()["impact"])
+        try:
+            error = cast("str", response.json()["error"])
+        except Exception:
+            error = response.text
+        logger.error(f"Error generating impact candidates: {response.status_code} - {error}")
+        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        console.rule()
+        return ""
+
 
 class LocalAiServiceClient(AiServiceClient):
     """Client for interacting with the local AI service."""

codeflash/api/aiservice.py

@@ -12,11 +12,12 @@
 
 def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]:
     """Extract the single dependent function from the code context excluding the main function."""
-    ast_tree = ast.parse(code_context.testgen_context_code)
-
-    dependent_functions = {
-        node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
-    }
+    dependent_functions = set()
+    for code_string in code_context.testgen_context.code_strings:
+        ast_tree = ast.parse(code_string.code)
+        dependent_functions.update(
+            {node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))}
+        )
 
     if main_function in dependent_functions:
         dependent_functions.discard(main_function)

codeflash/code_utils/coverage_utils.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import subprocess
 import tempfile
 import time
@@ -9,15 +8,12 @@
 from typing import TYPE_CHECKING, Optional
 
 import git
-from filelock import FileLock
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.compat import codeflash_cache_dir
 from codeflash.code_utils.git_utils import check_running_in_git_repo, git_root_dir
 
 if TYPE_CHECKING:
-    from typing import Any
-
     from git import Repo
 
 
@@ -100,71 +96,24 @@ def get_patches_dir_for_project() -> Path:
     return Path(patches_dir / project_id)
 
 
-def get_patches_metadata() -> dict[str, Any]:
-    project_patches_dir = get_patches_dir_for_project()
-    meta_file = project_patches_dir / "metadata.json"
-    if meta_file.exists():
-        with meta_file.open("r", encoding="utf-8") as f:
-            return json.load(f)
-    return {"id": get_git_project_id() or "", "patches": []}
-
-
-def save_patches_metadata(patch_metadata: dict) -> dict:
-    project_patches_dir = get_patches_dir_for_project()
-    meta_file = project_patches_dir / "metadata.json"
-    lock_file = project_patches_dir / "metadata.json.lock"
-
-    # we are not supporting multiple concurrent optimizations within the same process, but keep that in case we decide to do so in the future.
-    with FileLock(lock_file, timeout=10):
-        metadata = get_patches_metadata()
-
-        patch_metadata["id"] = time.strftime("%Y%m%d-%H%M%S")
-        metadata["patches"].append(patch_metadata)
-
-        meta_file.write_text(json.dumps(metadata, indent=2))
-
-    return patch_metadata
-
-
-def overwrite_patch_metadata(patches: list[dict]) -> bool:
-    project_patches_dir = get_patches_dir_for_project()
-    meta_file = project_patches_dir / "metadata.json"
-    lock_file = project_patches_dir / "metadata.json.lock"
-
-    with FileLock(lock_file, timeout=10):
-        metadata = get_patches_metadata()
-        metadata["patches"] = patches
-        meta_file.write_text(json.dumps(metadata, indent=2))
-    return True
-
-
 def create_diff_patch_from_worktree(
-    worktree_dir: Path,
-    files: list[str],
-    fto_name: Optional[str] = None,
-    metadata_input: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
+    worktree_dir: Path, files: list[str], fto_name: Optional[str] = None
+) -> Optional[Path]:
     repository = git.Repo(worktree_dir, search_parent_directories=True)
     uni_diff_text = repository.git.diff(None, "HEAD", *files, ignore_blank_lines=True, ignore_space_at_eol=True)
 
     if not uni_diff_text:
         logger.warning("No changes found in worktree.")
-        return {}
+        return None
 
     if not uni_diff_text.endswith("\n"):
         uni_diff_text += "\n"
 
     project_patches_dir = get_patches_dir_for_project()
     project_patches_dir.mkdir(parents=True, exist_ok=True)
 
-    final_function_name = fto_name or metadata_input.get("fto_name", "unknown")
-    patch_path = project_patches_dir / f"{worktree_dir.name}.{final_function_name}.patch"
+    patch_path = project_patches_dir / f"{worktree_dir.name}.{fto_name}.patch"
     with patch_path.open("w", encoding="utf8") as f:
         f.write(uni_diff_text)
 
-    final_metadata = {"patch_path": str(patch_path)}
-    if metadata_input:
-        final_metadata.update(metadata_input)
-        final_metadata = save_patches_metadata(final_metadata)
-
-    return final_metadata
+    return patch_path

codeflash/code_utils/git_worktree_utils.py

@@ -114,32 +114,32 @@ def get_code_optimization_context(
             read_only_context_code = ""
 
     # Extract code context for testgen
-    testgen_code_markdown = extract_code_string_context_from_files(
+    testgen_context = extract_code_markdown_context_from_files(
         helpers_of_fto_dict,
         helpers_of_helpers_dict,
         project_root_path,
         remove_docstrings=False,
         code_context_type=CodeContextType.TESTGEN,
     )
-    testgen_context_code = testgen_code_markdown.code
-    testgen_context_code_tokens = encoded_tokens_len(testgen_context_code)
-    if testgen_context_code_tokens > testgen_token_limit:
-        testgen_code_markdown = extract_code_string_context_from_files(
+    testgen_markdown_code = testgen_context.markdown
+    testgen_code_token_length = encoded_tokens_len(testgen_markdown_code)
+    if testgen_code_token_length > testgen_token_limit:
+        testgen_context = extract_code_markdown_context_from_files(
             helpers_of_fto_dict,
             helpers_of_helpers_dict,
             project_root_path,
             remove_docstrings=True,
             code_context_type=CodeContextType.TESTGEN,
         )
-        testgen_context_code = testgen_code_markdown.code
-        testgen_context_code_tokens = encoded_tokens_len(testgen_context_code)
-        if testgen_context_code_tokens > testgen_token_limit:
+        testgen_markdown_code = testgen_context.markdown
+        testgen_code_token_length = encoded_tokens_len(testgen_markdown_code)
+        if testgen_code_token_length > testgen_token_limit:
             raise ValueError("Testgen code context has exceeded token limit, cannot proceed")
     code_hash_context = hashing_code_context.markdown
     code_hash = hashlib.sha256(code_hash_context.encode("utf-8")).hexdigest()
 
     return CodeOptimizationContext(
-        testgen_context_code=testgen_context_code,
+        testgen_context=testgen_context,
         read_writable_code=final_read_writable_code,
         read_only_context_code=read_only_context_code,
         hashing_code_context=code_hash_context,

codeflash/context/code_context_extractor.py

@@ -279,7 +279,7 @@ def _expand_qualified_functions(self) -> set[str]:
         # Find class methods and add their containing classes and dunder methods
         for qualified_name in list(self.qualified_function_names):
             if "." in qualified_name:
-                class_name, method_name = qualified_name.split(".", 1)
+                class_name, _method_name = qualified_name.split(".", 1)
 
                 # Add the class itself
                 expanded.add(class_name)
@@ -511,7 +511,7 @@ def revert_unused_helper_functions(
     if not unused_helpers:
         return
 
-    logger.info(f"Reverting {len(unused_helpers)} unused helper function(s) to original definitions")
+    logger.debug(f"Reverting {len(unused_helpers)} unused helper function(s) to original definitions")
 
     # Group unused helpers by file path
     unused_helpers_by_file = defaultdict(list)

codeflash/context/unused_definition_remover.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import ast
+import enum
 import hashlib
 import os
 import pickle
@@ -11,12 +12,11 @@
 import unittest
 from collections import defaultdict
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import TYPE_CHECKING, Callable, Optional, final
 
 if TYPE_CHECKING:
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 
-import pytest
 from pydantic.dataclasses import dataclass
 from rich.panel import Panel
 from rich.text import Text
@@ -35,6 +35,22 @@
     from codeflash.verification.verification_utils import TestConfig
 
 
+@final
+class PytestExitCode(enum.IntEnum):  # don't need to import entire pytest just for this
+    #: Tests passed.
+    OK = 0
+    #: Tests failed.
+    TESTS_FAILED = 1
+    #: pytest was interrupted.
+    INTERRUPTED = 2
+    #: An internal error got in the way.
+    INTERNAL_ERROR = 3
+    #: pytest was misused.
+    USAGE_ERROR = 4
+    #: pytest couldn't find tests.
+    NO_TESTS_COLLECTED = 5
+
+
 @dataclass(frozen=True)
 class TestFunction:
     function_name: str
@@ -401,6 +417,7 @@ def discover_tests_pytest(
         with tmp_pickle_path.open(mode="rb") as f:
             exitcode, tests, pytest_rootdir = pickle.load(f)
     except Exception as e:
+        tests, pytest_rootdir = [], None
         logger.exception(f"Failed to discover tests: {e}")
         exitcode = -1
     finally:
@@ -412,15 +429,15 @@ def discover_tests_pytest(
             error_section = match.group(1) if match else result.stdout
 
             logger.warning(
-                f"Failed to collect tests. Pytest Exit code: {exitcode}={pytest.ExitCode(exitcode).name}\n {error_section}"
+                f"Failed to collect tests. Pytest Exit code: {exitcode}={PytestExitCode(exitcode).name}\n {error_section}"
             )
             if "ModuleNotFoundError" in result.stdout:
                 match = ImportErrorPattern.search(result.stdout).group()
                 panel = Panel(Text.from_markup(f"⚠️  {match} ", style="bold red"), expand=False)
                 console.print(panel)
 
         elif 0 <= exitcode <= 5:
-            logger.warning(f"Failed to collect tests. Pytest Exit code: {exitcode}={pytest.ExitCode(exitcode).name}")
+            logger.warning(f"Failed to collect tests. Pytest Exit code: {exitcode}={PytestExitCode(exitcode).name}")
         else:
             logger.warning(f"Failed to collect tests. Pytest Exit code: {exitcode}")
         console.rule()