Interruption issue

[hung-phan]

Hi there,

I am trying to use the library but encounter some issues. The interruption workflow i have is working well with memory checkpointer, but when i switch to AsyncRedisCheck i don't see interruption message get return by langgraph any more.

My code is sth like

def get_checkpointer(checkpointer_type: CheckpointerType) -> Callable[[], Coroutine[None, None, Checkpointer]]:
    return lambda: _get_checkpointer(checkpointer_type)


def _get_redis_client():
    return Redis(...)


@cached(cache=Cache.MEMORY)
async def _get_checkpointer(checkpointer_type: CheckpointerType) -> Checkpointer:
    match checkpointer_type:
        case CheckpointerType.SHALLOW:
            async with AsyncShallowRedisSaver(redis_client=_get_redis_client()) as checkpointer:
                await checkpointer.asetup()
                return checkpointer

        case CheckpointerType.DEEP:
            async with AsyncRedisSaver(redis_client=_get_redis_client()) as checkpointer:
                await checkpointer.asetup()
                return checkpointer

        case CheckpointerType.MEMORY:
            if not is_local():
                raise ValueError("Memory checkpointer is not supported in production environment due to memory leak.")

            return MemorySaver()

@dataclass
class WorkflowBuilder:
    name: str
    workflow_state_class: type[WorkflowState] | None = None
    checkpointer_func: Callable[..., Coroutine[None, None, Checkpointer]] | None = None

    async def build(self) -> "Workflow":
        self.__validate_nodes()

        builder = StateGraph(self.workflow_state_class)

        self.__add_nodes(builder)

        try:
            self.__add_start_edges(builder)
            self.__add_intra_edges(builder)
            self.__add_end_edges(builder)

            return Workflow(
                name=self.name,
                compiled_workflow=builder.compile(
                    checkpointer=(await self.checkpointer_func() if callable(self.checkpointer_func) else MemorySaver())
                ),
            )

@dataclass
class Workflow:
    compiled_workflow: CompiledStateGraph

    async def stream_using_workflow_state(
        self,
        workflow_state: MutableMapping[str, Any],
        *,
        config: RunnableConfig | None = None
    ) -> AsyncGenerator[WorkflowStreamEvent, Any]:
        config = config or RunnableConfig()
        config.setdefault("configurable", {}).setdefault("thread_id", str(ULID()))

        thread_id: str = config["configurable"]["thread_id"]
        input_data: Command = Command(update=workflow_state)

        try:
            while True:
                async for event in self.compiled_workflow.astream(
                    input=input_data,
                    config=config,
                    stream_mode=["custom", "updates"],
                    subgraphs=include_nested_workflows,
                ):
                  print(event)

Then i build a simple graph, but some how it doesn't return anything when i start interrupt

[bsbodden]

@hung-phan Could you create a minimal example showing the issue? That way, I can build an integration test. Thanks!

[hung-phan]

The example here won't execute with langgraph-redis-checkpoint. It will hang and waiting for sth.

cc @bsbodden

    from langchain_core.messages import HumanMessage
    from langgraph.checkpoint.redis import AsyncRedisSaver
    from langgraph.constants import END, START
    from langgraph.graph import MessagesState, StateGraph
    from redis.asyncio import Redis

    def func1(_: MessagesState) -> MessagesState:
        return {"messages": [HumanMessage(id="1", content="func1", role="user")]}

    def func2(_: MessagesState) -> MessagesState:
        value = interrupt({"message": "Interrupt from func2"})

        return {"messages": [HumanMessage(id="2", content=f"func2+{value}", role="user")]}

    builder = StateGraph(MessagesState)

    builder.add_node("func1", func1)
    builder.add_node("func2", func2)

    builder.add_edge(START, "func1")
    builder.add_edge("func1", "func2")
    builder.add_edge("func2", END)

    async with (
        Redis(host="localhost", port=6379, ssl=False) as client,
        AsyncRedisSaver(redis_client=client) as checkpointer,
    ):
        await checkpointer.asetup()

        compiled_graph = builder.compile(checkpointer=checkpointer)

        config = RunnableConfig(configurable={"thread_id": "my-thread-id"})

        data = await compiled_graph.ainvoke({"messages": []}, config=config)
        assert data == {"messages": [HumanMessage(id="1", content="func1", role="user")]}

        data = await compiled_graph.ainvoke(Command(resume="resume from command"), config=config)
        assert data == {
            "messages": [
                HumanMessage(id="1", content="func1", role="user"),
                HumanMessage(id="2", content="func2+resume from command", role="user"),
            ]
        }

[abrookins]

Thanks, @hung-phan. That code didn't work as-is, but I was able to modify it and then reduce the problem further. It looks like something is going wrong during 'shutdown' if a graph raises any exception:

import asyncio
import sys
import uuid
from typing import Optional
from typing_extensions import TypedDict

from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command

from langgraph.checkpoint.redis import RedisSaver
from langgraph.checkpoint.redis.aio import AsyncRedisSaver


class State(TypedDict):
    """The graph state."""

    foo: str
    human_value: Optional[str]
    """Human value will be updated using an interrupt."""


async def node(state: State):
    # answer = interrupt(
    #     # This value will be sent to the client
    #     # as part of the interrupt information.
    #     "what is your age?"
    # )
    # print(f"> Received an input from the interrupt: {answer}")
    # return {"human_value": answer}
    raise RuntimeError("test")


builder = StateGraph(State)
builder.add_node("node", node)
builder.add_edge(START, "node")


async def main_redis():
    async with AsyncRedisSaver(redis_url="redis://localhost:6379") as checkpointer:
        graph = builder.compile(checkpointer=checkpointer)

        config = {
            "configurable": {
                "thread_id": uuid.uuid4(),
            }
        }

        async for chunk in graph.astream({"foo": "abc"}, config):
            print(chunk)

        command = Command(resume="some input from a human!!!")

        async for chunk in graph.astream(command, config):
            print(chunk)


async def main_memory():
    graph = builder.compile(checkpointer=MemorySaver())

    config = {
        "configurable": {
            "thread_id": uuid.uuid4(),
        }
    }

    async for chunk in graph.astream({"foo": "abc"}, config):
        print(chunk)

    command = Command(resume="some input from a human!!!")

    async for chunk in graph.astream(command, config):
        print(chunk)


if __name__ == "__main__":
    which_saver = sys.argv[1] if len(sys.argv) > 1 else "memory"
    if which_saver == "memory":
        asyncio.run(main_memory())
    elif which_saver == "redis":
        asyncio.run(main_redis())
    else:
        raise ValueError(f"Invalid saver: {which_saver}")

We'll look into it ASAP!

[bsbodden]

@hung-phan could you test and let us know if the issue is now gone?