From dff7535b1001233e0a0b62a2162f36f45930657f Mon Sep 17 00:00:00 2001
From: vasanthasaikalluri
 <165021735+vasanthasaikalluri@users.noreply.github.com>
Date: Tue, 17 Sep 2024 11:30:01 +0000
Subject: [PATCH] removed __Entity__ labels

---
 backend/src/chunkid_entities.py | 8 ++++++++
 backend/src/shared/constants.py | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/backend/src/chunkid_entities.py b/backend/src/chunkid_entities.py
index 30b1c15a9..53a0544f6 100644
--- a/backend/src/chunkid_entities.py
+++ b/backend/src/chunkid_entities.py
@@ -19,10 +19,18 @@ def process_records(records):
                 relationship = element['relationship']
 
                 if start_node['element_id'] not in seen_nodes:
+                    if "labels" in start_node.keys():
+                        labels = set(start_node["labels"])
+                        labels.discard("__Entity__")
+                        start_node["labels"] = list(labels)
                     nodes.append(start_node)
                     seen_nodes.add(start_node['element_id'])
 
                 if end_node['element_id'] not in seen_nodes:
+                    if "labels" in end_node.keys():
+                        labels = set(end_node["labels"])
+                        labels.discard("__Entity__")
+                        end_node["labels"] = list(labels)
                     nodes.append(end_node)
                     seen_nodes.add(end_node['element_id'])
 
diff --git a/backend/src/shared/constants.py b/backend/src/shared/constants.py
index 4ec88e8a5..8a1d2bba2 100644
--- a/backend/src/shared/constants.py
+++ b/backend/src/shared/constants.py
@@ -74,7 +74,7 @@
 MATCH (chunk)-[:PART_OF]->(d:Document)
 CALL {WITH chunk
 MATCH (chunk)-[:HAS_ENTITY]->(e) 
-MATCH path=(e)(()-[rels:!HAS_ENTITY&!PART_OF]-()){0,2}(:!Chunk &! Document) 
+MATCH path=(e)(()-[rels:!HAS_ENTITY&!PART_OF]-()){0,2}(:!Chunk &! Document &! `__Community__`) 
 UNWIND rels as r
 RETURN collect(distinct r) as rels
 }