mem0/mem0/graphs/neptune/base.py

import logging
from abc import ABC, abstractmethod

from mem0.memory.utils import format_entities

try:
    from rank_bm25 import BM25Okapi
except ImportError:
    raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")

from mem0.graphs.tools import (
    DELETE_MEMORY_STRUCT_TOOL_GRAPH,
    DELETE_MEMORY_TOOL_GRAPH,
    EXTRACT_ENTITIES_STRUCT_TOOL,
    EXTRACT_ENTITIES_TOOL,
    RELATIONS_STRUCT_TOOL,
    RELATIONS_TOOL,
)
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
from mem0.utils.factory import EmbedderFactory, LlmFactory, VectorStoreFactory

logger = logging.getLogger(__name__)


class NeptuneBase(ABC):
    """
    Abstract base class for neptune (neptune analytics and neptune db) calls using OpenCypher
    to store/retrieve data
    """

    @staticmethod
    def _create_embedding_model(config):
        """
        :return: the Embedder model used for memory store
        """
        return EmbedderFactory.create(
            config.embedder.provider,
            config.embedder.config,
            {"enable_embeddings": True},
        )

    @staticmethod
    def _create_llm(config, llm_provider):
        """
        :return: the llm model used for memory store
        """
        return LlmFactory.create(llm_provider, config.llm.config)

    @staticmethod
    def _create_vector_store(vector_store_provider, config):
        """
        :param vector_store_provider: name of vector store
        :param config: the vector_store configuration
        :return:
        """
        return VectorStoreFactory.create(vector_store_provider, config.vector_store.config)

    def add(self, data, filters):
        """
        Adds data to the graph.

        Args:
            data (str): The data to add to the graph.
            filters (dict): A dictionary containing filters to be applied during the addition.
        """
        entity_type_map = self._retrieve_nodes_from_data(data, filters)
        to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
        search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
        to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)

        deleted_entities = self._delete_entities(to_be_deleted, filters["user_id"])
        added_entities = self._add_entities(to_be_added, filters["user_id"], entity_type_map)

        return {"deleted_entities": deleted_entities, "added_entities": added_entities}

    def _retrieve_nodes_from_data(self, data, filters):
        """
        Extract all entities mentioned in the query.
        """
        _tools = [EXTRACT_ENTITIES_TOOL]
        if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
            _tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
        search_results = self.llm.generate_response(
            messages=[
                {
                    "role": "system",
                    "content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
                },
                {"role": "user", "content": data},
            ],
            tools=_tools,
        )

        entity_type_map = {}

        try:
            for tool_call in search_results["tool_calls"]:
                if tool_call["name"] != "extract_entities":
                    continue
                for item in tool_call.get("arguments", {}).get("entities", []):
                    entity_type_map[item["entity"]] = item["entity_type"]
        except Exception as e:
            logger.exception(
                f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
            )

        entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
        return entity_type_map

    def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
        """
        Establish relations among the extracted nodes.
        """
        if self.config.graph_store.custom_prompt:
            messages = [
                {
                    "role": "system",
                    "content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]).replace(
                        "CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}"
                    ),
                },
                {"role": "user", "content": data},
            ]
        else:
            messages = [
                {
                    "role": "system",
                    "content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]),
                },
                {
                    "role": "user",
                    "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}",
                },
            ]

        _tools = [RELATIONS_TOOL]
        if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
            _tools = [RELATIONS_STRUCT_TOOL]

        extracted_entities = self.llm.generate_response(
            messages=messages,
            tools=_tools,
        )

        entities = []
        if extracted_entities["tool_calls"]:
            entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])

        entities = self._remove_spaces_from_entities(entities)
        logger.debug(f"Extracted entities: {entities}")
        return entities

    def _remove_spaces_from_entities(self, entity_list):
        for item in entity_list:
            item["source"] = item["source"].lower().replace(" ", "_")
            item["relationship"] = item["relationship"].lower().replace(" ", "_")
            item["destination"] = item["destination"].lower().replace(" ", "_")
        return entity_list

    def _get_delete_entities_from_search_output(self, search_output, data, filters):
        """
        Get the entities to be deleted from the search output.
        """

        search_output_string = format_entities(search_output)
        system_prompt, user_prompt = get_delete_messages(search_output_string, data, filters["user_id"])

        _tools = [DELETE_MEMORY_TOOL_GRAPH]
        if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
            _tools = [
                DELETE_MEMORY_STRUCT_TOOL_GRAPH,
            ]

        memory_updates = self.llm.generate_response(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            tools=_tools,
        )

        to_be_deleted = []
        for item in memory_updates["tool_calls"]:
            if item["name"] == "delete_graph_memory":
                to_be_deleted.append(item["arguments"])
        # in case if it is not in the correct format
        to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
        logger.debug(f"Deleted relationships: {to_be_deleted}")
        return to_be_deleted

    def _delete_entities(self, to_be_deleted, user_id):
        """
        Delete the entities from the graph.
        """

        results = []
        for item in to_be_deleted:
            source = item["source"]
            destination = item["destination"]
            relationship = item["relationship"]

            # Delete the specific relationship between nodes
            cypher, params = self._delete_entities_cypher(source, destination, relationship, user_id)
            result = self.graph.query(cypher, params=params)
            results.append(result)
        return results

    @abstractmethod
    def _delete_entities_cypher(self, source, destination, relationship, user_id):
        """
        Returns the OpenCypher query and parameters for deleting entities in the graph DB
        """

        pass

    def _add_entities(self, to_be_added, user_id, entity_type_map):
        """
        Add the new entities to the graph. Merge the nodes if they already exist.
        """

        results = []
        for item in to_be_added:
            # entities
            source = item["source"]
            destination = item["destination"]
            relationship = item["relationship"]

            # types
            source_type = entity_type_map.get(source, "__User__")
            destination_type = entity_type_map.get(destination, "__User__")

            # embeddings
            source_embedding = self.embedding_model.embed(source)
            dest_embedding = self.embedding_model.embed(destination)

            # search for the nodes with the closest embeddings
            source_node_search_result = self._search_source_node(source_embedding, user_id, threshold=self.threshold)
            destination_node_search_result = self._search_destination_node(dest_embedding, user_id, threshold=self.threshold)

            cypher, params = self._add_entities_cypher(
                source_node_search_result,
                source,
                source_embedding,
                source_type,
                destination_node_search_result,
                destination,
                dest_embedding,
                destination_type,
                relationship,
                user_id,
            )
            result = self.graph.query(cypher, params=params)
            results.append(result)
        return results

    def _add_entities_cypher(
        self,
        source_node_list,
        source,
        source_embedding,
        source_type,
        destination_node_list,
        destination,
        dest_embedding,
        destination_type,
        relationship,
        user_id,
    ):
        """
        Returns the OpenCypher query and parameters for adding entities in the graph DB
        """
        if not destination_node_list and source_node_list:
            return self._add_entities_by_source_cypher(
                source_node_list,
                destination,
                dest_embedding,
                destination_type,
                relationship,
                user_id)
        elif destination_node_list and not source_node_list:
            return self._add_entities_by_destination_cypher(
                source,
                source_embedding,
                source_type,
                destination_node_list,
                relationship,
                user_id)
        elif source_node_list and destination_node_list:
            return self._add_relationship_entities_cypher(
                source_node_list,
                destination_node_list,
                relationship,
                user_id)
        # else source_node_list and destination_node_list are empty
        return self._add_new_entities_cypher(
            source,
            source_embedding,
            source_type,
            destination,
            dest_embedding,
            destination_type,
            relationship,
            user_id)

    @abstractmethod
    def _add_entities_by_source_cypher(
            self,
            source_node_list,
            destination,
            dest_embedding,
            destination_type,
            relationship,
            user_id,
    ):
        pass

    @abstractmethod
    def _add_entities_by_destination_cypher(
            self,
            source,
            source_embedding,
            source_type,
            destination_node_list,
            relationship,
            user_id,
    ):
        pass

    @abstractmethod
    def _add_relationship_entities_cypher(
            self,
            source_node_list,
            destination_node_list,
            relationship,
            user_id,
    ):
        pass

    @abstractmethod
    def _add_new_entities_cypher(
            self,
            source,
            source_embedding,
            source_type,
            destination,
            dest_embedding,
            destination_type,
            relationship,
            user_id,
    ):
        pass

    def search(self, query, filters, limit=100):
        """
        Search for memories and related graph data.

        Args:
            query (str): Query to search for.
            filters (dict): A dictionary containing filters to be applied during the search.
            limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.

        Returns:
            dict: A dictionary containing:
                - "contexts": List of search results from the base data store.
                - "entities": List of related graph data based on the query.
        """

        entity_type_map = self._retrieve_nodes_from_data(query, filters)
        search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)

        if not search_output:
            return []

        search_outputs_sequence = [
            [item["source"], item["relationship"], item["destination"]] for item in search_output
        ]
        bm25 = BM25Okapi(search_outputs_sequence)

        tokenized_query = query.split(" ")
        reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)

        search_results = []
        for item in reranked_results:
            search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})

        return search_results

    def _search_source_node(self, source_embedding, user_id, threshold=0.9):
        cypher, params = self._search_source_node_cypher(source_embedding, user_id, threshold)
        result = self.graph.query(cypher, params=params)
        return result

    @abstractmethod
    def _search_source_node_cypher(self, source_embedding, user_id, threshold):
        """
        Returns the OpenCypher query and parameters to search for source nodes
        """
        pass

    def _search_destination_node(self, destination_embedding, user_id, threshold=0.9):
        cypher, params = self._search_destination_node_cypher(destination_embedding, user_id, threshold)
        result = self.graph.query(cypher, params=params)
        return result

    @abstractmethod
    def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
        """
        Returns the OpenCypher query and parameters to search for destination nodes
        """
        pass

    def delete(self, data, filters):
        """
        Delete graph entities associated with the given memory text.

        Extracts entities and relationships from the memory text using the same
        pipeline as add(), then deletes the matching relationships in the graph.

        Args:
            data (str): The memory text whose graph entities should be removed.
            filters (dict): Scope filters (user_id, agent_id, run_id).
        """
        try:
            entity_type_map = self._retrieve_nodes_from_data(data, filters)
            if not entity_type_map:
                logger.debug("No entities found in memory text, skipping graph cleanup")
                return
            to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
            if to_be_deleted:
                self._delete_entities(to_be_deleted, filters["user_id"])
        except Exception as e:
            logger.error(f"Error during graph cleanup for memory delete: {e}")

    def delete_all(self, filters):
        cypher, params = self._delete_all_cypher(filters)
        self.graph.query(cypher, params=params)

    @abstractmethod
    def _delete_all_cypher(self, filters):
        """
        Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
        """
        pass

    def get_all(self, filters, limit=100):
        """
        Retrieves all nodes and relationships from the graph database based on filtering criteria.

        Args:
            filters (dict): A dictionary containing filters to be applied during the retrieval.
            limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
        Returns:
            list: A list of dictionaries, each containing:
                - 'contexts': The base data store response for each memory.
                - 'entities': A list of strings representing the nodes and relationships
        """

        # return all nodes and relationships
        query, params = self._get_all_cypher(filters, limit)
        results = self.graph.query(query, params=params)

        final_results = []
        for result in results:
            final_results.append(
                {
                    "source": result["source"],
                    "relationship": result["relationship"],
                    "target": result["target"],
                }
            )

        logger.debug(f"Retrieved {len(final_results)} relationships")

        return final_results

    @abstractmethod
    def _get_all_cypher(self, filters, limit):
        """
        Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
        """
        pass

    def _search_graph_db(self, node_list, filters, limit=100):
        """
        Search similar nodes among and their respective incoming and outgoing relations.
        """
        result_relations = []

        for node in node_list:
            n_embedding = self.embedding_model.embed(node)
            cypher_query, params = self._search_graph_db_cypher(n_embedding, filters, limit)
            ans = self.graph.query(cypher_query, params=params)
            result_relations.extend(ans)

        return result_relations

    @abstractmethod
    def _search_graph_db_cypher(self, n_embedding, filters, limit):
        """
        Returns the OpenCypher query and parameters to search for similar nodes in the memory store
        """
        pass

    # Reset is not defined in base.py
    def reset(self):
        """
        Reset the graph by clearing all nodes and relationships.

        link: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/neptune-graph/client/reset_graph.html
        """

        logger.warning("Clearing graph...")
        graph_id = self.graph.graph_identifier
        self.graph.client.reset_graph(
            graphIdentifier=graph_id,
            skipSnapshot=True,
        )
        waiter = self.graph.client.get_waiter("graph_available")
        waiter.wait(graphIdentifier=graph_id, WaiterConfig={"Delay": 10, "MaxAttempts": 60})