Merge branch 'main' of github.com:codeflash-ai/codeflash into lsp/demo-optimization

Author: mohammedahmed18

.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/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/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,
     ) -> str:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -327,6 +330,7 @@ def get_new_explanation(  # noqa: D417
         - original_throughput: str | None - throughput for the baseline code (operations per second)
         - optimized_throughput: str | None - throughput for the optimized code (operations per second)
         - throughput_improvement: str | None - throughput improvement percentage
+        - function_references: str | None - where the function is called in the codebase
 
         Returns
         -------
@@ -349,6 +353,7 @@ def get_new_explanation(  # noqa: D417
             "original_throughput": original_throughput,
             "optimized_throughput": optimized_throughput,
             "throughput_improvement": throughput_improvement,
+            "function_references": function_references,
         }
         logger.info("loading|Generating explanation")
         console.rule()
@@ -373,7 +378,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.
 
@@ -382,6 +392,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
         -------
@@ -394,6 +405,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()
@@ -594,6 +606,7 @@ def get_optimization_review(
             "optimized_runtime": humanize_runtime(explanation.best_runtime_ns),
             "original_runtime": humanize_runtime(explanation.original_runtime_ns),
             "calling_fn_details": calling_fn_details,
+            "python_version": platform.python_version(),
         }
         console.rule()
         try:

codeflash/api/cfapi.py

@@ -1,10 +1,7 @@
-"""Module for interacting with the Codeflash API."""
-
 from __future__ import annotations
 
 import json
 import os
-import sys
 from dataclasses import dataclass
 from functools import lru_cache
 from pathlib import Path
@@ -101,12 +98,13 @@ 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
 
@@ -127,19 +125,21 @@ 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
 
-    # Handle 403 (Invalid API key) - exit with error message
     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 = (
-            "Invalid Codeflash API key. The API key you provided is not valid.\n"
+            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"

codeflash/cli_cmds/cmd_init.py

@@ -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:

codeflash/code_utils/code_extractor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import time
 from dataclasses import dataclass
 from itertools import chain
 from pathlib import Path
@@ -1138,6 +1139,7 @@ def find_specific_function_in_file(
 def get_fn_references_jedi(
     source_code: str, file_path: Path, project_root: Path, target_function: str, target_class: str | None
 ) -> list[Path]:
+    start_time = time.perf_counter()
     function_position: CodePosition = find_specific_function_in_file(
         source_code, file_path, target_function, target_class
     )
@@ -1146,6 +1148,8 @@ def get_fn_references_jedi(
         # Get references to the function
         references = script.get_references(line=function_position.line_no, column=function_position.col_no)
         # Collect unique file paths where references are found
+        end_time = time.perf_counter()
+        logger.debug(f"Jedi for function references ran in {end_time - start_time:.2f} seconds")
         reference_files = set()
         for ref in references:
             if ref.module_path:
@@ -1163,6 +1167,7 @@ def get_fn_references_jedi(
 def get_opt_review_metrics(
     source_code: str, file_path: Path, qualified_name: str, project_root: Path, tests_root: Path
 ) -> str:
+    start_time = time.perf_counter()
     try:
         qualified_name_split = qualified_name.rsplit(".", maxsplit=1)
         if len(qualified_name_split) == 1:
@@ -1176,4 +1181,6 @@ def get_opt_review_metrics(
     except Exception as e:
         calling_fns_details = ""
         logger.debug(f"Investigate {e}")
+    end_time = time.perf_counter()
+    logger.debug(f"Got function references in {end_time - start_time:.2f} seconds")
     return calling_fns_details

codeflash/code_utils/code_utils.py

@@ -17,6 +17,7 @@
 
 from codeflash.cli_cmds.console import logger, paneled_text
 from codeflash.code_utils.config_parser import find_pyproject_toml, get_all_closest_config_files
+from codeflash.lsp.helpers import is_LSP_enabled
 
 ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
 
@@ -352,6 +353,10 @@ def restore_conftest(path_to_content_map: dict[Path, str]) -> None:
 
 
 def exit_with_message(message: str, *, error_on_exit: bool = False) -> None:
+    """Don't Call it inside the lsp process, it will terminate the lsp server."""
+    if is_LSP_enabled():
+        logger.error(message)
+        return
     paneled_text(message, panel_args={"style": "red"})
 
     sys.exit(1 if error_on_exit else 0)

codeflash/code_utils/env_utils.py

@@ -13,6 +13,7 @@
 from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.formatter import format_code
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config, save_api_key_to_rc
+from codeflash.lsp.helpers import is_LSP_enabled
 
 
 def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool = True) -> bool:  # noqa
@@ -70,7 +71,10 @@ def get_codeflash_api_key() -> str:
         except Exception as e:
             logger.debug(f"Failed to automatically save API key to shell config: {e}")
 
-    api_key = env_api_key or shell_api_key
+    # Prefer the shell configuration over environment variables for lsp,
+    # as the API key may change in the RC file during lsp runtime. Since the LSP client (extension) can restart
+    # within the same process, the environment variable could become outdated.
+    api_key = shell_api_key or env_api_key if is_LSP_enabled() else env_api_key or shell_api_key
 
     api_secret_docs_message = "For more information, refer to the documentation at [https://docs.codeflash.ai/getting-started/codeflash-github-actions#add-your-api-key-to-your-repository-secrets]."  # noqa
     if not api_key: