node_merger.py

import os, sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))


from langchain_community.graphs import MemgraphGraph
import re
from typing import Union
from langchain_community.graphs.graph_document import GraphDocument
from data_processing.data_loading.retriever import TemporalGraphRetriever
import logging 


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class MergeMode: 
    
    EXACT_MATCH = "exact_match"
    VECTOR_SIMILARITY = "vector_similarity"


class NodeMerger:
    """Merge similar nodes in the knowledge graph."""
    
    def __init__(self, graph: MemgraphGraph):
        self.graph = graph
        self.retriever = TemporalGraphRetriever()

    def merge_entities(self, mode: Union[MergeMode.EXACT_MATCH, MergeMode.VECTOR_SIMILARITY], graph_doc: GraphDocument) : # type: ignore

        if mode == MergeMode.EXACT_MATCH:
            self.merge_by_name(graph_doc)
        else : 
            self.merge_by_vector_similarity(graph_doc)


    def merge_by_name(self, graph_doc: GraphDocument) -> int:
        """
        Merge nodes with the same name by exact match.
        """
        
        simplified_query = """
        MATCH (a:Entity), (b:Entity)
        WHERE a.id <> b.id
        AND a.name IS NOT NULL AND b.name IS NOT NULL
        AND toLower(trim(a.name)) = toLower(trim(b.name))
        
        OPTIONAL MATCH (a)-[r1:MERGED_WITH]->(b)
        OPTIONAL MATCH (b)-[r2:MERGED_WITH]->(a)
        WHERE r1 IS NULL AND r2 IS NULL
 
        SET a += b
        SET a.merged_ids = CASE WHEN a.merged_ids IS NULL THEN [b.id] ELSE a.merged_ids + [b.id] END
        SET a.merged_count = CASE WHEN a.merged_count IS NULL THEN 1 ELSE a.merged_count + 1 END

        CREATE (b)-[:MERGED_WITH]->(a)
        RETURN count(b) as merged_count
        """

        try:
            result = self.graph.query(simplified_query)
            merged_count = result[0]["merged_count"] if result else 0
            return merged_count
        except Exception as e2:
            print(f"Error with simplified merging: {e2}")
            return 0
    

    def merge_by_vector_similarity(self, graph_doc: GraphDocument):
        """
        Merge nodes with similar vector representations.
        """
    
        for node in graph_doc.nodes:

            related_node_id = self.retriever.node_retrieval([node.properties.get("name", "")], limit=1, threshold=0.95)

            if not related_node_id:
                continue

            if related_node_id == node.id:
                continue

            try: 
                check_merged_query = """
                MATCH (a {id: $node_id})-[:MERGED_WITH]-(b {id: $related_id})
                RETURN count(*) > 0 as already_merged
                """

                result = self.graph.query(
                    check_merged_query, 
                    {
                        "node_id": node.id,
                        "related_id": related_node_id
                    }
                )
                already_merged = result[0]["already_merged"] if result else False
                if already_merged:
                    continue

                get_related_query = """
                MATCH (n {id: $related_id})
                RETURN n.name as name, n.description as description, 
                       labels(n) as labels, properties(n) as properties
                """

                related_result = self.graph.query(get_related_query, {"related_id": related_node_id})
                
                if not related_result:
                    continue

                related_info = related_result[0]

                merged_properties = dict(node.properties)
                related_properties = dict(related_info["properties"])

                curr_desc = merged_properties.get("description", "")
                related_desc = related_properties.get("description", "")

                if curr_desc and related_desc:
                    merged_properties["description"] = f"{curr_desc} {related_desc}".strip()
                elif related_desc:
                    merged_properties["description"] = related_desc.strip()
                elif curr_desc:
                    merged_properties["description"] = curr_desc.strip()
                else:
                    merged_properties["description"] = ""

                for key, value in related_properties.items():
                    if key not in merged_properties:
                        merged_properties[key] = value
                    elif isinstance(merged_properties[key], list) and isinstance(value, list):
                        merged_properties[key].extend(value)
                    elif isinstance(merged_properties[key], str) and isinstance(value, str):
                        merged_properties[key] += f" {value}"

                merged_ids = merged_properties.get("merged_ids", [])
                if isinstance(merged_ids, str):
                    merged_ids = [merged_ids]
                else:
                    merged_ids = []

                if related_node_id not in merged_ids:
                    merged_ids.append(related_node_id)

                merged_properties["merged_ids"] = merged_ids

                update_node_query = """
                MATCH (n {id: $node_id})
                SET n += $properties
                """
                
                self.graph.query(update_node_query, {
                    "node_id": node.id,
                    "properties": merged_properties
                })
                
                # Create merge relationship
                create_merge_rel_query = """
                MATCH (main {id: $node_id}), (related {id: $related_id})
                CREATE (related)-[:MERGED_WITH {merged_at: timestamp()}]->(main)
                """
                
                self.graph.query(create_merge_rel_query, {
                    "node_id": node.id,
                    "related_id": related_node_id
                })

                self._redirect_relationships(node.id, related_node_id)


            except Exception as e:
                logger.error(f"Error checking merged nodes: {e}")
            


    def _redirect_relationships(self, from_node_id: str, to_node_id: str): 
        """
        Redirect relationships from one node to another.
        """

        try: 
            redirect_incoming_query = """ 
                MATCH (source)-[r]->(target {id: $from_id})
                WHERE source.id <> $to_id
                WITH source, r, target, type(r) as rel_type, properties(r) as rel_props
                
                OPTIONAL MATCH (source)-[existing_rel]->(new_target {id: $to_id})
                WHERE type(existing_rel) = rel_type
                
                FOREACH (x IN CASE WHEN existing_rel IS NULL THEN [1] ELSE [] END |
                    MERGE (source)-[new_r:`{}`]->(new_target {id: $to_id})
                    SET new_r = rel_props
                )
                DELETE r
            """.format("` + rel_type + `")

            self.graph.query(
                redirect_incoming_query, 
                {
                    "from_id": from_node_id,
                    "to_id": to_node_id
                }
            )

            redirect_outgoing_query = """
                MATCH (source {id: $from_id})-[r]->(target)
                WHERE target.id <> $to_id
                WITH source, r, target, type(r) as rel_type, properties(r) as rel_props
                
                OPTIONAL MATCH (new_source {id: $to_id})-[existing_rel]->(target)
                WHERE type(existing_rel) = rel_type
                
                FOREACH (x IN CASE WHEN existing_rel IS NULL THEN [1] ELSE [] END |
                    MERGE (new_source {id: $to_id})-[new_r:`{}`]->(target)
                    SET new_r = rel_props
                )
                DELETE r
            """.format("` + rel_type + `")

            self.graph.query(
                redirect_outgoing_query, 
                {
                    "from_id": from_node_id,
                    "to_id": to_node_id
                }
            )

        except Exception as e:
            logger.error(f"Error redirecting relationships: {e}")