Update graph_rag.py

graph_rag.py

@@ -13,28 +13,36 @@ logger = logging.getLogger(__name__)
 load_dotenv()
 
 def unzip_folder(zip_path, extract_path):
-    output_dir = os.path.join(extract_path, "ragtest")
-    os.makedirs(output_dir, exist_ok=True)
-    
     with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-        zip_ref.extractall(output_dir)
+        zip_ref.extractall(extract_path)
     
-    logger.info(f"Extracted contents to {output_dir}")
-    logger.info(f"Contents of {output_dir}:")
-    for file in os.listdir(output_dir):
-        logger.info(os.path.join(output_dir, file))
+    logger.info(f"Extracted contents to {extract_path}")
+    logger.info(f"Contents of {extract_path}:")
+    for root, dirs, files in os.walk(extract_path):
+        for file in files:
+            logger.info(os.path.join(root, file))
     
-    return output_dir
+    return extract_path
+
+def find_data_dir(root_dir):
+    for root, dirs, files in os.walk(root_dir):
+        if 'output' in dirs:
+            return os.path.join(root, 'output')
+    return root_dir
 
-def run_graphrag_query(query):
+def run_graphrag_query(query, root_dir):
+    data_dir = find_data_dir(root_dir)
+    
     # Define the command
     command = [
         "python", "-m", "graphrag.query",
-        "--root", "./ragtest",
-        "--method", "local",
+        "--root", data_dir,
+        "--method", "global",
         query
     ]
     
+    logger.info(f"Running command: {' '.join(command)}")
+    
     # Run the command
     result = subprocess.run(command, capture_output=True, text=True)
     
@@ -42,22 +50,20 @@ def run_graphrag_query(query):
     if result.returncode == 0:
         return result.stdout
     else:
+        logger.error(f"GraphRAG error: {result.stderr}")
         return result.stderr
 
 def qa_tool_graph_rag(user_question):
     try:
         # Paths from environment variables
         zip_path = os.getenv('ZIP_PATH', '/home/user/app/ragtest.zip')
-        extract_path = os.getenv('EXTRACT_PATH', '/home/user/app')
+        extract_path = os.getenv('EXTRACT_PATH', '/home/user/app/ragtest')
         
         # Unzip the folder
         output_dir = unzip_folder(zip_path, extract_path)
         
-        # Change to the directory containing the ragtest folder
-        os.chdir(extract_path)
-        
         # Run the GraphRAG query
-        answer = run_graphrag_query(user_question)
+        answer = run_graphrag_query(user_question, output_dir)
         
         logger.info(f"GraphRAG answer generated: {answer}")
         
@@ -73,8 +79,5 @@ def qa_tool_graph_rag(user_question):
     
     finally:
         # Clean up: remove the extracted directory
-        if 'output_dir' in locals():
-            shutil.rmtree(output_dir)
-        
-        # Change back to the original directory
-        os.chdir('/home/user/app')
+        if os.path.exists(extract_path):
+            shutil.rmtree(extract_path)