INTPYTHON-567 Fix AsyncMongoDBSaver setup (#109)

Author: blink1073

libs/langgraph-checkpoint-mongodb/langgraph/checkpoint/mongodb/aio.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import builtins
 import sys
@@ -9,6 +11,8 @@
 from langchain_core.runnables import RunnableConfig
 from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
 from pymongo import UpdateOne
+from pymongo.asynchronous.database import AsyncDatabase
+from pymongo.asynchronous.mongo_client import AsyncMongoClient
 from pymongo.driver_info import DriverInfo
 
 from langgraph.checkpoint.base import (
@@ -68,12 +72,12 @@ class AsyncMongoDBSaver(BaseCheckpointSaver):
         input=3, output=4
     """
 
-    client: AsyncIOMotorClient
-    db: AsyncIOMotorDatabase
+    client: AsyncIOMotorClient | AsyncMongoClient
+    db: AsyncIOMotorDatabase | AsyncDatabase
 
     def __init__(
         self,
-        client: AsyncIOMotorClient,
+        client: AsyncIOMotorClient | AsyncMongoClient,
         db_name: str = "checkpointing_db",
         checkpoint_collection_name: str = "checkpoints_aio",
         writes_collection_name: str = "checkpoint_writes_aio",
@@ -92,12 +96,24 @@ async def _setup(self):
         if self._setup_future is not None:
             return await self._setup_future
         self._setup_future = asyncio.Future()
-        if len(await self.checkpoint_collection.list_indexes().to_list()) < 2:
+        if isinstance(self.client, AsyncMongoClient):
+            num_indexes = len(
+                await (await self.checkpoint_collection.list_indexes()).to_list()
+            )
+        else:
+            num_indexes = len(await self.checkpoint_collection.list_indexes().to_list())
+        if num_indexes < 2:
             await self.checkpoint_collection.create_index(
                 keys=[("thread_id", 1), ("checkpoint_ns", 1), ("checkpoint_id", -1)],
                 unique=True,
             )
-        if len(await self.writes_collection.list_indexes().to_list()) < 2:
+        if isinstance(self.client, AsyncMongoClient):
+            num_indexes = len(
+                await (await self.writes_collection.list_indexes()).to_list()
+            )
+        else:
+            num_indexes = len(await self.writes_collection.list_indexes().to_list())
+        if num_indexes < 2:
             await self.writes_collection.create_index(
                 keys=[
                     ("thread_id", 1),
@@ -119,7 +135,7 @@ async def from_conn_string(
         checkpoint_collection_name: str = "checkpoints_aio",
         writes_collection_name: str = "checkpoint_writes_aio",
         **kwargs: Any,
-    ) -> AsyncIterator["AsyncMongoDBSaver"]:
+    ) -> AsyncIterator[AsyncMongoDBSaver]:
         """Create asynchronous checkpointer
 
         This includes creation of collections and indexes if they don't exist

libs/langgraph-checkpoint-mongodb/pyproject.toml

@@ -10,8 +10,8 @@ readme = "README.md"
 requires-python = ">=3.9"
 dependencies = [
     "langgraph-checkpoint>=2.0.23,<3.0.0",
-    "pymongo>=4.9,<4.12",
-    "motor>3.5.0",
+    "pymongo>=4.10,<4.12",
+    "motor>3.6.0",
 ]
 
 [dependency-groups]