diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index 2619f828a..28fea5e36 100644 --- a/examples/draft-revise-loop.py +++ b/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 diff --git a/examples/rag.py b/examples/rag.py index 64fe5306c..14a5a5c50 100644 --- a/examples/rag.py +++ b/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 = ( diff --git a/permchain/__init__.py b/permchain/__init__.py index 2f7bb6275..c4637331c 100644 --- a/permchain/__init__.py +++ b/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"] diff --git a/permchain/checkpoint/base.py b/permchain/checkpoint/base.py index c735d32c8..3b8252d4b 100644 --- a/permchain/checkpoint/base.py +++ b/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): diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 49ad344fc..6fad69399 100644 --- a/permchain/pregel/__init__.py +++ b/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(): diff --git a/permchain/pregel/reserved.py b/permchain/pregel/reserved.py new file mode 100644 index 000000000..aa138d092 --- /dev/null +++ b/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.""" diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py index 927a78021..344630f67 100644 --- a/permchain/pregel/validate.py +++ b/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) diff --git a/permchain/utils.py b/permchain/utils.py new file mode 100644 index 000000000..ef5e2c13d --- /dev/null +++ b/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 diff --git a/tests/test_pregel.py b/tests/test_pregel.py index d129c016e..dd1730747 100644 --- a/tests/test_pregel.py +++ b/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") diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 091a11234..95b4c776e 100644 --- a/tests/test_pregel_async.py +++ b/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")