Update chunk_python_code.py

chunk_python_code.py

@@ -1,65 +1,75 @@
 import ast
 from langchain.schema import Document 
 
-def chunk_python_code_with_metadata(source_code, source):
+def chunk_python_code_with_metadata(python_code, file_path):
     """
     Entry point method to process the Python file.
     It invokes the iterate_ast function.
     """
     documents = []
-   #print(f"Processing file: {source}")
+   #print(f"Processing file: {file_path}")
   
-    iterate_ast(source_code, documents, source)
-    # Determine usage based on the source path
-    if source.startswith("kadi_apy/lib/"):
+    _iterate_ast(python_code, documents, file_path)
+    # Determine usage based on the file_path
+    if file_path.startswith("kadi_apy/lib/"):
         usage = "kadi-apy python library"
-    elif source.startswith("kadi_apy/cli/"):
+    elif file_path.startswith("kadi_apy/cli/"):
         usage = "kadi-apy python cli library"
     else:
         usage = "undefined"
 
     # Add metadata for usage to all documents
     for doc in documents:
-        doc.metadata["source"] = source
+        doc.metadata["source"] = file_path
         doc.metadata["usage"] = usage  # Add the determined usage metadata
         #print(doc)
     return documents
 
 
-def iterate_ast(source_code, documents, source):
+def _iterate_ast(python_code, documents, file_path):
     """
     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=source)
+    documents = []
+    tree = ast.parse(python_code, filename=file_path)
 
     first_level_nodes = list(ast.iter_child_nodes(tree))
 
     # Check if there are no first-level nodes
     if not first_level_nodes:
-        handle_no_first_level_node_found(documents, source_code, source)
+        documents.extend(
+            _chunk_nodeless_python_code(documents, python_code, file_path))
         return
 
     all_imports = all(isinstance(node, (ast.Import, ast.ImportFrom)) for node in first_level_nodes)
     if all_imports:
-        handle_first_level_imports_only(documents, source_code, source)
+        docuements.extend(
+            _chunk_import_only_python_code(documents, python_code, file_path))
     
     # 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, documents, source_code)
+            documents.extend(
+                = handle_first_level_class(first_level_node, documents, python_code))
         elif isinstance(first_level_node, ast.FunctionDef):
-            handle_first_level_func(first_level_node, documents, source_code)
+            documents.extend(
+                _chunk_first_level_func_node(first_level_node, documents, python_code))
         elif isinstance(first_level_node, ast.Assign):
-            handle_first_level_assign(first_level_node, documents, source_code)
-
-
-def handle_first_level_imports_only(documents, source_code, source):
+            documents.extend(
+                _chunk_first_level_assign_node(first_level_node, documents, python_code))
+        else: 
+            documents.extend(
+                handle_notdefined_case(python_code))
+             
+
+def _chunk_import_only_python_code(python_code, file_path):
     """
     Handles cases where the first-level nodes are only imports.
     """
-    if source.endswith("__init__.py"):
+    documents = []
+    if file_path.endswith("__init__.py"):
         type = "__init__-file"
     else:
         type = "undefined"
@@ -69,17 +79,46 @@ def handle_first_level_imports_only(documents, source_code, source):
 
     # Create and store a Document with the full source code
     doc = Document(
-        page_content=source_code,
+        page_content=python_code,
         metadata=metadata
     )
     documents.append(doc)
+    return documents
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 
-def handle_no_first_level_node_found(documents, source_code, source):
+def handle_notdefined_case(python_code):
+    documents = []
+    documents.extend(
+        chunk_python_code_by_character)
+    
+    return documents
+    
+
+def chunk_python_code_by_character(python_code):
+    documents = []
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=512,         
+        chunk_overlap=128,      
+        separators=[]
+        
+    chunks = text_splitter.split_text(text)
+    doc = Document(
+        page_content=python_code
+    )
+    
+    documents.append(doc)
+
+    return documents
+    
+
+def _chunk_nodeless_python_code(python_code, file_path):
     """
     Handles cases where no top-level nodes are found in the AST.
     """
-    if source.endswith("__init__.py"):
+    documents = []
+    if file_path.endswith("__init__.py"):
         type = "__init__-file"
     else:
         type = "undefined"
@@ -89,19 +128,22 @@ def handle_no_first_level_node_found(documents, source_code, source):
 
     # Create and store a Document with the full source code
     doc = Document(
-        page_content=source_code,
+        page_content=python_code,
         metadata=metadata
     )
     documents.append(doc)
 
+    return documents
 
-def handle_first_level_assign(assign_node, documents, source_code):
+def _chunk_first_level_assign_node(ast_node, python_code):
+    
     """
     Handles assignment statements at the first level of the AST.
     """
-    assign_start_line = assign_node.lineno
-    assign_end_line = assign_node.end_lineno
-    assign_source = '\n'.join(source_code.splitlines()[assign_start_line-1:assign_end_line])
+    documents = []
+    assign_start_line = ast_node.lineno
+    assign_end_line = ast_node.end_lineno
+    assign_source = '\n'.join(python_code.splitlines()[assign_start_line-1:assign_end_line])
 
     # Create metadata without imports
     metadata = {"type": "Assign"}
@@ -113,20 +155,21 @@ def handle_first_level_assign(assign_node, documents, source_code):
     )
     documents.append(doc)
 
+    return documents
 
-def handle_first_level_class(class_node, documents, source_code):
+def handle_first_level_class(ast_node , python_code):
     """
     Handles classes at the first level of the AST.
     """
-    class_start_line = class_node.lineno
-    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
-    class_source = '\n'.join(source_code.splitlines()[class_start_line-1:class_end_line])
+    documents = []
+    class_start_line = ast_node.lineno
+    class_body_lines = [child.lineno for child in ast_node.body if isinstance(child, ast.FunctionDef)]
+    class_end_line = min(class_body_lines, default=ast_node.end_lineno) - 1
+    class_source = '\n'.join(python_code.splitlines()[class_start_line-1:class_end_line])
 
-    # Create metadata without imports
     metadata = {
         "type": "class",
-        "class": class_node.name,
+        "class": ast_node.name,
         "visibility": "public"
     }
 
@@ -138,14 +181,14 @@ def handle_first_level_class(class_node, documents, source_code):
     documents.append(doc)
 
     # Handle methods within the class
-    for second_level_node in ast.iter_child_nodes(class_node):
+    for second_level_node in ast.iter_child_nodes(ast_node):
         if isinstance(second_level_node, ast.FunctionDef):
             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])
+            method_source = '\n'.join(python_code.splitlines()[method_start_line-1:method_end_line])
 
             visibility = "internal" if second_level_node.name.startswith("_") else "public"
 
@@ -155,28 +198,30 @@ def handle_first_level_class(class_node, documents, source_code):
                     "type": "method",
                     "method": second_level_node.name,
                     "visibility": visibility,
-                    "class": class_node.name
+                    "class": ast_node.name
                 }
             )
             documents.append(doc)
 
+    return documents
 
-def handle_first_level_func(function_node, documents, source_code):
+def _chunk_first_level_func_node(ast_node, python_code):
     """
     Handles functions at the first level of the AST.
     """
+    documents = []
     function_start_line = (
-        function_node.decorator_list[0].lineno
-        if function_node.decorator_list else function_node.lineno
+        ast_node.decorator_list[0].lineno
+        if ast_node.decorator_list else ast_node.lineno
     )
-    function_end_line = function_node.end_lineno
-    function_source = '\n'.join(source_code.splitlines()[function_start_line-1:function_end_line])
+    function_end_line = ast_node.end_lineno
+    function_source = '\n'.join(python_code.splitlines()[function_start_line-1:function_end_line])
 
-    visibility = "internal" if function_node.name.startswith("_") else "public"
+    visibility = "internal" if ast_node.name.startswith("_") else "public"
 
     is_command = any(
         decorator.id == "apy_command"
-        for decorator in function_node.decorator_list
+        for decorator in ast_node.decorator_list
         if hasattr(decorator, "id")
     )
 
@@ -185,12 +230,14 @@ def handle_first_level_func(function_node, documents, source_code):
         "visibility": visibility
     }
     if is_command:
-        metadata["command"] = function_node.name
+        metadata["command"] = ast_node.name
     else:
-        metadata["method"] = function_node.name
+        metadata["method"] = ast_node.name
 
     doc = Document(
         page_content=function_source,
         metadata=metadata
     )
-    documents.append(doc)
+    documents.append(doc)
+
+    return documents