WIP Subgraphs

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

@@ -10,6 +10,7 @@ from psycopg.types.json import Jsonb
 from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
@@ -292,9 +293,14 @@ class AsyncPostgresSaver(BasePostgresSaver):
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
         """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
         async with self._cursor(pipeline=True) as cur:
             await cur.executemany(
-                self.UPSERT_CHECKPOINT_WRITES_SQL,
+                query,
                 await asyncio.to_thread(
                     self._dump_writes,
                     config["configurable"]["thread_id"],

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -108,6 +108,15 @@ UPSERT_CHECKPOINTS_SQL = """
 """
 
 UPSERT_CHECKPOINT_WRITES_SQL = """
+    INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
+    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+    ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
+        channel = EXCLUDED.channel,
+        type = EXCLUDED.type,
+        blob = EXCLUDED.blob;
+"""
+
+INSERT_CHECKPOINT_WRITES_SQL = """
     INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
     VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
     ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
@@ -120,6 +129,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
     UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
     UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
     UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
+    INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
 
     jsonplus_serde = JsonPlusSerializer()
 

libs/langgraph/langgraph/pregel/loop.py

@@ -242,6 +242,14 @@ class PregelLoop:
     ) -> bool:
         """Execute a single iteration of the Pregel loop.
         Returns True if more iterations are needed."""
+        print(
+            "tick",
+            self.config.get("configurable", {}).get("checkpoint_ns"),
+            self.step,
+            self.status,
+            self.input is INPUT_RESUMING,
+            self.input is INPUT_DONE,
+        )
 
         if self.status != "pending":
             raise RuntimeError("Cannot tick when status is no longer 'pending'")
@@ -345,7 +353,13 @@ class PregelLoop:
                     continue
                 if task := self.tasks.get(tid):
                     if k == SCHEDULED:
-                        self.tasks[tid] = task._replace(scheduled=True)
+                        if v == max(
+                            self.checkpoint["versions_seen"]
+                            .get(INTERRUPT, {})
+                            .values(),
+                            default=None,
+                        ):
+                            self.tasks[tid] = task._replace(scheduled=True)
                     else:
                         task.writes.append((k, v))
             # print output for any tasks we applied previous writes to

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -147,4 +147,5 @@ class KafkaExecutor(AbstractAsyncContextManager):
         task_id: str,
         writes: list[tuple[str, Any]],
     ) -> None:
+        print("put_writes", task_id, writes)
         return submit(self.graph.checkpointer.aput_writes, config, writes, task_id)

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -6,7 +6,7 @@ import aiokafka
 from langchain_core.runnables import ensure_config
 
 import langgraph.scheduler.kafka.serde as serde
-from langgraph.constants import CONFIG_KEY_DEDUPE_TASKS, SCHEDULED
+from langgraph.constants import CONFIG_KEY_DEDUPE_TASKS, INTERRUPT, SCHEDULED
 from langgraph.pregel import Pregel
 from langgraph.pregel.loop import INPUT_RESUMING, AsyncPregelLoop
 from langgraph.pregel.types import RetryPolicy
@@ -145,4 +145,19 @@ class KafkaOrchestrator(AbstractAsyncContextManager):
                     await asyncio.gather(*futures)
                     # mark as scheduled
                     for task in new_tasks:
-                        loop.put_writes(task.id, [(SCHEDULED, None)])
+                        loop.put_writes(
+                            task.id,
+                            [
+                                (
+                                    SCHEDULED,
+                                    max(
+                                        loop.checkpoint["versions_seen"]
+                                        .get(INTERRUPT, {})
+                                        .values(),
+                                        default=None,
+                                    ),
+                                )
+                            ],
+                        )
+            else:
+                pass

libs/scheduler-kafka/tests/test_subgraph.py

@@ -0,0 +1,245 @@
+import asyncio
+import functools
+import re
+from typing import Callable, Literal, Optional, ParamSpec, TypeVar, Union, cast
+
+import anyio
+import pytest
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+from langchain_core.language_models.fake_chat_models import (
+    FakeMessagesListChatModel,
+)
+from langchain_core.messages import AIMessage, HumanMessage, ToolCall
+from langchain_core.tools import tool
+
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.constants import END, START
+from langgraph.graph import MessagesState
+from langgraph.graph.state import StateGraph
+from langgraph.pregel import Pregel
+from langgraph.scheduler.kafka import serde
+from langgraph.scheduler.kafka.executor import KafkaExecutor
+from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+
+pytestmark = pytest.mark.anyio
+C = ParamSpec("C")
+R = TypeVar("R")
+
+
+class AnyStr(str):
+    def __init__(self, prefix: Union[str, re.Pattern] = "") -> None:
+        super().__init__()
+        self.prefix = prefix
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, str) and (
+            other.startswith(self.prefix)
+            if isinstance(self.prefix, str)
+            else self.prefix.match(other)
+        )
+
+    def __hash__(self) -> int:
+        return hash((str(self), self.prefix))
+
+
+def timeout(delay: int):
+    def decorator(func: Callable[C, R]) -> Callable[C, R]:
+        @functools.wraps(func)
+        async def new_func(*args: C.args, **kwargs: C.kwargs) -> R:
+            async with asyncio.timeout(delay):
+                return await func(*args, **kwargs)
+
+        return new_func
+
+    return decorator
+
+
+def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel:
+    # copied from test_weather_subgraph
+
+    # setup subgraph
+
+    @tool
+    def get_weather(city: str):
+        """Get the weather for a specific city"""
+        return f"I'ts sunny in {city}!"
+
+    weather_model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="",
+                tool_calls=[
+                    ToolCall(
+                        id="tool_call123",
+                        name="get_weather",
+                        args={"city": "San Francisco"},
+                    )
+                ],
+            )
+        ]
+    )
+
+    class SubGraphState(MessagesState):
+        city: str
+
+    def model_node(state: SubGraphState):
+        result = weather_model.invoke(state["messages"])
+        return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]}
+
+    def weather_node(state: SubGraphState):
+        result = get_weather.invoke({"city": state["city"]})
+        return {"messages": [{"role": "assistant", "content": result}]}
+
+    subgraph = StateGraph(SubGraphState)
+    subgraph.add_node(model_node)
+    subgraph.add_node(weather_node)
+    subgraph.add_edge(START, "model_node")
+    subgraph.add_edge("model_node", "weather_node")
+    subgraph.add_edge("weather_node", END)
+    subgraph = subgraph.compile(interrupt_before=["weather_node"])
+
+    # setup main graph
+
+    class RouterState(MessagesState):
+        route: Literal["weather", "other"]
+
+    router_model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="",
+                tool_calls=[
+                    ToolCall(
+                        id="tool_call123",
+                        name="router",
+                        args={"dest": "weather"},
+                    )
+                ],
+            )
+        ]
+    )
+
+    def router_node(state: RouterState):
+        system_message = "Classify the incoming query as either about weather or not."
+        messages = [{"role": "system", "content": system_message}] + state["messages"]
+        route = router_model.invoke(messages)
+        return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]}
+
+    def normal_llm_node(state: RouterState):
+        return {"messages": [AIMessage("Hello!")]}
+
+    def route_after_prediction(state: RouterState):
+        if state["route"] == "weather":
+            return "weather_graph"
+        else:
+            return "normal_llm_node"
+
+    async def weather_graph(state: RouterState):
+        return await subgraph.ainvoke(state)
+
+    graph = StateGraph(RouterState)
+    graph.add_node(router_node)
+    graph.add_node(normal_llm_node)
+    graph.add_node("weather_graph", weather_graph)
+    graph.add_edge(START, "router_node")
+    graph.add_conditional_edges("router_node", route_after_prediction)
+    graph.add_edge("normal_llm_node", END)
+    graph.add_edge("weather_graph", END)
+
+    return graph.compile(checkpointer=checkpointer)
+
+
+@timeout(10)
+async def test_subgraph_w_interrupt(
+    topics: Topics, checkpointer: BaseCheckpointSaver
+) -> None:
+    input = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
+    config = {"configurable": {"thread_id": "1"}}
+    graph = mk_weather_graph(checkpointer)
+    n_orch_msgs = 0
+    n_exec_msgs = 0
+    errors = []
+    scope: Optional[anyio.CancelScope] = None
+
+    async def orchestrator(expected: int) -> None:
+        nonlocal n_orch_msgs
+        async with KafkaOrchestrator(graph, topics) as orch:
+            async for msgs in orch:
+                n_orch_msgs += len(msgs)
+                print("orch", n_orch_msgs, msgs)
+                if n_orch_msgs == expected:
+                    break
+
+    async def executor(expected: int) -> None:
+        nonlocal n_exec_msgs
+        async with KafkaExecutor(graph, topics) as exec:
+            async for msgs in exec:
+                n_exec_msgs += len(msgs)
+                print("exec", n_exec_msgs, msgs)
+                if n_exec_msgs == expected:
+                    break
+
+    async def error_consumer() -> None:
+        async with AIOKafkaConsumer(topics.error) as consumer:
+            async for msg in consumer:
+                errors.append(msg)
+                if scope:
+                    scope.cancel()
+
+    # start error consumer
+    error_task = asyncio.create_task(error_consumer(), name="error_consumer")
+
+    # start a new run
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=input, config=config),
+        )
+
+    # run the orchestrator and executor
+    async with anyio.create_task_group() as tg:
+        scope = tg.cancel_scope
+        tg.start_soon(orchestrator, 4, name="orchestrator")
+        tg.start_soon(executor, 3, name="executor")
+
+    # check no errors
+    assert not errors
+
+    # check interrupted state
+    state = await graph.aget_state(config)
+    assert n_orch_msgs == 4
+    assert n_exec_msgs == 3
+    assert state.next == ("weather_graph",)
+    assert state.values == {
+        "messages": [HumanMessage(id=AnyStr(), content="what's the weather in sf")],
+        "route": "weather",
+    }
+
+    # resume the thread
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=None, config=config),
+        )
+
+    # run the orchestrator and executor
+    async with anyio.create_task_group() as tg:
+        scope = tg.cancel_scope
+        tg.start_soon(orchestrator, 6, name="orchestrator")
+        tg.start_soon(executor, 4, name="executor")
+
+    # check no errors
+    assert not errors
+
+    # check final state
+    state = await graph.aget_state(config)
+    assert n_orch_msgs == 6
+    assert n_exec_msgs == 4
+    assert state.next == ()
+    assert state.values == {
+        "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+        "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+        "query": "analyzed: query: analyzed: query: what is weather in sf",
+    }
+
+    error_task.cancel()