diff --git a/examples/persistence_mongodb.ipynb b/examples/persistence_mongodb.ipynb index 0eec3667c..99ce73057 100644 --- a/examples/persistence_mongodb.ipynb +++ b/examples/persistence_mongodb.ipynb @@ -1,1013 +1,917 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# How to create a custom checkpointer using MongoDB\n", - "\n", - "When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions.\n", - "\n", - "This example shows how to use `MongoDB` as the backend for persisting checkpoint state.\n", - "\n", - "NOTE: this is just an example implementation. You can implement your own checkpointer using a different database or modify this one as long as it conforms to the `BaseCheckpointSaver` interface." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Checkpointer implementation" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%capture --no-stderr\n", - "%pip install -U langgraph pymongo" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [], - "source": [ - "import pickle\n", - "from contextlib import AbstractContextManager\n", - "from types import TracebackType\n", - "from typing import Any, Dict, Iterator, Optional\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 pymongo import MongoClient\n", - "\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", - "\n", - "class MongoDBSaver(AbstractContextManager, BaseCheckpointSaver):\n", - " \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database.\n", - "\n", - " Args:\n", - " client (pymongo.MongoClient): The 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 pymongo import MongoClient\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 = MongoClient(\"mongodb://localhost:27017/\")\n", - " >>> memory = MongoDBSaver(client, \"checkpoints\", \"checkpoints\")\n", - " >>> graph = builder.compile(checkpointer=memory)\n", - " >>> config = {\"configurable\": {\"thread_id\": \"1\"}}\n", - " >>> graph.get_state(config)\n", - " >>> result = graph.invoke(3, config)\n", - " >>> graph.get_state(config)\n", - " StateSnapshot(values=4, next=(), config={'configurable': {'thread_id': '1', 'thread_ts': '2024-05-04T06:32:42.235444+00:00'}}, parent_config=None)\n", - " \"\"\"\n", - "\n", - " serde = JsonPlusSerializerCompat()\n", - "\n", - " client: MongoClient\n", - " db_name: str\n", - " collection_name: str\n", - "\n", - " def __init__(\n", - " self,\n", - " client: MongoClient,\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", - " def get_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", - " 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", - " def list(\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", - " ) -> Iterator[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", - " Iterator[CheckpointTuple]: An 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", - " 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", - " def put(\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", - " 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": [ - "## MongoDB connection" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ - "MONGO_URI = \"mongodb://localhost:27017/\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Basic example using graph" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "StateSnapshot(values=4, next=(), config={'configurable': {'thread_id': '123'}}, metadata={'source': 'loop', 'step': 1, 'writes': {'add_one': 4}}, created_at='2024-07-09T15:56:06.885848+00:00', parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3e0bc-09c1-6c26-8000-b9e1d26417ff'}})" - ] - }, - "execution_count": 3, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from langgraph.graph import StateGraph, START, END\n", - "\n", - "checkpointer = MongoDBSaver(\n", - " MongoClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\"\n", - ")\n", - "builder = StateGraph(int)\n", - "builder.add_node(\"add_one\", lambda x: x + 1)\n", - "builder.add_edge(START, \"add_one\")\n", - "builder.add_edge(\"add_one\", END)\n", - "graph = builder.compile(checkpointer=checkpointer)\n", - "config = {\"configurable\": {\"thread_id\": \"123\"}}\n", - "graph.get_state(config)\n", - "result = graph.invoke(3, config)\n", - "graph.get_state(config)" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "4" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "result" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'v': 1,\n", - " 'ts': '2024-07-09T15:56:06.885848+00:00',\n", - " 'id': '1ef3e0bc-09d1-6a75-8001-8f750e9a0782',\n", - " 'channel_values': {'__root__': 4, 'add_one': 'add_one'},\n", - " 'channel_versions': {'__start__': 2,\n", - " '__root__': 3,\n", - " 'start:add_one': 3,\n", - " 'add_one': 3},\n", - " 'versions_seen': {'__start__': {'__start__': 1},\n", - " 'add_one': {'start:add_one': 2}},\n", - " 'pending_sends': []}" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "checkpointer.get(config)" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3e0bc-09d1-6a75-8001-8f750e9a0782'}}, checkpoint={'v': 1, 'ts': '2024-07-09T15:56:06.885848+00:00', 'id': '1ef3e0bc-09d1-6a75-8001-8f750e9a0782', '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={'source': 'loop', 'step': 1, 'writes': {'add_one': 4}}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3e0bc-09c1-6c26-8000-b9e1d26417ff'}})\n", - "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3e0bc-09c1-6c26-8000-b9e1d26417ff'}}, checkpoint={'v': 1, 'ts': '2024-07-09T15:56:06.878338+00:00', 'id': '1ef3e0bc-09c1-6c26-8000-b9e1d26417ff', '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={'source': 'loop', 'step': 0, 'writes': None}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3e0bc-09bc-6e04-bfff-5342ac1ccc10'}})\n", - "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3e0bc-09bc-6e04-bfff-5342ac1ccc10'}}, checkpoint={'v': 1, 'ts': '2024-07-09T15:56:06.877337+00:00', 'id': '1ef3e0bc-09bc-6e04-bfff-5342ac1ccc10', 'channel_values': {'__start__': 3}, 'channel_versions': {'__start__': 1}, 'versions_seen': {}, 'pending_sends': []}, metadata={'source': 'input', 'step': -1, 'writes': 3}, parent_config=None)\n" - ] - } - ], - "source": [ - "list = checkpointer.list(config, limit=3)\n", - "for item in list:\n", - " print(item)" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "CheckpointTuple(config={'configurable': {'thread_id': '123'}}, checkpoint={'v': 1, 'ts': '2024-07-09T13:22:19.610402+00:00', 'id': '1ef3df64-4ba9-6b58-8001-ab084cc01a30', '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={'source': 'loop', 'step': 1, 'writes': {'add_one': 4}}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3df64-4ba2-660c-8000-569999697ff3'}})" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "checkpointer.get_tuple(config)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Setup environment" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "metadata": {}, - "outputs": [], - "source": [ - "import getpass\n", - "import os\n", - "\n", - "\n", - "def _set_env(var: str):\n", - " if not os.environ.get(var):\n", - " os.environ[var] = getpass.getpass(f\"{var}: \")\n", - "\n", - "\n", - "_set_env(\"OPENAI_API_KEY\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Setup model and tools for the graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%pip install langchain_openai" - ] - }, - { - "cell_type": "code", - "execution_count": 15, - "metadata": {}, - "outputs": [], - "source": [ - "from typing import Literal\n", - "from langchain_core.runnables import ConfigurableField\n", - "from langchain_core.tools import tool\n", - "from langchain_openai import ChatOpenAI\n", - "from langgraph.prebuilt import create_react_agent\n", - "\n", - "\n", - "@tool\n", - "def get_weather(city: Literal[\"nyc\", \"sf\"]):\n", - " \"\"\"Use this to get weather information.\"\"\"\n", - " if city == \"nyc\":\n", - " return \"It might be cloudy in nyc\"\n", - " elif city == \"sf\":\n", - " return \"It's always sunny in sf\"\n", - " else:\n", - " raise AssertionError(\"Unknown city\")\n", - "\n", - "\n", - "tools = [get_weather]\n", - "model = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)" - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "metadata": {}, - "outputs": [], - "source": [ - "graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", - "config = {\"configurable\": {\"thread_id\": \"1\"}}\n", - "res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)" - ] - }, - { - "cell_type": "code", - "execution_count": 17, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'messages': [HumanMessage(content=\"what's the weather in sf\", id='a624d383-13c6-499c-8f03-31ed11fa0cfb'),\n", - " AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_wapm4s91KQUQqE9y1L53QmmE', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 58, 'total_tokens': 72}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-614bc54a-ad37-4f17-9047-b80752bdf66e-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_wapm4s91KQUQqE9y1L53QmmE'}]),\n", - " ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='e58dd97d-b50d-4b0a-9492-7155106c975a', tool_call_id='call_wapm4s91KQUQqE9y1L53QmmE'),\n", - " AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 86, 'total_tokens': 96}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-e9984eca-f132-46d0-94ba-41d8ba5b7046-0')]}" - ] - }, - "execution_count": 17, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "res" - ] - }, - { - "cell_type": "code", - "execution_count": 18, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "CheckpointTuple(config={'configurable': {'thread_id': '1'}}, checkpoint={'v': 1, 'ts': '2024-07-09T13:22:49.794047+00:00', 'id': '1ef3df65-6b84-63fd-8003-888bcef289e3', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='a624d383-13c6-499c-8f03-31ed11fa0cfb'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_wapm4s91KQUQqE9y1L53QmmE', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 58, 'total_tokens': 72}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-614bc54a-ad37-4f17-9047-b80752bdf66e-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_wapm4s91KQUQqE9y1L53QmmE'}]), ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='e58dd97d-b50d-4b0a-9492-7155106c975a', tool_call_id='call_wapm4s91KQUQqE9y1L53QmmE'), AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 86, 'total_tokens': 96}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-e9984eca-f132-46d0-94ba-41d8ba5b7046-0')], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 86, 'total_tokens': 96}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-e9984eca-f132-46d0-94ba-41d8ba5b7046-0')]}}}, parent_config={'configurable': {'thread_id': '1', 'thread_ts': '1ef3df65-6252-6b30-8002-2c9e1e68364b'}})" - ] - }, - "execution_count": 18, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "checkpointer.get_tuple(config)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Checkpoints saved in MongoDB" - ] - }, - { - "cell_type": "code", - "execution_count": 19, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{'_id': ObjectId('668d398bb975d3e766de42ce'), 'thread_id': '123', 'thread_ts': '1ef3df64-4b98-68b4-bfff-592f97570cf6', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:19.603371+00:00\", \"id\": \"1ef3df64-4b98-68b4-bfff-592f97570cf6\", \"channel_values\": {\"__start__\": 3}, \"channel_versions\": {\"__start__\": 1}, \"versions_seen\": {}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"input\", \"step\": -1, \"writes\": 3}'}\n", - "{'_id': ObjectId('668d398bb975d3e766de42cf'), 'thread_id': '123', 'thread_ts': '1ef3df64-4ba2-660c-8000-569999697ff3', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:19.607399+00:00\", \"id\": \"1ef3df64-4ba2-660c-8000-569999697ff3\", \"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': '1ef3df64-4b98-68b4-bfff-592f97570cf6'}\n", - "{'_id': ObjectId('668d398bb975d3e766de42d0'), 'thread_id': '123', 'thread_ts': '1ef3df64-4ba9-6b58-8001-ab084cc01a30', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:19.610402+00:00\", \"id\": \"1ef3df64-4ba9-6b58-8001-ab084cc01a30\", \"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': '1ef3df64-4ba2-660c-8000-569999697ff3'}\n", - "{'_id': ObjectId('668d39a7b975d3e766de42d1'), 'thread_id': '1', 'thread_ts': '1ef3df65-585c-6abb-bfff-634fac15b6b3', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:47.785541+00:00\", \"id\": \"1ef3df65-585c-6abb-bfff-634fac15b6b3\", \"channel_values\": {\"messages\": [], \"__start__\": {\"messages\": [[\"human\", \"what\\'s the weather in sf\"]]}}, \"channel_versions\": {\"__start__\": 1}, \"versions_seen\": {}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"input\", \"step\": -1, \"writes\": {\"messages\": [[\"human\", \"what\\'s the weather in sf\"]]}}'}\n", - "{'_id': ObjectId('668d39a7b975d3e766de42d2'), 'thread_id': '1', 'thread_ts': '1ef3df65-5863-6fbc-8000-d45480f1ccc0', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:47.788537+00:00\", \"id\": \"1ef3df65-5863-6fbc-8000-d45480f1ccc0\", \"channel_values\": {\"messages\": [{\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"HumanMessage\"], \"kwargs\": {\"content\": \"what\\'s the weather in sf\", \"type\": \"human\", \"id\": \"a624d383-13c6-499c-8f03-31ed11fa0cfb\"}}], \"start:agent\": \"__start__\"}, \"channel_versions\": {\"__start__\": 2, \"messages\": 2, \"start:agent\": 2}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"agent\": {}, \"tools\": {}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 0, \"writes\": null}', 'parent_ts': '1ef3df65-585c-6abb-bfff-634fac15b6b3'}\n", - "{'_id': ObjectId('668d39a8b975d3e766de42d3'), 'thread_id': '1', 'thread_ts': '1ef3df65-6248-6ee8-8001-d5eef3fad087', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:48.826032+00:00\", \"id\": \"1ef3df65-6248-6ee8-8001-d5eef3fad087\", \"channel_values\": {\"messages\": [{\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"HumanMessage\"], \"kwargs\": {\"content\": \"what\\'s the weather in sf\", \"type\": \"human\", \"id\": \"a624d383-13c6-499c-8f03-31ed11fa0cfb\"}}, {\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"AIMessage\"], \"kwargs\": {\"content\": \"\", \"additional_kwargs\": {\"tool_calls\": [{\"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\", \"function\": {\"arguments\": \"{\\\\\"city\\\\\":\\\\\"sf\\\\\"}\", \"name\": \"get_weather\"}, \"type\": \"function\"}]}, \"response_metadata\": {\"token_usage\": {\"completion_tokens\": 14, \"prompt_tokens\": 58, \"total_tokens\": 72}, \"model_name\": \"gpt-3.5-turbo\", \"system_fingerprint\": null, \"finish_reason\": \"tool_calls\", \"logprobs\": null}, \"type\": \"ai\", \"id\": \"run-614bc54a-ad37-4f17-9047-b80752bdf66e-0\", \"tool_calls\": [{\"name\": \"get_weather\", \"args\": {\"city\": \"sf\"}, \"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\"}], \"invalid_tool_calls\": []}}], \"agent\": \"agent\", \"branch:agent:should_continue:tools\": \"agent\"}, \"channel_versions\": {\"__start__\": 2, \"messages\": 3, \"start:agent\": 3, \"agent\": 3, \"branch:agent:should_continue:tools\": 3}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"agent\": {\"start:agent\": 2}, \"tools\": {}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 1, \"writes\": {\"agent\": {\"messages\": [{\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"AIMessage\"], \"kwargs\": {\"content\": \"\", \"additional_kwargs\": {\"tool_calls\": [{\"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\", \"function\": {\"arguments\": \"{\\\\\"city\\\\\":\\\\\"sf\\\\\"}\", \"name\": \"get_weather\"}, \"type\": \"function\"}]}, \"response_metadata\": {\"token_usage\": {\"completion_tokens\": 14, \"prompt_tokens\": 58, \"total_tokens\": 72}, \"model_name\": \"gpt-3.5-turbo\", \"system_fingerprint\": null, \"finish_reason\": \"tool_calls\", \"logprobs\": null}, \"type\": \"ai\", \"id\": \"run-614bc54a-ad37-4f17-9047-b80752bdf66e-0\", \"tool_calls\": [{\"name\": \"get_weather\", \"args\": {\"city\": \"sf\"}, \"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\"}], \"invalid_tool_calls\": []}}]}}}', 'parent_ts': '1ef3df65-5863-6fbc-8000-d45480f1ccc0'}\n", - "{'_id': ObjectId('668d39a8b975d3e766de42d4'), 'thread_id': '1', 'thread_ts': '1ef3df65-6252-6b30-8002-2c9e1e68364b', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:48.830033+00:00\", \"id\": \"1ef3df65-6252-6b30-8002-2c9e1e68364b\", \"channel_values\": {\"messages\": [{\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"HumanMessage\"], \"kwargs\": {\"content\": \"what\\'s the weather in sf\", \"type\": \"human\", \"id\": \"a624d383-13c6-499c-8f03-31ed11fa0cfb\"}}, {\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"AIMessage\"], \"kwargs\": {\"content\": \"\", \"additional_kwargs\": {\"tool_calls\": [{\"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\", \"function\": {\"arguments\": \"{\\\\\"city\\\\\":\\\\\"sf\\\\\"}\", \"name\": \"get_weather\"}, \"type\": \"function\"}]}, \"response_metadata\": {\"token_usage\": {\"completion_tokens\": 14, \"prompt_tokens\": 58, \"total_tokens\": 72}, \"model_name\": \"gpt-3.5-turbo\", \"system_fingerprint\": null, \"finish_reason\": \"tool_calls\", \"logprobs\": null}, \"type\": \"ai\", \"id\": \"run-614bc54a-ad37-4f17-9047-b80752bdf66e-0\", \"tool_calls\": [{\"name\": \"get_weather\", \"args\": {\"city\": \"sf\"}, \"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\"}], \"invalid_tool_calls\": []}}, {\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"ToolMessage\"], \"kwargs\": {\"content\": \"It\\'s always sunny in sf\", \"type\": \"tool\", \"name\": \"get_weather\", \"id\": \"e58dd97d-b50d-4b0a-9492-7155106c975a\", \"tool_call_id\": \"call_wapm4s91KQUQqE9y1L53QmmE\"}}], \"tools\": \"tools\"}, \"channel_versions\": {\"__start__\": 2, \"messages\": 4, \"start:agent\": 3, \"agent\": 4, \"branch:agent:should_continue:tools\": 4, \"tools\": 4}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"agent\": {\"start:agent\": 2}, \"tools\": {\"branch:agent:should_continue:tools\": 3}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 2, \"writes\": {\"tools\": {\"messages\": [{\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"ToolMessage\"], \"kwargs\": {\"content\": \"It\\'s always sunny in sf\", \"type\": \"tool\", \"name\": \"get_weather\", \"id\": \"e58dd97d-b50d-4b0a-9492-7155106c975a\", \"tool_call_id\": \"call_wapm4s91KQUQqE9y1L53QmmE\"}}]}}}', 'parent_ts': '1ef3df65-6248-6ee8-8001-d5eef3fad087'}\n", - "{'_id': ObjectId('668d39a9b975d3e766de42d5'), 'thread_id': '1', 'thread_ts': '1ef3df65-6b84-63fd-8003-888bcef289e3', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-09T13:22:49.794047+00:00\", \"id\": \"1ef3df65-6b84-63fd-8003-888bcef289e3\", \"channel_values\": {\"messages\": [{\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"HumanMessage\"], \"kwargs\": {\"content\": \"what\\'s the weather in sf\", \"type\": \"human\", \"id\": \"a624d383-13c6-499c-8f03-31ed11fa0cfb\"}}, {\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"AIMessage\"], \"kwargs\": {\"content\": \"\", \"additional_kwargs\": {\"tool_calls\": [{\"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\", \"function\": {\"arguments\": \"{\\\\\"city\\\\\":\\\\\"sf\\\\\"}\", \"name\": \"get_weather\"}, \"type\": \"function\"}]}, \"response_metadata\": {\"token_usage\": {\"completion_tokens\": 14, \"prompt_tokens\": 58, \"total_tokens\": 72}, \"model_name\": \"gpt-3.5-turbo\", \"system_fingerprint\": null, \"finish_reason\": \"tool_calls\", \"logprobs\": null}, \"type\": \"ai\", \"id\": \"run-614bc54a-ad37-4f17-9047-b80752bdf66e-0\", \"tool_calls\": [{\"name\": \"get_weather\", \"args\": {\"city\": \"sf\"}, \"id\": \"call_wapm4s91KQUQqE9y1L53QmmE\"}], \"invalid_tool_calls\": []}}, {\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"ToolMessage\"], \"kwargs\": {\"content\": \"It\\'s always sunny in sf\", \"type\": \"tool\", \"name\": \"get_weather\", \"id\": \"e58dd97d-b50d-4b0a-9492-7155106c975a\", \"tool_call_id\": \"call_wapm4s91KQUQqE9y1L53QmmE\"}}, {\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"AIMessage\"], \"kwargs\": {\"content\": \"The weather in San Francisco is always sunny!\", \"response_metadata\": {\"token_usage\": {\"completion_tokens\": 10, \"prompt_tokens\": 86, \"total_tokens\": 96}, \"model_name\": \"gpt-3.5-turbo\", \"system_fingerprint\": null, \"finish_reason\": \"stop\", \"logprobs\": null}, \"type\": \"ai\", \"id\": \"run-e9984eca-f132-46d0-94ba-41d8ba5b7046-0\", \"tool_calls\": [], \"invalid_tool_calls\": []}}], \"agent\": \"agent\"}, \"channel_versions\": {\"__start__\": 2, \"messages\": 5, \"start:agent\": 3, \"agent\": 5, \"branch:agent:should_continue:tools\": 4, \"tools\": 5}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"agent\": {\"start:agent\": 3, \"tools\": 4}, \"tools\": {\"branch:agent:should_continue:tools\": 3}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 3, \"writes\": {\"agent\": {\"messages\": [{\"lc\": 1, \"type\": \"constructor\", \"id\": [\"langchain\", \"schema\", \"messages\", \"AIMessage\"], \"kwargs\": {\"content\": \"The weather in San Francisco is always sunny!\", \"response_metadata\": {\"token_usage\": {\"completion_tokens\": 10, \"prompt_tokens\": 86, \"total_tokens\": 96}, \"model_name\": \"gpt-3.5-turbo\", \"system_fingerprint\": null, \"finish_reason\": \"stop\", \"logprobs\": null}, \"type\": \"ai\", \"id\": \"run-e9984eca-f132-46d0-94ba-41d8ba5b7046-0\", \"tool_calls\": [], \"invalid_tool_calls\": []}}]}}}', 'parent_ts': '1ef3df65-6252-6b30-8002-2c9e1e68364b'}\n" - ] - } - ], - "source": [ - "client = MongoClient(MONGO_URI)\n", - "database = client[\"checkpoints_db\"]\n", - "collection = database[\"checkpoints_collection\"]\n", - "\n", - "for doc in collection.find():\n", - " print(doc)\n", - "\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", - "\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", - "\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, START\n", - "\n", - "checkpointer = MongoDBSaver(\n", - " AsyncIOMotorClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\"\n", - ")\n", - "builder = StateGraph(int)\n", - "builder.add_node(\"add_one\", lambda x: x + 1)\n", - "builder.add_edge(START, \"add_one\")\n", - "builder.add_edge(\"add_one\", END)\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", - "\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": { - "kernelspec": { - "display_name": "myenv", - "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.11.9" - } + "cells": [ + { + "cell_type": "markdown", + "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53", + "metadata": {}, + "source": [ + "# How to create a custom checkpointer using MongoDB\n", + "\n", + "When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions. \n", + "\n", + "This reference implementation shows how to use MongoDB as the backend for persisting checkpoint state. Make sure that you have MongoDB running on port `27017` for going through this guide.\n", + "\n", + "NOTE: this is just an reference implementation. You can implement your own checkpointer using a different database or modify this one as long as it conforms to the `BaseCheckpointSaver` interface." + ] }, - "nbformat": 4, - "nbformat_minor": 2 + { + "cell_type": "markdown", + "id": "456fa19c-93a5-4750-a410-f2d810b964ad", + "metadata": {}, + "source": [ + "## Setup environment" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "faadfb1b-cebe-4dcf-82fd-34044c380bc4", + "metadata": {}, + "outputs": [], + "source": [ + "%%capture --no-stderr\n", + "%pip install -U pymongo motor langgraph" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "eca9aafb-a155-407a-8036-682a2f1297d7", + "metadata": {}, + "outputs": [ + { + "name": "stdin", + "output_type": "stream", + "text": [ + "OPENAI_API_KEY: ········\n" + ] + } + ], + "source": [ + "import getpass\n", + "import os\n", + "\n", + "\n", + "def _set_env(var: str):\n", + " if not os.environ.get(var):\n", + " os.environ[var] = getpass.getpass(f\"{var}: \")\n", + "\n", + "\n", + "_set_env(\"OPENAI_API_KEY\")" + ] + }, + { + "cell_type": "markdown", + "id": "ecb23436-f238-4f8c-a2b7-67c7956121e2", + "metadata": {}, + "source": [ + "## Checkpointer implementation" + ] + }, + { + "cell_type": "markdown", + "id": "922822a8-f7d2-41ce-bada-206fc125c20c", + "metadata": {}, + "source": [ + "### MongoDBSaver" + ] + }, + { + "cell_type": "markdown", + "id": "c216852b-8318-4927-9000-1361d3ca81e8", + "metadata": {}, + "source": [ + "Below is an implementation of MongoDBSaver (for synchronous use of graph, i.e. `.invoke()`, `.stream()`). MongoDBSaver implements four methods that are required for any checkpointer:\n", + "\n", + "- `.put` - Store a checkpoint with its configuration and metadata.\n", + "- `.put_writes` - Store intermediate writes linked to a checkpoint (i.e. pending writes).\n", + "- `.get_tuple` - Fetch a checkpoint tuple using for a given configuration (`thread_id` and `checkpoint_id`).\n", + "- `.list` - List checkpoints that match a given configuration and filter criteria." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "98c8d65e-eb95-4cbd-8975-d33a52351d03", + "metadata": {}, + "outputs": [], + "source": [ + "from contextlib import asynccontextmanager, contextmanager\n", + "from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple\n", + "\n", + "from langchain_core.runnables import RunnableConfig\n", + "from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase\n", + "from pymongo import MongoClient, UpdateOne\n", + "from pymongo.database import Database as MongoDatabase\n", + "\n", + "from langgraph.checkpoint.base import (\n", + " BaseCheckpointSaver,\n", + " ChannelVersions,\n", + " Checkpoint,\n", + " CheckpointMetadata,\n", + " CheckpointTuple,\n", + " get_checkpoint_id,\n", + ")\n", + "\n", + "\n", + "class MongoDBSaver(BaseCheckpointSaver):\n", + " \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database.\"\"\"\n", + "\n", + " client: MongoClient\n", + " db: MongoDatabase\n", + "\n", + " def __init__(\n", + " self,\n", + " client: MongoClient,\n", + " db_name: str,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.client = client\n", + " self.db = self.client[db_name]\n", + "\n", + " @classmethod\n", + " @contextmanager\n", + " def from_conn_info(\n", + " cls, *, host: str, port: int, db_name: str\n", + " ) -> Iterator[\"MongoDBSaver\"]:\n", + " client = None\n", + " try:\n", + " client = MongoClient(host=host, port=port)\n", + " yield MongoDBSaver(client, db_name)\n", + " finally:\n", + " if client:\n", + " client.close()\n", + "\n", + " def get_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 \"checkpoint_id\" key, the checkpoint with\n", + " the matching thread ID and checkpoint ID 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", + " thread_id = config[\"configurable\"][\"thread_id\"]\n", + " checkpoint_ns = config[\"configurable\"].get(\"checkpoint_ns\", \"\")\n", + " if checkpoint_id := get_checkpoint_id(config):\n", + " query = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " }\n", + " else:\n", + " query = {\"thread_id\": thread_id, \"checkpoint_ns\": checkpoint_ns}\n", + "\n", + " result = self.db[\"checkpoints\"].find(query).sort(\"checkpoint_id\", -1).limit(1)\n", + " for doc in result:\n", + " config_values = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": doc[\"checkpoint_id\"],\n", + " }\n", + " checkpoint = self.serde.loads_typed((doc[\"type\"], doc[\"checkpoint\"]))\n", + " serialized_writes = self.db[\"checkpoint_writes\"].find(config_values)\n", + " pending_writes = [\n", + " (\n", + " doc[\"task_id\"],\n", + " doc[\"channel\"],\n", + " self.serde.loads_typed((doc[\"type\"], doc[\"value\"])),\n", + " )\n", + " for doc in serialized_writes\n", + " ]\n", + " return CheckpointTuple(\n", + " {\"configurable\": config_values},\n", + " checkpoint,\n", + " self.serde.loads(doc[\"metadata\"]),\n", + " (\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": doc[\"parent_checkpoint_id\"],\n", + " }\n", + " }\n", + " if doc.get(\"parent_checkpoint_id\")\n", + " else None\n", + " ),\n", + " pending_writes,\n", + " )\n", + "\n", + " def list(\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", + " ) -> Iterator[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 checkpoint ID in descending order (newest first).\n", + "\n", + " Args:\n", + " config (RunnableConfig): The config to use for listing the checkpoints.\n", + " filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.\n", + " before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.\n", + " limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.\n", + "\n", + " Yields:\n", + " Iterator[CheckpointTuple]: An iterator of checkpoint tuples.\n", + " \"\"\"\n", + " query = {}\n", + " if config is not None:\n", + " query = {\n", + " \"thread_id\": config[\"configurable\"][\"thread_id\"],\n", + " \"checkpoint_ns\": config[\"configurable\"].get(\"checkpoint_ns\", \"\"),\n", + " }\n", + "\n", + " if filter:\n", + " for key, value in filter.items():\n", + " query[f\"metadata.{key}\"] = value\n", + "\n", + " if before is not None:\n", + " query[\"checkpoint_id\"] = {\"$lt\": before[\"configurable\"][\"checkpoint_id\"]}\n", + "\n", + " result = self.db[\"checkpoints\"].find(query).sort(\"checkpoint_id\", -1)\n", + "\n", + " if limit is not None:\n", + " result = result.limit(limit)\n", + " for doc in result:\n", + " checkpoint = self.serde.loads_typed((doc[\"type\"], doc[\"checkpoint\"]))\n", + " yield CheckpointTuple(\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": doc[\"thread_id\"],\n", + " \"checkpoint_ns\": doc[\"checkpoint_ns\"],\n", + " \"checkpoint_id\": doc[\"checkpoint_id\"],\n", + " }\n", + " },\n", + " checkpoint,\n", + " self.serde.loads(doc[\"metadata\"]),\n", + " (\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": doc[\"thread_id\"],\n", + " \"checkpoint_ns\": doc[\"checkpoint_ns\"],\n", + " \"checkpoint_id\": doc[\"parent_checkpoint_id\"],\n", + " }\n", + " }\n", + " if doc.get(\"parent_checkpoint_id\")\n", + " else None\n", + " ),\n", + " )\n", + "\n", + " def put(\n", + " self,\n", + " config: RunnableConfig,\n", + " checkpoint: Checkpoint,\n", + " metadata: CheckpointMetadata,\n", + " new_versions: ChannelVersions,\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 (CheckpointMetadata): Additional metadata to save with the checkpoint.\n", + " new_versions (ChannelVersions): New channel versions as of this write.\n", + "\n", + " Returns:\n", + " RunnableConfig: Updated configuration after storing the checkpoint.\n", + " \"\"\"\n", + " thread_id = config[\"configurable\"][\"thread_id\"]\n", + " checkpoint_ns = config[\"configurable\"][\"checkpoint_ns\"]\n", + " checkpoint_id = checkpoint[\"id\"]\n", + " type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)\n", + " doc = {\n", + " \"parent_checkpoint_id\": config[\"configurable\"].get(\"checkpoint_id\"),\n", + " \"type\": type_,\n", + " \"checkpoint\": serialized_checkpoint,\n", + " \"metadata\": self.serde.dumps(metadata),\n", + " }\n", + " upsert_query = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " }\n", + " # Perform your operations here\n", + " self.db[\"checkpoints\"].update_one(upsert_query, {\"$set\": doc}, upsert=True)\n", + " return {\n", + " \"configurable\": {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " }\n", + " }\n", + "\n", + " def put_writes(\n", + " self,\n", + " config: RunnableConfig,\n", + " writes: Sequence[Tuple[str, Any]],\n", + " task_id: str,\n", + " ) -> None:\n", + " \"\"\"Store intermediate writes linked to a checkpoint.\n", + "\n", + " This method saves intermediate writes associated with a checkpoint to the MongoDB database.\n", + "\n", + " Args:\n", + " config (RunnableConfig): Configuration of the related checkpoint.\n", + " writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.\n", + " task_id (str): Identifier for the task creating the writes.\n", + " \"\"\"\n", + " thread_id = config[\"configurable\"][\"thread_id\"]\n", + " checkpoint_ns = config[\"configurable\"][\"checkpoint_ns\"]\n", + " checkpoint_id = config[\"configurable\"][\"checkpoint_id\"]\n", + " operations = []\n", + " for idx, (channel, value) in enumerate(writes):\n", + " upsert_query = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " \"task_id\": task_id,\n", + " \"idx\": idx,\n", + " }\n", + " type_, serialized_value = self.serde.dumps_typed(value)\n", + " operations.append(\n", + " UpdateOne(\n", + " upsert_query,\n", + " {\n", + " \"$set\": {\n", + " \"channel\": channel,\n", + " \"type\": type_,\n", + " \"value\": serialized_value,\n", + " }\n", + " },\n", + " upsert=True,\n", + " )\n", + " )\n", + " self.db[\"checkpoint_writes\"].bulk_write(operations)" + ] + }, + { + "cell_type": "markdown", + "id": "ec21ff00-75a7-4789-b863-93fffcc0b32d", + "metadata": {}, + "source": [ + "### AsyncMongoDBSaver" + ] + }, + { + "cell_type": "markdown", + "id": "9e5ad763-12ab-4918-af40-0be85678e35b", + "metadata": {}, + "source": [ + "Below is a reference implementation of AsyncMongoDBSaver (for asynchronous use of graph, i.e. `.ainvoke()`, `.astream()`). AsyncMongoDBSaver implements four methods that are required for any async checkpointer:\n", + "\n", + "- `.aput` - Store a checkpoint with its configuration and metadata.\n", + "- `.aput_writes` - Store intermediate writes linked to a checkpoint (i.e. pending writes).\n", + "- `.aget_tuple` - Fetch a checkpoint tuple using for a given configuration (`thread_id` and `checkpoint_id`).\n", + "- `.alist` - List checkpoints that match a given configuration and filter criteria." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "888302ee-c201-498f-b6e3-69ec5f1a039c", + "metadata": {}, + "outputs": [], + "source": [ + "class AsyncMongoDBSaver(BaseCheckpointSaver):\n", + " \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database asynchronously.\"\"\"\n", + "\n", + " client: AsyncIOMotorClient\n", + " db: AsyncIOMotorDatabase\n", + "\n", + " def __init__(\n", + " self,\n", + " client: AsyncIOMotorClient,\n", + " db_name: str,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.client = client\n", + " self.db = self.client[db_name]\n", + "\n", + " @classmethod\n", + " @asynccontextmanager\n", + " async def from_conn_info(\n", + " cls, *, host: str, port: int, db_name: str\n", + " ) -> AsyncIterator[\"AsyncMongoDBSaver\"]:\n", + " client = None\n", + " try:\n", + " client = AsyncIOMotorClient(host=host, port=port)\n", + " yield AsyncMongoDBSaver(client, db_name)\n", + " finally:\n", + " if client:\n", + " client.close()\n", + "\n", + " async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n", + " \"\"\"Get a checkpoint tuple from the database asynchronously.\n", + "\n", + " This method retrieves a checkpoint tuple from the MongoDB database based on the\n", + " provided config. If the config contains a \"checkpoint_id\" key, the checkpoint with\n", + " the matching thread ID and checkpoint ID 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", + " thread_id = config[\"configurable\"][\"thread_id\"]\n", + " checkpoint_ns = config[\"configurable\"].get(\"checkpoint_ns\", \"\")\n", + " if checkpoint_id := get_checkpoint_id(config):\n", + " query = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " }\n", + " else:\n", + " query = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " }\n", + "\n", + " result = self.db[\"checkpoints\"].find(query).sort(\"checkpoint_id\", -1).limit(1)\n", + " async for doc in result:\n", + " config_values = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": doc[\"checkpoint_id\"],\n", + " }\n", + " checkpoint = self.serde.loads_typed((doc[\"type\"], doc[\"checkpoint\"]))\n", + " serialized_writes = self.db[\"checkpoint_writes\"].find(config_values)\n", + " pending_writes = [\n", + " (\n", + " doc[\"task_id\"],\n", + " doc[\"channel\"],\n", + " self.serde.loads_typed((doc[\"type\"], doc[\"value\"])),\n", + " )\n", + " async for doc in serialized_writes\n", + " ]\n", + " return CheckpointTuple(\n", + " {\"configurable\": config_values},\n", + " checkpoint,\n", + " self.serde.loads(doc[\"metadata\"]),\n", + " (\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": doc[\"parent_checkpoint_id\"],\n", + " }\n", + " }\n", + " if doc.get(\"parent_checkpoint_id\")\n", + " else None\n", + " ),\n", + " pending_writes,\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 asynchronously.\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 checkpoint ID in descending order (newest first).\n", + "\n", + " Args:\n", + " config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.\n", + " filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.\n", + " before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.\n", + " limit (Optional[int]): Maximum number of checkpoints to return.\n", + "\n", + " Yields:\n", + " AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.\n", + " \"\"\"\n", + " query = {}\n", + " if config is not None:\n", + " query = {\n", + " \"thread_id\": config[\"configurable\"][\"thread_id\"],\n", + " \"checkpoint_ns\": config[\"configurable\"].get(\"checkpoint_ns\", \"\"),\n", + " }\n", + "\n", + " if filter:\n", + " for key, value in filter.items():\n", + " query[f\"metadata.{key}\"] = value\n", + "\n", + " if before is not None:\n", + " query[\"checkpoint_id\"] = {\"$lt\": before[\"configurable\"][\"checkpoint_id\"]}\n", + "\n", + " result = self.db[\"checkpoints\"].find(query).sort(\"checkpoint_id\", -1)\n", + "\n", + " if limit is not None:\n", + " result = result.limit(limit)\n", + " async for doc in result:\n", + " checkpoint = self.serde.loads_typed((doc[\"type\"], doc[\"checkpoint\"]))\n", + " yield CheckpointTuple(\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": doc[\"thread_id\"],\n", + " \"checkpoint_ns\": doc[\"checkpoint_ns\"],\n", + " \"checkpoint_id\": doc[\"checkpoint_id\"],\n", + " }\n", + " },\n", + " checkpoint,\n", + " self.serde.loads(doc[\"metadata\"]),\n", + " (\n", + " {\n", + " \"configurable\": {\n", + " \"thread_id\": doc[\"thread_id\"],\n", + " \"checkpoint_ns\": doc[\"checkpoint_ns\"],\n", + " \"checkpoint_id\": doc[\"parent_checkpoint_id\"],\n", + " }\n", + " }\n", + " if doc.get(\"parent_checkpoint_id\")\n", + " else None\n", + " ),\n", + " )\n", + "\n", + " async def aput(\n", + " self,\n", + " config: RunnableConfig,\n", + " checkpoint: Checkpoint,\n", + " metadata: CheckpointMetadata,\n", + " new_versions: ChannelVersions,\n", + " ) -> RunnableConfig:\n", + " \"\"\"Save a checkpoint to the database asynchronously.\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 (CheckpointMetadata): Additional metadata to save with the checkpoint.\n", + " new_versions (ChannelVersions): New channel versions as of this write.\n", + "\n", + " Returns:\n", + " RunnableConfig: Updated configuration after storing the checkpoint.\n", + " \"\"\"\n", + " thread_id = config[\"configurable\"][\"thread_id\"]\n", + " checkpoint_ns = config[\"configurable\"][\"checkpoint_ns\"]\n", + " checkpoint_id = checkpoint[\"id\"]\n", + " type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)\n", + " doc = {\n", + " \"parent_checkpoint_id\": config[\"configurable\"].get(\"checkpoint_id\"),\n", + " \"type\": type_,\n", + " \"checkpoint\": serialized_checkpoint,\n", + " \"metadata\": self.serde.dumps(metadata),\n", + " }\n", + " upsert_query = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " }\n", + " # Perform your operations here\n", + " await self.db[\"checkpoints\"].update_one(\n", + " upsert_query, {\"$set\": doc}, upsert=True\n", + " )\n", + " return {\n", + " \"configurable\": {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " }\n", + " }\n", + "\n", + " async def aput_writes(\n", + " self,\n", + " config: RunnableConfig,\n", + " writes: Sequence[Tuple[str, Any]],\n", + " task_id: str,\n", + " ) -> None:\n", + " \"\"\"Store intermediate writes linked to a checkpoint asynchronously.\n", + "\n", + " This method saves intermediate writes associated with a checkpoint to the database.\n", + "\n", + " Args:\n", + " config (RunnableConfig): Configuration of the related checkpoint.\n", + " writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.\n", + " task_id (str): Identifier for the task creating the writes.\n", + " \"\"\"\n", + " thread_id = config[\"configurable\"][\"thread_id\"]\n", + " checkpoint_ns = config[\"configurable\"][\"checkpoint_ns\"]\n", + " checkpoint_id = config[\"configurable\"][\"checkpoint_id\"]\n", + " operations = []\n", + " for idx, (channel, value) in enumerate(writes):\n", + " upsert_query = {\n", + " \"thread_id\": thread_id,\n", + " \"checkpoint_ns\": checkpoint_ns,\n", + " \"checkpoint_id\": checkpoint_id,\n", + " \"task_id\": task_id,\n", + " \"idx\": idx,\n", + " }\n", + " type_, serialized_value = self.serde.dumps_typed(value)\n", + " operations.append(\n", + " UpdateOne(\n", + " upsert_query,\n", + " {\n", + " \"$set\": {\n", + " \"channel\": channel,\n", + " \"type\": type_,\n", + " \"value\": serialized_value,\n", + " }\n", + " },\n", + " upsert=True,\n", + " )\n", + " )\n", + " await self.db[\"checkpoint_writes\"].bulk_write(operations)\n" + ] + }, + { + "cell_type": "markdown", + "id": "e26b3204-cca2-414c-800e-7e09032445ae", + "metadata": {}, + "source": [ + "## Setup model and tools for the graph" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "e5213193-5a7d-43e7-aeba-fe732bb1cd7a", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Literal\n", + "from langchain_core.runnables import ConfigurableField\n", + "from langchain_core.tools import tool\n", + "from langchain_openai import ChatOpenAI\n", + "from langgraph.prebuilt import create_react_agent\n", + "\n", + "\n", + "@tool\n", + "def get_weather(city: Literal[\"nyc\", \"sf\"]):\n", + " \"\"\"Use this to get weather information.\"\"\"\n", + " if city == \"nyc\":\n", + " return \"It might be cloudy in nyc\"\n", + " elif city == \"sf\":\n", + " return \"It's always sunny in sf\"\n", + " else:\n", + " raise AssertionError(\"Unknown city\")\n", + "\n", + "\n", + "tools = [get_weather]\n", + "model = ChatOpenAI(model_name=\"gpt-4o-mini\", temperature=0)" + ] + }, + { + "cell_type": "markdown", + "id": "e9342c62-dbb4-40f6-9271-7393f1ca48c4", + "metadata": {}, + "source": [ + "## Use sync connection" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "5fe54e79-9eaf-44e2-b2d9-1e0284b984d0", + "metadata": {}, + "outputs": [], + "source": [ + "with MongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n", + " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", + " config = {\"configurable\": {\"thread_id\": \"1\"}}\n", + " res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)\n", + "\n", + " latest_checkpoint = checkpointer.get(config)\n", + " latest_checkpoint_tuple = checkpointer.get_tuple(config)\n", + " checkpoint_tuples = list(checkpointer.list(config))" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "c298e627-115a-4b4c-ae17-520ca9a640cd", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'v': 1,\n", + " 'ts': '2024-08-09T16:19:39.102711+00:00',\n", + " 'id': '1ef566b2-d2a8-6cdc-8003-cc4d1980d188',\n", + " 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='f4227353-e0e5-43a9-984a-e4b9e2d8e7b8'),\n", + " AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-cd1d3187-470f-4ebd-938f-527a61824045-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n", + " ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='2d124101-696d-450f-bc9f-d8fdcc564101', tool_call_id='call_Y7PzHb7LrIdiTnO5UiSfelt3'),\n", + " AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-87c76dd2-33f4-433e-986a-9405cfe88c88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})],\n", + " 'agent': 'agent'},\n", + " 'channel_versions': {'__start__': 2,\n", + " 'messages': 5,\n", + " 'start:agent': 3,\n", + " 'agent': 5,\n", + " 'branch:agent:should_continue:tools': 4,\n", + " 'tools': 5},\n", + " 'versions_seen': {'__input__': {},\n", + " '__start__': {'__start__': 1},\n", + " 'agent': {'start:agent': 2, 'tools': 4},\n", + " 'tools': {'branch:agent:should_continue:tools': 3}},\n", + " 'pending_sends': [],\n", + " 'current_tasks': {}}" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "latest_checkpoint" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "922f9406-0f68-418a-9cb4-e0e29de4b5f9", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "CheckpointTuple(config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-d2a8-6cdc-8003-cc4d1980d188'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:39.102711+00:00', 'id': '1ef566b2-d2a8-6cdc-8003-cc4d1980d188', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='f4227353-e0e5-43a9-984a-e4b9e2d8e7b8'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-cd1d3187-470f-4ebd-938f-527a61824045-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}), ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='2d124101-696d-450f-bc9f-d8fdcc564101', tool_call_id='call_Y7PzHb7LrIdiTnO5UiSfelt3'), AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-87c76dd2-33f4-433e-986a-9405cfe88c88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-87c76dd2-33f4-433e-986a-9405cfe88c88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}}, 'step': 3}, parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-cdf7-6b98-8002-997748cc5052'}}, pending_writes=[])" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "latest_checkpoint_tuple" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "b2ce743b-5896-443b-9ec0-a655b065895c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[CheckpointTuple(config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-d2a8-6cdc-8003-cc4d1980d188'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:39.102711+00:00', 'id': '1ef566b2-d2a8-6cdc-8003-cc4d1980d188', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='f4227353-e0e5-43a9-984a-e4b9e2d8e7b8'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-cd1d3187-470f-4ebd-938f-527a61824045-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}), ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='2d124101-696d-450f-bc9f-d8fdcc564101', tool_call_id='call_Y7PzHb7LrIdiTnO5UiSfelt3'), AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-87c76dd2-33f4-433e-986a-9405cfe88c88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in San Francisco is always sunny!', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-87c76dd2-33f4-433e-986a-9405cfe88c88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}}, 'step': 3}, parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-cdf7-6b98-8002-997748cc5052'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-cdf7-6b98-8002-997748cc5052'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:38.610752+00:00', 'id': '1ef566b2-cdf7-6b98-8002-997748cc5052', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='f4227353-e0e5-43a9-984a-e4b9e2d8e7b8'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-cd1d3187-470f-4ebd-938f-527a61824045-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}), ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='2d124101-696d-450f-bc9f-d8fdcc564101', tool_call_id='call_Y7PzHb7LrIdiTnO5UiSfelt3')], 'tools': 'tools'}, 'channel_versions': {'__start__': 2, 'messages': 4, 'start:agent': 3, 'agent': 4, 'branch:agent:should_continue:tools': 4, 'tools': 4}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'tools': {'messages': [ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='2d124101-696d-450f-bc9f-d8fdcc564101', tool_call_id='call_Y7PzHb7LrIdiTnO5UiSfelt3')]}}, 'step': 2}, parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-cde3-6c60-8001-28d4cc36978d'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-cde3-6c60-8001-28d4cc36978d'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:38.602590+00:00', 'id': '1ef566b2-cde3-6c60-8001-28d4cc36978d', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='f4227353-e0e5-43a9-984a-e4b9e2d8e7b8'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-cd1d3187-470f-4ebd-938f-527a61824045-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71})], 'agent': 'agent', 'branch:agent:should_continue:tools': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 3, 'start:agent': 3, 'agent': 3, 'branch:agent:should_continue:tools': 3}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-cd1d3187-470f-4ebd-938f-527a61824045-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_Y7PzHb7LrIdiTnO5UiSfelt3', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71})]}}, 'step': 1}, parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-c72c-6fca-8000-aac6e4f4b809'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-c72c-6fca-8000-aac6e4f4b809'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:37.898584+00:00', 'id': '1ef566b2-c72c-6fca-8000-aac6e4f4b809', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='f4227353-e0e5-43a9-984a-e4b9e2d8e7b8')], 'start:agent': '__start__'}, 'channel_versions': {'__start__': 2, 'messages': 2, 'start:agent': 2}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': None, 'step': 0}, parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-c72a-6af4-bfff-919b9dc6abfe'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b2-c72a-6af4-bfff-919b9dc6abfe'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:37.897642+00:00', 'id': '1ef566b2-c72a-6af4-bfff-919b9dc6abfe', 'channel_values': {'messages': [], '__start__': {'messages': [['human', \"what's the weather in sf\"]]}}, 'channel_versions': {'__start__': 1}, 'versions_seen': {'__input__': {}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'input', 'writes': {'messages': [['human', \"what's the weather in sf\"]]}, 'step': -1}, parent_config=None, pending_writes=None)]" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "checkpoint_tuples" + ] + }, + { + "cell_type": "markdown", + "id": "c0a47d3e-e588-48fc-a5d4-2145dff17e77", + "metadata": {}, + "source": [ + "## Use async connection" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "6a39d1ff-ca37-4457-8b52-07d33b59c36e", + "metadata": {}, + "outputs": [], + "source": [ + "async with AsyncMongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n", + " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", + " config = {\"configurable\": {\"thread_id\": \"2\"}}\n", + " res = await graph.ainvoke({\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config)\n", + "\n", + " latest_checkpoint = await checkpointer.aget(config)\n", + " latest_checkpoint_tuple = await checkpointer.aget_tuple(config)\n", + " checkpoint_tuples = [c async for c in checkpointer.alist(config)]" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "51125ef1-bdb6-454e-82cc-4ae19a113606", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'v': 1,\n", + " 'ts': '2024-08-09T16:19:48.212051+00:00',\n", + " 'id': '1ef566b3-2988-664c-8003-5974c59c6bda',\n", + " 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1ae4b12f-b1cb-4d55-a754-42cf1c2fbcd5'),\n", + " AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b5da58b5-8f75-485d-af29-bfdeb09b0d94-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}),\n", + " ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='56d4e46b-6cb3-4efe-b369-27b666e62348', tool_call_id='call_IJvXEELx7Ir3kASCqr9dbvhU'),\n", + " AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-dcacbc70-b213-4ddc-ac08-c0d17b2766d8-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})],\n", + " 'agent': 'agent'},\n", + " 'channel_versions': {'__start__': 2,\n", + " 'messages': 5,\n", + " 'start:agent': 3,\n", + " 'agent': 5,\n", + " 'branch:agent:should_continue:tools': 4,\n", + " 'tools': 5},\n", + " 'versions_seen': {'__input__': {},\n", + " '__start__': {'__start__': 1},\n", + " 'agent': {'start:agent': 2, 'tools': 4},\n", + " 'tools': {'branch:agent:should_continue:tools': 3}},\n", + " 'pending_sends': [],\n", + " 'current_tasks': {}}" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "latest_checkpoint" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "97f8a87b-8423-41c6-a76b-9a6b30904e73", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-2988-664c-8003-5974c59c6bda'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:48.212051+00:00', 'id': '1ef566b3-2988-664c-8003-5974c59c6bda', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1ae4b12f-b1cb-4d55-a754-42cf1c2fbcd5'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b5da58b5-8f75-485d-af29-bfdeb09b0d94-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='56d4e46b-6cb3-4efe-b369-27b666e62348', tool_call_id='call_IJvXEELx7Ir3kASCqr9dbvhU'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-dcacbc70-b213-4ddc-ac08-c0d17b2766d8-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-dcacbc70-b213-4ddc-ac08-c0d17b2766d8-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}, 'step': 3}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-23c9-64ea-8002-036c32979035'}}, pending_writes=[])" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "latest_checkpoint_tuple" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "2b6d73ca-519e-45f7-90c2-1b8596624505", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-2988-664c-8003-5974c59c6bda'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:48.212051+00:00', 'id': '1ef566b3-2988-664c-8003-5974c59c6bda', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1ae4b12f-b1cb-4d55-a754-42cf1c2fbcd5'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b5da58b5-8f75-485d-af29-bfdeb09b0d94-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='56d4e46b-6cb3-4efe-b369-27b666e62348', tool_call_id='call_IJvXEELx7Ir3kASCqr9dbvhU'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-dcacbc70-b213-4ddc-ac08-c0d17b2766d8-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-dcacbc70-b213-4ddc-ac08-c0d17b2766d8-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}, 'step': 3}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-23c9-64ea-8002-036c32979035'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-23c9-64ea-8002-036c32979035'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:47.609498+00:00', 'id': '1ef566b3-23c9-64ea-8002-036c32979035', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1ae4b12f-b1cb-4d55-a754-42cf1c2fbcd5'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b5da58b5-8f75-485d-af29-bfdeb09b0d94-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='56d4e46b-6cb3-4efe-b369-27b666e62348', tool_call_id='call_IJvXEELx7Ir3kASCqr9dbvhU')], 'tools': 'tools'}, 'channel_versions': {'__start__': 2, 'messages': 4, 'start:agent': 3, 'agent': 4, 'branch:agent:should_continue:tools': 4, 'tools': 4}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'tools': {'messages': [ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='56d4e46b-6cb3-4efe-b369-27b666e62348', tool_call_id='call_IJvXEELx7Ir3kASCqr9dbvhU')]}}, 'step': 2}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-23b5-6de6-8001-a39c8ce6fd93'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-23b5-6de6-8001-a39c8ce6fd93'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:47.601527+00:00', 'id': '1ef566b3-23b5-6de6-8001-a39c8ce6fd93', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1ae4b12f-b1cb-4d55-a754-42cf1c2fbcd5'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b5da58b5-8f75-485d-af29-bfdeb09b0d94-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})], 'agent': 'agent', 'branch:agent:should_continue:tools': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 3, 'start:agent': 3, 'agent': 3, 'branch:agent:should_continue:tools': 3}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}, 'agent': {'start:agent': 2}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b5da58b5-8f75-485d-af29-bfdeb09b0d94-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_IJvXEELx7Ir3kASCqr9dbvhU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})]}}, 'step': 1}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-1d6c-6a02-8000-158f156ffce3'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-1d6c-6a02-8000-158f156ffce3'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:46.942389+00:00', 'id': '1ef566b3-1d6c-6a02-8000-158f156ffce3', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1ae4b12f-b1cb-4d55-a754-42cf1c2fbcd5')], 'start:agent': '__start__'}, 'channel_versions': {'__start__': 2, 'messages': 2, 'start:agent': 2}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': 1}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'loop', 'writes': None, 'step': 0}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-1d67-61e2-bfff-d91abbcc3a09'}}, pending_writes=None),\n", + " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef566b3-1d67-61e2-bfff-d91abbcc3a09'}}, checkpoint={'v': 1, 'ts': '2024-08-09T16:19:46.940133+00:00', 'id': '1ef566b3-1d67-61e2-bfff-d91abbcc3a09', 'channel_values': {'messages': [], '__start__': {'messages': [['human', \"what's the weather in nyc\"]]}}, 'channel_versions': {'__start__': 1}, 'versions_seen': {'__input__': {}}, 'pending_sends': [], 'current_tasks': {}}, metadata={'source': 'input', 'writes': {'messages': [['human', \"what's the weather in nyc\"]]}, 'step': -1}, parent_config=None, pending_writes=None)]" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "checkpoint_tuples" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "langgraph", + "language": "python", + "name": "langgraph" + }, + "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.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 } diff --git a/examples/persistence_redis.ipynb b/examples/persistence_redis.ipynb index b63c36f1d..3b2ad1170 100644 --- a/examples/persistence_redis.ipynb +++ b/examples/persistence_redis.ipynb @@ -7,9 +7,9 @@ "source": [ "# How to create a custom checkpointer using Redis\n", "\n", - "When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions. Make sure that you have Redis running on port 6379 for going through this guide.\n", + "When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions.\n", "\n", - "This reference implementation shows how to use Redis as the backend for persisting checkpoint state.\n", + "This reference implementation shows how to use Redis as the backend for persisting checkpoint state. Make sure that you have Redis running on port `6379` for going through this guide.\n", "\n", "NOTE: this is just an reference implementation. You can implement your own checkpointer using a different database or modify this one as long as it conforms to the `BaseCheckpointSaver` interface." ] @@ -93,7 +93,6 @@ "source": [ "\"\"\"Implementation of a langgraph checkpoint saver using Redis.\"\"\"\n", "from contextlib import asynccontextmanager, contextmanager\n", - "from hashlib import md5\n", "from typing import (\n", " Any,\n", " AsyncGenerator,\n", @@ -108,12 +107,10 @@ "\n", "from langgraph.checkpoint.base import (\n", " BaseCheckpointSaver,\n", - " ChannelProtocol,\n", " ChannelVersions,\n", " Checkpoint,\n", " CheckpointMetadata,\n", " CheckpointTuple,\n", - " EmptyChannelError,\n", " PendingWrite,\n", " get_checkpoint_id,\n", ")\n", @@ -230,22 +227,6 @@ " return writes\n", "\n", "\n", - "def _get_next_version(\n", - " serde: SerializerProtocol, current: Optional[str], channel: ChannelProtocol\n", - ") -> str:\n", - " \"\"\"Implementation of get_next_version for the checkpointers\"\"\"\n", - " if current is None:\n", - " current_v = 0\n", - " else:\n", - " current_v = int(current.split(\".\")[0])\n", - " next_v = current_v + 1\n", - " try:\n", - " next_h = md5(serde.dumps_typed(channel.checkpoint())[1]).hexdigest()\n", - " except EmptyChannelError:\n", - " next_h = \"\"\n", - " return f\"{next_v:032}.{next_h}\"\n", - "\n", - "\n", "def _parse_redis_checkpoint_data(\n", " serde: SerializerProtocol,\n", " key: str,\n", @@ -492,20 +473,6 @@ " if data and b\"checkpoint\" in data and b\"metadata\" in data:\n", " yield _parse_redis_checkpoint_data(self.serde, key.decode(), data)\n", "\n", - " def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:\n", - " \"\"\"Generate the next version ID for a channel.\n", - "\n", - " This method creates a new version identifier for a channel based on its current version.\n", - "\n", - " Args:\n", - " current (Optional[str]): The current version identifier of the channel.\n", - " channel (BaseChannel): The channel being versioned.\n", - "\n", - " Returns:\n", - " str: The next version identifier, which is guaranteed to be monotonically increasing.\n", - " \"\"\"\n", - " return _get_next_version(self.serde, current, channel)\n", - "\n", " def _get_checkpoint_key(\n", " self, conn, thread_id: str, checkpoint_ns: str, checkpoint_id: Optional[str]\n", " ) -> Optional[str]:\n", @@ -736,20 +703,6 @@ " if data and b\"checkpoint\" in data and b\"metadata\" in data:\n", " yield _parse_redis_checkpoint_data(self.serde, key.decode(), data)\n", "\n", - " def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:\n", - " \"\"\"Generate the next version ID for a channel.\n", - "\n", - " This method creates a new version identifier for a channel based on its current version.\n", - "\n", - " Args:\n", - " current (Optional[str]): The current version identifier of the channel.\n", - " channel (BaseChannel): The channel being versioned.\n", - "\n", - " Returns:\n", - " str: The next version identifier, which is guaranteed to be monotonically increasing.\n", - " \"\"\"\n", - " return _get_next_version(self.serde, current, channel)\n", - "\n", " async def _aget_checkpoint_key(\n", " self, conn, thread_id: str, checkpoint_ns: str, checkpoint_id: Optional[str]\n", " ) -> Optional[str]:\n",