Update graph_rag.py

graph_rag.py

@@ -65,22 +65,25 @@ def clean_response(response):
     search_str = "SUCCESS: Global Search Response:"
     start_index = response.find(search_str)
     
-    # If the search string is found, process the response from that point onward
+    # If the search string is found, return the substring starting from after this string
     if start_index != -1:
-        # Extract the relevant part of the response
-        cleaned_response = response[start_index + len(search_str):].strip()
-        
-        # Remove "[Data: Reports (XXX)]" references
-        pattern = r'\[Data: Reports \(\d+\)\]'
-        cleaned_response = re.sub(pattern, '', cleaned_response)
-        
-        # Remove extra spaces that might result from the substitution
-        cleaned_response = re.sub(r'\s+', ' ', cleaned_response).strip()
-        
-        return cleaned_response
+        # 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 remove_data_reports(response):
+    # Regular expression pattern to find "[Data: Reports (XXX)]" and similar patterns
+    pattern = r'\[Data: Reports \(\d+\)\]'
+    
+    # Substitute the pattern with an empty string
+    cleaned_response = re.sub(pattern, '', response)
+    
+    # Remove extra spaces that might result from the substitution
+    cleaned_response = re.sub(r'\s+', ' ', cleaned_response).strip()
+    
+    return cleaned_response
         
 def qa_tool_graph_rag(user_question):
     original_dir = os.getcwd()  # Store the original directory
@@ -100,6 +103,8 @@ def qa_tool_graph_rag(user_question):
         
         # Clean the response to remove everything before "SUCCESS: Global Search Response:"
         answer = clean_response(raw_answer)
+
+        answer = remove_data_reports(answer)
         
         logger.info(f"GraphRAG answer generated: {answer}")