diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 95eb46438..caff6da15 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -6,11 +6,14 @@ from collections import defaultdict, deque from typing import ( Any, AsyncIterator, + Awaitable, + Callable, Iterator, Mapping, Optional, Sequence, Type, + Union, cast, overload, ) @@ -26,7 +29,7 @@ from langchain.schema.runnable import ( RunnablePassthrough, RunnableSerializable, ) -from langchain.schema.runnable.base import RunnableLike, coerce_to_runnable +from langchain.schema.runnable.base import Input, Output, coerce_to_runnable from langchain.schema.runnable.config import ( RunnableConfig, get_executor_for_config, @@ -54,6 +57,19 @@ from permchain.pregel.reserved import ReservedChannels from permchain.pregel.validate import validate_chains_channels from permchain.pregel.write import ChannelWrite +WriteValue = Union[ + Runnable[Input, Output], + Callable[[Input], Output], + Callable[[Input], Awaitable[Output]], + Any, +] + + +def _coerce_write_value(value: WriteValue) -> Runnable[Input, Output]: + if not isinstance(value, Runnable) and not callable(value): + return coerce_to_runnable(lambda _: value) + return coerce_to_runnable(value) + class Channel: @overload @@ -94,13 +110,13 @@ class Channel: def write_to( cls, *channels: str, - **kwargs: RunnableLike, + **kwargs: WriteValue, ) -> ChannelWrite: """Writes to channels the result of the lambda, or None to skip writing.""" return ChannelWrite( channels=( [(c, RunnablePassthrough()) for c in channels] - + [(k, coerce_to_runnable(v)) for k, v in kwargs.items()] + + [(k, _coerce_write_value(v)) for k, v in kwargs.items()] ) ) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index dd1730747..a2777eec0 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -50,6 +50,29 @@ def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) - assert app.invoke(2) == 3 +def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = ( + Channel.subscribe_to("input") + | add_one + | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + ) + + app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"]) + + assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.output_schema.schema() == { + "title": "PregelOutput", + "type": "object", + "properties": { + "output": {"title": "Output"}, + "fixed": {"title": "Fixed"}, + "output_plus_one": {"title": "Output Plus One"}, + }, + } + assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} + + def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1}) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 95b4c776e..402dd62be 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -51,6 +51,29 @@ async def test_invoke_single_process_in_out_implicit_channels( assert await app.ainvoke(2) == 3 +async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = ( + Channel.subscribe_to("input") + | add_one + | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + ) + + app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"]) + + assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.output_schema.schema() == { + "title": "PregelOutput", + "type": "object", + "properties": { + "output": {"title": "Output"}, + "fixed": {"title": "Fixed"}, + "output_plus_one": {"title": "Output Plus One"}, + }, + } + assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} + + async def test_invoke_single_process_in_out_reserved_is_last( mocker: MockerFixture ) -> None: