|
11 | 11 | from jedi.api.classes import Name |
12 | 12 | from libcst import CSTNode |
13 | 13 |
|
14 | | -from codeflash.cli_cmds.console import logger |
| 14 | +from codeflash.cli_cmds.console import code_print, logger |
15 | 15 | from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects |
16 | 16 | from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages |
17 | 17 | from codeflash.discovery.functions_to_optimize import FunctionToOptimize |
@@ -73,6 +73,7 @@ def get_code_optimization_context( |
73 | 73 |
|
74 | 74 | # Handle token limits |
75 | 75 | tokenizer = tiktoken.encoding_for_model("gpt-4o") |
| 76 | + code_print(final_read_writable_code) |
76 | 77 | final_read_writable_tokens = len(tokenizer.encode(final_read_writable_code)) |
77 | 78 | if final_read_writable_tokens > optim_token_limit: |
78 | 79 | raise ValueError("Read-writable code has exceeded token limit, cannot proceed") |
@@ -356,6 +357,7 @@ def get_function_to_optimize_as_function_source( |
356 | 357 | name.type == "function" |
357 | 358 | and name.full_name |
358 | 359 | and name.name == function_to_optimize.function_name |
| 360 | + and name.full_name.startswith(name.module_name) |
359 | 361 | and get_qualified_name(name.module_name, name.full_name) == function_to_optimize.qualified_name |
360 | 362 | ): |
361 | 363 | function_source = FunctionSource( |
@@ -410,6 +412,7 @@ def get_function_sources_from_jedi( |
410 | 412 | and definition.full_name |
411 | 413 | and definition.type == "function" |
412 | 414 | and not belongs_to_function_qualified(definition, qualified_function_name) |
| 415 | + and definition.full_name.startswith(definition.module_name) |
413 | 416 | # Avoid nested functions or classes. Only class.function is allowed |
414 | 417 | and len((qualified_name := get_qualified_name(definition.module_name, definition.full_name)).split(".")) <= 2 |
415 | 418 | ): |
|
0 commit comments