Update process_python_code

process_python_code

@@ -1,25 +1,28 @@
 import ast
-from langchain.schema import Document 
+from langchain.schema import Document  # Assuming "Document" is imported from LangChain
 
-def chunk_python_source_code(source_code, references):
+
+def chunk_python_code_with_metadata(source_code, references):
     """
     Entry point method to process the Python file.
     It invokes the iterate_ast function.
     """
-    print(f"Processing file: {references}")
-
-    for source_code, reference in zip(source_code, references):
-        print(f"Processing text: {source_code} with reference: {reference}")
-        iterate_ast(source_code, reference)
-    
+    documents = []
+    print(f"Processing file: {file_path}")
+    iterate_ast(source_code, documents)
+    for doc in documents:
+        print(f"Stored Document:\n")
+        print(doc)
+    print(len(documents)) 
+    return documents
 
-def iterate_ast(source_code, reference):
+def iterate_ast(source_code, documents):
     """
     Parses the AST of the given Python file and delegates
     handling to specific methods based on node types.
     """
     # Parse the source code into an abstract syntax tree (AST)
-    tree = ast.parse(source_code, filename=reference)
+    tree = ast.parse(source_code, filename=file_path)
 
     # Gather all top-level imports for later use
     imports_dict = extract_imports(tree)
@@ -27,9 +30,9 @@ def iterate_ast(source_code, reference):
     # Iterate over first-level nodes
     for first_level_node in ast.iter_child_nodes(tree):
         if isinstance(first_level_node, ast.ClassDef):
-            handle_first_level_class(first_level_node, source_code, imports_dict)
+            handle_first_level_class(first_level_node, documents, source_code, imports_dict)
         elif isinstance(first_level_node, ast.FunctionDef):
-            handle_first_level_func(first_level_node, source_code, imports_dict)
+            handle_first_level_func(first_level_node, documents, source_code, imports_dict)
 
 def extract_imports(tree):
     """
@@ -58,101 +61,114 @@ def analyze_imports(node, imports_dict):
             relevant_imports.add(imports_dict[sub_node.id])
     return list(relevant_imports)
 
-def handle_first_level_class(class_node, source_code, imports_dict):
+def handle_first_level_class(class_node, documents, source_code, imports_dict):
+
     """
-    Handles classes at the first level of the AST by processing
-    the class and its methods. Stores each class method in a Document object.
+    Handles classes at the first level of the AST by storing them
+    in a Document object with metadata, including relevant imports.
     """
-    print(f"Class detected: {class_node.name}")
+    # Extract class source code
+    class_start_line = class_node.lineno
 
+    # Find the line where the first function (def) starts or the next top-level node
+    class_body_lines = [child.lineno for child in class_node.body if isinstance(child, ast.FunctionDef)]
+    class_end_line = min(class_body_lines, default=class_node.end_lineno) - 1  # Use `-1` to exclude the next node
+
+    # Generate the class source code
+    class_source = '\n'.join(source_code.splitlines()[class_start_line-1:class_end_line])
+    
     # Extract relevant imports for this class
     class_imports = analyze_imports(class_node, imports_dict)
     
-    # Extract the class source code
-    class_start_line = class_node.lineno
-    class_end_line = max(
-        [n.end_lineno for n in ast.walk(class_node) if hasattr(n, "end_lineno")], default=class_node.lineno
-    )
-    class_source = '\n'.join(source_code.splitlines()[class_start_line - 1:class_end_line])
-
-    # Store the class-level Document
-    class_doc = Document(
+    # Create and store Document for the class
+    doc = Document(
         page_content=class_source,
         metadata={
             "type": "class",
-            "class_name": class_node.name,
-            "imports": class_imports
+            "class": class_node.name,
+            "visibility": "public",
+            "imports": class_imports  # Add class-specific imports
         }
     )
-    print(f"Stored Class Document: {class_doc}\n")
-    
-    # Process methods within the class
+    documents.append(doc)
+
+
+    # Handle methods within the class
     for second_level_node in ast.iter_child_nodes(class_node):
         if isinstance(second_level_node, ast.FunctionDef):
-            handle_class_method(second_level_node, class_node.name, source_code, imports_dict)
+            # Extract method source code
+            method_start_line = (
+                second_level_node.decorator_list[0].lineno
+                if second_level_node.decorator_list else second_level_node.lineno
+            )
+            method_end_line = second_level_node.end_lineno
+            method_source = '\n'.join(source_code.splitlines()[method_start_line-1:method_end_line])
+            
+            # Determine visibility metadata
+            visibility = "internal" if second_level_node.name.startswith("_") else "public"
+            # Extract relevant imports for this method
+            method_imports = analyze_imports(second_level_node, imports_dict)
+            
+            # Create and store Document
 
-def handle_class_method(method_node, class_name, source_code, imports_dict):
-    """
-    Handles methods within a class by storing them in a Document object.
-    """
-    print(f"Method detected: {method_node.name} in class {class_name}")
-    
-    # Extract method source code
-    method_start_line = (
-        method_node.decorator_list[0].lineno
-        if method_node.decorator_list else method_node.lineno
-    )
-    method_end_line = method_node.end_lineno
-    method_source = '\n'.join(source_code.splitlines()[method_start_line - 1:method_end_line])
-    
-    # Determine visibility metadata
-    visibility = "internal" if method_node.name.startswith("_") else "public"
-    
-    # Extract relevant imports for this method
-    method_imports = analyze_imports(method_node, imports_dict)
-    
-    # Store the method-level Document
-    method_doc = Document(
-        page_content=method_source,
-        metadata={
-            "type": "method",
-            "class_name": class_name,
-            "method_name": method_node.name,
-            "visibility": visibility,
-            "imports": method_imports
-        }
-    )
-    print(f"Stored Method Document: {method_doc}\n")
+            doc = Document(
+                page_content=method_source,
+                metadata={
+                    "type": "method",
+                    "method": second_level_node.name,
+                    "visibility": "visibility",
+                    "imports": method_imports,  
+                    "class": class_node.name
+                }
+            )
+            documents.append(doc)
 
-def handle_first_level_func(function_node, source_code, imports_dict):
+def handle_first_level_func(function_node, documents, source_code, imports_dict):
     """
     Handles functions at the first level of the AST by storing them
     in a Document object with metadata, including relevant imports.
     """
-    print(f"Function detected: {function_node.name}")
-
     # Extract function source code
     function_start_line = (
         function_node.decorator_list[0].lineno
         if function_node.decorator_list else function_node.lineno
     )
     function_end_line = function_node.end_lineno
-    function_source = '\n'.join(source_code.splitlines()[function_start_line - 1:function_end_line])
+    function_source = '\n'.join(source_code.splitlines()[function_start_line-1:function_end_line])
     
     # Determine visibility metadata
     visibility = "internal" if function_node.name.startswith("_") else "public"
     
+    # Check if the function is a CLI command (e.g., decorated with `@apy_command`)
+    is_command = any(
+        decorator.id == "apy_command"  # Check decorator name
+        for decorator in function_node.decorator_list
+        if hasattr(decorator, "id")  # Ensure the decorator has an identifier
+    )
+    
     # Extract relevant imports for this function
     function_imports = analyze_imports(function_node, imports_dict)
     
-    # Store the function-level Document
-    function_doc = Document(
-        page_content=function_source,
-        metadata={
-            "type": "function",
-            "function_name": function_node.name,
-            "visibility": visibility,
-            "imports": function_imports
-        }
-    )
-    print(f"Stored Function Document: {function_doc}\n")
+    # Create and store Document
+    
+    if is_command:
+        doc = Document(
+            page_content=function_source,
+            metadata={
+                "type": "command",
+                "command": function_node.name,
+                "visibility": "public",
+                "imports": function_imports
+            }
+        )
+    else:
+          doc = Document(
+            page_content=function_source,
+            metadata={
+                "type": "function",
+                "method": function_node.name,
+                "visibility": visibility,
+                "imports": function_imports
+            }
+        )          
+    documents.append(doc)