Merge remote-tracking branch 'origin/main' into generated-tests-markdown

Author: Codeflash Bot

.github/workflows/unit-tests.yaml

@@ -28,8 +28,8 @@ jobs:
       - name: install dependencies
         run: uv sync
 
-      - name: Install test-only dependencies (Python 3.13)
-        if: matrix.python-version == '3.13'
+      - name: Install test-only dependencies (Python 3.9 and 3.13)
+        if: matrix.python-version == '3.9' || matrix.python-version == '3.13'
         run: uv sync --group tests
 
       - name: Unit tests

code_to_optimize/bubble_sort_method.py

@@ -15,3 +15,27 @@ def sorter(self, arr):
                     arr[j + 1] = temp
         print("stderr test", file=sys.stderr)
         return arr
+
+    @classmethod
+    def sorter_classmethod(cls, arr):
+        print("codeflash stdout : BubbleSorter.sorter_classmethod() called")
+        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("stderr test classmethod", file=sys.stderr)
+        return arr
+
+    @staticmethod
+    def sorter_staticmethod(arr):
+        print("codeflash stdout : BubbleSorter.sorter_staticmethod() called")
+        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("stderr test staticmethod", file=sys.stderr)
+        return arr

code_to_optimize/tests/pytest/test_topological_sort.py

@@ -10,7 +10,7 @@ def test_topological_sort():
     g.addEdge(2, 3)
     g.addEdge(3, 1)
 
-    assert g.topologicalSort() == [5, 4, 2, 3, 1, 0]
+    assert g.topologicalSort()[0] == [5, 4, 2, 3, 1, 0]
 
 
 def test_topological_sort_2():
@@ -20,15 +20,15 @@ def test_topological_sort_2():
         for j in range(i + 1, 10):
             g.addEdge(i, j)
 
-    assert g.topologicalSort() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    assert g.topologicalSort()[0] == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
 
     g = Graph(10)
 
     for i in range(10):
         for j in range(i + 1, 10):
             g.addEdge(i, j)
 
-    assert g.topologicalSort() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    assert g.topologicalSort()[0] == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
 
 
 def test_topological_sort_3():
@@ -38,4 +38,4 @@ def test_topological_sort_3():
         for j in range(i + 1, 1000):
             g.addEdge(j, i)
 
-    assert g.topologicalSort() == list(reversed(range(1000)))
+    assert g.topologicalSort()[0] == list(reversed(range(1000)))

codeflash-benchmark/codeflash_benchmark/plugin.py

@@ -52,10 +52,11 @@ 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)
+    # Only add benchmark options if pytest-benchmark is not installed for backward compatibility with existing pytest-benchmark setup
+    if not PYTEST_BENCHMARK_INSTALLED:
+        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

codeflash/api/aiservice.py

@@ -255,6 +255,8 @@ def optimize_python_code_refinement(self, request: list[AIServiceRefinerRequest]
                 "optimized_code_runtime": opt.optimized_code_runtime,
                 "speedup": opt.speedup,
                 "trace_id": opt.trace_id,
+                "function_references": opt.function_references,
+                "python_version": platform.python_version(),
             }
             for opt in request
         ]
@@ -308,6 +310,7 @@ def get_new_explanation(  # noqa: D417
         original_throughput: str | None = None,
         optimized_throughput: str | None = None,
         throughput_improvement: str | None = None,
+        function_references: str | None = None,
         codeflash_version: str = codeflash_version,
     ) -> str:
         """Optimize the given python code for performance by making a request to the Django endpoint.
@@ -329,6 +332,7 @@ def get_new_explanation(  # noqa: D417
         - optimized_throughput: str | None - throughput for the optimized code (operations per second)
         - throughput_improvement: str | None - throughput improvement percentage
         - current codeflash version
+        - function_references: str | None - where the function is called in the codebase
 
         Returns
         -------
@@ -351,6 +355,7 @@ def get_new_explanation(  # noqa: D417
             "original_throughput": original_throughput,
             "optimized_throughput": optimized_throughput,
             "throughput_improvement": throughput_improvement,
+            "function_references": function_references,
             "codeflash_version": codeflash_version,
         }
         logger.info("loading|Generating explanation")
@@ -376,7 +381,12 @@ def get_new_explanation(  # noqa: D417
         return ""
 
     def generate_ranking(  # noqa: D417
-        self, trace_id: str, diffs: list[str], optimization_ids: list[str], speedups: list[float]
+        self,
+        trace_id: str,
+        diffs: list[str],
+        optimization_ids: list[str],
+        speedups: list[float],
+        function_references: str | None = None,
     ) -> list[int] | None:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -385,6 +395,7 @@ def generate_ranking(  # noqa: D417
         - trace_id : unique uuid of function
         - diffs : list of unified diff strings of opt candidates
         - speedups : list of speedups of opt candidates
+        - function_references : where the function is called in the codebase
 
         Returns
         -------
@@ -397,6 +408,7 @@ def generate_ranking(  # noqa: D417
             "speedups": speedups,
             "optimization_ids": optimization_ids,
             "python_version": platform.python_version(),
+            "function_references": function_references,
         }
         logger.info("loading|Generating ranking")
         console.rule()
@@ -598,6 +610,7 @@ def get_optimization_review(
             "original_runtime": humanize_runtime(explanation.original_runtime_ns),
             "codeflash_version": codeflash_version,
             "calling_fn_details": calling_fn_details,
+            "python_version": platform.python_version(),
         }
         console.rule()
         try:

codeflash/api/cfapi.py

@@ -2,7 +2,7 @@
 
 import json
 import os
-import sys
+from dataclasses import dataclass
 from functools import lru_cache
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
@@ -13,6 +13,7 @@
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
+from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
 from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name
 from codeflash.github.PrComment import FileDiffContent, PrComment
@@ -26,14 +27,24 @@
 
 from packaging import version
 
-if os.environ.get("CODEFLASH_CFAPI_SERVER", "prod").lower() == "local":
-    CFAPI_BASE_URL = "http://localhost:3001"
-    CFWEBAPP_BASE_URL = "http://localhost:3000"
-    logger.info(f"Using local CF API at {CFAPI_BASE_URL}.")
-    console.rule()
-else:
-    CFAPI_BASE_URL = "https://app.codeflash.ai"
-    CFWEBAPP_BASE_URL = "https://app.codeflash.ai"
+
+@dataclass
+class BaseUrls:
+    cfapi_base_url: Optional[str] = None
+    cfwebapp_base_url: Optional[str] = None
+
+
+@lru_cache(maxsize=1)
+def get_cfapi_base_urls() -> BaseUrls:
+    if os.environ.get("CODEFLASH_CFAPI_SERVER", "prod").lower() == "local":
+        cfapi_base_url = "http://localhost:3001"
+        cfwebapp_base_url = "http://localhost:3000"
+        logger.info(f"Using local CF API at {cfapi_base_url}.")
+        console.rule()
+    else:
+        cfapi_base_url = "https://app.codeflash.ai"
+        cfwebapp_base_url = "https://app.codeflash.ai"
+    return BaseUrls(cfapi_base_url=cfapi_base_url, cfwebapp_base_url=cfwebapp_base_url)
 
 
 def make_cfapi_request(
@@ -53,8 +64,9 @@ def make_cfapi_request(
     :param suppress_errors: If True, suppress error logging for HTTP errors.
     :return: The response object from the API.
     """
-    url = f"{CFAPI_BASE_URL}/cfapi{endpoint}"
-    cfapi_headers = {"Authorization": f"Bearer {api_key or get_codeflash_api_key()}"}
+    url = f"{get_cfapi_base_urls().cfapi_base_url}/cfapi{endpoint}"
+    final_api_key = api_key or get_codeflash_api_key()
+    cfapi_headers = {"Authorization": f"Bearer {final_api_key}"}
     if extra_headers:
         cfapi_headers.update(extra_headers)
     try:
@@ -86,16 +98,22 @@ def make_cfapi_request(
 
 
 @lru_cache(maxsize=1)
-def get_user_id(api_key: Optional[str] = None) -> Optional[str]:
+def get_user_id(api_key: Optional[str] = None) -> Optional[str]:  # noqa: PLR0911
     """Retrieve the user's userid by making a request to the /cfapi/cli-get-user endpoint.
 
+    :param api_key: The API key to use. If None, uses get_codeflash_api_key().
     :return: The userid or None if the request fails.
     """
+    lsp_enabled = is_LSP_enabled()
     if not api_key and not ensure_codeflash_api_key():
         return None
 
     response = make_cfapi_request(
-        endpoint="/cli-get-user", method="GET", extra_headers={"cli_version": __version__}, api_key=api_key
+        endpoint="/cli-get-user",
+        method="GET",
+        extra_headers={"cli_version": __version__},
+        api_key=api_key,
+        suppress_errors=True,
     )
     if response.status_code == 200:
         if "min_version" not in response.text:
@@ -107,15 +125,31 @@ def get_user_id(api_key: Optional[str] = None) -> Optional[str]:
             if min_version and version.parse(min_version) > version.parse(__version__):
                 msg = "Your Codeflash CLI version is outdated. Please update to the latest version using `pip install --upgrade codeflash`."
                 console.print(f"[bold red]{msg}[/bold red]")
-                if is_LSP_enabled():
+                if lsp_enabled:
                     logger.debug(msg)
                     return f"Error: {msg}"
-                sys.exit(1)
+                exit_with_message(msg, error_on_exit=True)
             return userid
 
         logger.error("Failed to retrieve userid from the response.")
         return None
 
+    if response.status_code == 403:
+        error_title = "Invalid Codeflash API key. The API key you provided is not valid."
+        if lsp_enabled:
+            return f"Error: {error_title}"
+        msg = (
+            f"{error_title}\n"
+            "Please generate a new one at https://app.codeflash.ai/app/apikeys ,\n"
+            "then set it as a CODEFLASH_API_KEY environment variable.\n"
+            "For more information, refer to the documentation at \n"
+            "https://docs.codeflash.ai/optimizing-with-codeflash/codeflash-github-actions#manual-setup\n"
+            "or\n"
+            "https://docs.codeflash.ai/optimizing-with-codeflash/codeflash-github-actions#automated-setup-recommended"
+        )
+        exit_with_message(msg, error_on_exit=True)
+
+    # For other errors, log and return None (backward compatibility)
     logger.error(f"Failed to look up your userid; is your CF API key valid? ({response.reason})")
     return None
 

codeflash/cli_cmds/cmd_init.py

@@ -167,8 +167,8 @@ def ask_run_end_to_end_test(args: Namespace) -> None:
     console.rule()
 
     if run_tests:
-        bubble_sort_path, bubble_sort_test_path = create_bubble_sort_file_and_test(args)
-        run_end_to_end_test(args, bubble_sort_path, bubble_sort_test_path)
+        file_path = create_find_common_tags_file(args, "find_common_tags.py")
+        run_end_to_end_test(args, file_path)
 
 
 def is_valid_pyproject_toml(pyproject_toml_path: Path) -> tuple[bool, dict[str, Any] | None, str]:  # noqa: PLR0911
@@ -264,7 +264,7 @@ def get_valid_subdirs(current_dir: Optional[Path] = None) -> list[str]:
     ]
 
 
-def get_suggestions(section: str) -> tuple(list[str], Optional[str]):
+def get_suggestions(section: str) -> tuple[list[str], Optional[str]]:
     valid_subdirs = get_valid_subdirs()
     if section == CommonSections.module_root:
         return [d for d in valid_subdirs if d != "tests"], None
@@ -391,7 +391,7 @@ def collect_setup_info() -> CLISetupInfo:
     tests_root_answer = tests_answers["tests_root"]
 
     if tests_root_answer == create_for_me_option:
-        tests_root = Path(curdir) / default_tests_subdir
+        tests_root = Path(curdir) / (default_tests_subdir or "tests")
         tests_root.mkdir()
         click.echo(f"✅ Created directory {tests_root}{os.path.sep}{LF}")
     elif tests_root_answer == custom_dir_option:
@@ -1207,6 +1207,35 @@ def enter_api_key_and_save_to_rc() -> None:
     os.environ["CODEFLASH_API_KEY"] = api_key
 
 
+def create_find_common_tags_file(args: Namespace, file_name: str) -> Path:
+    find_common_tags_content = """def find_common_tags(articles: list[dict[str, list[str]]]) -> set[str]:
+    if not articles:
+        return set()
+
+    common_tags = articles[0]["tags"]
+    for article in articles[1:]:
+        common_tags = [tag for tag in common_tags if tag in article["tags"]]
+    return set(common_tags)
+"""
+
+    file_path = Path(args.module_root) / file_name
+    lsp_enabled = is_LSP_enabled()
+    if file_path.exists() and not lsp_enabled:
+        from rich.prompt import Confirm
+
+        overwrite = Confirm.ask(
+            f"🤔 {file_path} already exists. Do you want to overwrite it?", default=True, show_default=False
+        )
+        if not overwrite:
+            apologize_and_exit()
+        console.rule()
+
+    file_path.write_text(find_common_tags_content, encoding="utf8")
+    logger.info(f"Created demo optimization file: {file_path}")
+
+    return file_path
+
+
 def create_bubble_sort_file_and_test(args: Namespace) -> tuple[str, str]:
     bubble_sort_content = """from typing import Union, List
 def sorter(arr: Union[List[int],List[float]]) -> Union[List[int],List[float]]:
@@ -1276,7 +1305,7 @@ def test_sort():
     return str(bubble_sort_path), str(bubble_sort_test_path)
 
 
-def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test_path: str) -> None:
+def run_end_to_end_test(args: Namespace, find_common_tags_path: Path) -> None:
     try:
         check_formatter_installed(args.formatter_cmds)
     except Exception:
@@ -1285,7 +1314,7 @@ def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test
         )
         return
 
-    command = ["codeflash", "--file", "bubble_sort.py", "--function", "sorter"]
+    command = ["codeflash", "--file", "find_common_tags.py", "--function", "find_common_tags"]
     if args.no_pr:
         command.append("--no-pr")
     if args.verbose:
@@ -1316,10 +1345,8 @@ def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test
         console.rule()
         # Delete the bubble_sort.py file after the test
         logger.info("🧹 Cleaning up…")
-        for path in [bubble_sort_path, bubble_sort_test_path]:
-            console.rule()
-            Path(path).unlink(missing_ok=True)
-            logger.info(f"🗑️  Deleted {path}")
+        find_common_tags_path.unlink(missing_ok=True)
+        logger.info(f"🗑️  Deleted {find_common_tags_path}")
 
 
 def ask_for_telemetry() -> bool:

codeflash/cli_cmds/console.py

@@ -80,9 +80,16 @@ def paneled_text(
     console.print(panel)
 
 
-def code_print(code_str: str, file_name: Optional[str] = None, function_name: Optional[str] = None) -> None:
+def code_print(
+    code_str: str,
+    file_name: Optional[str] = None,
+    function_name: Optional[str] = None,
+    lsp_message_id: Optional[str] = None,
+) -> None:
     if is_LSP_enabled():
-        lsp_log(LspCodeMessage(code=code_str, file_name=file_name, function_name=function_name))
+        lsp_log(
+            LspCodeMessage(code=code_str, file_name=file_name, function_name=function_name, message_id=lsp_message_id)
+        )
         return
     """Print code with syntax highlighting."""
     from rich.syntax import Syntax