mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 11:19:54 +02:00
Implement dict input, add more tests
This commit is contained in:
Nuno Campos
@@ -36,8 +36,8 @@ Check `examples` for more examples.
|
||||
- [x] Finish updating tests to new API
|
||||
- [x] Implement input_schema and output_schema in Pregel
|
||||
- [ ] More tests
|
||||
- [ ] Test different input and output types (str, str sequence, None)
|
||||
- [ ] Add tests for Stream, UniqueInbox
|
||||
- [x] Test different input and output types (str, str sequence)
|
||||
- [x] Add tests for Stream, UniqueInbox
|
||||
- [ ] Add tests for subscribe_to_each().join()
|
||||
- [ ] Implement checkpointing
|
||||
- [ ] Save checkpoints at end of each step
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
from permchain.channels.base import Channel, EmptyChannelError, InvalidUpdateError
|
||||
from permchain.channels.binop import BinaryOperatorAggregate
|
||||
from permchain.channels.context import ContextManager
|
||||
from permchain.channels.inbox import Inbox, UniqueInbox
|
||||
from permchain.channels.last_value import LastValue
|
||||
from permchain.channels.binop import BinaryOperatorAggregate
|
||||
from permchain.channels.stream import Set, Stream
|
||||
from permchain.channels.context import ContextManager
|
||||
|
||||
__all__ = [
|
||||
"Channel",
|
||||
|
||||
+33
-10
@@ -288,7 +288,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
|
||||
|
||||
output: str | Sequence[str]
|
||||
|
||||
input: str | None
|
||||
input: str | Sequence[str]
|
||||
|
||||
step_timeout: Optional[float] = None
|
||||
|
||||
@@ -300,7 +300,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
|
||||
*chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch,
|
||||
channels: Mapping[str, Channel],
|
||||
output: str | Sequence[str],
|
||||
input: str | None = None,
|
||||
input: str | Sequence[str],
|
||||
step_timeout: Optional[float] = None,
|
||||
) -> None:
|
||||
chains_flat: list[PregelInvoke | PregelBatch] = []
|
||||
@@ -402,8 +402,13 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
|
||||
processes,
|
||||
channels,
|
||||
deque((self.input, chunk) for chunk in input)
|
||||
if self.input is not None
|
||||
else deque((k, v) for chunk in input for k, v in chunk.items()),
|
||||
if isinstance(self.input, str)
|
||||
else deque(
|
||||
(k, v)
|
||||
for chunk in input
|
||||
for k, v in chunk.items()
|
||||
if k in self.input
|
||||
),
|
||||
)
|
||||
|
||||
def read(chan: str) -> Any:
|
||||
@@ -495,8 +500,15 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
|
||||
processes,
|
||||
channels,
|
||||
deque([(self.input, chunk) async for chunk in input])
|
||||
if self.input is not None
|
||||
else deque([(k, v) async for chunk in input for k, v in chunk.items()]),
|
||||
if isinstance(self.input, str)
|
||||
else deque(
|
||||
[
|
||||
(k, v)
|
||||
async for chunk in input
|
||||
for k, v in chunk.items()
|
||||
if k in self.input
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
def read(chan: str) -> Any:
|
||||
@@ -695,7 +707,7 @@ def _apply_writes_and_prepare_next_tasks(
|
||||
def validate_chains_channels(
|
||||
chains: Sequence[PregelInvoke | PregelBatch],
|
||||
channels: Mapping[str, Channel],
|
||||
input: str | None,
|
||||
input: str | Sequence[str],
|
||||
output: str | Sequence[str],
|
||||
) -> None:
|
||||
subscribed_channels = set[str]()
|
||||
@@ -709,13 +721,24 @@ def validate_chains_channels(
|
||||
f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()"
|
||||
)
|
||||
|
||||
if input is not None and input not in subscribed_channels:
|
||||
raise ValueError(f"Input channel {input} is not subscribed to by any chain")
|
||||
|
||||
for chan in subscribed_channels:
|
||||
if chan not in channels:
|
||||
raise ValueError(f"Channel {chan} is subscribed to, but not initialized")
|
||||
|
||||
if isinstance(input, str):
|
||||
if input not in subscribed_channels:
|
||||
raise ValueError(f"Input channel {input} is not subscribed to by any chain")
|
||||
else:
|
||||
for chan in input:
|
||||
if chan not in subscribed_channels:
|
||||
raise ValueError(
|
||||
f"Input channel {chan} is not subscribed to by any chain"
|
||||
)
|
||||
|
||||
if isinstance(output, str):
|
||||
if output not in channels:
|
||||
raise ValueError(f"Output channel {output} is not initialized")
|
||||
else:
|
||||
for chan in output:
|
||||
if chan not in channels:
|
||||
raise ValueError(f"Output channel {chan} is not initialized")
|
||||
|
||||
+1
-1
@@ -67,6 +67,6 @@ asyncio_mode = "auto"
|
||||
#
|
||||
# https://github.com/tophat/syrupy
|
||||
# --snapshot-warn-unused Prints a warning on unused snapshots rather than fail the test suite.
|
||||
addopts = "--full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused"
|
||||
addopts = "-x --full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused"
|
||||
# Registering custom markers.
|
||||
# https://docs.pytest.org/en/7.1.x/example/markers.html#registering-markers
|
||||
|
||||
+95
-11
@@ -1,5 +1,7 @@
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator
|
||||
|
||||
import pytest
|
||||
from langchain.schema.runnable import RunnablePassthrough
|
||||
@@ -51,12 +53,39 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
assert app.invoke(2) == {"output": 3}
|
||||
|
||||
|
||||
def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
|
||||
|
||||
app = Pregel(
|
||||
(chain,),
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
"output": channels.LastValue(int),
|
||||
},
|
||||
input=["input"],
|
||||
output=["output"],
|
||||
)
|
||||
|
||||
assert app.input_schema.schema() == {
|
||||
"title": "PregelInput",
|
||||
"type": "object",
|
||||
"properties": {"input": {"title": "Input", "type": "integer"}},
|
||||
}
|
||||
assert app.output_schema.schema() == {
|
||||
"title": "PregelOutput",
|
||||
"type": "object",
|
||||
"properties": {"output": {"title": "Output", "type": "integer"}},
|
||||
}
|
||||
assert app.invoke({"input": 2}) == {"output": 3}
|
||||
|
||||
|
||||
def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
|
||||
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
|
||||
|
||||
pubsub = Pregel(
|
||||
app = Pregel(
|
||||
[chain_one, chain_two],
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
@@ -67,8 +96,26 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
output="output",
|
||||
)
|
||||
|
||||
# Then invoke pubsub
|
||||
assert pubsub.invoke(2) == 4
|
||||
assert app.invoke(2) == 4
|
||||
|
||||
|
||||
def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
|
||||
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
|
||||
|
||||
app = Pregel(
|
||||
[chain_one, chain_two],
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
"output": channels.LastValue(int),
|
||||
"inbox": channels.Inbox(int),
|
||||
},
|
||||
input=["input", "inbox"],
|
||||
output="output",
|
||||
)
|
||||
|
||||
assert [*app.stream({"input": 2, "inbox": 12})] == [13, 4] # [12 + 1, 2 + 1 + 1]
|
||||
|
||||
|
||||
def test_batch_two_processes_in_out() -> None:
|
||||
@@ -83,7 +130,7 @@ def test_batch_two_processes_in_out() -> None:
|
||||
Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output")
|
||||
)
|
||||
|
||||
pubsub = Pregel(
|
||||
app = Pregel(
|
||||
chain_one,
|
||||
chain_two,
|
||||
channels={
|
||||
@@ -95,8 +142,7 @@ def test_batch_two_processes_in_out() -> None:
|
||||
output="output",
|
||||
)
|
||||
|
||||
# Then invoke pubsub
|
||||
assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
|
||||
assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
|
||||
|
||||
|
||||
def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
@@ -151,7 +197,6 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
app = Pregel(*chains, channels=chans, input="input", output="output")
|
||||
|
||||
for _ in range(10):
|
||||
# Then invoke pubsub
|
||||
assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
|
||||
2 + test_size,
|
||||
1 + test_size,
|
||||
@@ -244,7 +289,7 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
|
||||
assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100
|
||||
|
||||
|
||||
def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None:
|
||||
def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
|
||||
|
||||
@@ -283,7 +328,6 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None:
|
||||
output="output",
|
||||
)
|
||||
|
||||
# Then invoke pubsub
|
||||
for _ in range(10):
|
||||
assert app.invoke([2, 3]) == 27
|
||||
|
||||
@@ -313,7 +357,6 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
|
||||
output="output",
|
||||
)
|
||||
|
||||
# Then invoke pubsub
|
||||
assert [c for c in app.stream(2)] == [3, 4]
|
||||
|
||||
|
||||
@@ -334,7 +377,6 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
||||
output="output",
|
||||
)
|
||||
|
||||
# Then invoke pubsub
|
||||
# It finishes executing (once no more messages being published)
|
||||
# but returns nothing, as nothing was published to OUT topic
|
||||
assert app.invoke(2) is None
|
||||
@@ -358,3 +400,45 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
|
||||
input="input",
|
||||
output="output",
|
||||
)
|
||||
|
||||
|
||||
def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
|
||||
setup = mocker.Mock()
|
||||
cleanup = mocker.Mock()
|
||||
|
||||
@contextmanager
|
||||
def an_int() -> Generator[int, None, None]:
|
||||
setup()
|
||||
try:
|
||||
yield 5
|
||||
finally:
|
||||
cleanup()
|
||||
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
|
||||
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
|
||||
|
||||
app = Pregel(
|
||||
[chain_one, chain_two],
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
"output": channels.LastValue(int),
|
||||
"inbox": channels.Inbox(int),
|
||||
"ctx": channels.ContextManager(an_int, typ=int),
|
||||
},
|
||||
input="input",
|
||||
output=["inbox", "output"],
|
||||
)
|
||||
|
||||
assert setup.call_count == 0
|
||||
assert cleanup.call_count == 0
|
||||
for i, chunk in enumerate(app.stream(2)):
|
||||
assert setup.call_count == 1, "Expected setup to be called once"
|
||||
assert cleanup.call_count == 0, "Expected cleanup to not be called yet"
|
||||
if i == 0:
|
||||
assert chunk == {"inbox": (3,)}
|
||||
elif i == 1:
|
||||
assert chunk == {"output": 4}
|
||||
else:
|
||||
assert False, "Expected only two chunks"
|
||||
assert cleanup.call_count == 1, "Expected cleanup to be called once"
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import Any, AsyncGenerator, AsyncIterator, Generator
|
||||
|
||||
import pytest
|
||||
from langchain.schema.runnable import RunnablePassthrough
|
||||
@@ -25,6 +27,56 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
assert await app.ainvoke(2) == 3
|
||||
|
||||
|
||||
async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
|
||||
|
||||
app = Pregel(
|
||||
(chain,),
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
"output": channels.LastValue(int),
|
||||
},
|
||||
input="input",
|
||||
output=["output"],
|
||||
)
|
||||
|
||||
assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
|
||||
assert app.output_schema.schema() == {
|
||||
"title": "PregelOutput",
|
||||
"type": "object",
|
||||
"properties": {"output": {"title": "Output", "type": "integer"}},
|
||||
}
|
||||
assert await app.ainvoke(2) == {"output": 3}
|
||||
|
||||
|
||||
async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
|
||||
|
||||
app = Pregel(
|
||||
(chain,),
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
"output": channels.LastValue(int),
|
||||
},
|
||||
input=["input"],
|
||||
output=["output"],
|
||||
)
|
||||
|
||||
assert app.input_schema.schema() == {
|
||||
"title": "PregelInput",
|
||||
"type": "object",
|
||||
"properties": {"input": {"title": "Input", "type": "integer"}},
|
||||
}
|
||||
assert app.output_schema.schema() == {
|
||||
"title": "PregelOutput",
|
||||
"type": "object",
|
||||
"properties": {"output": {"title": "Output", "type": "integer"}},
|
||||
}
|
||||
assert await app.ainvoke({"input": 2}) == {"output": 3}
|
||||
|
||||
|
||||
async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
|
||||
@@ -44,6 +96,26 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
assert await app.ainvoke(2) == 4
|
||||
|
||||
|
||||
async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
|
||||
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
|
||||
|
||||
pubsub = Pregel(
|
||||
[chain_one, chain_two],
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
"output": channels.LastValue(int),
|
||||
"inbox": channels.Inbox(int),
|
||||
},
|
||||
input=["input", "inbox"],
|
||||
output="output",
|
||||
)
|
||||
|
||||
# [12 + 1, 2 + 1 + 1]
|
||||
assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [13, 4]
|
||||
|
||||
|
||||
async def test_batch_two_processes_in_out() -> None:
|
||||
async def add_one_with_delay(inp: int) -> int:
|
||||
await asyncio.sleep(inp / 10)
|
||||
@@ -318,3 +390,68 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
||||
# It finishes executing (once no more messages being published)
|
||||
# but returns nothing, as nothing was published to OUT topic
|
||||
assert await app.ainvoke(2) is None
|
||||
|
||||
|
||||
async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
|
||||
setup_sync = mocker.Mock()
|
||||
cleanup_sync = mocker.Mock()
|
||||
setup_async = mocker.Mock()
|
||||
cleanup_async = mocker.Mock()
|
||||
|
||||
@contextmanager
|
||||
def an_int() -> Generator[int, None, None]:
|
||||
setup_sync()
|
||||
try:
|
||||
yield 5
|
||||
finally:
|
||||
cleanup_sync()
|
||||
|
||||
@asynccontextmanager
|
||||
async def an_int_async() -> AsyncGenerator[int, None]:
|
||||
setup_async()
|
||||
try:
|
||||
yield 5
|
||||
finally:
|
||||
cleanup_async()
|
||||
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
|
||||
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
|
||||
|
||||
app = Pregel(
|
||||
[chain_one, chain_two],
|
||||
channels={
|
||||
"input": channels.LastValue(int),
|
||||
"output": channels.LastValue(int),
|
||||
"inbox": channels.Inbox(int),
|
||||
"ctx": channels.ContextManager(an_int, an_int_async, typ=int),
|
||||
},
|
||||
input="input",
|
||||
output=["inbox", "output"],
|
||||
)
|
||||
|
||||
async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]:
|
||||
i = 0
|
||||
async for chunk in aiter:
|
||||
yield i, chunk
|
||||
i += 1
|
||||
|
||||
assert setup_sync.call_count == 0
|
||||
assert cleanup_sync.call_count == 0
|
||||
assert setup_async.call_count == 0
|
||||
assert cleanup_async.call_count == 0
|
||||
async for i, chunk in aenumerate(app.astream(2)):
|
||||
assert setup_sync.call_count == 0, "Sync context manager should not be used"
|
||||
assert cleanup_sync.call_count == 0, "Sync context manager should not be used"
|
||||
assert setup_async.call_count == 1, "Expected setup to be called once"
|
||||
assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet"
|
||||
if i == 0:
|
||||
assert chunk == {"inbox": (3,)}
|
||||
elif i == 1:
|
||||
assert chunk == {"output": 4}
|
||||
else:
|
||||
assert False, "Expected only two chunks"
|
||||
assert setup_sync.call_count == 0
|
||||
assert cleanup_sync.call_count == 0
|
||||
assert setup_async.call_count == 1, "Expected setup to be called once"
|
||||
assert cleanup_async.call_count == 1, "Expected cleanup to be called once"
|
||||
|
||||
Reference in New Issue
Block a user