Update utils.py

utils.py

@@ -3,48 +3,29 @@ import glob
 import os
 
 
-def parse_file_content(string: str) -> tuple[str | None, str | None]:
-    """
-    Parses the content of a file and returns the action and description.
-    :param string: The content of a file.
-    :return: A tuple containing the action and description.
-    """
+def parse_file_content(string: str):
     first_break = string.find("---")
     last_break = string.rfind("---")
     if first_break == -1 and last_break == -1 or first_break == last_break:
         return None, None
-    
-    # Find the newline after the last separator
-    nl_after = string.find("\n", last_break) + 1
-    description = string[nl_after:].strip() if nl_after > last_break else ""
+    nl_after = string.find("\n", last_break)
+    description = string[nl_after:]
+    return string[first_break + 4 : last_break], description.strip("\n")
 
-    return string[first_break + 4 : last_break], description
 
-
-def parse_action(string: str) -> tuple[str, str | None]:
-    """
-    Parses the action from a string.
-    :param string: The string to parse the action from.
-    :return: A tuple containing the action and action input.
-    """
+def parse_action(string: str):
     assert string.startswith("action:")
     idx = string.find("action_input=")
-    action = string[8: idx - 1] if idx > 8 else string[8:]
-    action_input = string[idx + 13 :].strip("'").strip('"') if idx > 8 else None
-    return action, action_input
+    if idx == -1:
+        return string[8:], None
+    return string[8 : idx - 1], string[idx + 13 :].strip("'").strip('"')
 
 
-def extract_imports(file_contents: str) -> tuple[list[str], list[ast.FunctionDef], list[ast.ClassDef]]:
-    """
-    Extracts imports, functions, and classes from a file's contents.
-    :param file_contents: The contents of a file.
-    :return: A tuple containing the imports, functions, and classes.
-    """
+def extract_imports(file_contents):
     module_ast = ast.parse(file_contents)
     imports = []
     functions = [n for n in module_ast.body if isinstance(n, ast.FunctionDef)]
     classes = [n for n in module_ast.body if isinstance(n, ast.ClassDef)]
-
     for node in ast.walk(module_ast):
         if isinstance(node, ast.Import):
             for alias in node.names:
@@ -57,16 +38,10 @@ def extract_imports(file_contents: str) -> tuple[list[str], list[ast.FunctionDef
                     imports.append(f"{module_name}.{name}")
                 else:
                     imports.append(name)
-
     return imports, functions, classes
 
 
-def read_python_module_structure(path: str) -> tuple[str, dict[str, str], dict[str, list[str]]]:
-    """
-    Reads the structure of a Python module and returns a prompt, content, and internal imports map.
-    :param path: The path to the Python module.
-    :return: A tuple containing the structure prompt, content, and internal imports map.
-    """
+def read_python_module_structure(path):
     file_types = ["*.py"]
     code = []
     for file_type in file_types:
@@ -75,7 +50,7 @@ def read_python_module_structure(path: str) -> tuple[str, dict[str, str], dict[s
     structure_prompt = "Files:\n"
     structure_prompt += "(listing all files and their functions and classes)\n\n"
 
-    def get_file_name(i: str) -> str:
+    def get_file_name(i):
         return "./{}.py".format(i.replace(".", "/"))
 
     content = {}
@@ -85,19 +60,20 @@ def read_python_module_structure(path: str) -> tuple[str, dict[str, str], dict[s
             continue
         with open(fn, "r") as f:
             content[fn] = f.read()
-
-        imports, functions, classes = extract_imports(content[fn])
-        internal_imports = [
-            ".".join(i.split(".")[:-1])
-            for i in imports
-            if i.startswith("app.")
-        ]
-        internal_imports_map[fn] = [
-            get_file_name(i) for i in set(internal_imports)
-        ]
-
-        structure_prompt += f"{fn}\n"
-        for function in functions:
-            structure_prompt += f"  {function.name}({function.args})\n"
-
-    return structure_prompt, content, internal_imports_map
+            imports, functions, classes = extract_imports(content[fn])
+            internal_imports = list(
+                {".".join(i.split(".")[:-1]) for i in imports if i.startswith("app.")}
+            )
+            internal_imports_map[fn] = [get_file_name(i) for i in internal_imports]
+            structure_prompt += f"{fn}\n"
+            for function in functions:
+                structure_prompt += "  {}()\n".format(function.name, function.args.args)
+
+            for class_ in classes:
+                structure_prompt += "  {}\n".format(class_.name)
+                methods = [n for n in class_.body if isinstance(n, ast.FunctionDef)]
+                for method in methods:
+                    structure_prompt += "  {}.{}()\n".format(class_.name, method.name)
+            structure_prompt += "\n"
+
+    return structure_prompt, content, internal_imports_map