Update graph_rag.py

graph_rag.py

@@ -17,17 +17,25 @@ def unzip_folder(zip_path, extract_path):
     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))
+    # Adjust to handle the extra ragtest folder inside the zip
+    actual_output_dir = os.path.join(output_dir, "ragtest")
     
-    return output_dir
+    if os.path.exists(actual_output_dir):
+        logger.info(f"Extracted contents to {actual_output_dir}")
+        logger.info(f"Contents of {actual_output_dir}:")
+        for file in os.listdir(actual_output_dir):
+            logger.info(os.path.join(actual_output_dir, file))
+    else:
+        logger.error(f"Expected directory {actual_output_dir} does not exist. Check the structure of the zip file.")
+        actual_output_dir = output_dir  # fallback in case the structure is not as expected
+    
+    return actual_output_dir
 
 def run_graphrag_query(query, ragtest_dir):
-    # Check and log the contents of the ragtest_dir
+    # Log the directory and its contents
+    logger.info(f"Running GraphRAG query with root: {ragtest_dir}")
     logger.info(f"Contents of {ragtest_dir}:")
     for file in os.listdir(ragtest_dir):
         logger.info(os.path.join(ragtest_dir, file))
@@ -47,9 +55,9 @@ def run_graphrag_query(query, ragtest_dir):
     if result.returncode == 0:
         return result.stdout
     else:
+        logger.error(f"GraphRAG query failed with error: {result.stderr}")
         return result.stderr
 
-
 def qa_tool_graph_rag(user_question):
     original_dir = os.getcwd()  # Store the original directory
     try:
@@ -77,3 +85,4 @@ def qa_tool_graph_rag(user_question):
             shutil.rmtree(output_dir)
         
         os.chdir(original_dir)  # Return to the original directory
+