mem0/examples/graph-db-demo/neptune-example.ipynb

Neptune as Graph Memory

In this notebook, we will be connecting using a Amazon Neptune Analytics instance as our memory graph storage for Mem0.

The Graph Memory storage persists memories in a graph or relationship form when performing m.add memory operations. It then uses vector distance algorithms to find related memories during a m.search operation. Relationships are returned in the result, and add context to the memories.

Reference: Vector Similarity using Neptune Analytics

Prerequisites

1. Install Mem0 with Graph Memory support

To use Mem0 with Graph Memory support, install it using pip:

pip install "mem0ai[graph]"

This command installs Mem0 along with the necessary dependencies for graph functionality.

2. Connect to Neptune

To connect to Amazon Neptune Analytics, you need to configure Neptune with your Amazon profile credentials. The best way to do this is to declare environment variables with IAM permission to your Neptune Analytics instance. The graph-identifier for the instance to persist memories needs to be defined in the Mem0 configuration under "graph_store", with the "neptune" provider. Note that the Neptune Analytics instance needs to have vector-search-configuration defined to meet the needs of the llm model's vector dimensions, see: https://docs.aws.amazon.com/neptune-analytics/latest/userguide/vector-index.html.

embedding_dimensions = 1536
graph_identifier = "<MY-GRAPH>" # graph with 1536 dimensions for vector search
config = {
    "embedder": {
        "provider": "openai",
        "config": {
            "model": "text-embedding-3-large",
            "embedding_dims": embedding_dimensions
        },
    },
    "graph_store": {
        "provider": "neptune",
        "config": {
            "endpoint": f"neptune-graph://{graph_identifier}",
        },
    },
}

3. Configure OpenSearch

We're going to use OpenSearch as our vector store. You can run OpenSearch from docker image:

docker pull opensearchproject/opensearch:2

And verify that it's running with a <custom-admin-password>:

 docker run -d -p 9200:9200 -p 9600:9600 -e "discovery.type=single-node" -e "OPENSEARCH_INITIAL_ADMIN_PASSWORD=<custom-admin-password>" opensearchproject/opensearch:latest

 curl https://localhost:9200 -ku admin:<custom-admin-password>

We're going to connect OpenSearch using the python client:

pip install "opensearch-py"

Configuration

Do all the imports and configure OpenAI (enter your OpenAI API key):

In [1]:

from mem0 import Memory
import os
import logging
import sys

logging.getLogger("mem0.graphs.neptune.main").setLevel(logging.DEBUG)
logging.getLogger("mem0.graphs.neptune.base").setLevel(logging.DEBUG)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

logging.basicConfig(
    format="%(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    stream=sys.stdout,  # Explicitly set output to stdout
)

Setup the Mem0 configuration using:

• openai as the embedder
• Amazon Neptune Analytics instance as a graph store
• OpenSearch as the vector store

In [2]:

graph_identifier = os.environ.get("GRAPH_ID")
opensearch_username = os.environ.get("OS_USERNAME")
opensearch_password = os.environ.get("OS_PASSWORD")
config = {
    "embedder": {
        "provider": "openai",
        "config": {"model": "text-embedding-3-large", "embedding_dims": 1536},
    },
    "graph_store": {
        "provider": "neptune",
        "config": {
            "endpoint": f"neptune-graph://{graph_identifier}",
        },
    },
    "vector_store": {
        "provider": "opensearch",
        "config": {
            "collection_name": "vector_store",
            "host": "localhost",
            "port": 9200,
            "user": opensearch_username,
            "password": opensearch_password,
            "use_ssl": False,
            "verify_certs": False,
        },
    },
}

Graph Memory initializiation

Initialize Memgraph as a Graph Memory store:

In [3]:

m = Memory.from_config(config_dict=config)

app_id = "movies"
user_id = "alice"

m.delete_all(user_id=user_id)

Out [3]:

WARNING - Creating index vector_store, it might take 1-2 minutes...
WARNING - Creating index mem0migrations, it might take 1-2 minutes...
DEBUG - delete_all query=
        MATCH (n  {user_id: $user_id})
        DETACH DELETE n
        

{'message': 'Memories deleted successfully!'}

Store memories

Create memories and store one at a time:

In [4]:

messages = [
    {
        "role": "user",
        "content": "I'm planning to watch a movie tonight. Any recommendations?",
    },
]

# Store inferred memories (default behavior)
result = m.add(messages, user_id=user_id, metadata={"category": "movie_recommendations"})

all_results = m.get_all(user_id=user_id)
for n in all_results["results"]:
    print(f"node \"{n['memory']}\": [hash: {n['hash']}]")

for e in all_results["relations"]:
    print(f"edge \"{e['source']}\" --{e['relationship']}--> \"{e['target']}\"")

DEBUG - Extracted entities: [{'source': 'alice', 'relationship': 'plans_to_watch', 'destination': 'movie'}]
DEBUG - _search_graph_db
  query=
            MATCH (n )
            WHERE n.user_id = $user_id
            WITH n, $n_embedding as n_embedding
            CALL neptune.algo.vectors.distanceByEmbedding(
                n_embedding,
                n,
                {metric:"CosineSimilarity"}
            ) YIELD distance
            WITH n, distance as similarity
            WHERE similarity >= $threshold
            CALL {
                WITH n
                MATCH (n)-[r]->(m) 
                RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id
                UNION ALL
                WITH n
                MATCH (m)-[r]->(n) 
                RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id
            }
            WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
            RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
            ORDER BY similarity DESC
            LIMIT $limit
            
DEBUG - Deleted relationships: []
DEBUG - _search_source_node
  query=
            MATCH (source_candidate )
            WHERE source_candidate.user_id = $user_id 

            WITH source_candidate, $source_embedding as v_embedding
            CALL neptune.algo.vectors.distanceByEmbedding(
                v_embedding,
                source_candidate,
                {metric:"CosineSimilarity"}
            ) YIELD distance
            WITH source_candidate, distance AS cosine_similarity
            WHERE cosine_similarity >= $threshold

            WITH source_candidate, cosine_similarity
            ORDER BY cosine_similarity DESC
            LIMIT 1

            RETURN id(source_candidate), cosine_similarity
            
DEBUG - _search_destination_node
  query=
                MATCH (destination_candidate )
                WHERE destination_candidate.user_id = $user_id
                
                WITH destination_candidate, $destination_embedding as v_embedding
                CALL neptune.algo.vectors.distanceByEmbedding(
                    v_embedding,
                    destination_candidate, 
                    {metric:"CosineSimilarity"}
                ) YIELD distance
                WITH destination_candidate, distance AS cosine_similarity
                WHERE cosine_similarity >= $threshold

                WITH destination_candidate, cosine_similarity
                ORDER BY cosine_similarity DESC
                LIMIT 1
    
                RETURN id(destination_candidate), cosine_similarity
                
DEBUG - _add_entities:
  destination_node_search_result=[]
  source_node_search_result=[]
  query=
                    MERGE (n :`__User__` {name: $source_name, user_id: $user_id})
                    ON CREATE SET n.created = timestamp(),
                                  n.mentions = 1
                                  
                    ON MATCH SET n.mentions = coalesce(n.mentions, 0) + 1
                    WITH n, $source_embedding as source_embedding
                    CALL neptune.algo.vectors.upsert(n, source_embedding)
                    WITH n
                    MERGE (m :`entertainment` {name: $dest_name, user_id: $user_id})
                    ON CREATE SET m.created = timestamp(),
                                  m.mentions = 1
                                  
                    ON MATCH SET m.mentions = coalesce(m.mentions, 0) + 1
                    WITH n, m, $dest_embedding as dest_embedding
                    CALL neptune.algo.vectors.upsert(m, dest_embedding)
                    WITH n, m
                    MERGE (n)-[rel:plans_to_watch]->(m)
                    ON CREATE SET rel.created = timestamp(), rel.mentions = 1
                    ON MATCH SET rel.mentions = coalesce(rel.mentions, 0) + 1
                    RETURN n.name AS source, type(rel) AS relationship, m.name AS target
                    
DEBUG - Retrieved 1 relationships
node "Planning to watch a movie tonight": [hash: bf55418607cfdca4afa311b5fd8496bd]
edge "alice" --plans_to_watch--> "movie"

In [6]:

messages = [
    {
        "role": "assistant",
        "content": "How about a thriller movies? They can be quite engaging.",
    },
]

# Store inferred memories (default behavior)
result = m.add(messages, user_id=user_id, metadata={"category": "movie_recommendations"})

all_results = m.get_all(user_id=user_id)
for n in all_results["results"]:
    print(f"node \"{n['memory']}\": [hash: {n['hash']}]")

for e in all_results["relations"]:
    print(f"edge \"{e['source']}\" --{e['relationship']}--> \"{e['target']}\"")

DEBUG - Extracted entities: [{'source': 'thriller_movies', 'relationship': 'is_engaging', 'destination': 'thriller_movies'}]
DEBUG - _search_graph_db
  query=
            MATCH (n )
            WHERE n.user_id = $user_id
            WITH n, $n_embedding as n_embedding
            CALL neptune.algo.vectors.distanceByEmbedding(
                n_embedding,
                n,
                {metric:"CosineSimilarity"}
            ) YIELD distance
            WITH n, distance as similarity
            WHERE similarity >= $threshold
            CALL {
                WITH n
                MATCH (n)-[r]->(m) 
                RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id
                UNION ALL
                WITH n
                MATCH (m)-[r]->(n) 
                RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id
            }
            WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
            RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
            ORDER BY similarity DESC
            LIMIT $limit
            
DEBUG - Deleted relationships: []
DEBUG - _search_source_node
  query=
            MATCH (source_candidate )
            WHERE source_candidate.user_id = $user_id 

            WITH source_candidate, $source_embedding as v_embedding
            CALL neptune.algo.vectors.distanceByEmbedding(
                v_embedding,
                source_candidate,
                {metric:"CosineSimilarity"}
            ) YIELD distance
            WITH source_candidate, distance AS cosine_similarity
            WHERE cosine_similarity >= $threshold

            WITH source_candidate, cosine_similarity
            ORDER BY cosine_similarity DESC
            LIMIT 1

            RETURN id(source_candidate), cosine_similarity
            
DEBUG - _search_destination_node
  query=
                MATCH (destination_candidate )
                WHERE destination_candidate.user_id = $user_id
                
                WITH destination_candidate, $destination_embedding as v_embedding
                CALL neptune.algo.vectors.distanceByEmbedding(
                    v_embedding,
                    destination_candidate, 
                    {metric:"CosineSimilarity"}
                ) YIELD distance
                WITH destination_candidate, distance AS cosine_similarity
                WHERE cosine_similarity >= $threshold

                WITH destination_candidate, cosine_similarity
                ORDER BY cosine_similarity DESC
                LIMIT 1
    
                RETURN id(destination_candidate), cosine_similarity
                
DEBUG - _add_entities:
  destination_node_search_result=[{'id(destination_candidate)': '67c49d52-e305-47fe-9fce-2cd5adc5d83c0', 'cosine_similarity': 0.999999}]
  source_node_search_result=[{'id(source_candidate)': '67c49d52-e305-47fe-9fce-2cd5adc5d83c0', 'cosine_similarity': 0.999999}]
  query=
                    MATCH (source)
                    WHERE id(source) = $source_id
                    SET source.mentions = coalesce(source.mentions, 0) + 1
                    WITH source
                    MATCH (destination)
                    WHERE id(destination) = $destination_id
                    SET destination.mentions = coalesce(destination.mentions) + 1
                    MERGE (source)-[r:is_engaging]->(destination)
                    ON CREATE SET 
                        r.created_at = timestamp(),
                        r.updated_at = timestamp(),
                        r.mentions = 1
                    ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
                    RETURN source.name AS source, type(r) AS relationship, destination.name AS target
                    
DEBUG - Retrieved 3 relationships
node "Planning to watch a movie tonight": [hash: bf55418607cfdca4afa311b5fd8496bd]
edge "thriller_movies" --is_a_type_of--> "movie"
edge "alice" --plans_to_watch--> "movie"
edge "thriller_movies" --is_engaging--> "thriller_movies"

In [ ]:

messages = [
    {
        "role": "user",
        "content": "I'm not a big fan of thriller movies but I love sci-fi movies.",
    },
]

# Store inferred memories (default behavior)
result = m.add(messages, user_id=user_id, metadata={"category": "movie_recommendations"})

all_results = m.get_all(user_id=user_id)
for n in all_results["results"]:
    print(f"node \"{n['memory']}\": [hash: {n['hash']}]")

for e in all_results["relations"]:
    print(f"edge \"{e['source']}\" --{e['relationship']}--> \"{e['target']}\"")

In [ ]:

messages = [
    {
        "role": "assistant",
        "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future.",
    },
]

# Store inferred memories (default behavior)
result = m.add(messages, user_id=user_id, metadata={"category": "movie_recommendations"})

all_results = m.get_all(user_id=user_id)
for n in all_results["results"]:
    print(f"node \"{n['memory']}\": [hash: {n['hash']}]")

for e in all_results["relations"]:
    print(f"edge \"{e['source']}\" --{e['relationship']}--> \"{e['target']}\"")

Search memories

In [ ]:

search_results = m.search("what does alice love?", user_id=user_id)
for result in search_results["results"]:
    print(f"\"{result['memory']}\" [score: {result['score']}]")
for relation in search_results["relations"]:
    print(f"{relation}")

In [ ]:

m.delete_all("user_id")
m.reset()