{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neptune as Graph Memory\n",
    "\n",
    "In this notebook, we will be connecting using a Amazon Neptune Analytics instance as our memory graph storage for Mem0.\n",
    "\n",
    "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.\n",
    "\n",
    "Reference: [Vector Similarity using Neptune Analytics](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/vector-similarity.html)"
   ]
  },
  {
   "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,extras]\"\n",
    "```\n",
    "\n",
    "This command installs Mem0 along with the necessary dependencies for graph functionality (`graph`) and other Amazon dependencies (`extras`).\n",
    "\n",
    "### 2. Connect to Amazon services\n",
    "\n",
    "For this sample notebook, configure `mem0ai` with [Amazon Neptune Analytics](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html) as the vector and graph store, and [Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) for generating embeddings.\n",
    "\n",
    "Use the following guide for setup details: [Setup AWS Bedrock, AOSS, and Neptune](https://docs.mem0.ai/examples/aws_example#aws-bedrock-and-aoss)\n",
    "\n",
    "The Neptune Analytics instance must be created using the same vector dimensions as the embedding model creates. See: https://docs.aws.amazon.com/neptune-analytics/latest/userguide/vector-index.html\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",
    "            \"embedding_dims\": 1024\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\": \"neptune\",\n",
    "        \"config\": {\n",
    "            \"endpoint\": f\"neptune-graph://my-graph-identifier\",\n",
    "        },\n",
    "    },\n",
    "    \"graph_store\": {\n",
    "        \"provider\": \"neptune\",\n",
    "        \"config\": {\n",
    "            \"endpoint\": f\"neptune-graph://my-graph-identifier\",\n",
    "        },\n",
    "    },\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Import all packages and setup logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mem0 import Memory\n",
    "import os\n",
    "import logging\n",
    "import sys\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "logging.getLogger(\"mem0.graphs.neptune.main\").setLevel(logging.INFO)\n",
    "logging.getLogger(\"mem0.graphs.neptune.base\").setLevel(logging.INFO)\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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup the Mem0 configuration using:\n",
    "- Amazon Bedrock as the embedder\n",
    "- Amazon Neptune Analytics instance as a vector / graph store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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",
    "graph_identifier = os.environ.get(\"GRAPH_ID\")\n",
    "\n",
    "config = {\n",
    "    \"embedder\": {\n",
    "        \"provider\": \"aws_bedrock\",\n",
    "        \"config\": {\n",
    "            \"model\": bedrock_embedder_model,\n",
    "            \"embedding_dims\": embedding_model_dims\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\": \"neptune\",\n",
    "        \"config\": {\n",
    "            \"endpoint\": f\"neptune-graph://{graph_identifier}\",\n",
    "        },\n",
    "    },\n",
    "    \"graph_store\": {\n",
    "        \"provider\": \"neptune\",\n",
    "        \"config\": {\n",
    "            \"endpoint\": f\"neptune-graph://{graph_identifier}\",\n",
    "        },\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Memory initializiation\n",
    "\n",
    "Initialize Memgraph as a Graph Memory store:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Store memories\n",
    "\n",
    "Create memories and store one at a time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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']}\\\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Explorer Visualization\n",
    "\n",
    "You can visualize the graph using a Graph Explorer connection to Neptune Analytics 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 analytics 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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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']}\\\"\")"
   ]
  },
  {
   "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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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']}\\\"\")"
   ]
  },
  {
   "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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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']}\\\"\")"
   ]
  },
  {
   "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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.delete_all(user_id)\n",
    "m.reset()"
   ]
  },
  {
   "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. Neptune Analytics can store the text chunks in a graph format with relationship entities."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
