Update graph_rag.py

graph_rag.py

@@ -1,5 +1,62 @@
-import gradio as gr
+import os
+import zipfile
+import tempfile
+import shutil
+from graphrag.query import query_engine
+import logging
+from dotenv import load_dotenv
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+load_dotenv()
+
+def unzip_folder(zip_path, extract_path):
+    output_dir = os.path.join(extract_path, "output_diagrams")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(output_dir)
+    
+    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))
+    
+    return output_dir
 
 def qa_tool_graph_rag(user_question):
-    # Placeholder function for Graph-RAG
-    return "Graph-RAG functionality not implemented yet.", [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    try:
+        # Paths from environment variables
+        zip_path = os.getenv('ZIP_PATH', '/home/user/app/output_diagrams.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
+        
+        # Perform the query
+        result = query_engine(user_question, method="global")
+        
+        # Process the result
+        answer = result.get('answer', "No answer found.")
+        
+        # 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:
+        logger.error(f"Error in GraphRAG processing: {str(e)}")
+        return f"An error occurred: {str(e)}", [], gr.update(visible=False), gr.update(visible=False)
+    
+    finally:
+        # Clean up: remove the extracted directory
+        if 'output_dir' in locals():
+            shutil.rmtree(output_dir)