Update MemoryReadingAtomicFlow.py

MemoryReadingAtomicFlow.py

@@ -1,8 +1,9 @@
 from typing import Dict, Any
 import os
-
+import ast
+import importlib
+import inspect
 from flows.base_flows import AtomicFlow
-from flows.utils.general_helpers import get_pyfile_functions_metadata_from_file
 
 
 class MemoryReadingAtomicFlow(AtomicFlow):
@@ -34,18 +35,77 @@ class MemoryReadingAtomicFlow(AtomicFlow):
             content = file.read()
         return content
 
+    def _get_pyfile_functions_metadata_from_file(self, file_location):
+        def python_file_path_to_module_name(file_path):
+            return os.path.basename(file_path).replace('.py', '')
+
+        def extract_top_level_function_names(python_file_path):
+            with open(python_file_path, 'r') as file:
+                file_content = file.read()
+                tree = ast.parse(file_content)
+            functions = filter(lambda node: isinstance(node, ast.FunctionDef), ast.iter_child_nodes(tree))
+            return [node.name for node in functions]
+
+        def load_module_from_file(file_path):
+            module_name = python_file_path_to_module_name(file_path)
+            spec = importlib.util.spec_from_file_location(module_name, file_path)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            return module
+
+        def get_function_from_name(function_name, module):
+            return getattr(module, function_name)
+
+        def function_to_dict(function):
+            if not callable(function):
+                raise ValueError("Provided object is not a function.")
+
+            function_dict = {
+                "name": function.__name__,
+                "doc": function.__doc__,
+                "args": []
+            }
+
+            signature = inspect.signature(function)
+            for name, param in signature.parameters.items():
+                arg_info = {
+                    "name": name,
+                    "default": param.default if param.default is not inspect.Parameter.empty else None,
+                    "type": str(param.annotation) if param.annotation is not inspect.Parameter.empty else "unknown"
+                }
+                function_dict["args"].append(arg_info)
+
+            return function_dict
+
+        function_names = extract_top_level_function_names(file_location)
+        module = load_module_from_file(file_location)
+        functions = [get_function_from_name(name, module) for name in function_names]
+        return [function_to_dict(function) for function in functions]
+
+    def _format_metadata(self, metadata):
+        lines = []
+        for function_data in metadata:
+            lines.append(f"Function: {function_data['name']}")
+            lines.append(f"Documentation: {function_data['doc']}")
+
+            args = function_data.get('args', [])
+            if args:
+                lines.append("Arguments:")
+                for arg in args:
+                    default = f" (default: {arg['default']})" if arg['default'] is not None else ""
+                    lines.append(f"  - {arg['name']} (type: {arg['type']}){default}")
+            else:
+                lines.append("Arguments: None")
+
+            lines.append("#########")
+
+        return '\n'.join(lines)
     def _read_py_code_library(self, file_location):
-        metadata = get_pyfile_functions_metadata_from_file(file_location)
+        metadata = self._get_pyfile_functions_metadata_from_file(file_location)
         if len(metadata) == 0:
             return "No functions yet."
-        descriptions = []
-        for function in metadata:
-            description = f"Function Name: {function['name']}\n"
-            description += f"Documentation: {function.get('docstring', 'No documentation')}\n"
-            description += f"Parameters: {', '.join(function.get('parameters', []))}\n"
-            description += f"Default Values: {function.get('defaults', 'None')}"
-            descriptions.append(description)
-        return "\n\n".join(descriptions)
+        formatted_metadata = self._format_metadata(metadata)
+        return formatted_metadata
 
     def run(
             self,