Update graph_rag.py

graph_rag.py

@@ -59,6 +59,19 @@ def run_graphrag_query(query, ragtest_dir):
         logger.error(f"GraphRAG query failed with error: {result.stderr}")
         return result.stderr
 
+def clean_response(response):
+    # Find the position of "SUCCESS: Global Search Response:"
+    search_str = "SUCCESS: Global Search Response:"
+    start_index = response.find(search_str)
+    
+    # If the search string is found, return the substring starting from after this string
+    if start_index != -1:
+        # Add the length of search_str to start_index to begin after the search string
+        return response[start_index + len(search_str):].strip()
+    else:
+        # If the search string is not found, return the original response
+        return response
+        
 def qa_tool_graph_rag(user_question):
     original_dir = os.getcwd()  # Store the original directory
     try:
@@ -73,7 +86,10 @@ def qa_tool_graph_rag(user_question):
         combined_input = f"{patient_info}\n\n{user_question}"
         
         # Run the GraphRAG query with the combined input
-        answer = run_graphrag_query(combined_input, output_dir)
+        raw_answer = run_graphrag_query(combined_input, output_dir)
+        
+        # Clean the response to remove everything before "SUCCESS: Global Search Response:"
+        answer = clean_response(raw_answer)
         
         logger.info(f"GraphRAG answer generated: {answer}")