diff --git a/examples/persistence_mongodb.ipynb b/examples/persistence_mongodb.ipynb index 7c3560258..bec75dea0 100644 --- a/examples/persistence_mongodb.ipynb +++ b/examples/persistence_mongodb.ipynb @@ -278,7 +278,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 4, "metadata": {}, "outputs": [], "source": [ @@ -574,6 +574,410 @@ "\n", "#The checkpoints from both the examples have been saved in the database." ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Asynchronous implementation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#Async package for MongoDB\n", + "%pip install motor" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import pickle\n", + "from contextlib import AbstractContextManager\n", + "from types import TracebackType\n", + "from typing import Any, Dict,Optional,AsyncIterator\n", + "\n", + "from langchain_core.runnables import RunnableConfig\n", + "from typing_extensions import Self\n", + "\n", + "from langgraph.checkpoint.base import (\n", + " BaseCheckpointSaver,\n", + " Checkpoint,\n", + " CheckpointMetadata,\n", + " CheckpointTuple,\n", + " SerializerProtocol,\n", + ")\n", + "from langgraph.serde.jsonplus import JsonPlusSerializer\n", + "from motor.motor_asyncio import AsyncIOMotorClient\n", + "\n", + "class JsonPlusSerializerCompat(JsonPlusSerializer):\n", + " \"\"\"A serializer that supports loading pickled checkpoints for backwards compatibility.\n", + "\n", + " This serializer extends the JsonPlusSerializer and adds support for loading pickled\n", + " checkpoints. If the input data starts with b\"\\x80\" and ends with b\".\", it is treated\n", + " as a pickled checkpoint and loaded using pickle.loads(). Otherwise, the default\n", + " JsonPlusSerializer behavior is used.\n", + "\n", + " Examples:\n", + " >>> import pickle\n", + " >>> from langgraph.checkpoint.sqlite import JsonPlusSerializerCompat\n", + " >>>\n", + " >>> serializer = JsonPlusSerializerCompat()\n", + " >>> pickled_data = pickle.dumps({\"key\": \"value\"})\n", + " >>> loaded_data = serializer.loads(pickled_data)\n", + " >>> print(loaded_data) # Output: {\"key\": \"value\"}\n", + " >>>\n", + " >>> json_data = '{\"key\": \"value\"}'.encode(\"utf-8\")\n", + " >>> loaded_data = serializer.loads(json_data)\n", + " >>> print(loaded_data) # Output: {\"key\": \"value\"}\n", + " \"\"\"\n", + "\n", + " def loads(self, data: bytes) -> Any:\n", + " if data.startswith(b\"\\x80\") and data.endswith(b\".\"):\n", + " return pickle.loads(data)\n", + " return super().loads(data)\n", + "\n", + "class MongoDBSaver(AbstractContextManager, BaseCheckpointSaver):\n", + " \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database.\n", + "\n", + " Args:\n", + " client (AsyncIOMotorClient): The Async MongoDB client.\n", + " db_name (str): The name of the database to use.\n", + " collection_name (str): The name of the collection to use.\n", + " serde (Optional[SerializerProtocol]): The serializer to use for serializing and deserializing checkpoints. Defaults to JsonPlusSerializerCompat.\n", + "\n", + " Examples:\n", + "\n", + " >>> from motor.motor_asyncio import AsyncIOMotorClient\n", + " >>> from langgraph.checkpoint.mongodb import MongoDBSaver\n", + " >>> from langgraph.graph import StateGraph\n", + " >>>\n", + " >>> builder = StateGraph(int)\n", + " >>> builder.add_node(\"add_one\", lambda x: x + 1)\n", + " >>> builder.set_entry_point(\"add_one\")\n", + " >>> builder.set_finish_point(\"add_one\")\n", + " >>> client = AsyncIOMotorClient(\"mongodb://localhost:27017/\")\n", + " >>> memory = MongoDBSaver(client, \"checkpoints\", \"checkpoints\")\n", + " >>> graph = builder.compile(checkpointer=memory)\n", + " >>> config = {\"configurable\": {\"thread_id\": \"1\"}}\n", + " >>> result = graph.ainvoke(3, config)\n", + " \"\"\"\n", + "\n", + " serde = JsonPlusSerializerCompat()\n", + "\n", + " client: AsyncIOMotorClient\n", + " db_name: str\n", + " collection_name: str\n", + "\n", + " def __init__(\n", + " self,\n", + " client: AsyncIOMotorClient,\n", + " db_name: str,\n", + " collection_name: str,\n", + " *,\n", + " serde: Optional[SerializerProtocol] = None,\n", + " ) -> None:\n", + " super().__init__(serde=serde)\n", + " self.client = client\n", + " self.db_name = db_name\n", + " self.collection_name = collection_name\n", + " self.collection = client[db_name][collection_name]\n", + "\n", + " def __enter__(self) -> Self:\n", + " return self\n", + "\n", + " def __exit__(\n", + " self,\n", + " __exc_type: Optional[type[BaseException]],\n", + " __exc_value: Optional[BaseException],\n", + " __traceback: Optional[TracebackType],\n", + " ) -> Optional[bool]:\n", + " return True\n", + "\n", + " async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n", + " \"\"\"Get a checkpoint tuple from the database.\n", + "\n", + " This method retrieves a checkpoint tuple from the MongoDB database based on the\n", + " provided config. If the config contains a \"thread_ts\" key, the checkpoint with\n", + " the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint\n", + " for the given thread ID is retrieved.\n", + "\n", + " Args:\n", + " config (RunnableConfig): The config to use for retrieving the checkpoint.\n", + "\n", + " Returns:\n", + " Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.\n", + " \"\"\"\n", + " if config[\"configurable\"].get(\"thread_ts\"):\n", + " query = {\n", + " \"thread_id\": config[\"configurable\"][\"thread_id\"],\n", + " \"thread_ts\": config[\"configurable\"][\"thread_ts\"],\n", + " }\n", + " else:\n", + " query = {\"thread_id\": config[\"configurable\"][\"thread_id\"]}\n", + " result = self.collection.find(query).sort(\"thread_ts\", -1).limit(1)\n", + " async for doc in result:\n", + " return CheckpointTuple(\n", + " config,\n", + " self.serde.loads(doc[\"checkpoint\"]),\n", + " self.serde.loads(doc[\"metadata\"]),\n", + " (\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": doc[\"thread_id\"],\n", + " \"thread_ts\": doc[\"parent_ts\"],\n", + " }\n", + " }\n", + " if doc.get(\"parent_ts\")\n", + " else None\n", + " ),\n", + " )\n", + "\n", + " async def alist(\n", + " self,\n", + " config: Optional[RunnableConfig],\n", + " *,\n", + " filter: Optional[Dict[str, Any]] = None,\n", + " before: Optional[RunnableConfig] = None,\n", + " limit: Optional[int] = None,\n", + " ) -> AsyncIterator[CheckpointTuple]:\n", + " \"\"\"List checkpoints from the database.\n", + "\n", + " This method retrieves a list of checkpoint tuples from the MongoDB database based\n", + " on the provided config. The checkpoints are ordered by timestamp in descending order.\n", + "\n", + " Args:\n", + " config (RunnableConfig): The config to use for listing the checkpoints.\n", + " before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.\n", + " limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.\n", + "\n", + " Yields:\n", + " AsyncIterator[CheckpointTuple]: An Async iterator of checkpoint tuples.\n", + " \"\"\"\n", + " query = {}\n", + " if config is not None:\n", + " query[\"thread_id\"] = config[\"configurable\"][\"thread_id\"]\n", + " if filter:\n", + " for key, value in filter.items():\n", + " query[f\"metadata.{key}\"] = value\n", + " if before is not None:\n", + " query[\"thread_ts\"] = {\"$lt\": before[\"configurable\"][\"thread_ts\"]}\n", + " result = self.collection.find(query).sort(\"thread_ts\", -1).limit(limit)\n", + " if limit is not None:\n", + " result = result.limit(limit)\n", + " async for doc in result:\n", + " yield CheckpointTuple(\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": doc[\"thread_id\"],\n", + " \"thread_ts\": doc[\"thread_ts\"],\n", + " }\n", + " },\n", + " self.serde.loads(doc[\"checkpoint\"]),\n", + " self.serde.loads(doc[\"metadata\"]),\n", + " (\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": doc[\"thread_id\"],\n", + " \"thread_ts\": doc[\"parent_ts\"],\n", + " }\n", + " }\n", + " if doc.get(\"parent_ts\")\n", + " else None\n", + " ),\n", + " )\n", + "\n", + " async def aput(\n", + " self,\n", + " config: RunnableConfig,\n", + " checkpoint: Checkpoint,\n", + " metadata: CheckpointMetadata,\n", + " ) -> RunnableConfig:\n", + " \"\"\"Save a checkpoint to the database.\n", + "\n", + " This method saves a checkpoint to the MongoDB database. The checkpoint is associated\n", + " with the provided config and its parent config (if any).\n", + "\n", + " Args:\n", + " config (RunnableConfig): The config to associate with the checkpoint.\n", + " checkpoint (Checkpoint): The checkpoint to save.\n", + " metadata (Optional[dict[str, Any]]): Additional metadata to save with the checkpoint. Defaults to None.\n", + "\n", + " Returns:\n", + " RunnableConfig: The updated config containing the saved checkpoint's timestamp.\n", + " \"\"\"\n", + " doc = {\n", + " \"thread_id\": config[\"configurable\"][\"thread_id\"],\n", + " \"thread_ts\": checkpoint[\"id\"],\n", + " \"checkpoint\": self.serde.dumps(checkpoint),\n", + " \"metadata\": self.serde.dumps(metadata),\n", + " }\n", + " if config[\"configurable\"].get(\"thread_ts\"):\n", + " doc[\"parent_ts\"] = config[\"configurable\"][\"thread_ts\"]\n", + " await self.collection.insert_one(doc)\n", + " return {\n", + " \"configurable\": {\n", + " \"thread_id\": config[\"configurable\"][\"thread_id\"],\n", + " \"thread_ts\": checkpoint[\"id\"],\n", + " }\n", + " }" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example with basic graph" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.graph import StateGraph\n", + "checkpointer = MongoDBSaver(AsyncIOMotorClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\")\n", + "builder = StateGraph(int)\n", + "builder.add_node(\"add_one\", lambda x: x + 1)\n", + "builder.set_entry_point(\"add_one\")\n", + "builder.set_finish_point(\"add_one\")\n", + "graph = builder.compile(checkpointer=checkpointer)\n", + "config = {\"configurable\": {\"thread_id\": \"123\"}}\n", + "res = await graph.ainvoke(3,config)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "4" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "res" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'v': 1,\n", + " 'ts': '2024-07-10T11:34:28.485660+00:00',\n", + " 'id': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb',\n", + " 'channel_values': {'__root__': 4, 'add_one': 'add_one'},\n", + " 'channel_versions': {'__start__': 5,\n", + " '__root__': 6,\n", + " 'start:add_one': 6,\n", + " 'add_one': 6},\n", + " 'versions_seen': {'__start__': {'__start__': 4},\n", + " 'add_one': {'start:add_one': 5}},\n", + " 'pending_sends': []}" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "await checkpointer.aget(config)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "CheckpointTuple(config={'configurable': {'thread_id': '123'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.485660+00:00', 'id': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb', 'channel_values': {'__root__': 4, 'add_one': 'add_one'}, 'channel_versions': {'__start__': 5, '__root__': 6, 'start:add_one': 6, 'add_one': 6}, 'versions_seen': {'__start__': {'__start__': 4}, 'add_one': {'start:add_one': 5}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 4, 'writes': {'add_one': 4}}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}})" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "await checkpointer.aget_tuple(config)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.485660+00:00', 'id': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb', 'channel_values': {'__root__': 4, 'add_one': 'add_one'}, 'channel_versions': {'__start__': 5, '__root__': 6, 'start:add_one': 6, 'add_one': 6}, 'versions_seen': {'__start__': {'__start__': 4}, 'add_one': {'start:add_one': 5}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 4, 'writes': {'add_one': 4}}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}})\n", + "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.477660+00:00', 'id': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1', 'channel_values': {'__root__': 3, 'start:add_one': '__start__'}, 'channel_versions': {'__start__': 5, '__root__': 5, 'start:add_one': 5, 'add_one': 4}, 'versions_seen': {'__start__': {'__start__': 4}, 'add_one': {'start:add_one': 2}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': None}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141'}})\n", + "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.476662+00:00', 'id': '1ef3eb05-e0bb-659e-8002-de83b4764141', 'channel_values': {'__root__': 4, '__start__': 3}, 'channel_versions': {'__start__': 4, '__root__': 3, 'start:add_one': 3, 'add_one': 4}, 'versions_seen': {'__start__': {'__start__': 1}, 'add_one': {'start:add_one': 2}}, 'pending_sends': []}, metadata={'source': 'input', 'step': 2, 'writes': 3}, parent_config=None)\n" + ] + } + ], + "source": [ + "list = checkpointer.alist(config, limit=3)\n", + "async for item in list:\n", + " print(item)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Checkpoints saved in MongoDB" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'_id': ObjectId('668e57930f55bbe62f358531'), 'thread_id': '123', 'thread_ts': '1ef3ea0c-18a5-67a6-bfff-0d85b77e4a09', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T09:42:43.453328+00:00\", \"id\": \"1ef3ea0c-18a5-67a6-bfff-0d85b77e4a09\", \"channel_values\": {\"__start__\": 3}, \"channel_versions\": {\"__start__\": 1}, \"versions_seen\": {}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"input\", \"step\": -1, \"writes\": 3}'}\n", + "{'_id': ObjectId('668e57930f55bbe62f358532'), 'thread_id': '123', 'thread_ts': '1ef3ea0c-18a7-6ea3-8000-9a52ba553d0c', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T09:42:43.454326+00:00\", \"id\": \"1ef3ea0c-18a7-6ea3-8000-9a52ba553d0c\", \"channel_values\": {\"__root__\": 3, \"start:add_one\": \"__start__\"}, \"channel_versions\": {\"__start__\": 2, \"__root__\": 2, \"start:add_one\": 2}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"add_one\": {}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 0, \"writes\": null}', 'parent_ts': '1ef3ea0c-18a5-67a6-bfff-0d85b77e4a09'}\n", + "{'_id': ObjectId('668e57930f55bbe62f358533'), 'thread_id': '123', 'thread_ts': '1ef3ea0c-18bc-6b54-8001-ef5781939492', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T09:42:43.462843+00:00\", \"id\": \"1ef3ea0c-18bc-6b54-8001-ef5781939492\", \"channel_values\": {\"__root__\": 4, \"add_one\": \"add_one\"}, \"channel_versions\": {\"__start__\": 2, \"__root__\": 3, \"start:add_one\": 3, \"add_one\": 3}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"add_one\": {\"start:add_one\": 2}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 1, \"writes\": {\"add_one\": 4}}', 'parent_ts': '1ef3ea0c-18a7-6ea3-8000-9a52ba553d0c'}\n", + "{'_id': ObjectId('668e71c4171972a41a226373'), 'thread_id': '123', 'thread_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T11:34:28.476662+00:00\", \"id\": \"1ef3eb05-e0bb-659e-8002-de83b4764141\", \"channel_values\": {\"__root__\": 4, \"__start__\": 3}, \"channel_versions\": {\"__start__\": 4, \"__root__\": 3, \"start:add_one\": 3, \"add_one\": 4}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"add_one\": {\"start:add_one\": 2}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"input\", \"step\": 2, \"writes\": 3}'}\n", + "{'_id': ObjectId('668e71c4171972a41a226374'), 'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T11:34:28.477660+00:00\", \"id\": \"1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1\", \"channel_values\": {\"__root__\": 3, \"start:add_one\": \"__start__\"}, \"channel_versions\": {\"__start__\": 5, \"__root__\": 5, \"start:add_one\": 5, \"add_one\": 4}, \"versions_seen\": {\"__start__\": {\"__start__\": 4}, \"add_one\": {\"start:add_one\": 2}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 3, \"writes\": null}', 'parent_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141'}\n", + "{'_id': ObjectId('668e71c4171972a41a226375'), 'thread_id': '123', 'thread_ts': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T11:34:28.485660+00:00\", \"id\": \"1ef3eb05-e0d1-651b-8004-15f129f5f4fb\", \"channel_values\": {\"__root__\": 4, \"add_one\": \"add_one\"}, \"channel_versions\": {\"__start__\": 5, \"__root__\": 6, \"start:add_one\": 6, \"add_one\": 6}, \"versions_seen\": {\"__start__\": {\"__start__\": 4}, \"add_one\": {\"start:add_one\": 5}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 4, \"writes\": {\"add_one\": 4}}', 'parent_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}\n" + ] + } + ], + "source": [ + "from pymongo import MongoClient\n", + "client = MongoClient(MONGO_URI)\n", + "database = client[\"checkpoints_db\"]\n", + "collection = database[\"checkpoints_collection\"]\n", + "\n", + "for doc in collection.find():\n", + " print(doc)" + ] } ], "metadata": {