Add concept of reserved channels (ie. channels updated by Pregel), Add a reserved channel for is_last_step

examples/draft-revise-loop.py

@@ -6,7 +6,6 @@ from langchain.prompts import SystemMessagePromptTemplate
 from langchain.schema.output_parser import StrOutputParser
 
 from permchain import Channel, Pregel
-from permchain.channels import LastValue
 
 # prompts
 

examples/rag.py

@@ -1,6 +1,11 @@
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.vectorstores import FAISS
+from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain.schema.messages import AIMessage, AnyMessage, FunctionMessage
+from langchain.vectorstores import FAISS
+
+from permchain import Channel, Pregel
+from permchain.channels import Topic
 
 texts = ["harrison went to kensho"]
 embeddings = OpenAIEmbeddings()
@@ -8,19 +13,11 @@ db = FAISS.from_texts(texts, embeddings)
 
 retriever = db.as_retriever()
 
-from langchain.prompts import PromptTemplate
 
 prompt = PromptTemplate.from_template(
     """Answer the question "{question}"  based on the following context: {context}"""
 )
 
-from langchain.schema.messages import AIMessage, AnyMessage, FunctionMessage
-
-from permchain import Channel, Pregel
-from permchain.channels import Topic
-
-from langchain.chat_models import ChatOpenAI
-
 model = ChatOpenAI()
 
 chain = (

permchain/__init__.py

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

permchain/checkpoint/base.py

@@ -1,5 +1,4 @@
 import asyncio
-import enum
 from abc import ABC, abstractmethod
 from typing import Any, Mapping, Sequence
 
@@ -7,12 +6,7 @@ from langchain.load.serializable import Serializable
 from langchain.schema.runnable import RunnableConfig
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
-
-# Before Python 3.11 native StrEnum is not available
-class StrEnum(str, enum.Enum):
-    """A string enum."""
-
-    pass
+from permchain.utils import StrEnum
 
 
 class CheckpointAt(StrEnum):

permchain/pregel/__init__.py

@@ -50,6 +50,7 @@ from permchain.pregel.debug import print_checkpoint, print_step_start
 from permchain.pregel.io import map_input, map_output
 from permchain.pregel.log import logger
 from permchain.pregel.read import ChannelBatch, ChannelInvoke
+from permchain.pregel.reserved import ReservedChannels
 from permchain.pregel.validate import validate_chains_channels
 from permchain.pregel.write import ChannelWrite
 
@@ -179,6 +180,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         run_manager: CallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> Iterator[dict[str, Any] | Any]:
+        if config["recursion_limit"] < 1:
+            raise ValueError("recursion_limit must be at least 1")
         processes = {**self.chains}
         checkpoint = (
             self.checkpoint.get(config) if self.checkpoint is not None else None
@@ -190,6 +193,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 processes,
                 channels,
                 deque(w for c in input for w in map_input(self.input, c)),
+                config,
+                0,
             )
 
             def read(chan: str) -> Any:
@@ -237,8 +242,9 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 _interrupt_or_proceed(done, inflight, step)
 
                 # apply writes to channels, decide on next step
+
                 next_tasks = _apply_writes_and_prepare_next_tasks(
-                    processes, channels, pending_writes
+                    processes, channels, pending_writes, config, step + 1
                 )
 
                 if self.debug:
@@ -274,6 +280,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         run_manager: AsyncCallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> AsyncIterator[dict[str, Any] | Any]:
+        if config["recursion_limit"] < 1:
+            raise ValueError("recursion_limit must be at least 1")
         processes = {**self.chains}
         checkpoint = (
             await self.checkpoint.aget(config) if self.checkpoint is not None else None
@@ -283,6 +291,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 processes,
                 channels,
                 deque([w async for c in input for w in map_input(self.input, c)]),
+                config,
+                0,
             )
 
             def read(chan: str) -> Any:
@@ -334,7 +344,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
                 # apply writes to channels, decide on next step
                 next_tasks = _apply_writes_and_prepare_next_tasks(
-                    processes, channels, pending_writes
+                    processes, channels, pending_writes, config, step + 1
                 )
 
                 if self.debug:
@@ -456,12 +466,23 @@ def _apply_writes_and_prepare_next_tasks(
     processes: Mapping[str, ChannelInvoke | ChannelBatch],
     channels: Mapping[str, BaseChannel],
     pending_writes: Sequence[tuple[str, Any]],
+    config: RunnableConfig,
+    for_step: int,
 ) -> list[tuple[Runnable, Any, str]]:
     pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
     # Group writes by channel
     for chan, val in pending_writes:
+        if chan in [c.value for c in ReservedChannels]:
+            raise ValueError(f"Can't write to reserved channel {chan}")
         pending_writes_by_channel[chan].append(val)
 
+    # Update reserved channels
+    pending_writes_by_channel[ReservedChannels.is_last_step] = [
+        for_step + 1 == config["recursion_limit"]
+    ]
+
+    print(for_step, config["recursion_limit"])
+
     updated_channels: set[str] = set()
     # Apply writes to channels
     for chan, vals in pending_writes_by_channel.items():

permchain/pregel/reserved.py

@@ -0,0 +1,8 @@
+from enum import StrEnum
+
+
+class ReservedChannels(StrEnum):
+    """Channels managed by the framework."""
+
+    is_last_step = "is_last_step"
+    """A channel that is True if the current step is the last step, False otherwise."""

permchain/pregel/validate.py

@@ -4,6 +4,7 @@ from permchain.channels.base import BaseChannel
 from permchain.channels.last_value import LastValue
 from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION
 from permchain.pregel.read import ChannelBatch, ChannelInvoke
+from permchain.pregel.reserved import ReservedChannels
 
 FORBIDDEN_CHANNEL_NAMES = {
     CHECKPOINT_KEY_TS,
@@ -57,3 +58,7 @@ def validate_chains_channels(
     for name in FORBIDDEN_CHANNEL_NAMES:
         if name in channels:
             raise ValueError(f"Channel name {name} is reserved")
+
+    for chan in ReservedChannels:
+        if chan not in channels:
+            channels[chan] = LastValue(Any)

permchain/utils.py

@@ -0,0 +1,8 @@
+import enum
+
+
+# Before Python 3.11 native StrEnum is not available
+class StrEnum(str, enum.Enum):
+    """A string enum."""
+
+    pass

tests/test_pregel.py

@@ -15,6 +15,7 @@ from permchain.channels.context import Context
 from permchain.channels.last_value import LastValue
 from permchain.channels.topic import Topic
 from permchain.checkpoint.memory import MemoryCheckpoint
+from permchain.pregel.reserved import ReservedChannels
 
 
 def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -49,6 +50,23 @@ def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -
     assert app.invoke(2) == 3
 
 
+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})
+
+    chain = (
+        Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
+        | add_one
+        | Channel.write_to("output")
+    )
+
+    app = Pregel(chains={"one": chain})
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {"title": "PregelOutput"}
+    assert app.invoke(2) == {"input": 3, "is_last_step": False}
+    assert app.invoke(2, {"recursion_limit": 1}) == {"input": 3, "is_last_step": True}
+
+
 def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")

tests/test_pregel_async.py

@@ -14,6 +14,7 @@ from permchain.channels.context import Context
 from permchain.channels.last_value import LastValue
 from permchain.channels.topic import Topic
 from permchain.checkpoint.memory import MemoryCheckpoint
+from permchain.pregel.reserved import ReservedChannels
 
 
 async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -50,6 +51,28 @@ async def test_invoke_single_process_in_out_implicit_channels(
     assert await app.ainvoke(2) == 3
 
 
+async 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})
+
+    chain = (
+        Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
+        | add_one
+        | Channel.write_to("output")
+    )
+
+    app = Pregel(chains={"one": chain})
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {"title": "PregelOutput"}
+    assert await app.ainvoke(2) == {"input": 3, "is_last_step": False}
+    assert await app.ainvoke(2, {"recursion_limit": 1}) == {
+        "input": 3,
+        "is_last_step": True,
+    }
+
+
 async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")