Allow passing values directly to Channel.write_toi

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()]
             )
         )
 

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})
 

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: