Initial test implementation.

Author: caseyclements

libs/langgraph-checkpoint-mongodb/tests/integration_tests/test_highlevel_graph.py

@@ -0,0 +1,225 @@
+"""
+Based on LangGraph's Benchmarking script,
+https://github.com/langchain-ai/langgraph/blob/main/libs/langgraph/bench/fanout_to_subgraph.py,
+this pattern of joke generation is used often in the examples.
+The fanout here is performed by the list comprehension of [class:~langgraph.types.Send] calls.
+The effect of this is a map (fanout) workflow where the graph invokes
+the same node multiple times in parallel.
+The node here is a subgraph.
+The subgraph is linear, with a conditional edge 'bump_loop' that repeatably calls
+the node 'bump' until a condition is met.
+This test can be used for benchmarking.
+It also demonstrates the high-level API of subgraphs, add_conditional_edges, and Send.
+"""
+
+import operator
+import os
+import time
+from collections.abc import Generator
+from typing import Annotated, TypedDict
+
+import langchain_core
+import pytest
+
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.checkpoint.mongodb import AsyncMongoDBSaver, MongoDBSaver
+from langgraph.constants import START, Send
+from langgraph.graph import END, StateGraph
+
+# --- Configuration ---
+MONGODB_URI = os.environ.get(
+    "MONGODB_URI", "mongodb://localhost:27017?directConnection=true"
+)
+DB_NAME = os.environ.get("DB_NAME", "langgraph-test")
+CHECKPOINT_CLXN_NAME = "fanout_checkpoints"
+WRITES_CLXN_NAME = "fanout_writes"
+
+N_SUBJECTS = 10  # increase for benchmarking
+
+
+class OverallState(TypedDict):
+    subjects: list[str]
+    jokes: Annotated[list[str], operator.add]
+
+
+class JokeInput(TypedDict):
+    subject: str
+
+
+class JokeOutput(TypedDict):
+    jokes: list[str]
+
+
+class JokeState(JokeInput, JokeOutput): ...
+
+
+@pytest.fixture
+def joke_subjects():
+    years = [str(2025 - 10 * i) for i in range(N_SUBJECTS)]
+    return {"subjects": years}
+
+
+@pytest.fixture(scope="function")
+def checkpointer_memory() -> Generator[InMemorySaver, None, None]:
+    yield InMemorySaver()
+
+
+@pytest.fixture(scope="function")
+def checkpointer_mongodb() -> Generator[MongoDBSaver, None, None]:
+    with MongoDBSaver.from_conn_string(
+        MONGODB_URI,
+        db_name=DB_NAME,
+        checkpoint_collection_name=CHECKPOINT_CLXN_NAME,
+        writes_collection_name=WRITES_CLXN_NAME,
+    ) as checkpointer:
+        checkpointer.checkpoint_collection.delete_many({})
+        checkpointer.writes_collection.delete_many({})
+        yield checkpointer
+        checkpointer.checkpoint_collection.drop()
+        checkpointer.writes_collection.drop()
+
+
+@pytest.fixture(scope="function")
+async def checkpointer_mongodb_async() -> Generator[AsyncMongoDBSaver, None, None]:
+    async with AsyncMongoDBSaver.from_conn_string(
+        MONGODB_URI,
+        db_name=DB_NAME,
+        checkpoint_collection_name=CHECKPOINT_CLXN_NAME + "_async",
+        writes_collection_name=WRITES_CLXN_NAME + "_async",
+    ) as checkpointer:
+        checkpointer.checkpoint_collection.delete_many({})
+        checkpointer.writes_collection.delete_many({})
+        yield checkpointer
+        checkpointer.checkpoint_collection.drop()
+        checkpointer.writes_collection.drop()
+
+
+@pytest.fixture(autouse=True)
+def disable_langsmith():
+    """Disable LangSmith tracing for all tests"""
+    os.environ["LANGCHAIN_TRACING_V2"] = "false"
+    os.environ["LANGCHAIN_API_KEY"] = ""
+
+
+def test_sync(
+    joke_subjects,
+    checkpointer_mongodb,
+    checkpointer_memory,
+) -> None:
+    checkpointers = {
+        "mongodb": checkpointer_mongodb,
+        "in_memory": checkpointer_memory,
+    }
+
+    def fanout_to_subgraph() -> StateGraph:
+        # Subgraph nodes create a joke
+        def edit(state: JokeInput):
+            return {"subject": f"{state["subject"]}, and cats"}
+
+        def generate(state: JokeInput):
+            return {"jokes": [f"Joke about the year {state['subject']}"]}
+
+        def bump(state: JokeOutput):
+            return {"jokes": [state["jokes"][0] + " and another"]}
+
+        def bump_loop(state: JokeOutput):
+            return END if state["jokes"][0].endswith(" and another" * 10) else "bump"
+
+        subgraph = StateGraph(JokeState, joke_subjects=JokeInput, output=JokeOutput)
+        subgraph.add_node("edit", edit)
+        subgraph.add_node("generate", generate)
+        subgraph.add_node("bump", bump)
+        subgraph.set_entry_point("edit")
+        subgraph.add_edge("edit", "generate")
+        subgraph.add_edge("generate", "bump")
+        subgraph.add_node("bump_loop", bump_loop)
+        subgraph.add_conditional_edges("bump", bump_loop)
+        subgraph.set_finish_point("generate")
+        subgraphc = subgraph.compile()
+
+        # parent graph maps the joke-generating subgraph
+        def fanout(state: OverallState):
+            return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
+
+        parentgraph = StateGraph(OverallState)
+        parentgraph.add_node("generate_joke", subgraphc)
+        parentgraph.add_conditional_edges(START, fanout)
+        parentgraph.add_edge("generate_joke", END)
+        return parentgraph
+
+    print("\n\nBegin test_sync")
+    for cname, checkpointer in checkpointers.items():
+        assert isinstance(checkpointer, BaseCheckpointSaver)
+
+        graphc = fanout_to_subgraph().compile(checkpointer=checkpointer)
+        assert isinstance(graphc.get_graph(), langchain_core.runnables.graph.Graph)
+        config = {"configurable": {"thread_id": cname}}
+        start = time.monotonic()
+        out = [c for c in graphc.stream(joke_subjects, config=config)]
+        assert len(out) == N_SUBJECTS
+        assert isinstance(out[0], dict)
+        assert out[0].keys() == {"generate_joke"}
+        assert set(out[0]["generate_joke"].keys()) == {"jokes"}
+        end = time.monotonic()
+        print(f"{cname}: {end - start:.4f} seconds")
+
+
+async def test_async(
+    joke_subjects, checkpointer_mongodb_async, checkpointer_memory
+) -> None:
+    checkpointers = {
+        "mongodb_async": checkpointer_mongodb_async,
+        "in_memory_async": checkpointer_memory,
+    }
+
+    async def fanout_to_subgraph() -> StateGraph:
+        # Subgraph nodes create a joke
+        async def edit(state: JokeInput):
+            subject = state["subject"]
+            return {"subject": f"{subject}, and cats"}
+
+        async def generate(state: JokeInput):
+            return {"jokes": [f"Joke about the year {state['subject']}"]}
+
+        async def bump(state: JokeOutput):
+            return {"jokes": [state["jokes"][0] + " and another"]}
+
+        async def bump_loop(state: JokeOutput):
+            return END if state["jokes"][0].endswith(" and another" * 10) else "bump"
+
+        subgraph = StateGraph(JokeState, joke_subjects=JokeInput, output=JokeOutput)
+        subgraph.add_node("edit", edit)
+        subgraph.add_node("generate", generate)
+        subgraph.add_node("bump", bump)
+        subgraph.set_entry_point("edit")
+        subgraph.add_edge("edit", "generate")
+        subgraph.add_edge("generate", "bump")
+        subgraph.add_conditional_edges("bump", bump_loop)
+        subgraph.set_finish_point("generate")
+        subgraphc = subgraph.compile()
+
+        # parent graph maps the joke-generating subgraph
+        async def fanout(state: OverallState):
+            return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
+
+        parentgraph = StateGraph(OverallState)
+        parentgraph.add_node("generate_joke", subgraphc)
+        parentgraph.add_conditional_edges(START, fanout)
+        parentgraph.add_edge("generate_joke", END)
+        return parentgraph
+
+    print("\n\nBegin test_async")
+    for cname, checkpointer in checkpointers.items():
+        assert isinstance(checkpointer, BaseCheckpointSaver)
+
+        graphc = (await fanout_to_subgraph()).compile(checkpointer=checkpointer)
+        config = {"configurable": {"thread_id": cname}}
+        start = time.monotonic()
+        out = [c async for c in graphc.astream(joke_subjects, config=config)]
+        assert len(out) == N_SUBJECTS
+        assert isinstance(out[0], dict)
+        assert out[0].keys() == {"generate_joke"}
+        assert set(out[0]["generate_joke"].keys()) == {"jokes"}
+        end = time.monotonic()
+        print(f"{cname}: {end - start:.4f} seconds")