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