WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-26 17:42:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove kafka

Browse Source
This commit is contained in:
Nuno Campos
2025-03-01 10:13:47 -08:00
parent e0fb56c6a3
commit 35c2e8a679
27 changed files with 0 additions and 5519 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/scheduler-kafka/LICENSE
-21
View File
@@ -1,21 +0,0 @@
MIT License
Copyright (c) 2024 LangChain, Inc.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
libs/scheduler-kafka/Makefile
-48
View File
@@ -1,48 +0,0 @@
.PHONY: test test_watch lint format
######################
# TESTING AND COVERAGE
######################
start-services:
docker compose -f tests/compose.yml up -V --force-recreate --wait --remove-orphans
stop-services:
docker compose -f tests/compose.yml down
TEST_PATH ?= .
test:
make start-services && poetry run pytest $(TEST_PATH); \
EXIT_CODE=$$?; \
make stop-services; \
exit $$EXIT_CODE
test_watch:
make start-services && poetry run ptw . -- -x $(TEST_PATH); \
EXIT_CODE=$$?; \
make stop-services; \
exit $$EXIT_CODE
######################
# LINTING AND FORMATTING
######################
# Define a variable for Python and notebook files.
PYTHON_FILES=.
MYPY_CACHE=.mypy_cache
lint format: PYTHON_FILES=.
lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --relative --diff-filter=d main . | grep -E '\.py$$|\.ipynb$$')
lint_package: PYTHON_FILES=langgraph
lint_tests: PYTHON_FILES=tests
lint_tests: MYPY_CACHE=.mypy_cache_test
lint lint_diff lint_package lint_tests:
poetry run ruff check .
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
format format_diff:
poetry run ruff format $(PYTHON_FILES)
poetry run ruff check --select I --fix $(PYTHON_FILES)
libs/scheduler-kafka/README.md
-139
View File
@@ -1,139 +0,0 @@
# LangGraph Scheduler for Kafka
This library implements a distributed scheduler for LangGraph using Kafka as the message broker.
## Architecture
![](./langgraph-distributed.png)
- Combination of Kafka (at least once) with a LangGraph Checkpointer provides exactly once semantics both for orchestrator and executor messages
- Checkpointer ensures writes for a given task are saved only once, even if the task is re-executed
- Checkpointer is used to record whether each task in each step has been successfully published to Kafka, to ensure tasks aren't lost, or published more than once
- Orchestrator and Executor manage commit of offsets manually to ensure tasks are marked as done only after finished processing
- Orchestrator and Executor pick up from the earliest message not yet consumed when restarted, to ensure no message is lost, and avoid processing messages more than once
- Orchestrator messages are keyed by thread ID and checkpoint NS, to ensure that no two consumers can process updates for same step of same thread concurrently
- Executor messages are not keyed, as they can be processed concurrently
- Orchestrator and Executor execute messages in configurable batches (up to N messages within space of X seconds), and dedupe messages intra-batch where appropriate (this is purely a performance optimization, with no impact on correctness whether applied or not)
## Basic Usage
Launch orchestrator and executor processes:
`orchestrator.py`
```python
import asyncio
import logging
import os
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
logger = logging.getLogger(__name__)
topics = Topics(
orchestrator=os.environ['KAFKA_TOPIC_ORCHESTRATOR'],
executor=os.environ['KAFKA_TOPIC_EXECUTOR'],
error=os.environ['KAFKA_TOPIC_ERROR'],
)
async def main():
async with AsyncKafkaOrchestrator(graph, topics) as orch:
async for msgs in orch:
logger.info('Procesed %d messages', len(msgs))
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO)
asyncio.run(main())
```
`executor.py`
```python
import asyncio
import logging
import os
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
logger = logging.getLogger(__name__)
topics = Topics(
orchestrator=os.environ['KAFKA_TOPIC_ORCHESTRATOR'],
executor=os.environ['KAFKA_TOPIC_EXECUTOR'],
error=os.environ['KAFKA_TOPIC_ERROR'],
)
async def main():
async with AsyncKafkaExecutor(graph, topics) as orch:
async for msgs in orch:
logger.info('Procesed %d messages', len(msgs))
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO)
asyncio.run(main())
```
```bash
export KAFKA_TOPIC_ORCHESTRATOR='orchestrator'
export KAFKA_TOPIC_EXECUTOR='executor'
export KAFKA_TOPIC_ERROR='error'
python orchestrator.py &
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.
- batch_max_ms (int): Maximum time in milliseconds to wait for messages to include in a batch. Default: 1000.
- retry_policy (langgraph.types.RetryPolicy): Controls which graph-level errors will be retried when processing messages. A good use for this is to retry database errors thrown by the checkpointer. Defaults to None.
### Connection settings
By default the orchestrator and executor will attempt to connect to a Kafka broker running on `localhost:9092`. You can change connection settings by passing any of the following values as `kwargs` to either `KafkaOrchestrator` or `KafkaExecutor`:
- bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
strings) that the consumer should contact to bootstrap initial
cluster metadata. This does not have to be the full node list.
It just needs to have at least one broker that will respond to
Metadata API Request. Default port is 9092. If no servers are
specified, will default to localhost:9092.
- client_id (str): a name for this client. This string is passed in
each request to servers and can be used to identify specific
server-side log entries that correspond to this client. Also
submitted to GroupCoordinator for logging with respect to
consumer group administration. Default: 'aiokafka-{ver}'
- request_timeout_ms (int): Client request timeout in milliseconds.
Default: 40000.
- metadata_max_age_ms (int): The period of time in milliseconds after
which we force a refresh of metadata even if we haven't seen
any partition leadership changes to proactively discover any
new brokers or partitions. Default: 300000
- retry_backoff_ms (int): Milliseconds to backoff when retrying on
errors. Default: 100.
- api_version (str): specify which kafka API version to use.
AIOKafka supports Kafka API versions >=0.9 only.
If set to 'auto', will attempt to infer the broker version by
probing various APIs. Default: auto
- security_protocol (str): Protocol used to communicate with brokers.
Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
Default: PLAINTEXT.
- ssl_context (ssl.SSLContext): pre-configured SSLContext for wrapping
socket connections. For more information see :ref:`ssl_auth`.
Default: None.
- 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.
libs/scheduler-kafka/langgraph-distributed.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 75 KiB

libs/scheduler-kafka/langgraph/scheduler/kafka/__init__.py
View File
libs/scheduler-kafka/langgraph/scheduler/kafka/default_async.py
-9
View File
@@ -1,9 +0,0 @@
import aiokafka
class DefaultAsyncConsumer(aiokafka.AIOKafkaConsumer):
pass
class DefaultAsyncProducer(aiokafka.AIOKafkaProducer):
pass
libs/scheduler-kafka/langgraph/scheduler/kafka/default_sync.py
-39
View File
@@ -1,39 +0,0 @@
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()
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
-473
View File
@@ -1,473 +0,0 @@
import asyncio
import concurrent.futures
from collections.abc import Sequence
from contextlib import (
AbstractAsyncContextManager,
AbstractContextManager,
AsyncExitStack,
ExitStack,
)
from functools import partial
from typing import Any, Optional
from uuid import UUID
import orjson
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
import langgraph.scheduler.kafka.serde as serde
from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR
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,
BackgroundExecutor,
Submit,
)
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.runner import PregelRunner
from langgraph.scheduler.kafka.retry import aretry, retry
from langgraph.scheduler.kafka.types import (
AsyncConsumer,
AsyncProducer,
Consumer,
ErrorMessage,
MessageToExecutor,
MessageToOrchestrator,
Producer,
Sendable,
Topics,
)
from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
class AsyncKafkaExecutor(AbstractAsyncContextManager):
consumer: AsyncConsumer
producer: AsyncProducer
def __init__(
self,
graph: Pregel,
topics: Topics,
*,
batch_max_n: int = 10,
batch_max_ms: int = 1000,
retry_policy: Optional[RetryPolicy] = None,
consumer: Optional[AsyncConsumer] = None,
producer: Optional[AsyncProducer] = None,
**kwargs: Any,
) -> None:
self.graph = graph
self.topics = topics
self.stack = AsyncExitStack()
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
async def __aenter__(self) -> Self:
loop = asyncio.get_running_loop()
self.subgraphs = {
k: v async for k, v in self.graph.aget_subgraphs(recurse=True)
}
if self.consumer is None:
from langgraph.scheduler.kafka.default_async import DefaultAsyncConsumer
self.consumer = await self.stack.enter_async_context(
DefaultAsyncConsumer(
self.topics.executor,
auto_offset_reset="earliest",
group_id="executor",
enable_auto_commit=False,
loop=loop,
**self.kwargs,
)
)
if self.producer is None:
from langgraph.scheduler.kafka.default_async import DefaultAsyncProducer
self.producer = await self.stack.enter_async_context(
DefaultAsyncProducer(
loop=loop,
**self.kwargs,
)
)
return self
async def __aexit__(self, *args: Any) -> None:
return await self.stack.__aexit__(*args)
def __aiter__(self) -> Self:
return self
async def __anext__(self) -> Sequence[MessageToExecutor]:
# wait for next batch
recs = await 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
await asyncio.gather(*(self.each(msg) for msg in msgs))
# commit offsets
await self.consumer.commit()
# return message
return msgs
async def each(self, msg: MessageToExecutor) -> None:
try:
await aretry(self.retry_policy, self.attempt, msg)
except CheckpointNotLatest:
pass
except GraphDelegate as exc:
for arg in exc.args:
fut = await self.producer.send(
self.topics.orchestrator,
value=serde.dumps(
MessageToOrchestrator(
config=arg["config"],
input=orjson.Fragment(
self.graph.checkpointer.serde.dumps(arg["input"])
),
finally_send=[
Sendable(topic=self.topics.executor, value=msg)
],
)
),
# use thread_id, checkpoint_ns as partition key
key=serde.dumps(
(
arg["config"]["configurable"]["thread_id"],
arg["config"]["configurable"].get("checkpoint_ns"),
)
),
)
await fut
except Exception as exc:
fut = await self.producer.send(
self.topics.error,
value=serde.dumps(
ErrorMessage(
topic=self.topics.executor,
msg=msg,
error=repr(exc),
)
),
)
await fut
async def attempt(self, msg: MessageToExecutor) -> None:
# find graph
if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
# remove task_ids from checkpoint_ns
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast} not found")
else:
graph = self.graph
# process message
saved = await self.graph.checkpointer.aget_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()
async with (
AsyncChannelsManager(
graph.channels,
saved.checkpoint,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed),
AsyncBackgroundExecutor(msg["config"]) as submit,
):
if task := await asyncio.to_thread(
prepare_single_task,
msg["task"]["path"],
msg["task"]["id"],
checkpoint=saved.checkpoint,
pending_writes=saved.pending_writes or [],
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,
store=self.graph.store,
):
# execute task, saving writes
runner = PregelRunner(
submit=submit,
put_writes=partial(self._put_writes, submit, msg["config"]),
schedule_task=self._schedule_task,
)
async for _ in runner.atick([task], reraise=False):
pass
else:
# task was not found
await self.graph.checkpointer.aput_writes(
msg["config"], [(ERROR, TaskNotFound())], str(UUID(int=0))
)
# notify orchestrator
fut = await self.producer.send(
self.topics.orchestrator,
value=serde.dumps(
MessageToOrchestrator(
input=None,
config=msg["config"],
finally_send=msg.get("finally_send"),
)
),
# use thread_id, checkpoint_ns as partition key
key=serde.dumps(
(
msg["config"]["configurable"]["thread_id"],
msg["config"]["configurable"].get("checkpoint_ns"),
)
),
)
await fut
def _schedule_task(
self,
task: PregelExecutableTask,
idx: int,
) -> None:
# will be scheduled by orchestrator when executor finishes
pass
def _put_writes(
self,
submit: Submit,
config: RunnableConfig,
task_id: str,
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_send=[
Sendable(topic=self.topics.executor, value=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 = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast} 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,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed),
BackgroundExecutor({}) as submit,
):
if task := prepare_single_task(
msg["task"]["path"],
msg["task"]["id"],
checkpoint=saved.checkpoint,
pending_writes=saved.pending_writes or [],
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"]),
schedule_task=self._schedule_task,
)
for _ in runner.tick([task], reraise=False):
pass
else:
# task was not found
self.graph.checkpointer.put_writes(
msg["config"], [(ERROR, TaskNotFound())], str(UUID(int=0))
)
# notify orchestrator
fut = self.producer.send(
self.topics.orchestrator,
value=serde.dumps(
MessageToOrchestrator(
input=None,
config=msg["config"],
finally_send=msg.get("finally_send"),
)
),
# 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 _schedule_task(
self,
task: PregelExecutableTask,
idx: int,
) -> None:
# will be scheduled by orchestrator when executor finishes
pass
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)
libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py
-408
View File
@@ -1,408 +0,0 @@
import asyncio
import concurrent.futures
from contextlib import (
AbstractAsyncContextManager,
AbstractContextManager,
AsyncExitStack,
ExitStack,
)
from typing import Any, Optional
from langchain_core.runnables import ensure_config
from typing_extensions import Self
import langgraph.scheduler.kafka.serde as serde
from langgraph.constants import (
CONFIG_KEY_DEDUPE_TASKS,
CONFIG_KEY_ENSURE_LATEST,
INTERRUPT,
SCHEDULED,
)
from langgraph.errors import CheckpointNotLatest, GraphInterrupt
from langgraph.pregel import Pregel
from langgraph.pregel.executor import BackgroundExecutor, Submit
from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
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.types import RetryPolicy
from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
consumer: AsyncConsumer
producer: AsyncProducer
def __init__(
self,
graph: Pregel,
topics: Topics,
batch_max_n: int = 10,
batch_max_ms: int = 1000,
retry_policy: Optional[RetryPolicy] = None,
consumer: Optional[AsyncConsumer] = None,
producer: Optional[AsyncProducer] = None,
**kwargs: Any,
) -> None:
self.graph = graph
self.topics = topics
self.stack = AsyncExitStack()
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
async def __aenter__(self) -> Self:
loop = asyncio.get_running_loop()
self.subgraphs = {
k: v async for k, v in self.graph.aget_subgraphs(recurse=True)
}
if self.consumer is None:
from langgraph.scheduler.kafka.default_async import DefaultAsyncConsumer
self.consumer = await self.stack.enter_async_context(
DefaultAsyncConsumer(
self.topics.orchestrator,
auto_offset_reset="earliest",
group_id="orchestrator",
enable_auto_commit=False,
loop=loop,
**self.kwargs,
)
)
if self.producer is None:
from langgraph.scheduler.kafka.default_async import DefaultAsyncProducer
self.producer = await self.stack.enter_async_context(
DefaultAsyncProducer(
loop=loop,
**self.kwargs,
)
)
return self
async def __aexit__(self, *args: Any) -> None:
return await self.stack.__aexit__(*args)
def __aiter__(self) -> Self:
return self
async def __anext__(self) -> list[MessageToOrchestrator]:
# wait for next batch
recs = await self.consumer.getmany(
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)
msgs: list[MessageToOrchestrator] = [serde.loads(msg) for msg in uniq]
# process batch
await asyncio.gather(*(self.each(msg) for msg in msgs))
# commit offsets
await self.consumer.commit()
# return message
return msgs
async def each(self, msg: MessageToOrchestrator) -> None:
try:
await aretry(self.retry_policy, self.attempt, msg)
except CheckpointNotLatest:
pass
except GraphInterrupt:
pass
except Exception as exc:
fut = await self.producer.send(
self.topics.error,
value=serde.dumps(
ErrorMessage(
topic=self.topics.orchestrator,
msg=msg,
error=repr(exc),
)
),
)
await fut
async def attempt(self, msg: MessageToOrchestrator) -> None:
# find graph
if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
# remove task_ids from checkpoint_ns
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast} not found")
else:
graph = self.graph
# process message
async with AsyncPregelLoop(
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,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
) as loop:
if loop.tick(input_keys=graph.input_channels):
# wait for checkpoint to be saved
if hasattr(loop, "_put_checkpoint_fut"):
await loop._put_checkpoint_fut
# schedule any new tasks
if new_tasks := [
t for t in loop.tasks.values() if not t.scheduled and not t.writes
]:
config = patch_configurable(
loop.config,
{
**loop.checkpoint_config["configurable"],
CONFIG_KEY_DEDUPE_TASKS: True,
CONFIG_KEY_ENSURE_LATEST: True,
},
)
# send messages to executor
futures = await asyncio.gather(
*(
self.producer.send(
self.topics.executor,
value=serde.dumps(
MessageToExecutor(
config=config,
task=ExecutorTask(id=task.id, path=task.path),
finally_send=msg.get("finally_send"),
)
),
)
for task in new_tasks
)
)
# wait for messages to be sent
await asyncio.gather(*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_send"):
# send any finally_send messages
futs = await asyncio.gather(
*(
self.producer.send(
m["topic"],
value=serde.dumps(m["value"]) if m.get("value") else None,
key=serde.dumps(m["key"]) if m.get("key") else None,
)
for m in msg["finally_send"]
)
)
# 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
)
# 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 = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast} 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,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
) as loop:
if loop.tick(input_keys=graph.input_channels):
# 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 and not t.writes
]:
config = patch_configurable(
loop.config,
{
**loop.checkpoint_config["configurable"],
CONFIG_KEY_DEDUPE_TASKS: True,
CONFIG_KEY_ENSURE_LATEST: True,
},
)
# send messages to executor
futures = [
self.producer.send(
self.topics.executor,
value=serde.dumps(
MessageToExecutor(
config=config,
task=ExecutorTask(id=task.id, path=task.path),
finally_send=msg.get("finally_send"),
)
),
)
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_send"):
# schedule any finally_send msgs
futs = [
self.producer.send(
m["topic"],
value=serde.dumps(m["value"]) if m.get("value") else None,
key=serde.dumps(m["key"]) if m.get("key") else None,
)
for m in msg["finally_send"]
]
# wait for messages to be sent
concurrent.futures.wait(futs)
libs/scheduler-kafka/langgraph/scheduler/kafka/py.typed
View File
libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py
-98
View File
@@ -1,98 +0,0 @@
import asyncio
import logging
import random
import time
from typing import Awaitable, Callable, Optional
from typing_extensions import ParamSpec
from langgraph.types import RetryPolicy
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]],
*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:
await 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,
)
await asyncio.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,
)
libs/scheduler-kafka/langgraph/scheduler/kafka/serde.py
-20
View File
@@ -1,20 +0,0 @@
from typing import Any
import orjson
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
SERIALIZER = JsonPlusSerializer()
def loads(v: bytes) -> Any:
return SERIALIZER.loads(v)
def dumps(v: Any) -> bytes:
return orjson.dumps(v, default=_default)
def _default(v: Any) -> Any:
# things we don't know how to serialize (eg. functions) ignore
return None
libs/scheduler-kafka/langgraph/scheduler/kafka/types.py
-98
View File
@@ -1,98 +0,0 @@
import asyncio
import concurrent.futures
from typing import Any, NamedTuple, Optional, Protocol, Sequence, TypedDict, Union
from langchain_core.runnables import RunnableConfig
class Topics(NamedTuple):
orchestrator: str
executor: str
error: str
class Sendable(TypedDict):
topic: str
value: Optional[Any]
key: Optional[Any]
class MessageToOrchestrator(TypedDict):
input: Optional[dict[str, Any]]
config: RunnableConfig
finally_send: Optional[Sequence[Sendable]]
class ExecutorTask(TypedDict):
id: Optional[str]
path: tuple[Union[str, int], ...]
class MessageToExecutor(TypedDict):
config: RunnableConfig
task: ExecutorTask
finally_send: Optional[Sequence[Sendable]]
class ErrorMessage(TypedDict):
topic: str
error: str
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[TopicPartition, Sequence[ConsumerRecord]]: ...
def commit(self) -> None: ...
class AsyncConsumer(Protocol):
async def getmany(
self, timeout_ms: int, max_records: int
) -> dict[TopicPartition, Sequence[ConsumerRecord]]: ...
async def commit(self) -> None: ...
class Producer(Protocol):
def send(
self,
topic: str,
*,
key: Optional[bytes] = None,
value: Optional[bytes] = None,
) -> concurrent.futures.Future: ...
class AsyncProducer(Protocol):
async def send(
self,
topic: str,
*,
key: Optional[bytes] = None,
value: Optional[bytes] = None,
) -> asyncio.Future: ...
libs/scheduler-kafka/poetry.lock
Generated
-1256
View File
File diff suppressed because it is too large Load Diff
libs/scheduler-kafka/pyproject.toml
-58
View File
@@ -1,58 +0,0 @@
[tool.poetry]
name = "langgraph-scheduler-kafka"
version = "1.0.0"
description = "Library with Kafka-based work scheduler."
authors = []
license = "MIT"
readme = "README.md"
repository = "https://www.github.com/langchain-ai/langgraph"
packages = [{ include = "langgraph" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
orjson = "^3.10.7"
crc32c = "^2.7.post1"
aiokafka = "^0.11.0"
langgraph = "^0.2.19"
[tool.poetry.group.dev.dependencies]
ruff = "^0.6.2"
codespell = "^2.2.0"
pytest = "^7.2.1"
pytest-mock = "^3.11.1"
pytest-watcher = "^0.4.1"
mypy = "^1.10.0"
langgraph = {path = "../langgraph", develop = true}
langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true}
langgraph-checkpoint = {path = "../checkpoint", develop = true}
kafka-python-ng = "^2.2.2"
[tool.pytest.ini_options]
# --strict-markers will raise errors on unknown marks.
# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
#
# https://docs.pytest.org/en/7.1.x/reference/reference.html
# --strict-config any warnings encountered while parsing the `pytest`
# section of the configuration file raise errors.
addopts = "--strict-markers --strict-config --durations=5 -vv"
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.ruff]
lint.select = [
"E", # pycodestyle
"F", # Pyflakes
"UP", # pyupgrade
"B", # flake8-bugbear
"I", # isort
]
lint.ignore = ["E501", "B008", "UP007", "UP006"]
[tool.pytest-watcher]
now = true
delay = 0.1
runner_args = ["--ff", "-v", "--tb", "short", "-s"]
patterns = ["*.py"]
libs/scheduler-kafka/tests/__init__.py
View File
libs/scheduler-kafka/tests/any.py
-63
View File
@@ -1,63 +0,0 @@
import re
from typing import Union
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))
class AnyDict(dict):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def __eq__(self, other: object) -> bool:
if not self and isinstance(other, dict):
return True
if not isinstance(other, dict) or len(self) != len(other):
return False
for k, v in self.items():
if kk := next((kk for kk in other if kk == k), None):
if v == other[kk]:
continue
else:
return False
else:
return True
class AnyList(list):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def __eq__(self, other: object) -> bool:
if not self and isinstance(other, list):
return True
if not isinstance(other, list) or len(self) != len(other):
return False
for i, v in enumerate(self):
if v == other[i]:
continue
else:
return False
else:
return True
class AnyInt(int):
def __init__(self) -> None:
super().__init__()
def __eq__(self, other: object) -> bool:
return isinstance(other, int)
libs/scheduler-kafka/tests/compose.yml
-21
View File
@@ -1,21 +0,0 @@
name: scheduler-kafka-tests
services:
broker:
image: apache/kafka:latest
ports:
- "9092:9092"
postgres:
image: postgres:16
ports:
- "5443:5432"
environment:
POSTGRES_DB: postgres
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
healthcheck:
test: pg_isready -U postgres
start_period: 10s
timeout: 1s
retries: 5
interval: 60s
start_interval: 1s
libs/scheduler-kafka/tests/conftest.py
-83
View File
@@ -1,83 +0,0 @@
from typing import AsyncIterator, Iterator
from uuid import uuid4
import kafka.admin
import pytest
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
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5443/"
@pytest.fixture
def anyio_backend():
return "asyncio"
@pytest.fixture
def topics() -> Iterator[Topics]:
o = f"test_o_{uuid4().hex[:16]}"
e = f"test_e_{uuid4().hex[:16]}"
z = f"test_z_{uuid4().hex[:16]}"
admin = kafka.admin.KafkaAdminClient()
# create topics
admin.create_topics(
[
kafka.admin.NewTopic(name=o, num_partitions=1, replication_factor=1),
kafka.admin.NewTopic(name=e, num_partitions=1, replication_factor=1),
kafka.admin.NewTopic(name=z, num_partitions=1, replication_factor=1),
]
)
# yield topics
yield Topics(orchestrator=o, executor=e, error=z)
# delete topics
admin.delete_topics([o, e, z])
admin.close()
@pytest.fixture
async def acheckpointer() -> AsyncIterator[AsyncPostgresSaver]:
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncConnectionPool(
DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
) as pool:
checkpointer = AsyncPostgresSaver(pool)
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
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}")
libs/scheduler-kafka/tests/drain.py
-168
View File
@@ -1,168 +0,0 @@
import asyncio
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, TypeVar
import anyio
from aiokafka import AIOKafkaConsumer
from typing_extensions import ParamSpec
from langgraph.pregel import Pregel
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(
topics: Topics, graph: Pregel, *, debug: bool = False
) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]:
scope: Optional[anyio.CancelScope] = None
orch_msgs = []
exec_msgs = []
errors = []
def done() -> bool:
return (
len(orch_msgs) > 0
and any(orch_msgs)
and len(exec_msgs) > 0
and any(exec_msgs)
and not orch_msgs[-1]
and not exec_msgs[-1]
)
async def orchestrator() -> None:
async with AsyncKafkaOrchestrator(graph, topics) as orch:
async for msgs in orch:
orch_msgs.append(msgs)
if debug:
print("\n---\norch", len(msgs), msgs)
if done():
scope.cancel()
async def executor() -> None:
async with AsyncKafkaExecutor(graph, topics) as exec:
async for msgs in exec:
exec_msgs.append(msgs)
if debug:
print("\n---\nexec", len(msgs), msgs)
if done():
scope.cancel()
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")
# run the orchestrator and executor until break_when
async with anyio.create_task_group() as tg:
tg.cancel_scope.deadline = anyio.current_time() + 20
scope = tg.cancel_scope
tg.start_soon(orchestrator, name="orchestrator")
tg.start_soon(executor, name="executor")
# cancel error consumer
error_task.cancel()
try:
await error_task
except asyncio.CancelledError:
pass
# 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]
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 any(orch_msgs)
and len(exec_msgs) > 0
and any(exec_msgs)
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():
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():
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]
libs/scheduler-kafka/tests/messages.py
-28
View File
@@ -1,28 +0,0 @@
"""Redefined messages as a work-around for pydantic issue with AnyStr.
The code below creates version of pydantic models
that will work in unit tests with AnyStr as id field
Please note that the `id` field is assigned AFTER the model is created
to workaround an issue with pydantic ignoring the __eq__ method on
subclassed strings.
"""
from typing import Any
from langchain_core.messages import AIMessage, HumanMessage
from tests.any import AnyStr
def _AnyIdAIMessage(**kwargs: Any) -> AIMessage:
"""Create ai message with an any id field."""
message = AIMessage(**kwargs)
message.id = AnyStr()
return message
def _AnyIdHumanMessage(**kwargs: Any) -> HumanMessage:
"""Create a human message with an any id field."""
message = HumanMessage(**kwargs)
message.id = AnyStr()
return message
libs/scheduler-kafka/tests/test_fanout.py
-274
View File
@@ -1,274 +0,0 @@
import asyncio
import operator
from typing import (
Annotated,
Sequence,
TypedDict,
Union,
)
import pytest
from aiokafka import AIOKafkaProducer
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.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.drain import drain_topics_async
pytestmark = pytest.mark.anyio
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]
async def rewrite_query(data: State) -> State:
return {"query": f'query: {data["query"]}'}
async def retriever_picker(data: State) -> list[str]:
return ["analyzer_one", "retriever_two"]
async def analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
await asyncio.sleep(0.1)
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
async 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)
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(acheckpointer)
# 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),
)
# drain topics
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check state
state = await graph.aget_state(config)
assert state.next == ()
assert (
state.values
== await graph.ainvoke(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 async for c in graph.aget_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_send": 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_send": None,
}
for c in reversed(history)
for t in c.tasks
]
async def test_fanout_graph_w_interrupt(
topics: Topics, acheckpointer: BaseCheckpointSaver
) -> None:
input = {"query": "what is weather in sf"}
config = {"configurable": {"thread_id": "1"}}
graph = mk_fanout_graph(acheckpointer, interrupt_before=["qa"])
# 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),
)
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check interrupted state
state = await graph.aget_state(config)
assert state.next == ("qa",)
assert (
state.values
== await graph.ainvoke(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 async for c in graph.aget_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_send": 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_send": None,
}
for c in reversed(history[1:]) # the last one wasn't executed
for t in c.tasks
]
# resume the thread
async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
await producer.send_and_wait(
topics.orchestrator,
MessageToOrchestrator(input=None, config=config),
)
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check final state
state = await graph.aget_state(config)
assert state.next == ()
assert (
state.values
== await graph.ainvoke(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",
}
)
libs/scheduler-kafka/tests/test_fanout_sync.py
-273
View File
@@ -1,273 +0,0 @@
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_send": 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_send": 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_send": 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_send": 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",
}
)
libs/scheduler-kafka/tests/test_push.py
-206
View File
@@ -1,206 +0,0 @@
import operator
from typing import (
Annotated,
Literal,
Union,
)
import pytest
from aiokafka import AIOKafkaProducer
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Command, Send
from tests.any import AnyDict
from tests.drain import drain_topics_async
pytestmark = pytest.mark.anyio
def mk_push_graph(
checkpointer: BaseCheckpointSaver,
) -> CompiledStateGraph:
# copied from test_send_dedupe_on_resume
class InterruptOnce:
ticks: int = 0
def __call__(self, state):
self.ticks += 1
if self.ticks == 1:
raise NodeInterrupt("Bahh")
return ["|".join(("flaky", str(state)))]
class Node:
def __init__(self, name: str):
self.name = name
self.ticks = 0
self.__name__ = name
def __call__(self, state):
self.ticks += 1
update = (
[self.name]
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, Command):
return state.copy(update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("flaky", 4))),
"3.1",
]
def route_to_three(state) -> Literal["3"]:
return "3"
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(Node("2"))
builder.add_node(Node("3"))
builder.add_node(Node("3.1"))
builder.add_node("flaky", InterruptOnce())
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
return builder.compile(checkpointer=checkpointer)
@pytest.mark.skip("TODO: re-enable in next PR")
async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
graph_compare = mk_push_graph(acheckpointer)
# 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),
)
# drain topics
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check state
state = await graph.aget_state(config)
assert all(not t.error for t in state.tasks)
assert state.next == ("flaky",)
assert (
state.values
== await graph_compare.ainvoke(input, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
]
)
# check history
history = [c async for c in graph.aget_state_history(config)]
assert len(history) == 2
# 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_send": 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": _convert_path(t.path),
},
"finally_send": None,
}
for c in reversed(history)
for t in c.tasks
]
# resume the thread
async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
await producer.send_and_wait(
topics.orchestrator,
MessageToOrchestrator(input=None, config=config),
)
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check final state
state = await graph.aget_state(config)
assert state.next == ()
assert (
state.values
== await graph_compare.ainvoke(None, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
]
)
# check history
history = [c async for c in graph.aget_state_history(config)]
assert len(history) == 4
# check executions
# node "2" doesn't get called again, as we recover writes saved before
assert graph.builder.nodes["2"].runnable.func.ticks == 3
# node "flaky" gets called again, as it was interrupted
assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
def _convert_path(
path: tuple[Union[str, int, tuple], ...],
) -> list[Union[str, int, list]]:
return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)
libs/scheduler-kafka/tests/test_push_sync.py
-208
View File
@@ -1,208 +0,0 @@
import operator
from typing import (
Annotated,
Literal,
Union,
)
import pytest
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Command, Send
from tests.any import AnyDict
from tests.drain import drain_topics
pytestmark = pytest.mark.anyio
def mk_push_graph(
checkpointer: BaseCheckpointSaver,
) -> CompiledStateGraph:
# copied from test_send_dedupe_on_resume
class InterruptOnce:
ticks: int = 0
def __call__(self, state):
self.ticks += 1
if self.ticks == 1:
raise NodeInterrupt("Bahh")
return ["|".join(("flaky", str(state)))]
class Node:
def __init__(self, name: str):
self.name = name
self.ticks = 0
self.__name__ = name
def __call__(self, state):
self.ticks += 1
update = (
[self.name]
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, Command):
return state.copy(update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("flaky", 4))),
"3.1",
]
def route_to_three(state) -> Literal["3"]:
return "3"
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(Node("2"))
builder.add_node(Node("3"))
builder.add_node(Node("3.1"))
builder.add_node("flaky", InterruptOnce())
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
return builder.compile(checkpointer=checkpointer)
@pytest.mark.skip("TODO: re-enable in next PR")
def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
graph_compare = mk_push_graph(acheckpointer)
# 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)
# check state
state = graph.get_state(config)
assert all(not t.error for t in state.tasks)
assert state.next == ("flaky",)
assert (
state.values
== graph_compare.invoke(input, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
]
)
# check history
history = [c for c in graph.get_state_history(config)]
assert len(history) == 2
# 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_send": 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": _convert_path(t.path),
},
"finally_send": None,
}
for c in reversed(history)
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_compare.invoke(None, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
]
)
# check history
history = [c for c in graph.get_state_history(config)]
assert len(history) == 4
# check executions
# node "2" doesn't get called again, as we recover writes saved before
assert graph.builder.nodes["2"].runnable.func.ticks == 3
# node "flaky" gets called again, as it was interrupted
assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
def _convert_path(
path: tuple[Union[str, int, tuple], ...],
) -> list[Union[str, int, list]]:
return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)
libs/scheduler-kafka/tests/test_subgraph.py
-765
View File
@@ -1,765 +0,0 @@
from typing import Literal, cast
import pytest
from aiokafka import AIOKafkaProducer
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.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.drain import drain_topics_async
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"
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)
async def test_subgraph_w_interrupt(
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(acheckpointer)
# 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),
)
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check interrupted state
state = await graph.aget_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 async for c in graph.aget_state_history(config)]
assert len(history) == 3
# check child history
child_history = [
c async for c in graph.aget_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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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": [
_AnyIdHumanMessage(content="what's the weather in sf")
],
"route": "weather",
},
"finally_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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
async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
await producer.send_and_wait(
topics.orchestrator,
MessageToOrchestrator(input=None, config=config),
)
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check final state
state = await graph.aget_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 async for c in graph.aget_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 async for c in graph.aget_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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_send": 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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_send": None,
"task": {
"id": history[1].tasks[0].id,
"path": list(history[1].tasks[0].path),
},
}
]
)
)
libs/scheduler-kafka/tests/test_subgraph_sync.py
-763
View File
@@ -1,763 +0,0 @@
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
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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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": [
_AnyIdHumanMessage(content="what's the weather in sf")
],
"route": "weather",
},
"finally_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_store": None,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_previous": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_store": None,
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_previous": None,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_send": 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_send": 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_call": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
"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_send": [
{
"topic": topics.executor,
"value": {
"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_send": 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_send": None,
"task": {
"id": history[1].tasks[0].id,
"path": list(history[1].tasks[0].path),
},
}
]
)
)