Update app.py

app.py

@@ -183,6 +183,116 @@ def process_directory(directory, partial_paths=None, file_paths=None):
     print(f"Print the text for testing broooo {all_texts}")
     return all_texts, file_references
 
+import ast
+
+def get_source_segment(source_lines, node):
+    start_line, start_col = node.lineno - 1, node.col_offset
+    end_line = node.end_lineno - 1 if hasattr(node, 'end_lineno') else node.lineno - 1
+    end_col = node.end_col_offset if hasattr(node, 'end_col_offset') else len(source_lines[end_line])
+
+    lines = source_lines[start_line:end_line + 1]
+    lines[0] = lines[0][start_col:]
+    lines[-1] = lines[-1][:end_col]
+
+    return ''.join(lines)
+
+from langchain.schema import Document
+
+def chunk_python_file_content(content, char_limit=1572):
+    source_lines = content.splitlines(keepends=True)
+    
+    # Parse the content into an abstract syntax tree (AST)
+    tree = ast.parse(content)
+
+    chunks = []
+    current_chunk = ""
+    current_chunk_size = 0
+    
+    # Find all class definitions and top-level functions in the AST
+    class_nodes = [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)]
+    function_nodes = [node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and not isinstance(node, ast.ClassDef)]
+
+    for class_node in class_nodes:
+        method_nodes = [node for node in class_node.body if isinstance(node, ast.FunctionDef)]
+
+        if method_nodes:
+            first_method_start_line = method_nodes[0].lineno - 1
+            class_def_lines = source_lines[class_node.lineno - 1:first_method_start_line]
+        else:
+            class_def_lines = source_lines[class_node.lineno - 1:class_node.end_lineno]
+
+        class_def = ''.join(class_def_lines)
+        class_def_size = len(class_def)
+        
+        # Add class definition to the current chunk if it fits
+        if current_chunk_size + class_def_size <= char_limit:
+            current_chunk += f"{class_def.strip()}\n"
+            current_chunk_size += class_def_size
+        else:
+            # Start a new chunk if the class definition exceeds the limit
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+                current_chunk = ""
+                current_chunk_size = 0
+            current_chunk += f"{class_def.strip()}\n"
+            current_chunk_size = class_def_size
+        
+        for method_node in method_nodes:
+            method_def = get_source_segment(source_lines, method_node)
+            method_def_size = len(method_def)
+            
+            # Add method definition to the current chunk if it fits
+            if current_chunk_size + method_def_size <= char_limit:
+                current_chunk += f"{method_def.strip()}\n"
+                current_chunk_size += method_def_size
+            else:
+                # Start a new chunk if the method definition exceeds the limit
+                if current_chunk:
+                    chunks.append(current_chunk.strip())
+                    current_chunk = ""
+                    current_chunk_size = 0
+                current_chunk += f"# This is a class method of class: {class_node.name}\n{method_def.strip()}\n"
+                current_chunk_size = method_def_size
+
+    for function_node in function_nodes:
+        function_def = get_source_segment(source_lines, function_node)
+        function_def_size = len(function_def)
+        
+        # Add function definition to the current chunk if it fits
+        if current_chunk_size + function_def_size <= char_limit:
+            current_chunk += f"{function_def.strip()}\n"
+            current_chunk_size += function_def_size
+        else:
+            # Start a new chunk if the function definition exceeds the limit
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+                current_chunk = ""
+                current_chunk_size = 0
+            current_chunk += f"{function_def.strip()}\n"
+            current_chunk_size = function_def_size
+
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
+    return chunks
+
+
+
+# Split python code into chunks
+def split_pythoncode_into_chunks(texts, references, chunk_size, chunk_overlap):
+    chunks = []
+    
+    for text, reference in zip(texts, references):
+        file_chunks = chunk_python_file_content(text, char_limit=chunk_size)
+        
+        for chunk in file_chunks:
+            document = Document(page_content=chunk, metadata={"source": reference})
+            chunks.append(document)
+    
+    print(f"Total number of chunks: {len(chunks)}")
+    return chunks
+
+
 # Split text into chunks
 def split_into_chunks(texts, references, chunk_size, chunk_overlap):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)