agentrun-mem0ai 0.0.11__py3-none-any.whl

agentrun_mem0/graphs/neptune/base.py

@@ -0,0 +1,497 @@
+import logging
+from abc import ABC, abstractmethod
+
+from agentrun_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 agentrun_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 agentrun_mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
+from agentrun_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["arguments"]["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]["arguments"]["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=0.9)
+            destination_node_search_result = self._search_destination_node(dest_embedding, user_id, threshold=0.9)
+
+            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_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})