From 9282ae368df842092f610c654bf907e45b609e7e Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 4 Apr 2025 14:14:04 -0700 Subject: [PATCH 1/2] found a case where jedi's full name did not start with the module name. we need this check so it doesn't error out --- codeflash/context/code_context_extractor.py | 5 ++++- codeflash/optimization/function_context.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/codeflash/context/code_context_extractor.py b/codeflash/context/code_context_extractor.py index 13c7ec1dd..bfd2de885 100644 --- a/codeflash/context/code_context_extractor.py +++ b/codeflash/context/code_context_extractor.py @@ -11,7 +11,7 @@ from jedi.api.classes import Name from libcst import CSTNode -from codeflash.cli_cmds.console import logger +from codeflash.cli_cmds.console import code_print, logger from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -73,6 +73,7 @@ def get_code_optimization_context( # Handle token limits tokenizer = tiktoken.encoding_for_model("gpt-4o") + code_print(final_read_writable_code) final_read_writable_tokens = len(tokenizer.encode(final_read_writable_code)) if final_read_writable_tokens > optim_token_limit: raise ValueError("Read-writable code has exceeded token limit, cannot proceed") @@ -356,6 +357,7 @@ def get_function_to_optimize_as_function_source( name.type == "function" and name.full_name and name.name == function_to_optimize.function_name + and name.full_name.startswith(name.module_name) and get_qualified_name(name.module_name, name.full_name) == function_to_optimize.qualified_name ): function_source = FunctionSource( @@ -410,6 +412,7 @@ def get_function_sources_from_jedi( and definition.full_name and definition.type == "function" and not belongs_to_function_qualified(definition, qualified_function_name) + and definition.full_name.startswith(definition.module_name) # Avoid nested functions or classes. Only class.function is allowed and len((qualified_name := get_qualified_name(definition.module_name, definition.full_name)).split(".")) <= 2 ): diff --git a/codeflash/optimization/function_context.py b/codeflash/optimization/function_context.py index 3c28a92db..d55aa2dec 100644 --- a/codeflash/optimization/function_context.py +++ b/codeflash/optimization/function_context.py @@ -31,7 +31,7 @@ def belongs_to_class(name: Name, class_name: str) -> bool: def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool: """Check if the given jedi Name is a direct child of the specified function, matched by qualified function name.""" try: - if get_qualified_name(name.module_name, name.full_name) == qualified_function_name: + if name.full_name.startswith(name.module_name) and get_qualified_name(name.module_name, name.full_name) == qualified_function_name: # Handles function definition and recursive function calls return False if name := name.parent(): From 64f79270fe805f427c414cb00763a69f8b2730f4 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 4 Apr 2025 14:15:40 -0700 Subject: [PATCH 2/2] fix --- codeflash/context/code_context_extractor.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/codeflash/context/code_context_extractor.py b/codeflash/context/code_context_extractor.py index bfd2de885..516f3c94e 100644 --- a/codeflash/context/code_context_extractor.py +++ b/codeflash/context/code_context_extractor.py @@ -11,7 +11,7 @@ from jedi.api.classes import Name from libcst import CSTNode -from codeflash.cli_cmds.console import code_print, logger +from codeflash.cli_cmds.console import logger from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -73,7 +73,6 @@ def get_code_optimization_context( # Handle token limits tokenizer = tiktoken.encoding_for_model("gpt-4o") - code_print(final_read_writable_code) final_read_writable_tokens = len(tokenizer.encode(final_read_writable_code)) if final_read_writable_tokens > optim_token_limit: raise ValueError("Read-writable code has exceeded token limit, cannot proceed")