Merge pull request #124 from codeflash-ai/jedi_ctx_fix

found a case where jedi's full name did not start with the module nam…

Author: alvin-r

codeflash/context/code_context_extractor.py

@@ -356,6 +356,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 +411,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
                     ):

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():