Update graph_rag.py

graph_rag.py

@@ -1,10 +1,10 @@
 import os
 import zipfile
-import tempfile
 import shutil
-from graphrag.query import query_engine
 import logging
+import subprocess
 from dotenv import load_dotenv
+import gradio as gr
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -13,7 +13,7 @@ logger = logging.getLogger(__name__)
 load_dotenv()
 
 def unzip_folder(zip_path, extract_path):
-    output_dir = os.path.join(extract_path, "output_diagrams")
+    output_dir = os.path.join(extract_path, "ragtest")
     os.makedirs(output_dir, exist_ok=True)
     
     with zipfile.ZipFile(zip_path, 'r') as zip_ref:
@@ -26,30 +26,45 @@ def unzip_folder(zip_path, extract_path):
     
     return output_dir
 
+def run_graphrag_query(query):
+    # Define the command
+    command = [
+        "python", "-m", "graphrag.query",
+        "--root", "./ragtest",
+        "--method", "local",
+        query
+    ]
+    
+    # Run the command
+    result = subprocess.run(command, capture_output=True, text=True)
+    
+    # Return the output or error message
+    if result.returncode == 0:
+        return result.stdout
+    else:
+        return result.stderr
+
 def qa_tool_graph_rag(user_question):
     try:
         # Paths from environment variables
-        zip_path = os.getenv('ZIP_PATH', '/home/user/app/output_diagrams.zip')
+        zip_path = os.getenv('ZIP_PATH', '/home/user/app/ragtest.zip')
         extract_path = os.getenv('EXTRACT_PATH', '/home/user/app')
         
         # Unzip the folder
         output_dir = unzip_folder(zip_path, extract_path)
         
-        # Set the GRAPHRAG_ROOT environment variable
-        os.environ['GRAPHRAG_ROOT'] = output_dir
+        # Change to the directory containing the ragtest folder
+        os.chdir(extract_path)
         
-        # Perform the query
-        result = query_engine(user_question, method="global")
+        # Run the GraphRAG query
+        answer = run_graphrag_query(user_question)
         
-        # Process the result
-        answer = result.get('answer', "No answer found.")
+        logger.info(f"GraphRAG answer generated: {answer}")
         
         # For simplicity, we're not handling images here. 
         # You may need to adjust this based on how GraphRAG returns results.
         images = []
         
-        logger.info(f"GraphRAG answer generated: {answer}")
-        
         return answer, images, gr.update(visible=True), gr.update(visible=True)
     
     except Exception as e:
@@ -59,4 +74,7 @@ def qa_tool_graph_rag(user_question):
     finally:
         # Clean up: remove the extracted directory
         if 'output_dir' in locals():
-            shutil.rmtree(output_dir)
+            shutil.rmtree(output_dir)
+        
+        # Change back to the original directory
+        os.chdir('/home/user/app')