From 1f79facd6ba1a65d2ec182d58f544f9d4e1bc827 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Mon, 23 Jun 2025 15:43:54 -0700 Subject: [PATCH 1/5] get installed modules --- codeflash/code_utils/code_utils.py | 39 ++++++++++++++++++++++++++++-- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 82a5b979..11a831a5 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import ast +import importlib.metadata import os import re import shutil @@ -10,6 +11,7 @@ from functools import lru_cache from pathlib import Path from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING import tomlkit @@ -19,8 +21,12 @@ ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE) +if TYPE_CHECKING: + from collections.abc import Generator + + @contextmanager -def custom_addopts() -> None: +def custom_addopts() -> Generator[None, None, None]: pyproject_file = find_pyproject_toml() original_content = None non_blacklist_plugin_args = "" @@ -58,7 +64,7 @@ def custom_addopts() -> None: @contextmanager -def add_addopts_to_pyproject() -> None: +def add_addopts_to_pyproject() -> Generator[None, None, None]: pyproject_file = find_pyproject_toml() original_content = None try: @@ -220,3 +226,32 @@ def exit_with_message(message: str, *, error_on_exit: bool = False) -> None: paneled_text(message, panel_args={"style": "red"}) sys.exit(1 if error_on_exit else 0) + + +blacklist_installed_pkgs = { + "codeflash", + "pytest", + "coverage", + "__", # this is for private packages or ones that contain "__" in order to mangle names i.e 3204bda914b7f2c6f497__mypyc + "setuptools", + "pip", + "wheel", + "importlib_metadata", + "importlib_resources", + "isort", + "black", + "tomlkit", + "stubs", +} + + +def get_installed_packages() -> set[str]: + pkgs = importlib.metadata.packages_distributions().keys() + return { + pkg + for pkg in pkgs + if not any(blacklisted in pkg for blacklisted in blacklist_installed_pkgs) and not pkg.startswith("_") + } + + +print(get_installed_packages()) From 1609d45152ebfb11b3fdb8276f124c1cebd94ba7 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Mon, 23 Jun 2025 17:53:55 -0700 Subject: [PATCH 2/5] implement aiservice side --- codeflash/api/aiservice.py | 4 ++++ codeflash/code_utils/code_utils.py | 9 +++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/codeflash/api/aiservice.py b/codeflash/api/aiservice.py index 81ab84d4..cc62a5ce 100644 --- a/codeflash/api/aiservice.py +++ b/codeflash/api/aiservice.py @@ -10,6 +10,7 @@ from pydantic.json import pydantic_encoder from codeflash.cli_cmds.console import console, logger +from codeflash.code_utils.code_utils import get_installed_packages 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.models.models import OptimizedCandidate @@ -27,6 +28,7 @@ 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.installed_packages = get_installed_packages() def get_aiservice_base_url(self) -> str: if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local": @@ -66,6 +68,8 @@ def make_ai_service_request( """ url = f"{self.base_url}/ai{endpoint}" if method.upper() == "POST": + if self.installed_packages: + payload["installed_packages"] = self.installed_packages json_payload = json.dumps(payload, indent=None, default=pydantic_encoder) headers = {**self.headers, "Content-Type": "application/json"} response = requests.post(url, data=json_payload, headers=headers, timeout=timeout) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 11a831a5..4756964d 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -245,13 +245,10 @@ def exit_with_message(message: str, *, error_on_exit: bool = False) -> None: } -def get_installed_packages() -> set[str]: +def get_installed_packages() -> list[str]: pkgs = importlib.metadata.packages_distributions().keys() - return { + return [ pkg for pkg in pkgs if not any(blacklisted in pkg for blacklisted in blacklist_installed_pkgs) and not pkg.startswith("_") - } - - -print(get_installed_packages()) + ] From bff5d232371ade01d811bf8c43fdb9081324750c Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Mon, 23 Jun 2025 18:04:57 -0700 Subject: [PATCH 3/5] make mypy happy --- codeflash/api/aiservice.py | 1 + 1 file changed, 1 insertion(+) diff --git a/codeflash/api/aiservice.py b/codeflash/api/aiservice.py index cc62a5ce..0c4bb252 100644 --- a/codeflash/api/aiservice.py +++ b/codeflash/api/aiservice.py @@ -68,6 +68,7 @@ def make_ai_service_request( """ url = f"{self.base_url}/ai{endpoint}" if method.upper() == "POST": + assert payload is not None, "Payload must be provided for POST requests" if self.installed_packages: payload["installed_packages"] = self.installed_packages json_payload = json.dumps(payload, indent=None, default=pydantic_encoder) From 9a7c3468969135cbee6db3a9471a4fbc4a72739c Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Mon, 23 Jun 2025 18:16:21 -0700 Subject: [PATCH 4/5] PLC0415 messes with import time optimizations --- codeflash/code_utils/code_utils.py | 11 +++++++++-- pyproject.toml | 3 ++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 4756964d..327f55ed 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -1,7 +1,6 @@ from __future__ import annotations import ast -import importlib.metadata import os import re import shutil @@ -246,7 +245,15 @@ def exit_with_message(message: str, *, error_on_exit: bool = False) -> None: def get_installed_packages() -> list[str]: - pkgs = importlib.metadata.packages_distributions().keys() + try: + import importlib.metadata as importlib_metadata + except ImportError: + try: + import importlib_metadata + except ImportError: + return [] + + pkgs = importlib_metadata.packages_distributions().keys() return [ pkg for pkg in pkgs diff --git a/pyproject.toml b/pyproject.toml index 3a990197..e056740f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -237,7 +237,8 @@ ignore = [ "S301", "D104", "PERF203", - "LOG015" + "LOG015", + "PLC0415" ] [tool.ruff.lint.flake8-type-checking] From 69a234042693f3fe77499bf51d3965901bf4a5a3 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Mon, 23 Jun 2025 18:22:36 -0700 Subject: [PATCH 5/5] pass 3.9 too --- codeflash/code_utils/code_utils.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 327f55ed..0c860298 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -246,16 +246,20 @@ def exit_with_message(message: str, *, error_on_exit: bool = False) -> None: def get_installed_packages() -> list[str]: try: - import importlib.metadata as importlib_metadata - except ImportError: try: - import importlib_metadata + import importlib.metadata as importlib_metadata except ImportError: - return [] + import importlib_metadata + except ImportError: + return [] + + try: + pkgs = importlib_metadata.packages_distributions().keys() + except AttributeError: + pkgs = [dist.metadata.get("Name", "") for dist in importlib_metadata.distributions()] - pkgs = importlib_metadata.packages_distributions().keys() return [ pkg for pkg in pkgs - if not any(blacklisted in pkg for blacklisted in blacklist_installed_pkgs) and not pkg.startswith("_") + if pkg and not pkg.startswith("_") and not any(blacklisted in pkg for blacklisted in blacklist_installed_pkgs) ]