Merge branch 'main' into chore/add-staging/docs

Author: HeshamHM28

README.md

@@ -65,6 +65,13 @@ For detailed installation and usage instructions, visit our documentation at [do
 
 https://github.com/user-attachments/assets/38f44f4e-be1c-4f84-8db9-63d5ee3e61e5
 
+- Optiming a workflow end to end automatically with `codeflash optimize`
+
+
+https://github.com/user-attachments/assets/355ba295-eb5a-453a-8968-7fb35c70d16c
+
+
+
 ## Support
 
 Join our community for support and discussions. If you have any questions, feel free to reach out to us using one of the following methods:

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/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/api/cfapi.py

@@ -101,6 +101,9 @@ def get_user_id() -> 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 console.quiet:  # lsp
+                    logger.debug(msg)
+                    return f"Error: {msg}"
                 sys.exit(1)
             return userid
 

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/env_utils.py

@@ -34,7 +34,12 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
 
 @lru_cache(maxsize=1)
 def get_codeflash_api_key() -> str:
-    api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+    if console.quiet:  # lsp
+        # prefer shell config over env var in lsp mode
+        api_key = read_api_key_from_shell_config()
+    else:
+        api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+
     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:
         msg = (

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/code_utils/shell_utils.py

@@ -42,7 +42,7 @@ def get_shell_rc_path() -> Path:
 
 
 def get_api_key_export_line(api_key: str) -> str:
-    return f"{SHELL_RC_EXPORT_PREFIX}{api_key}"
+    return f'{SHELL_RC_EXPORT_PREFIX}"{api_key}"'
 
 
 def save_api_key_to_rc(api_key: str) -> Result[str, str]:

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),
         )
     )