mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 18:27:52 +02:00
Implement LangGraph Scheduler for Kafka
- Orchestrator and Executor classes to run LangGraph in a distributed fashion using Kafka as a message bus for communication - Orchestrator and Executor run on-demand when a new message is published to the topic they listen to - Orchestrator is responsible for running the Pregel algorithm (deciding next tasks to run) and sending messages to the executor topic - Executor is responsible for executing each task (node), and sending messages to the orchestrator topic when done
This commit is contained in:
Nuno Campos
@@ -9,6 +9,7 @@ CONFIG_KEY_STREAM = "__pregel_stream"
|
||||
CONFIG_KEY_STORE = "__pregel_store"
|
||||
CONFIG_KEY_RESUMING = "__pregel_resuming"
|
||||
CONFIG_KEY_TASK_ID = "__pregel_task_id"
|
||||
CONFIG_KEY_DEDUPE_TASKS = "__pregel_dedupe_tasks"
|
||||
# this one part of public API so more readable
|
||||
CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map"
|
||||
INTERRUPT = "__interrupt__"
|
||||
@@ -30,6 +31,7 @@ RESERVED = {
|
||||
CONFIG_KEY_STORE,
|
||||
CONFIG_KEY_RESUMING,
|
||||
CONFIG_KEY_TASK_ID,
|
||||
CONFIG_KEY_DEDUPE_TASKS,
|
||||
INPUT,
|
||||
RUNTIME_PLACEHOLDER,
|
||||
}
|
||||
|
||||
@@ -39,6 +39,7 @@ from langgraph.checkpoint.base import (
|
||||
)
|
||||
from langgraph.constants import (
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_DEDUPE_TASKS,
|
||||
CONFIG_KEY_RESUMING,
|
||||
CONFIG_KEY_STREAM,
|
||||
CONFIG_KEY_TASK_ID,
|
||||
@@ -179,7 +180,10 @@ class PregelLoop:
|
||||
self.output_keys = output_keys
|
||||
self.stream_keys = stream_keys
|
||||
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get("configurable", {})
|
||||
self.skip_done_tasks = "checkpoint_id" not in config["configurable"]
|
||||
self.skip_done_tasks = (
|
||||
"checkpoint_id" not in config["configurable"]
|
||||
or CONFIG_KEY_DEDUPE_TASKS in config["configurable"]
|
||||
)
|
||||
self.debug = debug
|
||||
if CONFIG_KEY_STREAM in config["configurable"]:
|
||||
self.stream = DuplexStream(
|
||||
@@ -595,6 +599,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
|
||||
**self.config,
|
||||
**saved.config,
|
||||
"configurable": {
|
||||
"checkpoint_ns": "",
|
||||
**self.config.get("configurable", {}),
|
||||
**saved.config.get("configurable", {}),
|
||||
},
|
||||
@@ -697,6 +702,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
|
||||
**self.config,
|
||||
**saved.config,
|
||||
"configurable": {
|
||||
"checkpoint_ns": "",
|
||||
**self.config.get("configurable", {}),
|
||||
**saved.config.get("configurable", {}),
|
||||
},
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
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.
|
||||
@@ -0,0 +1,46 @@
|
||||
.PHONY: test test_watch lint format
|
||||
|
||||
######################
|
||||
# TESTING AND COVERAGE
|
||||
######################
|
||||
|
||||
start-services:
|
||||
docker compose -f testss/compose.yml up -V --force-recreate --wait --remove-orphans
|
||||
|
||||
stop-services:
|
||||
docker compose -f testss/compose.yml down
|
||||
|
||||
test:
|
||||
make start-services && poetry run pytest; \
|
||||
EXIT_CODE=$$?; \
|
||||
make stop-services; \
|
||||
exit $$EXIT_CODE
|
||||
|
||||
test_watch:
|
||||
make start-services && poetry run ptw .; \
|
||||
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)
|
||||
@@ -0,0 +1,3 @@
|
||||
# LangGraph Scheduler for Kafka
|
||||
|
||||
...
|
||||
@@ -0,0 +1,94 @@
|
||||
import asyncio
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from typing import Any
|
||||
|
||||
import aiokafka
|
||||
|
||||
import langgraph.scheduler.kafka.serde as serde
|
||||
from langgraph.constants import ERROR
|
||||
from langgraph.errors import TaskNotFound
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel.algo import prepare_single_task
|
||||
from langgraph.pregel.executor import AsyncBackgroundExecutor
|
||||
from langgraph.pregel.manager import AsyncChannelsManager
|
||||
from langgraph.pregel.runner import PregelRunner
|
||||
from langgraph.scheduler.kafka.types import (
|
||||
MessageToExecutor,
|
||||
MessageToOrchestrator,
|
||||
Topics,
|
||||
)
|
||||
|
||||
|
||||
class KafkaExecutor(AbstractAsyncContextManager):
|
||||
def __init__(self, graph: Pregel, topics: Topics, **kwargs: Any) -> None:
|
||||
self.graph = graph
|
||||
self.topics = topics
|
||||
self.consumer = aiokafka.AIOKafkaConsumer(
|
||||
topics.executor, value_deserializer=serde.loads, **kwargs
|
||||
)
|
||||
self.producer = aiokafka.AIOKafkaProducer(
|
||||
value_serializer=serde.dumps, **kwargs
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> "KafkaExecutor":
|
||||
await self.consumer.start()
|
||||
await self.producer.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.consumer.stop()
|
||||
await self.producer.stop()
|
||||
|
||||
def __aiter__(self) -> "KafkaExecutor":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> Any:
|
||||
# wait for next message
|
||||
try:
|
||||
rec = await self.consumer.getone()
|
||||
msg: MessageToExecutor = rec.value
|
||||
except aiokafka.ConsumerStoppedError:
|
||||
raise StopAsyncIteration from None
|
||||
# process message
|
||||
saved = await self.graph.checkpointer.aget_tuple(msg["config"])
|
||||
if saved is None:
|
||||
raise RuntimeError("Checkpoint not found")
|
||||
async with AsyncChannelsManager(
|
||||
self.graph.channels, saved.checkpoint, msg["config"], self.graph.store
|
||||
) as (channels, managed), AsyncBackgroundExecutor() as submit:
|
||||
|
||||
def put_writes(task_id: str, writes: list[tuple[str, Any]]) -> None:
|
||||
print("put_writes", task_id, writes)
|
||||
return submit(
|
||||
self.graph.checkpointer.aput_writes, msg["config"], writes, task_id
|
||||
)
|
||||
|
||||
if task := await asyncio.to_thread(
|
||||
prepare_single_task,
|
||||
msg["task"]["path"],
|
||||
msg["task"]["id"],
|
||||
checkpoint=saved.checkpoint,
|
||||
processes=self.graph.nodes,
|
||||
channels=channels,
|
||||
managed=managed,
|
||||
config=msg["config"],
|
||||
step=msg["task"]["step"],
|
||||
for_execution=True,
|
||||
is_resuming=msg["task"]["resuming"],
|
||||
):
|
||||
# execute task, saving writes
|
||||
runner = PregelRunner(submit=submit, put_writes=put_writes)
|
||||
async for _ in runner.atick([task]):
|
||||
pass
|
||||
else:
|
||||
# task was not found
|
||||
await self.graph.checkpointer.put_writes(
|
||||
msg["config"], [(ERROR, TaskNotFound())]
|
||||
)
|
||||
# notify orchestrator
|
||||
await self.producer.send(
|
||||
self.topics.orchestrator,
|
||||
value=MessageToOrchestrator(input=None, config=msg["config"]),
|
||||
)
|
||||
# return message
|
||||
return msg
|
||||
@@ -0,0 +1,94 @@
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from typing import Any
|
||||
|
||||
import aiokafka
|
||||
from langchain_core.runnables import ensure_config
|
||||
|
||||
import langgraph.scheduler.kafka.serde as serde
|
||||
from langgraph.constants import CONFIG_KEY_DEDUPE_TASKS, SCHEDULED
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel.loop import INPUT_RESUMING, AsyncPregelLoop
|
||||
from langgraph.scheduler.kafka.types import (
|
||||
ExecutorTask,
|
||||
MessageToExecutor,
|
||||
MessageToOrchestrator,
|
||||
Topics,
|
||||
)
|
||||
from langgraph.utils.config import patch_configurable
|
||||
|
||||
|
||||
class KafkaOrchestrator(AbstractAsyncContextManager):
|
||||
def __init__(self, graph: Pregel, topics: Topics, **kwargs: Any) -> None:
|
||||
self.graph = graph
|
||||
self.topics = topics
|
||||
self.consumer = aiokafka.AIOKafkaConsumer(
|
||||
topics.orchestrator, value_deserializer=serde.loads, **kwargs
|
||||
)
|
||||
self.producer = aiokafka.AIOKafkaProducer(
|
||||
value_serializer=serde.dumps, **kwargs
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> "KafkaOrchestrator":
|
||||
await self.consumer.start()
|
||||
await self.producer.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.consumer.stop()
|
||||
await self.producer.stop()
|
||||
|
||||
def __aiter__(self) -> "KafkaOrchestrator":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> Any:
|
||||
# wait for next message
|
||||
try:
|
||||
rec = await self.consumer.getone()
|
||||
msg: MessageToOrchestrator = rec.value
|
||||
except aiokafka.ConsumerStoppedError:
|
||||
raise StopAsyncIteration from None
|
||||
# process message
|
||||
async with AsyncPregelLoop(
|
||||
msg["input"],
|
||||
config=ensure_config(msg["config"]),
|
||||
stream=None,
|
||||
store=self.graph.store,
|
||||
checkpointer=self.graph.checkpointer,
|
||||
nodes=self.graph.nodes,
|
||||
specs=self.graph.channels,
|
||||
output_keys=self.graph.output_channels,
|
||||
stream_keys=self.graph.stream_channels,
|
||||
) as loop:
|
||||
if loop.tick(input_keys=self.graph.input_channels):
|
||||
if hasattr(loop, "_put_checkpoint_fut"):
|
||||
await loop._put_checkpoint_fut
|
||||
if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
|
||||
# send messages to executor
|
||||
for task in new_tasks:
|
||||
if task.scheduled:
|
||||
continue
|
||||
await self.producer.send(
|
||||
self.topics.executor,
|
||||
value=MessageToExecutor(
|
||||
config=patch_configurable(
|
||||
loop.config,
|
||||
{
|
||||
**loop.checkpoint_config["configurable"],
|
||||
CONFIG_KEY_DEDUPE_TASKS: True,
|
||||
},
|
||||
),
|
||||
task=ExecutorTask(
|
||||
id=task.id,
|
||||
path=task.path,
|
||||
step=loop.step,
|
||||
resuming=loop.input is INPUT_RESUMING,
|
||||
),
|
||||
),
|
||||
)
|
||||
# flush producer
|
||||
await self.producer.flush()
|
||||
# mark as scheduled
|
||||
for task in new_tasks:
|
||||
loop.put_writes(task.id, [(SCHEDULED, None)])
|
||||
# return message
|
||||
return msg
|
||||
@@ -0,0 +1,16 @@
|
||||
from typing import Any
|
||||
|
||||
import orjson
|
||||
|
||||
|
||||
def loads(v: bytes) -> Any:
|
||||
return orjson.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
|
||||
@@ -0,0 +1,25 @@
|
||||
from typing import Any, NamedTuple, Optional, TypedDict
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
|
||||
class Topics(NamedTuple):
|
||||
orchestrator: str
|
||||
executor: str
|
||||
|
||||
|
||||
class MessageToOrchestrator(TypedDict):
|
||||
input: Optional[dict[str, Any]]
|
||||
config: RunnableConfig
|
||||
|
||||
|
||||
class ExecutorTask(TypedDict):
|
||||
id: str
|
||||
path: tuple[str, ...]
|
||||
step: int
|
||||
resuming: bool
|
||||
|
||||
|
||||
class MessageToExecutor(TypedDict):
|
||||
config: RunnableConfig
|
||||
task: ExecutorTask
|
||||
Generated
+1182
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,57 @@
|
||||
[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"
|
||||
kafka-python-ng = "^2.2.2"
|
||||
|
||||
[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}
|
||||
|
||||
[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 = ["-s", "--ff", "-v", "--tb", "short"]
|
||||
patterns = ["*.py"]
|
||||
@@ -0,0 +1,20 @@
|
||||
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
|
||||
@@ -0,0 +1,58 @@
|
||||
from typing import AsyncIterator, Iterator
|
||||
from uuid import uuid4
|
||||
|
||||
import kafka.admin
|
||||
import pytest
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
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_{uuid4().hex[:16]}"
|
||||
e = f"test_{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),
|
||||
]
|
||||
)
|
||||
# yield topics
|
||||
yield Topics(orchestrator=o, executor=e)
|
||||
# delete topics
|
||||
admin.delete_topics([o, e])
|
||||
admin.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def checkpointer() -> 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 AsyncPostgresSaver.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + database
|
||||
) as checkpointer:
|
||||
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}")
|
||||
@@ -0,0 +1,124 @@
|
||||
import asyncio
|
||||
import operator
|
||||
from typing import Annotated, 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.executor import KafkaExecutor
|
||||
from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator
|
||||
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def mk_fanout_graph(checkpointer: BaseCheckpointSaver) -> 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:
|
||||
print("rewrite_query", data)
|
||||
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)
|
||||
|
||||
|
||||
async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -> None:
|
||||
graph = mk_fanout_graph(checkpointer)
|
||||
n_orch_msgs = 0
|
||||
n_exec_msgs = 0
|
||||
|
||||
async def orchestrator() -> None:
|
||||
nonlocal n_orch_msgs
|
||||
async with KafkaOrchestrator(graph, topics) as orch:
|
||||
async for msg in orch:
|
||||
n_orch_msgs += 1
|
||||
print("orch", msg)
|
||||
|
||||
async def executor() -> None:
|
||||
nonlocal n_exec_msgs
|
||||
async with KafkaExecutor(graph, topics) as exec:
|
||||
async for msg in exec:
|
||||
n_exec_msgs += 1
|
||||
print("exec", msg)
|
||||
|
||||
async with asyncio.TaskGroup() as tg:
|
||||
o = tg.create_task(orchestrator(), name="orchestrator")
|
||||
e = tg.create_task(executor(), name="executor")
|
||||
|
||||
# start a new run
|
||||
producer = AIOKafkaProducer(value_serializer=serde.dumps)
|
||||
await producer.start()
|
||||
await producer.send_and_wait(
|
||||
topics.orchestrator,
|
||||
MessageToOrchestrator(
|
||||
input={"query": "what is weather in sf"},
|
||||
config={"configurable": {"thread_id": "1"}},
|
||||
),
|
||||
)
|
||||
await producer.stop()
|
||||
|
||||
# wait for the run to finish
|
||||
await asyncio.sleep(5)
|
||||
o.cancel()
|
||||
e.cancel()
|
||||
|
||||
assert n_orch_msgs == 13
|
||||
assert n_exec_msgs == 12
|
||||
Reference in New Issue
Block a user