diff --git a/libs/scheduler-kafka/Makefile b/libs/scheduler-kafka/Makefile index 7462627e1..5d899a3ea 100644 --- a/libs/scheduler-kafka/Makefile +++ b/libs/scheduler-kafka/Makefile @@ -10,14 +10,16 @@ start-services: stop-services: docker compose -f tests/compose.yml down +TEST_PATH ?= . + test: - make start-services && poetry run pytest; \ + make start-services && poetry run pytest $(TEST_PATH); \ EXIT_CODE=$$?; \ make stop-services; \ exit $$EXIT_CODE test_watch: - make start-services && poetry run ptw .; \ + make start-services && poetry run ptw . -- $(TEST_PATH); \ EXIT_CODE=$$?; \ make stop-services; \ exit $$EXIT_CODE diff --git a/libs/scheduler-kafka/README.md b/libs/scheduler-kafka/README.md index 2d934d59f..637a337dd 100644 --- a/libs/scheduler-kafka/README.md +++ b/libs/scheduler-kafka/README.md @@ -26,7 +26,7 @@ import asyncio import logging import os -from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator +from langgraph.scheduler.kafka.orchestrator import AsyncKafkaOrchestrator from langgraph.scheduler.kafka.types import Topics from your_lib import graph # graph expected to be a compiled LangGraph graph @@ -40,7 +40,7 @@ topics = Topics( ) async def main(): - async with KafkaOrchestrator(graph, topics) as orch: + async with AsyncKafkaOrchestrator(graph, topics) as orch: async for msgs in orch: logger.info('Procesed %d messages', len(msgs)) @@ -56,7 +56,7 @@ import asyncio import logging import os -from langgraph.scheduler.kafka.executor import KafkaExecutor +from langgraph.scheduler.kafka.executor import AsyncKafkaExecutor from langgraph.scheduler.kafka.types import Topics from your_lib import graph # graph expected to be a compiled LangGraph graph @@ -70,7 +70,7 @@ topics = Topics( ) async def main(): - async with KafkaExecutor(graph, topics) as orch: + async with AsyncKafkaExecutor(graph, topics) as orch: async for msgs in orch: logger.info('Procesed %d messages', len(msgs)) @@ -89,6 +89,8 @@ python executor.py & ## Configuration +We offer sync and async versions of the orchestrator and executor, `KafkaOrchestrator` and `AsyncKafkaOrchestrator`, and `KafkaExecutor` and `AsyncKafkaExecutor` respectively. The async versions are recommended, especially if you want to process tasks in batches. With the async classes we recommend using `uvloop` for better performance. + You can pass any of the following values as `kwargs` to either `KafkaOrchestrator` or `KafkaExecutor` to configure the consumer: - batch_max_n (int): Maximum number of messages to include in a single batch. Default: 10. @@ -131,3 +133,7 @@ By default the orchestrator and executor will attempt to connect to a Kafka brok - connections_max_idle_ms (int): Close idle connections after the number of milliseconds specified by this config. Specifying `None` will disable idle checks. Default: 540000 (9 minutes). + +### Custom consumer/producer + +Both the orchestrator and executor accept a `consumer` and `producer` argument, which should implement the `Consumer` or `Producer` protocols respectively. We expect the consumer to have auto-commit disabled, and the producer and consumer to have no serializers/deserializers set. diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/default_async.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/default_async.py index ce6703a03..6b6e762bc 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/default_async.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/default_async.py @@ -1,15 +1,8 @@ -import dataclasses -from typing import Any, Sequence - import aiokafka class DefaultAsyncConsumer(aiokafka.AIOKafkaConsumer): - async def getmany( - self, timeout_ms: int, max_records: int - ) -> dict[str, Sequence[dict[str, Any]]]: - batch = await super().getmany(timeout_ms=timeout_ms, max_records=max_records) - return {t: [dataclasses.asdict(m) for m in msgs] for t, msgs in batch.items()} + pass class DefaultAsyncProducer(aiokafka.AIOKafkaProducer): diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/default_sync.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/default_sync.py new file mode 100644 index 000000000..30dda7bf1 --- /dev/null +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/default_sync.py @@ -0,0 +1,39 @@ +import concurrent.futures +from typing import Optional, Sequence + +from kafka import KafkaConsumer, KafkaProducer +from langgraph.scheduler.kafka.types import ConsumerRecord, TopicPartition + + +class DefaultConsumer(KafkaConsumer): + def getmany( + self, timeout_ms: int, max_records: int + ) -> dict[TopicPartition, Sequence[ConsumerRecord]]: + return self.poll(timeout_ms=timeout_ms, max_records=max_records) + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + +class DefaultProducer(KafkaProducer): + def send( + self, + topic: str, + *, + key: Optional[bytes] = None, + value: Optional[bytes] = None, + ) -> concurrent.futures.Future: + fut = concurrent.futures.Future() + kfut = super().send(topic, key=key, value=value) + kfut.add_callback(fut.set_result) + kfut.add_errback(fut.set_exception) + return fut + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index f7bceb86a..8e586c96e 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,5 +1,11 @@ import asyncio -from contextlib import AbstractAsyncContextManager, AsyncExitStack +import concurrent.futures +from contextlib import ( + AbstractAsyncContextManager, + AbstractContextManager, + AsyncExitStack, + ExitStack, +) from functools import partial from typing import Any, Optional, Sequence @@ -12,23 +18,29 @@ from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR, NS_END, NS_SEP from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound from langgraph.pregel import Pregel from langgraph.pregel.algo import prepare_single_task -from langgraph.pregel.executor import AsyncBackgroundExecutor, Submit -from langgraph.pregel.manager import AsyncChannelsManager +from langgraph.pregel.executor import ( + AsyncBackgroundExecutor, + BackgroundExecutor, + Submit, +) +from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.runner import PregelRunner from langgraph.pregel.types import RetryPolicy -from langgraph.scheduler.kafka.retry import aretry +from langgraph.scheduler.kafka.retry import aretry, retry from langgraph.scheduler.kafka.types import ( AsyncConsumer, AsyncProducer, + Consumer, ErrorMessage, MessageToExecutor, MessageToOrchestrator, + Producer, Topics, ) from langgraph.utils.config import patch_configurable -class KafkaExecutor(AbstractAsyncContextManager): +class AsyncKafkaExecutor(AbstractAsyncContextManager): consumer: AsyncConsumer producer: AsyncProducer @@ -93,7 +105,7 @@ class KafkaExecutor(AbstractAsyncContextManager): timeout_ms=self.batch_max_ms, max_records=self.batch_max_n ) msgs: list[MessageToExecutor] = [ - serde.loads(msg["value"]) for msgs in recs.values() for msg in msgs + serde.loads(msg.value) for msgs in recs.values() for msg in msgs ] # process batch await asyncio.gather(*(self.each(msg) for msg in msgs)) @@ -189,7 +201,7 @@ class KafkaExecutor(AbstractAsyncContextManager): pass else: # task was not found - await self.graph.checkpointer.put_writes( + await self.graph.checkpointer.aput_writes( msg["config"], [(ERROR, TaskNotFound())] ) # notify orchestrator @@ -220,3 +232,195 @@ class KafkaExecutor(AbstractAsyncContextManager): writes: list[tuple[str, Any]], ) -> None: return submit(self.graph.checkpointer.aput_writes, config, writes, task_id) + + +class KafkaExecutor(AbstractContextManager): + consumer: Consumer + + producer: Producer + + def __init__( + self, + graph: Pregel, + topics: Topics, + *, + batch_max_n: int = 10, + batch_max_ms: int = 1000, + retry_policy: Optional[RetryPolicy] = None, + consumer: Optional[Consumer] = None, + producer: Optional[Producer] = None, + **kwargs: Any, + ) -> None: + self.graph = graph + self.topics = topics + self.stack = ExitStack() + self.kwargs = kwargs + self.consumer = consumer + self.producer = producer + self.batch_max_n = batch_max_n + self.batch_max_ms = batch_max_ms + self.retry_policy = retry_policy + + def __enter__(self) -> Self: + self.subgraphs = dict(self.graph.get_subgraphs(recurse=True)) + self.submit = self.stack.enter_context(BackgroundExecutor({})) + if self.consumer is None: + from langgraph.scheduler.kafka.default_sync import DefaultConsumer + + self.consumer = self.stack.enter_context( + DefaultConsumer( + self.topics.executor, + auto_offset_reset="earliest", + group_id="executor", + enable_auto_commit=False, + **self.kwargs, + ) + ) + if self.producer is None: + from langgraph.scheduler.kafka.default_sync import DefaultProducer + + self.producer = self.stack.enter_context( + DefaultProducer( + **self.kwargs, + ) + ) + return self + + def __exit__(self, *args: Any) -> None: + return self.stack.__exit__(*args) + + def __iter__(self) -> Self: + return self + + def __next__(self) -> Sequence[MessageToExecutor]: + # wait for next batch + recs = self.consumer.getmany( + timeout_ms=self.batch_max_ms, max_records=self.batch_max_n + ) + msgs: list[MessageToExecutor] = [ + serde.loads(msg.value) for msgs in recs.values() for msg in msgs + ] + # process batch + concurrent.futures.wait(self.submit(self.each, msg) for msg in msgs) + # commit offsets + self.consumer.commit() + # return message + return msgs + + def each(self, msg: MessageToExecutor) -> None: + try: + retry(self.retry_policy, self.attempt, msg) + except CheckpointNotLatest: + pass + except GraphDelegate as exc: + for arg in exc.args: + fut = self.producer.send( + self.topics.orchestrator, + value=serde.dumps( + MessageToOrchestrator( + config=arg["config"], + input=orjson.Fragment( + self.graph.checkpointer.serde.dumps(arg["input"]) + ), + finally_executor=[msg], + ) + ), + # use thread_id, checkpoint_ns as partition key + key=serde.dumps( + ( + arg["config"]["configurable"]["thread_id"], + arg["config"]["configurable"].get("checkpoint_ns"), + ) + ), + ) + fut.result() + except Exception as exc: + fut = self.producer.send( + self.topics.error, + value=serde.dumps( + ErrorMessage( + topic=self.topics.executor, + msg=msg, + error=repr(exc), + ) + ), + ) + fut.result() + + def attempt(self, msg: MessageToExecutor) -> None: + # find graph + if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): + # remove task_ids from checkpoint_ns + recast_checkpoint_ns = NS_SEP.join( + part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) + ) + # find the subgraph with the matching name + if recast_checkpoint_ns in self.subgraphs: + graph = self.subgraphs[recast_checkpoint_ns] + else: + raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + else: + graph = self.graph + # process message + saved = self.graph.checkpointer.get_tuple( + patch_configurable(msg["config"], {"checkpoint_id": None}) + ) + if saved is None: + raise RuntimeError("Checkpoint not found") + if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: + raise CheckpointNotLatest() + with ChannelsManager( + graph.channels, saved.checkpoint, msg["config"], self.graph.store + ) as (channels, managed), BackgroundExecutor({}) as submit: + if task := prepare_single_task( + msg["task"]["path"], + msg["task"]["id"], + checkpoint=saved.checkpoint, + processes=graph.nodes, + channels=channels, + managed=managed, + config=patch_configurable(msg["config"], {CONFIG_KEY_DELEGATE: True}), + step=saved.metadata["step"] + 1, + for_execution=True, + checkpointer=self.graph.checkpointer, + ): + # execute task, saving writes + runner = PregelRunner( + submit=submit, + put_writes=partial(self._put_writes, submit, msg["config"]), + ) + for _ in runner.tick([task], reraise=False): + pass + else: + # task was not found + self.graph.checkpointer.put_writes( + msg["config"], [(ERROR, TaskNotFound())] + ) + # notify orchestrator + fut = self.producer.send( + self.topics.orchestrator, + value=serde.dumps( + MessageToOrchestrator( + input=None, + config=msg["config"], + finally_executor=msg.get("finally_executor"), + ) + ), + # use thread_id, checkpoint_ns as partition key + key=serde.dumps( + ( + msg["config"]["configurable"]["thread_id"], + msg["config"]["configurable"].get("checkpoint_ns"), + ) + ), + ) + fut.result() + + def _put_writes( + self, + submit: Submit, + config: RunnableConfig, + task_id: str, + writes: list[tuple[str, Any]], + ) -> None: + return submit(self.graph.checkpointer.put_writes, config, writes, task_id) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index e594d88ba..a04d49537 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -1,5 +1,11 @@ import asyncio -from contextlib import AbstractAsyncContextManager, AsyncExitStack +import concurrent.futures +from contextlib import ( + AbstractAsyncContextManager, + AbstractContextManager, + AsyncExitStack, + ExitStack, +) from typing import Any, Optional from langchain_core.runnables import ensure_config @@ -16,22 +22,25 @@ from langgraph.constants import ( ) from langgraph.errors import CheckpointNotLatest, GraphInterrupt from langgraph.pregel import Pregel -from langgraph.pregel.loop import AsyncPregelLoop +from langgraph.pregel.executor import BackgroundExecutor, Submit +from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop from langgraph.pregel.types import RetryPolicy -from langgraph.scheduler.kafka.retry import aretry +from langgraph.scheduler.kafka.retry import aretry, retry from langgraph.scheduler.kafka.types import ( AsyncConsumer, AsyncProducer, + Consumer, ErrorMessage, ExecutorTask, MessageToExecutor, MessageToOrchestrator, + Producer, Topics, ) from langgraph.utils.config import patch_configurable -class KafkaOrchestrator(AbstractAsyncContextManager): +class AsyncKafkaOrchestrator(AbstractAsyncContextManager): consumer: AsyncConsumer producer: AsyncProducer @@ -95,7 +104,7 @@ class KafkaOrchestrator(AbstractAsyncContextManager): timeout_ms=self.batch_max_ms, max_records=self.batch_max_n ) # dedupe messages, eg. if multiple nodes finish around same time - uniq = set(msg["value"] for msgs in recs.values() for msg in msgs) + uniq = set(msg.value for msgs in recs.values() for msg in msgs) msgs: list[MessageToOrchestrator] = [serde.loads(msg) for msg in uniq] # process batch await asyncio.gather(*(self.each(msg) for msg in msgs)) @@ -213,3 +222,183 @@ class KafkaOrchestrator(AbstractAsyncContextManager): ) # wait for messages to be sent await asyncio.gather(*futs) + + +class KafkaOrchestrator(AbstractContextManager): + consumer: Consumer + + producer: Producer + + submit: Submit + + def __init__( + self, + graph: Pregel, + topics: Topics, + batch_max_n: int = 10, + batch_max_ms: int = 1000, + retry_policy: Optional[RetryPolicy] = None, + consumer: Optional[Consumer] = None, + producer: Optional[Producer] = None, + **kwargs: Any, + ) -> None: + self.graph = graph + self.topics = topics + self.stack = ExitStack() + self.kwargs = kwargs + self.consumer = consumer + self.producer = producer + self.batch_max_n = batch_max_n + self.batch_max_ms = batch_max_ms + self.retry_policy = retry_policy + + def __enter__(self) -> Self: + self.subgraphs = dict(self.graph.get_subgraphs(recurse=True)) + self.submit = self.stack.enter_context(BackgroundExecutor({})) + if self.consumer is None: + from langgraph.scheduler.kafka.default_sync import DefaultConsumer + + self.consumer = self.stack.enter_context( + DefaultConsumer( + self.topics.orchestrator, + auto_offset_reset="earliest", + group_id="orchestrator", + enable_auto_commit=False, + **self.kwargs, + ) + ) + if self.producer is None: + from langgraph.scheduler.kafka.default_sync import DefaultProducer + + self.producer = self.stack.enter_context( + DefaultProducer( + **self.kwargs, + ) + ) + return self + + def __exit__(self, *args: Any) -> None: + return self.stack.__exit__(*args) + + def __iter__(self) -> Self: + return self + + def __next__(self) -> list[MessageToOrchestrator]: + # wait for next batch + recs = self.consumer.getmany( + timeout_ms=self.batch_max_ms, max_records=self.batch_max_n + ) + print("orch.__next__", recs) + # dedupe messages, eg. if multiple nodes finish around same time + uniq = set(msg.value for msgs in recs.values() for msg in msgs) + msgs: list[MessageToOrchestrator] = [serde.loads(msg) for msg in uniq] + # process batch + concurrent.futures.wait(self.submit(self.each, msg) for msg in msgs) + # commit offsets + self.consumer.commit() + # return message + return msgs + + def each(self, msg: MessageToOrchestrator) -> None: + try: + retry(self.retry_policy, self.attempt, msg) + except CheckpointNotLatest: + pass + except GraphInterrupt: + pass + except Exception as exc: + fut = self.producer.send( + self.topics.error, + value=serde.dumps( + ErrorMessage( + topic=self.topics.orchestrator, + msg=msg, + error=repr(exc), + ) + ), + ) + fut.result() + + def attempt(self, msg: MessageToOrchestrator) -> None: + # find graph + if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): + # remove task_ids from checkpoint_ns + recast_checkpoint_ns = NS_SEP.join( + part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) + ) + # find the subgraph with the matching name + if recast_checkpoint_ns in self.subgraphs: + graph = self.subgraphs[recast_checkpoint_ns] + else: + raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + else: + graph = self.graph + # process message + with SyncPregelLoop( + msg["input"], + config=ensure_config(msg["config"]), + stream=None, + store=self.graph.store, + checkpointer=self.graph.checkpointer, + nodes=graph.nodes, + specs=graph.channels, + output_keys=graph.output_channels, + stream_keys=graph.stream_channels, + ) as loop: + if loop.tick( + input_keys=graph.input_channels, + interrupt_after=graph.interrupt_after_nodes, + interrupt_before=graph.interrupt_before_nodes, + ): + # wait for checkpoint to be saved + if hasattr(loop, "_put_checkpoint_fut"): + loop._put_checkpoint_fut.result() + # schedule any new tasks + if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]: + # send messages to executor + futures = [ + self.producer.send( + self.topics.executor, + value=serde.dumps( + MessageToExecutor( + config=patch_configurable( + loop.config, + { + **loop.checkpoint_config["configurable"], + CONFIG_KEY_DEDUPE_TASKS: True, + CONFIG_KEY_ENSURE_LATEST: True, + }, + ), + task=ExecutorTask(id=task.id, path=task.path), + finally_executor=msg.get("finally_executor"), + ) + ), + ) + for task in new_tasks + ] + # wait for messages to be sent + concurrent.futures.wait(futures) + # mark as scheduled + for task in new_tasks: + loop.put_writes( + task.id, + [ + ( + SCHEDULED, + max( + loop.checkpoint["versions_seen"] + .get(INTERRUPT, {}) + .values(), + default=None, + ), + ) + ], + ) + elif loop.status == "done" and msg.get("finally_executor"): + # schedule any finally_executor tasks + futs = [ + self.producer.send(self.topics.executor, value=serde.dumps(m)) + for m in msg["finally_executor"] + ] + # wait for messages to be sent + concurrent.futures.wait(futs) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py index 6f3137e68..bb80047f8 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py @@ -1,6 +1,7 @@ import asyncio import logging import random +import time from typing import Awaitable, Callable, Optional from typing_extensions import ParamSpec @@ -11,6 +12,49 @@ logger = logging.getLogger(__name__) P = ParamSpec("P") +def retry( + retry_policy: Optional[RetryPolicy], + func: Callable[P, None], + *args: P.args, + **kwargs: P.kwargs, +) -> None: + """Run a task asynchronously with retries.""" + interval = retry_policy.initial_interval if retry_policy else 0 + attempts = 0 + while True: + try: + func(*args, **kwargs) + # if successful, end + break + except Exception as exc: + if retry_policy is None: + raise + # increment attempts + attempts += 1 + # check if we should retry + if callable(retry_policy.retry_on): + if not retry_policy.retry_on(exc): + raise + elif not isinstance(exc, retry_policy.retry_on): + raise + # check if we should give up + if attempts >= retry_policy.max_attempts: + raise + # sleep before retrying + interval = min( + retry_policy.max_interval, + interval * retry_policy.backoff_factor, + ) + time.sleep( + interval + random.uniform(0, 1) if retry_policy.jitter else interval + ) + # log the retry + logger.info( + f"Retrying function {func} with {args} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}", + exc_info=exc, + ) + + async def aretry( retry_policy: Optional[RetryPolicy], func: Callable[P, Awaitable[None]], diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py index 65385a3a0..2471b6eb6 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py @@ -34,10 +34,32 @@ class ErrorMessage(TypedDict): msg: Union[MessageToExecutor, MessageToOrchestrator] +class TopicPartition(Protocol): + topic: str + partition: int + + +class ConsumerRecord(Protocol): + topic: str + "The topic this record is received from" + partition: int + "The partition from which this record is received" + offset: int + "The position of this record in the corresponding Kafka partition." + timestamp: int + "The timestamp of this record" + timestamp_type: int + "The timestamp type of this record" + key: Optional[bytes] + "The key (or `None` if no key is specified)" + value: Optional[bytes] + "The value" + + class Consumer(Protocol): def getmany( self, timeout_ms: int, max_records: int - ) -> dict[str, Sequence[dict[str, Any]]]: ... + ) -> dict[TopicPartition, Sequence[ConsumerRecord]]: ... def commit(self) -> None: ... @@ -45,7 +67,7 @@ class Consumer(Protocol): class AsyncConsumer(Protocol): async def getmany( self, timeout_ms: int, max_records: int - ) -> dict[str, Sequence[dict[str, Any]]]: ... + ) -> dict[TopicPartition, Sequence[ConsumerRecord]]: ... async def commit(self) -> None: ... diff --git a/libs/scheduler-kafka/tests/conftest.py b/libs/scheduler-kafka/tests/conftest.py index 74b8a7800..331855d46 100644 --- a/libs/scheduler-kafka/tests/conftest.py +++ b/libs/scheduler-kafka/tests/conftest.py @@ -3,9 +3,10 @@ from uuid import uuid4 import kafka.admin import pytest -from psycopg import AsyncConnection -from psycopg_pool import AsyncConnectionPool +from psycopg import AsyncConnection, Connection +from psycopg_pool import AsyncConnectionPool, ConnectionPool +from langgraph.checkpoint.postgres import PostgresSaver from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.scheduler.kafka.types import Topics @@ -39,7 +40,7 @@ def topics() -> Iterator[Topics]: @pytest.fixture -async def checkpointer() -> AsyncIterator[AsyncPostgresSaver]: +async def acheckpointer() -> AsyncIterator[AsyncPostgresSaver]: database = f"test_{uuid4().hex[:16]}" # create unique db async with await AsyncConnection.connect( @@ -60,3 +61,23 @@ async def checkpointer() -> AsyncIterator[AsyncPostgresSaver]: DEFAULT_POSTGRES_URI, autocommit=True ) as conn: await conn.execute(f"DROP DATABASE {database}") + + +@pytest.fixture +def checkpointer() -> Iterator[PostgresSaver]: + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + with ConnectionPool( + DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} + ) as pool: + checkpointer = PostgresSaver(pool) + checkpointer.setup() + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") diff --git a/libs/scheduler-kafka/tests/drain.py b/libs/scheduler-kafka/tests/drain.py index 21fcd4d4c..0f0c22d44 100644 --- a/libs/scheduler-kafka/tests/drain.py +++ b/libs/scheduler-kafka/tests/drain.py @@ -1,4 +1,7 @@ import asyncio +import threading +import time +from concurrent.futures import ThreadPoolExecutor from typing import Optional, TypeVar import anyio @@ -6,15 +9,19 @@ from aiokafka import AIOKafkaConsumer from typing_extensions import ParamSpec from langgraph.pregel import Pregel -from langgraph.scheduler.kafka.executor import KafkaExecutor -from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator +from langgraph.scheduler.kafka.default_sync import DefaultConsumer +from langgraph.scheduler.kafka.executor import AsyncKafkaExecutor, KafkaExecutor +from langgraph.scheduler.kafka.orchestrator import ( + AsyncKafkaOrchestrator, + KafkaOrchestrator, +) from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics C = ParamSpec("C") R = TypeVar("R") -async def drain_topics( +async def drain_topics_async( topics: Topics, graph: Pregel, *, debug: bool = False ) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]: scope: Optional[anyio.CancelScope] = None @@ -31,7 +38,7 @@ async def drain_topics( ) async def orchestrator() -> None: - async with KafkaOrchestrator(graph, topics) as orch: + async with AsyncKafkaOrchestrator(graph, topics) as orch: async for msgs in orch: orch_msgs.append(msgs) if debug: @@ -40,7 +47,7 @@ async def drain_topics( scope.cancel() async def executor() -> None: - async with KafkaExecutor(graph, topics) as exec: + async with AsyncKafkaExecutor(graph, topics) as exec: async for msgs in exec: exec_msgs.append(msgs) if debug: @@ -77,3 +84,83 @@ async def drain_topics( assert not errors, errors return [m for mm in orch_msgs for m in mm], [m for mm in exec_msgs for m in mm] + + +def drain_topics( + topics: Topics, graph: Pregel, *, debug: bool = False +) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]: + orch_msgs = [] + exec_msgs = [] + errors = [] + event = threading.Event() + + def done() -> bool: + return ( + len(orch_msgs) > 0 + and len(exec_msgs) > 0 + and not orch_msgs[-1] + and not exec_msgs[-1] + ) + + def orchestrator() -> None: + try: + with KafkaOrchestrator(graph, topics) as orch: + for msgs in orch: + orch_msgs.append(msgs) + if debug: + print("\n---\norch", len(msgs), msgs) + if done(): + print("am i done? orchestrator") + event.set() + if event.is_set(): + break + except Exception as e: + errors.append(e) + event.set() + + def executor() -> None: + try: + with KafkaExecutor(graph, topics) as exec: + for msgs in exec: + exec_msgs.append(msgs) + if debug: + print("\n---\nexec", len(msgs), msgs) + if done(): + print("am i done? executor") + event.set() + if event.is_set(): + break + except Exception as e: + errors.append(e) + event.set() + + def error_consumer() -> None: + try: + with DefaultConsumer(topics.error) as consumer: + while not event.is_set(): + if msg := consumer.poll(timeout_ms=100): + errors.append(msg) + event.set() + except Exception as e: + errors.append(e) + event.set() + + with ThreadPoolExecutor() as pool: + # start error consumer + pool.submit(error_consumer) + + # run the orchestrator and executor until break_when + pool.submit(orchestrator) + pool.submit(executor) + + # timeout + start = time.time() + while not event.is_set(): + time.sleep(0.1) + if time.time() - start > 20: + event.set() + + # check no errors + assert not errors, errors + + return [m for mm in orch_msgs for m in mm], [m for mm in exec_msgs for m in mm] diff --git a/libs/scheduler-kafka/tests/test_fanout.py b/libs/scheduler-kafka/tests/test_fanout.py index 9f3726648..ee8719509 100644 --- a/libs/scheduler-kafka/tests/test_fanout.py +++ b/libs/scheduler-kafka/tests/test_fanout.py @@ -16,7 +16,7 @@ from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics from tests.any import AnyDict -from tests.drain import drain_topics +from tests.drain import drain_topics_async pytestmark = pytest.mark.anyio @@ -86,10 +86,10 @@ def mk_fanout_graph( return builder.compile(checkpointer, interrupt_before=interrupt_before) -async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -> None: +async def test_fanout_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None: input = {"query": "what is weather in sf"} config = {"configurable": {"thread_id": "1"}} - graph = mk_fanout_graph(checkpointer) + graph = mk_fanout_graph(acheckpointer) # start a new run async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: @@ -99,7 +99,7 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - ) # drain topics - orch_msgs, exec_msgs = await drain_topics(topics, graph) + orch_msgs, exec_msgs = await drain_topics_async(topics, graph) # check state state = await graph.aget_state(config) @@ -169,11 +169,11 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - async def test_fanout_graph_w_interrupt( - topics: Topics, checkpointer: BaseCheckpointSaver + topics: Topics, acheckpointer: BaseCheckpointSaver ) -> None: input = {"query": "what is weather in sf"} config = {"configurable": {"thread_id": "1"}} - graph = mk_fanout_graph(checkpointer, interrupt_before=["qa"]) + graph = mk_fanout_graph(acheckpointer, interrupt_before=["qa"]) # start a new run async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: @@ -182,7 +182,7 @@ async def test_fanout_graph_w_interrupt( MessageToOrchestrator(input=input, config=config), ) - orch_msgs, exec_msgs = await drain_topics(topics, graph) + orch_msgs, exec_msgs = await drain_topics_async(topics, graph) # check interrupted state state = await graph.aget_state(config) @@ -258,7 +258,7 @@ async def test_fanout_graph_w_interrupt( MessageToOrchestrator(input=None, config=config), ) - orch_msgs, exec_msgs = await drain_topics(topics, graph) + orch_msgs, exec_msgs = await drain_topics_async(topics, graph) # check final state state = await graph.aget_state(config) diff --git a/libs/scheduler-kafka/tests/test_fanout_sync.py b/libs/scheduler-kafka/tests/test_fanout_sync.py new file mode 100644 index 000000000..9829e5cee --- /dev/null +++ b/libs/scheduler-kafka/tests/test_fanout_sync.py @@ -0,0 +1,273 @@ +import operator +import time +from typing import ( + Annotated, + Sequence, + TypedDict, + Union, +) + +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.graph.state import StateGraph +from langgraph.pregel import Pregel +from langgraph.scheduler.kafka import serde +from langgraph.scheduler.kafka.default_sync import DefaultProducer +from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics +from tests.any import AnyDict +from tests.drain import drain_topics + + +def mk_fanout_graph( + checkpointer: BaseCheckpointSaver, interrupt_before: Sequence[str] = () +) -> Pregel: + # copied from test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + def retriever_picker(data: State) -> list[str]: + return ["analyzer_one", "retriever_two"] + + def analyzer_one(data: State) -> State: + return {"query": f'analyzed: {data["query"]}'} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + def decider(data: State) -> None: + return None + + def decider_cond(data: State) -> str: + if data["query"].count("analyzed") > 1: + return "qa" + else: + return "rewrite_query" + + builder = StateGraph(State) + + builder.add_node("rewrite_query", rewrite_query) + builder.add_node("analyzer_one", analyzer_one) + builder.add_node("retriever_one", retriever_one) + builder.add_node("retriever_two", retriever_two) + builder.add_node("decider", decider) + builder.add_node("qa", qa) + + builder.set_entry_point("rewrite_query") + builder.add_conditional_edges("rewrite_query", retriever_picker) + builder.add_edge("analyzer_one", "retriever_one") + builder.add_edge(["retriever_one", "retriever_two"], "decider") + builder.add_conditional_edges("decider", decider_cond) + builder.set_finish_point("qa") + + return builder.compile(checkpointer, interrupt_before=interrupt_before) + + +def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -> None: + input = {"query": "what is weather in sf"} + config = {"configurable": {"thread_id": "1"}} + graph = mk_fanout_graph(checkpointer) + + # start a new run + with DefaultProducer() as producer: + producer.send( + topics.orchestrator, + value=serde.dumps(MessageToOrchestrator(input=input, config=config)), + ) + producer.flush() + + # drain topics + orch_msgs, exec_msgs = drain_topics(topics, graph, debug=1) + + # check state + state = graph.get_state(config) + assert state.next == () + assert ( + state.values + == graph.invoke(input, {"configurable": {"thread_id": "2"}}) + == { + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + } + ) + + # check history + history = [c for c in graph.get_state_history(config)] + assert len(history) == 11 + + # check messages + assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": None, + } + for c in reversed(history) + for _ in c.tasks + ] + assert exec_msgs == [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": None, + } + for c in reversed(history) + for t in c.tasks + ] + + +def test_fanout_graph_w_interrupt( + topics: Topics, checkpointer: BaseCheckpointSaver +) -> None: + input = {"query": "what is weather in sf"} + config = {"configurable": {"thread_id": "1"}} + graph = mk_fanout_graph(checkpointer, interrupt_before=["qa"]) + + # start a new run + with DefaultProducer() as producer: + producer.send( + topics.orchestrator, + value=serde.dumps(MessageToOrchestrator(input=input, config=config)), + ) + producer.flush() + + orch_msgs, exec_msgs = drain_topics(topics, graph, debug=1) + + # check interrupted state + state = graph.get_state(config) + assert state.next == ("qa",) + assert ( + state.values + == graph.invoke(input, {"configurable": {"thread_id": "2"}}) + == { + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + "query": "analyzed: query: analyzed: query: what is weather in sf", + } + ) + + # check history + history = [c for c in graph.get_state_history(config)] + assert len(history) == 10 + + # check messages + assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": None, + } + for c in reversed(history[1:]) # the last one wasn't executed + # orchestrator messages appear only after tasks for that checkpoint + # finish executing, ie. after executor sends message to resume checkpoint + for _ in c.tasks + ] + assert exec_msgs == [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": None, + } + for c in reversed(history[1:]) # the last one wasn't executed + for t in c.tasks + ] + + # resume the thread + with DefaultProducer() as producer: + producer.send( + topics.orchestrator, + value=serde.dumps(MessageToOrchestrator(input=None, config=config)), + ) + producer.flush() + + orch_msgs, exec_msgs = drain_topics(topics, graph) + + # check final state + state = graph.get_state(config) + assert state.next == () + assert ( + state.values + == graph.invoke(None, {"configurable": {"thread_id": "2"}}) + == { + "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", + } + ) diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index c7ccb7e50..a9708b951 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -16,7 +16,7 @@ from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics from tests.any import AnyDict, AnyStr -from tests.drain import drain_topics +from tests.drain import drain_topics_async from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage pytestmark = pytest.mark.anyio @@ -117,11 +117,11 @@ def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel: async def test_subgraph_w_interrupt( - topics: Topics, checkpointer: BaseCheckpointSaver + topics: Topics, acheckpointer: BaseCheckpointSaver ) -> None: input = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} config = {"configurable": {"thread_id": "1"}} - graph = mk_weather_graph(checkpointer) + graph = mk_weather_graph(acheckpointer) # start a new run async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: @@ -130,7 +130,7 @@ async def test_subgraph_w_interrupt( MessageToOrchestrator(input=input, config=config), ) - orch_msgs, exec_msgs = await drain_topics(topics, graph) + orch_msgs, exec_msgs = await drain_topics_async(topics, graph) # check interrupted state state = await graph.aget_state(config) @@ -415,7 +415,7 @@ async def test_subgraph_w_interrupt( MessageToOrchestrator(input=None, config=config), ) - orch_msgs, exec_msgs = await drain_topics(topics, graph) + orch_msgs, exec_msgs = await drain_topics_async(topics, graph) # check final state state = await graph.aget_state(config) diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py new file mode 100644 index 000000000..a616049d9 --- /dev/null +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -0,0 +1,699 @@ +from typing import Literal, cast + +import pytest +from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, +) +from langchain_core.messages import AIMessage, 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.default_sync import DefaultProducer +from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics +from tests.any import AnyDict, AnyStr +from tests.drain import drain_topics +from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage + +pytestmark = pytest.mark.anyio + + +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" + + def weather_graph(state: RouterState): + return subgraph.invoke(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) + + +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) + + # start a new run + with DefaultProducer() as producer: + producer.send( + topics.orchestrator, + value=serde.dumps(MessageToOrchestrator(input=input, config=config)), + ) + producer.flush() + + orch_msgs, exec_msgs = drain_topics(topics, graph) + + # check interrupted state + state = graph.get_state(config) + assert state.next == ("weather_graph",) + assert state.values == { + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + } + + # check outer history + history = [c for c in graph.get_state_history(config)] + assert len(history) == 3 + + # check child history + child_history = [c for c in graph.get_state_history(history[0].tasks[0].state)] + assert len(child_history) == 3 + + # check messages + assert ( + orch_msgs + == ( + # initial message to outer graph + [MessageToOrchestrator(input=input, config=config)] + # outer graph messages, until interrupted + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": None, + } + for c in reversed(history[1:]) # the last one wasn't executed + # orchestrator messages appear only after tasks for that checkpoint + # finish executing, ie. after executor sends message to resume checkpoint + for _ in c.tasks + ] + # initial message to child graph + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "__pregel_task_id": history[0].tasks[0].id, + "checkpoint_id": None, + "checkpoint_map": { + "": history[0].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[0] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": { + "messages": [ + { + "id": [ + "langchain", + "schema", + "messages", + "HumanMessage", + ], + "kwargs": { + "content": "what's the weather in sf", + "id": AnyStr(), + "type": "human", + }, + "lc": 1, + "type": "constructor", + } + ], + "route": "weather", + }, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": False, + "checkpoint_id": history[0].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[0].tasks[0].id, + "path": list(history[0].tasks[0].path), + }, + } + ], + } + ] + # child graph messages, until interrupted + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "__pregel_task_id": history[0].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[0].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[0] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": False, + "checkpoint_id": history[0].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[0].tasks[0].id, + "path": list(history[0].tasks[0].path), + }, + } + ], + } + for c in reversed(child_history[1:]) # the last one wasn't executed + # orchestrator messages appear only after tasks for that checkpoint + # finish executing, ie. after executor sends message to resume checkpoint + for _ in c.tasks + ] + ) + ) + assert ( + exec_msgs + == ( + # outer graph tasks + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": None, + } + for c in reversed(history) + for t in c.tasks + ] + # child graph tasks + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "__pregel_task_id": history[0].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[0].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[0] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": False, + "checkpoint_id": history[0].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[0].tasks[0].id, + "path": list(history[0].tasks[0].path), + }, + } + ], + } + for c in reversed(child_history[1:]) # the last one wasn't executed + for t in c.tasks + ] + ) + ) + + # resume the thread + with DefaultProducer() as producer: + producer.send( + topics.orchestrator, + value=serde.dumps(MessageToOrchestrator(input=None, config=config)), + ) + producer.flush() + + orch_msgs, exec_msgs = drain_topics(topics, graph) + + # check final state + state = graph.get_state(config) + assert state.next == () + assert state.values == { + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="I'ts sunny in San Francisco!"), + ], + "route": "weather", + } + + # check outer history + history = [c for c in graph.get_state_history(config)] + assert len(history) == 4 + + # check child history + # accessing second to last checkpoint, since that's the one w/ subgraph task + child_history = [c for c in graph.get_state_history(history[1].tasks[0].state)] + assert len(child_history) == 4 + + # check messages + assert ( + orch_msgs + == ( + # initial message to outer graph + [MessageToOrchestrator(input=None, config=config)] + # initial message to child graph + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "__pregel_task_id": history[1].tasks[0].id, + "checkpoint_id": None, + "checkpoint_map": { + "": history[1].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[1] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ], + } + ] + # child graph messages, from previous last checkpoint onwards + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "__pregel_task_id": history[1].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[1].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[1] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ], + } + for c in reversed(child_history[:2]) + for _ in c.tasks + ] + # outer graph messages, from previous last checkpoint onwards + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": None, + } + for c in reversed(history[:2]) + for _ in c.tasks + ] + ) + ) + assert ( + exec_msgs + == ( + # outer graph tasks + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": None, + } + for c in reversed(history[:2]) + for t in c.tasks + ] + # child graph tasks + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "__pregel_task_id": history[1].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[1].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[1] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ], + } + for c in reversed(child_history[:2]) + for t in c.tasks + ] + # "finally" tasks + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ] + ) + )