Typing

Author: caseyclements

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

@@ -359,7 +359,7 @@ async def aput(
         type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
         metadata = metadata.copy()
         metadata.update(config.get("metadata", {}))
-        doc = {
+        doc: dict[str, Any] = {
             "parent_checkpoint_id": config["configurable"].get("checkpoint_id"),
             "type": type_,
             "checkpoint": serialized_checkpoint,
@@ -423,7 +423,7 @@ async def aput_writes(
 
             type_, serialized_value = self.serde.dumps_typed(value)
 
-            update_doc = {
+            update_doc: dict[str, Any] = {
                 "channel": channel,
                 "type": type_,
                 "value": serialized_value,

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

@@ -438,7 +438,7 @@ def put_writes(
 
             type_, serialized_value = self.serde.dumps_typed(value)
 
-            update_doc = {
+            update_doc: dict[str, Any] = {
                 "channel": channel,
                 "type": type_,
                 "value": serialized_value,

libs/langgraph-checkpoint-mongodb/tests/unit_tests/test_interrupt.py

@@ -3,9 +3,10 @@
 
 import pytest
 import pytest_asyncio
-from pymongo import AsyncMongoClient, MongoClient
+from langchain_core.runnables import RunnableConfig
+from pymongo import MongoClient
 
-from langgraph.checkpoint.mongodb import AsyncMongoDBSaver, MongoDBSaver
+from langgraph.checkpoint.mongodb import MongoDBSaver
 from langgraph.types import Interrupt
 
 MONGODB_URI = os.environ.get(
@@ -17,35 +18,23 @@
 TTL: int = 60 * 60
 
 
-@pytest_asyncio.fixture(params=["run_in_executor", "aio"])
+@pytest_asyncio.fixture()
 async def async_saver(request: pytest.FixtureRequest) -> AsyncGenerator:
-    if request.param == "aio":
-        # Use async client and checkpointer
-        aclient: AsyncMongoClient = AsyncMongoClient(MONGODB_URI)
-        adb = aclient[DB_NAME]
-        for clxn in await adb.list_collection_names():
-            await adb.drop_collection(clxn)
-        async with AsyncMongoDBSaver.from_conn_string(
-            MONGODB_URI, DB_NAME, COLLECTION_NAME, WRITES_COLLECTION_NAME, TTL
-        ) as checkpointer:
-            yield checkpointer
-        await aclient.close()
-    else:
-        # Use sync client and checkpointer with async methods run in executor
-        client: MongoClient = MongoClient(MONGODB_URI)
-        db = client[DB_NAME]
-        for clxn in db.list_collection_names():
-            db.drop_collection(clxn)
-        with MongoDBSaver.from_conn_string(
-            MONGODB_URI, DB_NAME, COLLECTION_NAME, WRITES_COLLECTION_NAME, TTL
-        ) as checkpointer:
-            yield checkpointer
-        client.close()
-
-
-def test_put_writes_on_interrupt(async_saver: MongoDBSaver):
+    # Use sync client and checkpointer with async methods run in executor
+    client: MongoClient = MongoClient(MONGODB_URI)
+    db = client[DB_NAME]
+    for clxn in db.list_collection_names():
+        db.drop_collection(clxn)
+    with MongoDBSaver.from_conn_string(
+        MONGODB_URI, DB_NAME, COLLECTION_NAME, WRITES_COLLECTION_NAME, TTL
+    ) as checkpointer:
+        yield checkpointer
+    client.close()
+
+
+async def test_put_writes_on_interrupt(async_saver: MongoDBSaver) -> None:
     """Test that no error is raised when interrupted workflow updates writes."""
-    config = {
+    config: RunnableConfig = {
         "configurable": {
             "checkpoint_id": "check1",
             "thread_id": "thread1",
@@ -61,13 +50,11 @@ def test_put_writes_on_interrupt(async_saver: MongoDBSaver):
             (
                 Interrupt(
                     value="please provide input",
-                    resumable=True,
-                    ns=["human_feedback:1b798da3"],
                 ),
             ),
         )
     ]
-    async_saver.aput_writes(config, writes1, task_id, task_path)
+    await async_saver.aput_writes(config, writes1, task_id, task_path)
 
     writes2 = [("__interrupt__", (Interrupt(value="please provide another input"),))]
-    async_saver.aput_writes(config, writes2, task_id, task_path)
+    await async_saver.aput_writes(config, writes2, task_id, task_path)