Update rag.py

rag.py

@@ -9,6 +9,7 @@ from llama_index.core import (
 )
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.core.llms import MockLLM
+from llama_index.readers.file import PyMuPDFReader
 
 # Set cache directory for Hugging Face models
 os.environ["HF_HOME"] = "/tmp/hf"
@@ -29,13 +30,17 @@ def initialize_index():
         index = load_index_from_storage(storage_context)
     else:
         os.makedirs(DATA_DIR, exist_ok=True)
-        # Load documents if folder is not empty
+        loader = SimpleDirectoryReader(
+            DATA_DIR, 
+            recursive=True, 
+            required_exts=[".pdf"],
+            file_extractor={".pdf": PyMuPDFReader()}
+        )
         if os.listdir(DATA_DIR):
-            documents = SimpleDirectoryReader(DATA_DIR).load_data()
+            documents = loader.load_data()
             index = VectorStoreIndex.from_documents(documents)
             index.storage_context.persist(persist_dir=PERSIST_DIR)
         else:
-            # Create an empty index if no data is found
             index = VectorStoreIndex.from_documents([])
     return index