{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neptune as Graph Memory\n",
    "\n",
    "In this notebook, we will be connecting using an Amazon Neptune DC Cluster instance as our memory graph storage for Mem0. Unlike other graph stores, Neptune DB doesn't store vectors itself. To detect vector similary in nodes, we store the node vectors in our defined vector store, and use vector search to retrieve similar nodes.\n",
    "\n",
    "For this reason, a vector store is required to configure neptune-db.\n",
    "\n",
    "The Graph Memory storage persists memories in a graph or relationship form when performing `m.add` memory operations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "### 1. Install Mem0 with Graph Memory support \n",
    "\n",
    "To use Mem0 with Graph Memory support (as well as other Amazon services), use pip install:\n",
    "\n",
    "```bash\n",
    "pip install \"mem0ai[graph,vector_stores,extras]\"\n",
    "```\n",
    "\n",
    "This command installs Mem0 along with the necessary dependencies for graph functionality (`graph`), vector stores, and other Amazon dependencies (`extras`).\n",
    "\n",
    "### 2. Connect to Amazon services\n",
    "\n",
    "For this sample notebook, configure `mem0ai` with [Amazon Neptune Database Cluster](https://docs.aws.amazon.com/neptune/latest/userguide/intro.html) as the graph store, [Amazon OpenSearch Serverless](https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-overview.html) as the vector store, and [Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) for generating embeddings.\n",
    "\n",
    "Your configuration should look similar to:\n",
    "\n",
    "```python\n",
    "config = {\n",
    "    \"embedder\": {\n",
    "        \"provider\": \"aws_bedrock\",\n",
    "        \"config\": {\n",
    "            \"model\": \"amazon.titan-embed-text-v2:0\"\n",
    "        }\n",
    "    },\n",
    "    \"llm\": {\n",
    "        \"provider\": \"aws_bedrock\",\n",
    "        \"config\": {\n",
    "            \"model\": \"us.anthropic.claude-3-7-sonnet-20250219-v1:0\",\n",
    "            \"temperature\": 0.1,\n",
    "            \"max_tokens\": 2000\n",
    "        }\n",
    "    },\n",
    "    \"vector_store\": {\n",
    "        \"provider\": \"opensearch\",\n",
    "        \"config\": {\n",
    "            \"collection_name\": \"mem0\",\n",
    "            \"host\": \"your-opensearch-domain.us-west-2.es.amazonaws.com\",\n",
    "            \"port\": 443,\n",
    "            \"http_auth\": auth,\n",
    "            \"connection_class\": RequestsHttpConnection,\n",
    "            \"pool_maxsize\": 20,\n",
    "            \"use_ssl\": True,\n",
    "            \"verify_certs\": True,\n",
    "            \"embedding_model_dims\": 1024,\n",
    "        }\n",
    "    },\n",
    "    \"graph_store\": {\n",
    "        \"provider\": \"neptunedb\",\n",
    "        \"config\": {\n",
    "            \"\": \"\",\n",
    "            \"endpoint\": f\"neptune-db://my-graph-host\",\n",
    "        },\n",
    "    },\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Import all packages and setup logging"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from mem0 import Memory\n",
    "import os\n",
    "import logging\n",
    "import sys\n",
    "import boto3\n",
    "from opensearchpy import RequestsHttpConnection, AWSV4SignerAuth\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "logging.getLogger(\"mem0.graphs.neptune.neptunedb\").setLevel(logging.DEBUG)\n",
    "logging.getLogger(\"mem0.graphs.neptune.base\").setLevel(logging.DEBUG)\n",
    "logger = logging.getLogger(__name__)\n",
    "logger.setLevel(logging.DEBUG)\n",
    "\n",
    "logging.basicConfig(\n",
    "    format=\"%(levelname)s - %(message)s\",\n",
    "    datefmt=\"%Y-%m-%d %H:%M:%S\",\n",
    "    stream=sys.stdout,  # Explicitly set output to stdout\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup the Mem0 configuration using:\n",
    "- Amazon Bedrock as the LLM and embedder\n",
    "- Amazon Neptune DB instance as a graph store with node vectors in OpenSearch (collection: `mem0ai_neptune_entities`)\n",
    "- OpenSearch as the text summaries vector store (collection: `mem0ai_text_summaries`)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "bedrock_embedder_model = \"amazon.titan-embed-text-v2:0\"\n",
    "bedrock_llm_model = \"us.anthropic.claude-3-7-sonnet-20250219-v1:0\"\n",
    "embedding_model_dims = 1024\n",
    "\n",
    "neptune_host = os.environ.get(\"GRAPH_HOST\")\n",
    "\n",
    "opensearch_host = os.environ.get(\"OS_HOST\")\n",
    "opensearch_port = 443\n",
    "\n",
    "credentials = boto3.Session().get_credentials()\n",
    "region = os.environ.get(\"AWS_REGION\")\n",
    "auth = AWSV4SignerAuth(credentials, region)\n",
    "\n",
    "config = {\n",
    "    \"embedder\": {\n",
    "        \"provider\": \"aws_bedrock\",\n",
    "        \"config\": {\n",
    "            \"model\": bedrock_embedder_model,\n",
    "        }\n",
    "    },\n",
    "    \"llm\": {\n",
    "        \"provider\": \"aws_bedrock\",\n",
    "        \"config\": {\n",
    "            \"model\": bedrock_llm_model,\n",
    "            \"temperature\": 0.1,\n",
    "            \"max_tokens\": 2000\n",
    "        }\n",
    "    },\n",
    "    \"vector_store\": {\n",
    "        \"provider\": \"opensearch\",\n",
    "        \"config\": {\n",
    "            \"collection_name\": \"mem0ai_text_summaries\",\n",
    "            \"host\": opensearch_host,\n",
    "            \"port\": opensearch_port,\n",
    "            \"http_auth\": auth,\n",
    "            \"embedding_model_dims\": embedding_model_dims,\n",
    "            \"use_ssl\": True,\n",
    "            \"verify_certs\": True,\n",
    "            \"connection_class\": RequestsHttpConnection,\n",
    "        },\n",
    "    },\n",
    "    \"graph_store\": {\n",
    "        \"provider\": \"neptunedb\",\n",
    "        \"config\": {\n",
    "            \"collection_name\": \"mem0ai_neptune_entities\",\n",
    "            \"endpoint\": f\"neptune-db://{neptune_host}\",\n",
    "        },\n",
    "    },\n",
    "}"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Memory initializiation\n",
    "\n",
    "Initialize Memgraph as a Graph Memory store:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "m = Memory.from_config(config_dict=config)\n",
    "\n",
    "app_id = \"movies\"\n",
    "user_id = \"alice\"\n",
    "\n",
    "m.delete_all(user_id=user_id)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Store memories\n",
    "\n",
    "Create memories and store one at a time:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"I'm planning to watch a movie tonight. Any recommendations?\",\n",
    "    },\n",
    "]\n",
    "\n",
    "# Store inferred memories (default behavior)\n",
    "result = m.add(messages, user_id=user_id, metadata={\"category\": \"movie_recommendations\"})\n",
    "\n",
    "all_results = m.get_all(user_id=user_id)\n",
    "for n in all_results[\"results\"]:\n",
    "    print(f\"node \\\"{n['memory']}\\\": [hash: {n['hash']}]\")\n",
    "\n",
    "for e in all_results[\"relations\"]:\n",
    "    print(f\"edge \\\"{e['source']}\\\" --{e['relationship']}--> \\\"{e['target']}\\\"\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Explorer Visualization\n",
    "\n",
    "You can visualize the graph using a Graph Explorer connection to Neptune-DB in Neptune Notebooks in the Amazon console.  See [Using Amazon Neptune with graph notebooks](https://docs.aws.amazon.com/neptune/latest/userguide/graph-notebooks.html) for instructions on how to setup a Neptune Notebook with Graph Explorer.\n",
    "\n",
    "Once the graph has been generated, you can open the visualization in the Neptune > Notebooks and click on Actions > Open Graph Explorer.  This will automatically connect to your neptune db graph that was provided in the notebook setup.\n",
    "\n",
    "Once in Graph Explorer, visit Open Connections and send all the available nodes and edges to Explorer. Visit Open Graph Explorer to see the nodes and edges in the graph.\n",
    "\n",
    "### Graph Explorer Visualization Example\n",
    "\n",
    "_Note that the visualization given below represents only a single example of the possible results generated by the LLM._\n",
    "\n",
    "Visualization for the relationship:\n",
    "```\n",
    "\"alice\" --plans_to_watch--> \"movie\"\n",
    "```\n",
    "\n",
    "![neptune-example-visualization-1.png](./neptune-example-visualization-1.png)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": \"How about a thriller movies? They can be quite engaging.\",\n",
    "    },\n",
    "]\n",
    "\n",
    "# Store inferred memories (default behavior)\n",
    "result = m.add(messages, user_id=user_id, metadata={\"category\": \"movie_recommendations\"})\n",
    "\n",
    "all_results = m.get_all(user_id=user_id)\n",
    "for n in all_results[\"results\"]:\n",
    "    print(f\"node \\\"{n['memory']}\\\": [hash: {n['hash']}]\")\n",
    "\n",
    "for e in all_results[\"relations\"]:\n",
    "    print(f\"edge \\\"{e['source']}\\\" --{e['relationship']}--> \\\"{e['target']}\\\"\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph Explorer Visualization Example\n",
    "\n",
    "_Note that the visualization given below represents only a single example of the possible results generated by the LLM._\n",
    "\n",
    "Visualization for the relationship:\n",
    "```\n",
    "\"alice\" --plans_to_watch--> \"movie\"\n",
    "\"thriller\" --type_of--> \"movie\"\n",
    "\"movie\" --can_be--> \"engaging\"\n",
    "```\n",
    "\n",
    "![neptune-example-visualization-2.png](./neptune-example-visualization-2.png)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"I'm not a big fan of thriller movies but I love sci-fi movies.\",\n",
    "    },\n",
    "]\n",
    "\n",
    "# Store inferred memories (default behavior)\n",
    "result = m.add(messages, user_id=user_id, metadata={\"category\": \"movie_recommendations\"})\n",
    "\n",
    "all_results = m.get_all(user_id=user_id)\n",
    "for n in all_results[\"results\"]:\n",
    "    print(f\"node \\\"{n['memory']}\\\": [hash: {n['hash']}]\")\n",
    "\n",
    "for e in all_results[\"relations\"]:\n",
    "    print(f\"edge \\\"{e['source']}\\\" --{e['relationship']}--> \\\"{e['target']}\\\"\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph Explorer Visualization Example\n",
    "\n",
    "_Note that the visualization given below represents only a single example of the possible results generated by the LLM._\n",
    "\n",
    "Visualization for the relationship:\n",
    "```\n",
    "\"alice\" --dislikes--> \"thriller_movies\"\n",
    "\"alice\" --loves--> \"sci-fi_movies\"\n",
    "\"alice\" --plans_to_watch--> \"movie\"\n",
    "\"thriller\" --type_of--> \"movie\"\n",
    "\"movie\" --can_be--> \"engaging\"\n",
    "```\n",
    "\n",
    "![neptune-example-visualization-3.png](./neptune-example-visualization-3.png)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": \"Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future.\",\n",
    "    },\n",
    "]\n",
    "\n",
    "# Store inferred memories (default behavior)\n",
    "result = m.add(messages, user_id=user_id, metadata={\"category\": \"movie_recommendations\"})\n",
    "\n",
    "all_results = m.get_all(user_id=user_id)\n",
    "for n in all_results[\"results\"]:\n",
    "    print(f\"node \\\"{n['memory']}\\\": [hash: {n['hash']}]\")\n",
    "\n",
    "for e in all_results[\"relations\"]:\n",
    "    print(f\"edge \\\"{e['source']}\\\" --{e['relationship']}--> \\\"{e['target']}\\\"\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph Explorer Visualization Example\n",
    "\n",
    "_Note that the visualization given below represents only a single example of the possible results generated by the LLM._\n",
    "\n",
    "Visualization for the relationship:\n",
    "```\n",
    "\"alice\" --recommends--> \"sci-fi\"\n",
    "\"alice\" --dislikes--> \"thriller_movies\"\n",
    "\"alice\" --loves--> \"sci-fi_movies\"\n",
    "\"alice\" --plans_to_watch--> \"movie\"\n",
    "\"alice\" --avoids--> \"thriller\"\n",
    "\"thriller\" --type_of--> \"movie\"\n",
    "\"movie\" --can_be--> \"engaging\"\n",
    "\"sci-fi\" --type_of--> \"movie\"\n",
    "```\n",
    "\n",
    "![neptune-example-visualization-4.png](./neptune-example-visualization-4.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Search memories\n",
    "\n",
    "Search all memories for \"what does alice love?\".  Since \"alice\" the user, this will search for a relationship that fits the users love of \"sci-fi\" movies and dislike of \"thriller\" movies."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "search_results = m.search(\"what does alice love?\", user_id=user_id)\n",
    "for result in search_results[\"results\"]:\n",
    "    print(f\"\\\"{result['memory']}\\\" [score: {result['score']}]\")\n",
    "for relation in search_results[\"relations\"]:\n",
    "    print(f\"{relation}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "m.delete_all(user_id)\n",
    "m.reset()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this example we demonstrated how an AWS tech stack can be used to store and retrieve memory context. Bedrock LLM models can be used to interpret given conversations.  OpenSearch can store text chunks with vector embeddings. Neptune Database can store the text entities in a graph format with relationship entities."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
