Merge branch 'main' into fix/revert-code-issue

Author: KRRT7

SECURITY.md

@@ -0,0 +1,19 @@
+# Security Policy
+
+This document outlines Codeflash's vulnerability disclosure policy. For more information about Codeflash's approach to security, please visit [codeflash.ai/security](https://www.codeflash.ai/security).
+
+## Supported Versions
+
+Since Codeflash is moving quickly, we can only commit to fixing security issues for the latest version of codeflash client.
+If a vulnerability is discovered in our backend, we will release the fix for all the users.
+
+## Reporting a Vulnerability
+
+
+Please do not report security vulnerabilities through public GitHub issues.
+
+Instead, please report them to our [GitHub Security page](https://github.com/codeflash-ai/codeflash/security). If you prefer to submit one without using GitHub, you can also email us at [email protected].
+
+We commit to acknowledging vulnerability reports immediately, and will work to fix active vulnerabilities as soon as we can. We will publish resolved vulnerabilities in the form of security advisories on our GitHub security page. Critical incidents will be communicated both on the GitHub security page and via email to all affected users.
+
+We appreciate your help in making Codeflash more secure for everyone. Thank you for your support and responsible disclosure.

code_to_optimize/code_directories/circular_deps/constants.py

@@ -1,8 +1,2 @@
 DEFAULT_API_URL = "https://api.galileo.ai/"
 DEFAULT_APP_URL = "https://app.galileo.ai/"
-
-
-# function_names: GalileoApiClient.get_console_url
-# module_abs_path : /home/mohammed/Work/galileo-python/src/galileo/api_client.py
-# preexisting_objects: {('GalileoApiClient', ()), ('_set_destination', ()), ('get_console_url', (FunctionParent(name='GalileoApiClient', type='ClassDef'),))}
-# project_root_path: /home/mohammed/Work/galileo-python/src

codeflash-benchmark/codeflash_benchmark/plugin.py

@@ -8,21 +8,54 @@
 
 PYTEST_BENCHMARK_INSTALLED = importlib.util.find_spec("pytest_benchmark") is not None
 
+benchmark_options = [
+    ("--benchmark-columns", "store", None, "Benchmark columns"),
+    ("--benchmark-group-by", "store", None, "Benchmark group by"),
+    ("--benchmark-name", "store", None, "Benchmark name pattern"),
+    ("--benchmark-sort", "store", None, "Benchmark sort column"),
+    ("--benchmark-json", "store", None, "Benchmark JSON output file"),
+    ("--benchmark-save", "store", None, "Benchmark save name"),
+    ("--benchmark-warmup", "store", None, "Benchmark warmup"),
+    ("--benchmark-warmup-iterations", "store", None, "Benchmark warmup iterations"),
+    ("--benchmark-min-time", "store", None, "Benchmark minimum time"),
+    ("--benchmark-max-time", "store", None, "Benchmark maximum time"),
+    ("--benchmark-min-rounds", "store", None, "Benchmark minimum rounds"),
+    ("--benchmark-timer", "store", None, "Benchmark timer"),
+    ("--benchmark-calibration-precision", "store", None, "Benchmark calibration precision"),
+    ("--benchmark-disable", "store_true", False, "Disable benchmarks"),
+    ("--benchmark-skip", "store_true", False, "Skip benchmarks"),
+    ("--benchmark-only", "store_true", False, "Only run benchmarks"),
+    ("--benchmark-verbose", "store_true", False, "Verbose benchmark output"),
+    ("--benchmark-histogram", "store", None, "Benchmark histogram"),
+    ("--benchmark-compare", "store", None, "Benchmark compare"),
+    ("--benchmark-compare-fail", "store", None, "Benchmark compare fail threshold"),
+]
+
 
 def pytest_configure(config: pytest.Config) -> None:
     """Register the benchmark marker and disable conflicting plugins."""
     config.addinivalue_line("markers", "benchmark: mark test as a benchmark that should be run with codeflash tracing")
 
-    if config.getoption("--codeflash-trace") and PYTEST_BENCHMARK_INSTALLED:
-        config.option.benchmark_disable = True
-        config.pluginmanager.set_blocked("pytest_benchmark")
-        config.pluginmanager.set_blocked("pytest-benchmark")
+    if config.getoption("--codeflash-trace"):
+        # When --codeflash-trace is used, ignore all benchmark options by resetting them to defaults
+        for option, _, default, _ in benchmark_options:
+            option_name = option.replace("--", "").replace("-", "_")
+            if hasattr(config.option, option_name):
+                setattr(config.option, option_name, default)
+
+        if PYTEST_BENCHMARK_INSTALLED:
+            config.pluginmanager.set_blocked("pytest_benchmark")
+            config.pluginmanager.set_blocked("pytest-benchmark")
 
 
 def pytest_addoption(parser: pytest.Parser) -> None:
     parser.addoption(
         "--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing for benchmarks"
     )
+    # These options are ignored when --codeflash-trace is used
+    for option, action, default, help_text in benchmark_options:
+        help_suffix = " (ignored when --codeflash-trace is used)"
+        parser.addoption(option, action=action, default=default, help=help_text + help_suffix)
 
 
 @pytest.fixture
@@ -37,7 +70,7 @@ def benchmark(request: pytest.FixtureRequest) -> object:
     # If pytest-benchmark is installed and --codeflash-trace is not enabled,
     # return the normal pytest-benchmark fixture
     if PYTEST_BENCHMARK_INSTALLED:
-        from pytest_benchmark.fixture import BenchmarkFixture as BSF  # noqa: N814
+        from pytest_benchmark.fixture import BenchmarkFixture as BSF  # pyright: ignore[reportMissingImports]  # noqa: I001, N814
 
         bs = getattr(config, "_benchmarksession", None)
         if bs and bs.skip:

codeflash-benchmark/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "codeflash-benchmark"
-version = "0.1.0"
+version = "0.2.0"
 description = "Pytest benchmarking plugin for codeflash.ai - automatic code performance optimization"
 authors = [{ name = "CodeFlash Inc.", email = "[email protected]" }]
 requires-python = ">=3.9"
@@ -25,8 +25,8 @@ Repository = "https://github.com/codeflash-ai/codeflash-benchmark"
 codeflash-benchmark = "codeflash_benchmark.plugin"
 
 [build-system]
-requires = ["setuptools>=45", "wheel", "setuptools_scm"]
+requires = ["setuptools>=45", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools]
-packages = ["codeflash_benchmark"]
+packages = ["codeflash_benchmark"]

codeflash/api/aiservice.py

@@ -13,7 +13,7 @@
 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 AIServiceRefinerRequest, OptimizedCandidate
+from codeflash.models.models import AIServiceRefinerRequest, CodeStringsMarkdown, OptimizedCandidate
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.version import __version__ as codeflash_version
 
@@ -136,7 +136,7 @@ def optimize_python_code(  # noqa: D417
             logger.debug(f"Generating optimizations took {end_time - start_time:.2f} seconds.")
             return [
                 OptimizedCandidate(
-                    source_code=opt["source_code"],
+                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"],
                 )
@@ -206,7 +206,7 @@ def optimize_python_code_line_profiler(  # noqa: D417
             console.rule()
             return [
                 OptimizedCandidate(
-                    source_code=opt["source_code"],
+                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"],
                 )
@@ -263,7 +263,7 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
             console.rule()
             return [
                 OptimizedCandidate(
-                    source_code=opt["source_code"],
+                    source_code=CodeStringsMarkdown.parse_markdown_code(opt["source_code"]),
                     explanation=opt["explanation"],
                     optimization_id=opt["optimization_id"][:-4] + "refi",
                 )

codeflash/code_utils/code_replacer.py

@@ -19,7 +19,7 @@
     from pathlib import Path
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-    from codeflash.models.models import CodeOptimizationContext, OptimizedCandidate, ValidCode
+    from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, OptimizedCandidate, ValidCode
 
 ASTNodeT = TypeVar("ASTNodeT", bound=ast.AST)
 
@@ -408,16 +408,17 @@ def replace_functions_and_add_imports(
 
 def replace_function_definitions_in_module(
     function_names: list[str],
-    optimized_code: str,
+    optimized_code: CodeStringsMarkdown,
     module_abspath: Path,
     preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")
+    code_to_apply = get_optimized_code_for_module(module_abspath.relative_to(project_root_path), optimized_code)
     new_code: str = replace_functions_and_add_imports(
-        add_global_assignments(optimized_code, source_code),
+        add_global_assignments(code_to_apply, source_code),
         function_names,
-        optimized_code,
+        code_to_apply,
         module_abspath,
         preexisting_objects,
         project_root_path,
@@ -428,6 +429,19 @@ def replace_function_definitions_in_module(
     return True
 
 
+def get_optimized_code_for_module(relative_path: Path, optimized_code: CodeStringsMarkdown) -> str:
+    file_to_code_context = optimized_code.file_to_path()
+    module_optimized_code = file_to_code_context.get(str(relative_path))
+    if module_optimized_code is None:
+        logger.warning(
+            f"Optimized code not found for {relative_path} In the context\n-------\n{optimized_code}\n-------\n"
+            "re-check your 'markdown code structure'"
+            f"existing files are {file_to_code_context.keys()}"
+        )
+        module_optimized_code = ""
+    return module_optimized_code
+
+
 def is_zero_diff(original_code: str, new_code: str) -> bool:
     return normalize_code(original_code) == normalize_code(new_code)
 

codeflash/code_utils/formatter.py

@@ -104,7 +104,7 @@ def is_diff_line(line: str) -> bool:
 def format_code(
     formatter_cmds: list[str],
     path: Union[str, Path],
-    optimized_function: str = "",
+    optimized_code: str = "",
     check_diff: bool = False,  # noqa
     print_status: bool = True,  # noqa
     exit_on_failure: bool = True,  # noqa
@@ -121,7 +121,7 @@ def format_code(
 
         if check_diff and original_code_lines > 50:
             # we dont' count the formatting diff for the optimized function as it should be well-formatted
-            original_code_without_opfunc = original_code.replace(optimized_function, "")
+            original_code_without_opfunc = original_code.replace(optimized_code, "")
 
             original_temp = Path(test_dir_str) / "original_temp.py"
             original_temp.write_text(original_code_without_opfunc, encoding="utf8")

codeflash/context/code_context_extractor.py

@@ -61,13 +61,14 @@ def get_code_optimization_context(
     )
 
     # Extract code context for optimization
-    final_read_writable_code = extract_code_string_context_from_files(
+    final_read_writable_code = extract_code_markdown_context_from_files(
         helpers_of_fto_dict,
         {},
         project_root_path,
         remove_docstrings=False,
         code_context_type=CodeContextType.READ_WRITABLE,
-    ).code
+    )
+
     read_only_code_markdown = extract_code_markdown_context_from_files(
         helpers_of_fto_dict,
         helpers_of_helpers_dict,
@@ -84,14 +85,14 @@ def get_code_optimization_context(
     )
 
     # Handle token limits
-    final_read_writable_tokens = encoded_tokens_len(final_read_writable_code)
+    final_read_writable_tokens = encoded_tokens_len(final_read_writable_code.markdown)
     if final_read_writable_tokens > optim_token_limit:
         raise ValueError("Read-writable code has exceeded token limit, cannot proceed")
 
     # Setup preexisting objects for code replacer
     preexisting_objects = set(
         chain(
-            find_preexisting_objects(final_read_writable_code),
+            *(find_preexisting_objects(codestring.code) for codestring in final_read_writable_code.code_strings),
             *(find_preexisting_objects(codestring.code) for codestring in read_only_code_markdown.code_strings),
         )
     )

codeflash/context/unused_definition_remover.py

@@ -3,16 +3,15 @@
 import ast
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from pathlib import Path
+from itertools import chain
+from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
 import libcst as cst
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_replacer import replace_function_definitions_in_module
+from codeflash.models.models import CodeString, CodeStringsMarkdown
 
 if TYPE_CHECKING:
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
@@ -530,7 +529,11 @@ def revert_unused_helper_functions(
                 helper_names = [helper.qualified_name for helper in helpers_in_file]
                 reverted_code = replace_function_definitions_in_module(
                     function_names=helper_names,
-                    optimized_code=original_code,  # Use original code as the "optimized" code to revert
+                    optimized_code=CodeStringsMarkdown(
+                        code_strings=[
+                            CodeString(code=original_code, file_path=Path(file_path).relative_to(project_root))
+                        ]
+                    ),  # Use original code as the "optimized" code to revert
                     module_abspath=file_path,
                     preexisting_objects=set(),  # Empty set since we're reverting
                     project_root_path=project_root,
@@ -609,7 +612,9 @@ def _analyze_imports_in_optimized_code(
 
 
 def detect_unused_helper_functions(
-    function_to_optimize: FunctionToOptimize, code_context: CodeOptimizationContext, optimized_code: str
+    function_to_optimize: FunctionToOptimize,
+    code_context: CodeOptimizationContext,
+    optimized_code: str | CodeStringsMarkdown,
 ) -> list[FunctionSource]:
     """Detect helper functions that are no longer called by the optimized entrypoint function.
 
@@ -622,6 +627,14 @@ def detect_unused_helper_functions(
         List of FunctionSource objects representing unused helper functions
 
     """
+    if isinstance(optimized_code, CodeStringsMarkdown) and len(optimized_code.code_strings) > 0:
+        return list(
+            chain.from_iterable(
+                detect_unused_helper_functions(function_to_optimize, code_context, code.code)
+                for code in optimized_code.code_strings
+            )
+        )
+
     try:
         # Parse the optimized code to analyze function calls and imports
         optimized_ast = ast.parse(optimized_code)

codeflash/discovery/pytest_new_process_discovery.py

@@ -41,7 +41,8 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s
 
     try:
         exitcode = pytest.main(
-            [tests_root, "-p no:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()]
+            [tests_root, "-p no:logging", "--collect-only", "-m", "not skip", "-p", "no:codeflash-benchmark"],
+            plugins=[PytestCollectionPlugin()],
         )
     except Exception as e:
         print(f"Failed to collect tests: {e!s}")