Update graph_rag.py

graph_rag.py

@@ -64,10 +64,19 @@ def clean_response(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 the search string is found, process the response from that point onward
     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()
+        # 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
     else:
         # If the search string is not found, return the original response
         return response