Update app.py

app.py

@@ -15,6 +15,8 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_groq import ChatGroq
 from dotenv import load_dotenv
 from langchain.docstore.document import Document
+from langchain.schema import Document
+from chunk_python_code import chunk_python_code_with_metadata
 
 # Load environment variables from .env file
 load_dotenv()
@@ -165,7 +167,6 @@ def get_file(temp_dir, file_path):
     return full_path
 
 
-#getFilesFromRepo
 def process_directory(directory, folder_paths, file_paths):
     all_texts = []
     file_references = []
@@ -198,22 +199,15 @@ def process_directory(directory, folder_paths, file_paths):
 
             print(f"Total number of files: {len(files)}")
             for file_path in files:
-                #print(f"Paths of files: {iles}")
+                
                 file_ext = os.path.splitext(file_path)[1]
-
                 if os.path.getsize(file_path) == 0:
                     print(f"Skipping an empty file: {file_path}")
                     continue
 
                 with open(file_path, 'rb') as f:
-                    if file_ext in ['.rst', '.md', '.txt', '.html', '.json', '.yaml', '.py']:
+                    if file_ext in ['.rst', '.py']:
                         text = f.read().decode('utf-8')
-                        
-                    elif file_ext in ['.svg']:
-                        text = f"SVG file content from {file_path}"
-                    elif file_ext in ['.png', '.ico']:
-                        text = f"Image metadata from {file_path}"
-                    else:
                         continue
 
                     all_texts.append(text)
@@ -223,177 +217,16 @@ def process_directory(directory, folder_paths, file_paths):
                     file_references.append(relative_path)
                     print("AAAAAAAAAAAAAAAAAAAAAAAAAAAAA: ", relative_path)
                     
-
-                    # Add this snippet after the 'all_texts.append(text)' and 'file_references.append(file_path)' lines.
-
-
     return all_texts, file_references
                 
- #   with tempfile.TemporaryDirectory() as tmpdirname:
-    # Unzip the file into the temporary directory
- #   with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
- #       zip_ref.extractall(tmpdirname)
-
 
-  #      unzipped_root = os.listdir(tmpdirname)
 
-            
-def process_directory5(directory, partial_paths=None, file_paths=None):
-    all_texts = []
-    file_references = []
-
-    zip_files = [file for file in os.listdir(directory) if file.endswith('.zip')]
-    
-    if not zip_files:
-        print("No zip file found in the directory.")
-        return all_texts, file_references
-
-    if len(zip_files) > 1:
-        print("More than one zip file found.")
-        return all_texts, file_references
-    else:
-        zip_file_path = os.path.join(directory, zip_files[0])
-
-        # Create a temporary directory for the zip file
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            # Unzip the file into the temporary directory
-            with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
-                zip_ref.extractall(tmpdirname)
-            
-            files = []
-            print("tmpdirname: " , tmpdirname)
-            unzipped_root = os.listdir(tmpdirname)
-            print("unzipped_root ", unzipped_root)
-            if len(unzipped_root) == 1 and os.path.isdir(os.path.join(tmpdirname, unzipped_root[0])):
-                tmpsubdirpath= os.path.join(tmpdirname, unzipped_root[0])
-                print("AYYYYYYY 11111")
-            else:
-                tmpsubdirpath = tmpdirname
-                print("AYYYYYYY 22222")
-                    
-            if not partial_paths and not file_paths:
-                for root, _, files_list in os.walk(tmpdirname):
-                    for file in files_list:
-                        files.append(os.path.join(root, file))
-            else: 
-                if partial_paths:
-                    for partial_path in partial_paths:
-                        files += get_all_files_in_folder(tmpsubdirpath, partial_path) 
-                if file_paths:
-                    files += [get_file(tmpsubdirpath, file_path) for file_path in file_paths] 
-
-            print(f"Total number of files: {len(files)}")
-            for file_path in files:
-                #print(f"Paths of files: {iles}")
-                file_ext = os.path.splitext(file_path)[1]
-
-                if os.path.getsize(file_path) == 0:
-                    print(f"Skipping an empty file: {file_path}")
-                    continue
-
-                with open(file_path, 'rb') as f:
-                    if file_ext in ['.rst', '.md', '.txt', '.html', '.json', '.yaml', '.py']:
-                        text = f.read().decode('utf-8')
-                    elif file_ext in ['.svg']:
-                        text = f"SVG file content from {file_path}"
-                    elif file_ext in ['.png', '.ico']:
-                        text = f"Image metadata from {file_path}"
-                    else:
-                        continue
-
-                    all_texts.append(text)
-                    file_references.append(file_path)
 
-    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)]
-
-    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"# This is a class method of class: {class_node.name}\n{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
-
-    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)
+def split_python_code_into_chunks(texts, file_paths):
+    chunks = [] 
+    for text, file_path in zip(texts, file_paths):
+        document_chunks = chunk_python_code_with_metadata(text, file_path)
+        chunks.extend(document_chunks)   
     return chunks
 
 
@@ -445,7 +278,7 @@ def rag_workflow(query):
     
     doc_context = "\n\n".join([doc_chunk for doc_chunk, _ in retrieved_doc_chunks])
     code_context = "\n\n".join([code_chunk for code_chunk, _ in retrieved_code_chunks]) 
-    
+
     doc_references = "\n".join([f"[{i+1}] {ref}" for i, (_, ref) in enumerate(retrieved_doc_chunks)])
     code_references = "\n".join([f"[{i+1}] {ref}" for i, (_, ref) in enumerate(retrieved_code_chunks)])
 
@@ -505,15 +338,15 @@ def initialize():
     doc_partial_paths = ['docs/source/setup/']
     doc_file_paths = ['docs/source/usage/lib.rst']
     
-    code_files, code_file_references = process_directory(DATA_DIR,code_partial_paths, code_file_paths)
+    code_texts, code_text_references = process_directory(DATA_DIR, code_partial_paths, code_file_paths)
     print("LEEEEEEEEEEEENGTH of code_files: ", len(code_files))
 
     
-    doc_files, doc_file_references = process_directory(DATA_DIR,doc_partial_paths, doc_file_paths)
+    doc_files, doc_file_references = process_directory(DATA_DIR, doc_partial_paths, doc_file_paths)
     print("LEEEEEEEEEEEENGTH of doc_files: ", len(doc_files))
     
-    code_chunks = split_pythoncode_into_chunks(code_files, code_file_references, 1500, 0)
-    doc_chunks = split_into_chunks(doc_files, doc_file_references, CHUNK_SIZE, CHUNK_OVERLAP)
+    code_chunks = split_python_code_into_chunks(code_texts, code_text_references)
+    doc_chunks = split_into_chunks(doc_file_texts, doc_file_references, CHUNK_SIZE, CHUNK_OVERLAP)
 
     print(f"Total number of code_chunks: {len(code_chunks)}")
     print(f"Total number of doc_chunks: {len(doc_chunks)}")