Merge pull request #20 from langchain-ai/nc/28nov

Update to new langchain release

permchain/__init__.py

@@ -1,4 +1,10 @@
+from permchain.checkpoint.base import BaseCheckpointAdapter, CheckpointAt
 from permchain.pregel import Channel, Pregel, ReservedChannels
-from permchain.pregel.read import ChannelRead
 
-__all__ = ["Channel", "Pregel", "ReservedChannels", "ChannelRead"]
+__all__ = [
+    "Channel",
+    "Pregel",
+    "ReservedChannels",
+    "BaseCheckpointAdapter",
+    "CheckpointAt",
+]

permchain/channels/base.py

@@ -54,7 +54,7 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
 
     @asynccontextmanager
     async def aempty(
-        self, checkpoint: Optional[str] = None
+        self, checkpoint: Optional[Checkpoint] = None
     ) -> AsyncGenerator[Self, None]:
         """Return a new identical channel, optionally initialized from a checkpoint."""
         with self.empty(checkpoint) as value:

permchain/channels/topic.py

@@ -43,7 +43,7 @@ class Topic(
         return Sequence[self.typ]  # type: ignore[name-defined]
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> Any:
         """The type of the update received by the channel."""
         return Union[self.typ, list[self.typ]]  # type: ignore[name-defined]
 

permchain/checkpoint/base.py

@@ -1,6 +1,6 @@
 import asyncio
 from abc import ABC, abstractmethod
-from typing import Any, Mapping, Sequence
+from typing import Any, Mapping
 
 from langchain.load.serializable import Serializable
 from langchain.schema.runnable import RunnableConfig
@@ -18,7 +18,7 @@ class BaseCheckpointAdapter(Serializable, ABC):
     at: CheckpointAt = CheckpointAt.END_OF_RUN
 
     @property
-    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+    def config_specs(self) -> list[ConfigurableFieldSpec]:
         return []
 
     @abstractmethod

permchain/checkpoint/memory.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Mapping, Sequence
+from typing import Any, Dict, Mapping
 
 from langchain.pydantic_v1 import Field
 from langchain.schema.runnable import RunnableConfig
@@ -11,7 +11,7 @@ class MemoryCheckpoint(BaseCheckpointAdapter):
     storage: Dict[str, Mapping[str, Any]] = Field(default_factory=dict)
 
     @property
-    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+    def config_specs(self) -> list[ConfigurableFieldSpec]:
         return [
             ConfigurableFieldSpec(
                 id="thread_id",
@@ -19,6 +19,7 @@ class MemoryCheckpoint(BaseCheckpointAdapter):
                 name="Thread ID",
                 description=None,
                 default="",
+                is_shared=True,
             ),
         ]
 

permchain/pregel/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import concurrent.futures
 from collections import defaultdict, deque
+from functools import partial
 from typing import (
     Any,
     AsyncIterator,
@@ -26,7 +27,6 @@ from langchain.globals import get_debug
 from langchain.pydantic_v1 import BaseModel, Field, create_model, root_validator
 from langchain.schema.runnable import (
     Runnable,
-    RunnablePassthrough,
     RunnableSerializable,
 )
 from langchain.schema.runnable.base import Input, Output, coerce_to_runnable
@@ -98,7 +98,8 @@ class Channel:
                 {key: channels}
                 if isinstance(channels, str)
                 else {chan: chan for chan in channels},
-            )
+            ),
+            triggers=[channels] if isinstance(channels, str) else channels,
         )
 
     @classmethod
@@ -115,7 +116,7 @@ class Channel:
         """Writes to channels the result of the lambda, or None to skip writing."""
         return ChannelWrite(
             channels=(
-                [(c, RunnablePassthrough()) for c in channels]
+                [(c, None) for c in channels]
                 + [(k, _coerce_write_value(v)) for k, v in kwargs.items()]
             )
         )
@@ -150,9 +151,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
     def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
         return get_unique_config_specs(
             [spec for chain in self.chains.values() for spec in chain.config_specs]
-            + self.checkpoint.config_specs
-            if self.checkpoint is not None
-            else []
+            + (self.checkpoint.config_specs if self.checkpoint is not None else [])
         )
 
     @property
@@ -213,11 +212,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 0,
             )
 
-            def read(chan: str) -> Any:
-                try:
-                    return channels[chan].get()
-                except EmptyChannelError:
-                    return None
+            if not next_tasks:
+                return
+
+            read = partial(_read_channel, channels)
 
             # Similarly to Bulk Synchronous Parallel / Pregel model
             # computation proceeds in steps, while there are channel updates
@@ -311,11 +309,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 0,
             )
 
-            def read(chan: str) -> Any:
-                try:
-                    return channels[chan].get()
-                except EmptyChannelError:
-                    return None
+            if not next_tasks:
+                return
+
+            read = partial(_read_channel, channels)
 
             # Similarly to Bulk Synchronous Parallel / Pregel model
             # computation proceeds in steps, while there are channel updates
@@ -478,6 +475,15 @@ def _interrupt_or_proceed(
         raise TimeoutError(f"Timed out at step {step}")
 
 
+def _read_channel(
+    channels: Mapping[str, BaseChannel], chan: str, catch: bool = True
+) -> Any:
+    try:
+        return channels[chan].get()
+    except EmptyChannelError:
+        return None
+
+
 def _apply_writes_and_prepare_next_tasks(
     processes: Mapping[str, ChannelInvoke | ChannelBatch],
     channels: Mapping[str, BaseChannel],
@@ -516,19 +522,24 @@ def _apply_writes_and_prepare_next_tasks(
     for name, proc in processes.items():
         if isinstance(proc, ChannelInvoke):
             # If any of the channels read by this process were updated
-            if any(chan in updated_channels for chan in proc.channels.values()):
-                # If all channels read by this process have been initialized
+            if any(chan in updated_channels for chan in proc.triggers):
+                # If all channels subscribed by this process have been initialized
                 try:
-                    val = {k: channels[chan].get() for k, chan in proc.channels.items()}
+                    val = {
+                        k: _read_channel(
+                            channels, chan, catch=chan not in proc.triggers
+                        )
+                        for k, chan in proc.channels.items()
+                    }
                 except EmptyChannelError:
                     continue
 
                 # Processes that subscribe to a single keyless channel get
                 # the value directly, instead of a dict
                 if list(proc.channels.keys()) == [None]:
-                    tasks.append((proc, val[None], name))
-                else:
-                    tasks.append((proc, val, name))
+                    val = val[None]
+
+                tasks.append((proc, val, name))
         elif isinstance(proc, ChannelBatch):
             # If the channel read by this process was updated
             if proc.channel in updated_channels:

permchain/pregel/read.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Mapping, Optional, Sequence
+from typing import Any, Callable, List, Mapping, Optional, Sequence
 
 from langchain.pydantic_v1 import Field
 from langchain.schema.runnable import (
@@ -67,6 +67,10 @@ default_bound = RunnablePassthrough()
 class ChannelInvoke(RunnableBindingBase):
     channels: Mapping[None, str] | Mapping[str, str]
 
+    triggers: List[str] = Field(default_factory=list)
+
+    skip: Optional[Callable[[Any], bool]] = None
+
     bound: Runnable[Any, Any] = Field(default=default_bound)
 
     kwargs: Mapping[str, Any] = Field(default_factory=dict)
@@ -74,6 +78,7 @@ class ChannelInvoke(RunnableBindingBase):
     def __init__(
         self,
         channels: Mapping[None, str] | Mapping[str, str],
+        triggers: Sequence[str],
         *,
         bound: Optional[Runnable[Any, Any]] = None,
         kwargs: Optional[Mapping[str, Any]] = None,
@@ -82,6 +87,7 @@ class ChannelInvoke(RunnableBindingBase):
     ) -> None:
         super().__init__(
             channels=channels,
+            triggers=triggers,
             bound=bound or default_bound,
             kwargs=kwargs or {},
             config=config,
@@ -92,13 +98,19 @@ class ChannelInvoke(RunnableBindingBase):
         assert isinstance(channels, list) or isinstance(
             channels, tuple
         ), "channels must be a list or tuple"
-        joiner = RunnablePassthrough.assign(
-            **{chan: ChannelRead(chan) for chan in channels}
+        assert all(
+            k is not None for k in self.channels.keys()
+        ), "all channels must be named when using .join()"
+        return ChannelInvoke(
+            channels={
+                **self.channels,
+                **{chan: chan for chan in channels},
+            },
+            triggers=self.triggers,
+            bound=self.bound,
+            kwargs=self.kwargs,
+            config=self.config,
         )
-        if self.bound is default_bound:
-            return ChannelInvoke(channels=self.channels, bound=joiner)
-        else:
-            return ChannelInvoke(channels=self.channels, bound=self.bound | joiner)
 
     def __or__(
         self,
@@ -108,11 +120,21 @@ class ChannelInvoke(RunnableBindingBase):
     ) -> ChannelInvoke:
         if self.bound is default_bound:
             return ChannelInvoke(
-                channels=self.channels, bound=coerce_to_runnable(other)
+                channels=self.channels,
+                triggers=self.triggers,
+                bound=coerce_to_runnable(other),
+                kwargs=self.kwargs,
+                config=self.config,
             )
         else:
-            # delegate to __or__ in self.bound
-            return ChannelInvoke(channels=self.channels, bound=self.bound | other)
+            return ChannelInvoke(
+                channels=self.channels,
+                triggers=self.triggers,
+                # delegate to __or__ in self.bound
+                bound=self.bound | other,
+                kwargs=self.kwargs,
+                config=self.config,
+            )
 
     def __ror__(
         self,

permchain/pregel/validate.py

@@ -14,7 +14,7 @@ FORBIDDEN_CHANNEL_NAMES = {
 
 def validate_chains_channels(
     chains: Mapping[str, ChannelInvoke | ChannelBatch],
-    channels: Mapping[str, BaseChannel],
+    channels: dict[str, BaseChannel],
     input: str | Sequence[str],
     output: str | Sequence[str],
 ) -> None:
@@ -31,17 +31,17 @@ def validate_chains_channels(
 
     for chan in subscribed_channels:
         if chan not in channels:
-            channels[chan] = LastValue(Any)
+            channels[chan] = LastValue(Any)  # type: ignore[arg-type]
 
     if isinstance(input, str):
         if input not in channels:
-            channels[input] = LastValue(Any)
+            channels[input] = LastValue(Any)  # type: ignore[arg-type]
         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 channels:
-                channels[chan] = LastValue(Any)
+                channels[chan] = LastValue(Any)  # type: ignore[arg-type]
         if all(chan not in subscribed_channels for chan in input):
             raise ValueError(
                 f"None of the input channels {input} are subscribed to by any chain"
@@ -49,11 +49,11 @@ def validate_chains_channels(
 
     if isinstance(output, str):
         if output not in channels:
-            channels[output] = LastValue(Any)
+            channels[output] = LastValue(Any)  # type: ignore[arg-type]
     else:
         for chan in output:
             if chan not in channels:
-                channels[chan] = LastValue(Any)
+                channels[chan] = LastValue(Any)  # type: ignore[arg-type]
 
     for name in FORBIDDEN_CHANNEL_NAMES:
         if name in channels:
@@ -61,4 +61,4 @@ def validate_chains_channels(
 
     for chan in ReservedChannels:
         if chan not in channels:
-            channels[chan] = LastValue(Any)
+            channels[chan] = LastValue(Any)  # type: ignore[arg-type]

permchain/pregel/write.py

@@ -15,7 +15,7 @@ TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 
 
 class ChannelWrite(RunnablePassthrough):
-    channels: Sequence[tuple[str, Runnable]]
+    channels: Sequence[tuple[str, Runnable | None]]
     """
     Mapping of write channels to Runnables that return the value to be written,
     or None to skip writing.
@@ -27,7 +27,7 @@ class ChannelWrite(RunnablePassthrough):
     def __init__(
         self,
         *,
-        channels: Sequence[tuple[str, Runnable]],
+        channels: Sequence[tuple[str, Runnable | None]],
     ):
         super().__init__(func=self._write, afunc=self._awrite, channels=channels)
 
@@ -44,12 +44,17 @@ class ChannelWrite(RunnablePassthrough):
         ]
 
     def _write(self, input: Any, config: RunnableConfig) -> None:
-        values = [(chan, r.invoke(input, config)) for chan, r in self.channels]
+        values = [
+            (chan, r.invoke(input, config) if r else input) for chan, r in self.channels
+        ]
 
         self.do_write(config, **dict(values))
 
     async def _awrite(self, input: Any, config: RunnableConfig) -> None:
-        values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels]
+        values = [
+            (chan, await r.ainvoke(input, config) if r else input)
+            for chan, r in self.channels
+        ]
 
         self.do_write(config, **dict(values))
 

tests/test_pregel.py

@@ -37,6 +37,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
     assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
     assert app.invoke(2) == 3
+    assert repr(app), "does not raise recursion error"
 
 
 def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None:
@@ -280,27 +281,39 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
         | raise_if_above_10
     )
 
+    memory = MemoryCheckpoint()
+
     app = Pregel(
         chains={"chain_one": chain_one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpoint=MemoryCheckpoint(),
+        checkpoint=memory,
     )
 
     # total starts out as 0, so output is 0+2=2
     assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 2
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 2
     # total is now 2, so output is 2+3=5
     assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
     with pytest.raises(ValueError):
         app.invoke(4, {"configurable": {"thread_id": "1"}})
     # checkpoint is not updated
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # on a new thread, total starts out as 0, so output is 0+5=5
     assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
-    assert app.checkpoint.get({"configurable": {"thread_id": "2"}}).get("total") == 5
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
+    checkpoint = memory.get({"configurable": {"thread_id": "2"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 5
 
 
 def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:

tests/test_pregel_async.py

@@ -294,37 +294,39 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
         | raise_if_above_10
     )
 
+    memory = MemoryCheckpoint()
+
     app = Pregel(
         chains={"chain_one": chain_one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpoint=MemoryCheckpoint(),
+        checkpoint=memory,
     )
 
     # total starts out as 0, so output is 0+2=2
     assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 2
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 2
     # total is now 2, so output is 2+3=5
     assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 7
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
     with pytest.raises(ValueError):
         await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
     # checkpoint is not updated
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 7
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # on a new thread, total starts out as 0, so output is 0+5=5
     assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 7
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "2"}})).get(
-        "total"
-    ) == 5
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
+    checkpoint = await memory.aget({"configurable": {"thread_id": "2"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 5
 
 
 async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: