diff --git a/examples/Untitled.ipynb b/examples/Untitled.ipynb new file mode 100644 index 000000000..36f1f3a76 --- /dev/null +++ b/examples/Untitled.ipynb @@ -0,0 +1,139 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "2589ed8b-a781-45cc-aeb6-0eebb6469100", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.hub import pull" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "de600f76-dd51-415f-afae-f8574cb6f99e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " \n" + ] + } + ], + "source": [ + "pull('homanp/superagent')" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b4831d38-7bb1-49a2-9efa-802adc24110f", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.load import loads" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "2ef796f9-7b34-462e-acb6-25cd116b9045", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/nuno/dev/langchain/libs/core/langchain_core/_api/beta_decorator.py:162: LangChainBetaWarning: The function `loads` is in beta. It is actively being worked on, so the API may change.\n", + " warn_beta(\n" + ] + }, + { + "data": { + "text/plain": [ + "{}" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "loads('{}')" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "4b6d8788-0b6f-4bd6-b6cf-e2e4e80f4707", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_core.beta.runnables.context import ContextGet" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "6975707d-58eb-4642-9618-76e37cdeb5e1", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/nuno/dev/langchain/libs/core/langchain_core/_api/beta_decorator.py:162: LangChainBetaWarning: The class `ContextGet` is in beta. It is actively being worked on, so the API may change.\n", + " warn_beta(\n" + ] + }, + { + "data": { + "text/plain": [ + "ContextGet(key='hello')" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ContextGet(key='hello')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1977eda2-4c4a-4bb0-886a-fe274599dc50", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/langgraph.ipynb b/examples/langgraph.ipynb index 505ee9162..b502b1e33 100644 --- a/examples/langgraph.ipynb +++ b/examples/langgraph.ipynb @@ -14,13 +14,18 @@ "id": "d642e6af-217a-4414-a78c-509b44155eca", "metadata": {}, "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "input_variables=['agent_scratchpad', 'input'] input_types={'chat_history': typing.List[typing.Union[langchain_core.messages.ai.AIMessage, langchain_core.messages.human.HumanMessage, langchain_core.messages.chat.ChatMessage, langchain_core.messages.system.SystemMessage, langchain_core.messages.function.FunctionMessage, langchain_core.messages.tool.ToolMessage]], 'agent_scratchpad': typing.List[typing.Union[langchain_core.messages.ai.AIMessage, langchain_core.messages.human.HumanMessage, langchain_core.messages.chat.ChatMessage, langchain_core.messages.system.SystemMessage, langchain_core.messages.function.FunctionMessage, langchain_core.messages.tool.ToolMessage]]} messages=[SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=[], template='You are a helpful assistant')), MessagesPlaceholder(variable_name='chat_history', optional=True), HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], template='{input}')), MessagesPlaceholder(variable_name='agent_scratchpad')]\n" + ] + }, { "name": "stderr", "output_type": "stream", "text": [ - "/Users/nuno/dev/permchain/.venv/lib/python3.11/site-packages/langchain_core/_api/beta_decorator.py:160: LangChainBetaWarning: The function `loads` is in beta. It is actively being worked on, so the API may change.\n", - " warn_beta(\n", - "/Users/nuno/dev/permchain/.venv/lib/python3.11/site-packages/langchain_core/_api/deprecation.py:189: LangChainDeprecationWarning: The class `ChatOpenAI` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use langchain_openai.ChatOpenAI instead.\n", + "/Users/nuno/dev/langchain/libs/core/langchain_core/_api/deprecation.py:191: LangChainDeprecationWarning: The class `ChatOpenAI` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use langchain_openai.ChatOpenAI instead.\n", " warn_deprecated(\n" ] } @@ -92,25 +97,10 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "c46bd262-9605-4449-9391-f6b6e0fe440e", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'input': 'what is the weather in sf',\n", - " 'intermediate_steps': [(AgentActionMessageLog(tool='tavily_search_results_json', tool_input={'query': 'current weather in San Francisco'}, log=\"\\nInvoking: `tavily_search_results_json` with `{'query': 'current weather in San Francisco'}`\\n\\n\\n\", message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'tavily_search_results_json', 'arguments': '{\"query\":\"current weather in San Francisco\"}'}})]),\n", - " [{'url': 'https://www.weather25.com/north-america/usa/california/san-francisco',\n", - " 'content': 'will give you an idea of weather trends in San Francisco. For example, the weather in San Francisco in January 2024. San Francisco 14 day weather The weather today in San Francisco San Francisco weather report the weather in San Francisco including humidity, wind, chance of rain and more on the San Francisco current weather Weather25.com provides all the information that you need to know about the weather in San Francisco, United States.The current temperature in San Francisco is ° F. You can find more detailed information about the weather in San Francisco including humidity, wind, chance of rain and more on the San Francisco current weather page. The weather in San Francisco'}])],\n", - " 'agent_outcome': AgentFinish(return_values={'output': 'The current temperature in San Francisco is not available, but you can find detailed information about the weather in San Francisco, including humidity, wind, and chance of rain on the San Francisco current weather page. You can visit [Weather25.com](https://www.weather25.com/north-america/usa/california/san-francisco) for more information.'}, log='The current temperature in San Francisco is not available, but you can find detailed information about the weather in San Francisco, including humidity, wind, and chance of rain on the San Francisco current weather page. You can visit [Weather25.com](https://www.weather25.com/north-america/usa/california/san-francisco) for more information.')}" - ] - }, - "execution_count": 2, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "chain.invoke({\"input\": \"what is the weather in sf\", \"intermediate_steps\": []})" ] @@ -125,19 +115,10 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "f6f96e81-4a20-4599-a625-8d18df6fa76d", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/nuno/dev/permchain/.venv/lib/python3.11/site-packages/langchain_core/_api/deprecation.py:189: LangChainDeprecationWarning: The class `ChatOpenAI` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use langchain_openai.ChatOpenAI instead.\n", - " warn_deprecated(\n" - ] - } - ], + "outputs": [], "source": [ "from langchain.agents import AgentExecutor, BaseMultiActionAgent, Tool\n", "from langchain.schema import AgentAction, AgentFinish\n", @@ -268,7 +249,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "7708fa95-547b-4bea-b126-3656de7d5873", "metadata": {}, "outputs": [], @@ -334,21 +315,10 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "d6cdd1cd-e480-4dd7-99b4-9018eb243b4d", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "AgentFinish(return_values={'output': 'The current weather in San Francisco can be accessed through various weather reporting services, which provide up-to-date temperature, humidity, wind, and chance of rain information [1]. Historically, San Francisco experiences a Mediterranean climate with average temperatures ranging from the low 50s to mid-60s Fahrenheit. The city is known for its microclimates, leading to significant weather variations across different neighborhoods. Summer temperatures are often cooler compared to other California areas due to the cold California Current and frequent fog, particularly in June and July. Winters are mild and moist, with most rainfall occurring between November and March, averaging around 23 inches annually. Wind is a prominent feature, especially in spring. For historical weather extremes and average wind speeds, additional specific data would be required.\\n\\nReferences:\\n[1] https://www.weather25.com/north-america/usa/california/san-francisco', 'initial_answer': \"The weather in San Francisco (SF) is characterized by a mild, Mediterranean-like climate with wet winters and dry summers. The city's unique topography and coastal location result in microclimates, where weather conditions can vary significantly from one neighborhood to another. Average temperatures typically range from the low 50s to the mid-60s Fahrenheit throughout the year. Summers in San Francisco are often cooler than in other parts of California due to the cold California Current offshore and the presence of fog, particularly in June and July. The fog usually burns off by the afternoon, leading to clearer skies and slightly warmer temperatures. Winters are mild and moist, with the majority of the city's rainfall occurring between November and March. Rainfall averages around 23 inches annually. Wind is also a notable feature of San Francisco's weather, with spring being the windiest season. Despite the general patterns, it's always advisable to dress in layers due to the potential for rapid weather changes.\"}, log='Reached max steps.')" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "workflow = Graph()\n", "\n", @@ -411,7 +381,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.6" + "version": "3.11.5" } }, "nbformat": 4, diff --git a/permchain/langgraph/__init__.py b/permchain/langgraph/__init__.py index cca29a9d1..f968b98b9 100644 --- a/permchain/langgraph/__init__.py +++ b/permchain/langgraph/__init__.py @@ -107,7 +107,7 @@ class Graph: if hasattr(self, "finish_point"): outgoing_edges[self.finish_point].append(END) - chains = { + nodes = { key: ( Channel.subscribe_to(key) | node @@ -118,10 +118,10 @@ class Graph: for key, branches in self.branches.items(): for branch in branches: - chains[key] |= branch.runnable + nodes[key] |= branch.runnable return Pregel( - chains=chains, + nodes=nodes, input=self.entry_point, output=END, ) diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index e97926b4c..d7444b5d8 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -61,7 +61,7 @@ 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.validate import validate_graph from permchain.pregel.write import ChannelWrite WriteValue = Union[ @@ -81,12 +81,22 @@ def _coerce_write_value(value: WriteValue) -> Runnable[Input, Output]: class Channel: @overload @classmethod - def subscribe_to(cls, channels: str, key: Optional[str] = None) -> ChannelInvoke: + def subscribe_to( + cls, + channels: str, + key: Optional[str] = None, + when: Callable[[Any], bool] | None = None, + ) -> ChannelInvoke: ... @overload @classmethod - def subscribe_to(cls, channels: Sequence[str], key: None = None) -> ChannelInvoke: + def subscribe_to( + cls, + channels: Sequence[str], + key: None = None, + when: Callable[[Any], bool] | None = None, + ) -> ChannelInvoke: ... @classmethod @@ -134,7 +144,7 @@ class Channel: class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): - chains: Mapping[str, ChannelInvoke | ChannelBatch] + nodes: Mapping[str, ChannelInvoke | ChannelBatch] channels: Mapping[str, BaseChannel] = Field(default_factory=dict) @@ -153,15 +163,15 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): @root_validator(skip_on_failure=True) def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]: - validate_chains_channels( - values["chains"], values["channels"], values["input"], values["output"] + validate_graph( + values["nodes"], values["channels"], values["input"], values["output"] ) return values @property def config_specs(self) -> list[ConfigurableFieldSpec]: return get_unique_config_specs( - [spec for chain in self.chains.values() for spec in chain.config_specs] + [spec for node in self.nodes.values() for spec in node.config_specs] + (self.saver.config_specs if self.saver is not None else []) ) @@ -208,8 +218,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): ) -> Iterator[tuple[dict[str, Any] | Any, CheckpointView]]: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") - # copy chains to ignore mutations during execution - processes = {**self.chains} + # copy nodes to ignore mutations during execution + processes = {**self.nodes} # get checkpoint from saver, or create an empty one checkpoint = self.saver.get(config) if self.saver else None checkpoint = checkpoint or empty_checkpoint() @@ -305,8 +315,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): ) -> AsyncIterator[tuple[dict[str, Any] | Any, CheckpointView]]: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") - # copy chains to ignore mutations during execution - processes = {**self.chains} + # copy nodes to ignore mutations during execution + processes = {**self.nodes} # get checkpoint from saver, or create an empty one checkpoint = await self.saver.aget(config) if self.saver else None checkpoint = checkpoint or empty_checkpoint() diff --git a/permchain/pregel/model.py b/permchain/pregel/model.py new file mode 100644 index 000000000..8607ed374 --- /dev/null +++ b/permchain/pregel/model.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, ClassVar, Dict +from langchain_core.pydantic_v1 import dataclasses + +from permchain.checkpoint.base import BaseCheckpointAdapter +from permchain.pregel import Pregel + +dataclasses.DataclassClassOrWrapper + +from langchain_core.messages import BaseMessage +from langchain_core.runnables import RunnableLambda + +from permchain.pregel.read import ChannelInvoke + + +def on_change(field_name: str): + def decorator(func): + return ChannelInvoke(bound=RunnableLambda(func), triggers=[field_name]) + + return decorator + + +class DataclassProtocol: + __dataclass_fields__: ClassVar[Dict[str, Any]] + + +def thread(data: type[DataclassProtocol], thread_id: str, saver: BaseCheckpointAdapter): + return Pregel( + chains={ + v.__name__: v + for v in data.__dict__.values() + if isinstance(v, ChannelInvoke) + } + ) + + +@dataclass +class Agent: + messages: list[BaseMessage] = field(default_factory=list) + actions: list[BaseMessage] = field(default_factory=list) + + @on_change("messages") + def plan(self): + ... + + @on_change("actions") + def execute(self): + ... diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py index 5c0dd6718..66a5c06b4 100644 --- a/permchain/pregel/validate.py +++ b/permchain/pregel/validate.py @@ -6,21 +6,21 @@ from permchain.pregel.read import ChannelBatch, ChannelInvoke from permchain.pregel.reserved import ReservedChannels -def validate_chains_channels( - chains: Mapping[str, ChannelInvoke | ChannelBatch], +def validate_graph( + nodes: Mapping[str, ChannelInvoke | ChannelBatch], channels: dict[str, BaseChannel], input: str | Sequence[str], output: str | Sequence[str], ) -> None: subscribed_channels = set[str]() - for chain in chains.values(): - if isinstance(chain, ChannelInvoke): - subscribed_channels.update(chain.channels.values()) - elif isinstance(chain, ChannelBatch): - subscribed_channels.add(chain.channel) + for node in nodes.values(): + if isinstance(node, ChannelInvoke): + subscribed_channels.update(node.channels.values()) + elif isinstance(node, ChannelBatch): + subscribed_channels.add(node.channel) else: raise TypeError( - f"Invalid chain type {type(chain)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()" + f"Invalid node type {type(node)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()" ) for chan in subscribed_channels: @@ -31,14 +31,14 @@ def validate_chains_channels( if input not in channels: 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") + raise ValueError(f"Input channel {input} is not subscribed to by any node") else: for chan in input: if chan not in channels: 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" + f"None of the input channels {input} are subscribed to by any node" ) if isinstance(output, str): diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 163c9ea11..1a1187367 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -23,7 +23,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={ + nodes={ "one": chain, }, channels={ @@ -51,7 +51,7 @@ def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) - add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel(chains={"one": chain}) + app = Pregel(nodes={"one": chain}) assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == {"title": "PregelOutput"} @@ -66,7 +66,7 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) ) - app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"]) + app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"]) assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == { @@ -90,7 +90,7 @@ def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) -> | Channel.write_to("output") ) - app = Pregel(chains={"one": chain}) + app = Pregel(nodes={"one": chain}) assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == {"title": "PregelOutput"} @@ -103,7 +103,7 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={ + nodes={ "one": chain, }, output=["output"], @@ -123,7 +123,7 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={ + nodes={ "one": chain, }, input=["input"], @@ -145,11 +145,11 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, ) assert app.invoke(2) == 4 @@ -233,13 +233,11 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_two = ( - Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") - ) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, channels={"inbox": Topic(int)}, input=["input", "inbox"], ) @@ -252,14 +250,10 @@ def test_batch_two_processes_in_out() -> None: time.sleep(inp / 10) return inp + 1 - chain_one = ( - Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") - ) - chain_two = ( - Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") - ) + one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") + two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") - app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) + app = Pregel(nodes={"one": one, "two": two}) assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] @@ -278,14 +272,14 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chains[str(i)] = ( + nodes[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains) + app = Pregel(nodes=nodes) for _ in range(10): assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size @@ -300,14 +294,14 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chains[str(i)] = ( + nodes[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains) + app = Pregel(nodes=nodes) for _ in range(3): assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ @@ -331,10 +325,10 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) + app = Pregel(nodes={"one": one, "two": two}) with pytest.raises(InvalidUpdateError): # LastValue channels can only be updated once per iteration @@ -344,11 +338,11 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, channels={"output": Topic(int)}, ) @@ -364,7 +358,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: raise ValueError("Input is too large") return input - chain_one = ( + one = ( Channel.subscribe_to(["input"]).join(["total"]) | add_one | Channel.write_to("output", "total") @@ -374,7 +368,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: memory = MemoryCheckpoint() app = Pregel( - chains={"chain_one": chain_one}, + nodes={"one": one}, channels={"total": BinaryOperatorAggregate(int, operator.add)}, saver=memory, ) @@ -410,15 +404,15 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") chain_four = ( Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output") ) app = Pregel( - chains={ - "chain_one": chain_one, + nodes={ + "one": one, "chain_three": chain_three, "chain_four": chain_four, }, @@ -440,17 +434,17 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) inner_app = Pregel( - chains={ + nodes={ "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") } ) - chain_one = ( + one = ( Channel.subscribe_to("input") | add_10_each | Channel.write_to("inbox_one").map() ) - chain_two = ( + two = ( Channel.subscribe_to("inbox_one") | inner_app.map() | sorted @@ -459,9 +453,9 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") app = Pregel( - chains={ - "chain_one": chain_one, - "chain_two": chain_two, + nodes={ + "one": one, + "two": two, "chain_three": chain_three, }, channels={"inbox_one": Topic(int)}, @@ -477,28 +471,24 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = ( + one = ( Channel.subscribe_to("input") | add_one | Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) - chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") + two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - ) + app = Pregel(nodes={"one": one, "two": two}) assert [c for c in app.stream(2)] == [3, 4] def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") - chain_two = Channel.subscribe_to("between") | add_one + one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") + two = Channel.subscribe_to("between") | add_one - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - ) + app = Pregel(nodes={"one": one, "two": two}) # It finishes executing (once no more messages being published) # but returns nothing, as nothing was published to OUT topic @@ -508,13 +498,11 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - chain_two = Channel.subscribe_to("between") | add_one + one = Channel.subscribe_to("between") | add_one | Channel.write_to("output") + two = Channel.subscribe_to("between") | add_one with pytest.raises(ValueError): - Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - ) + Pregel(nodes={"one": one, "two": two}) def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: @@ -530,13 +518,11 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: cleanup() add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_two = ( - Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") - ) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, channels={ "inbox": Topic(int), "ctx": Context(an_int, typ=int), diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 90e82d94c..879df9dfe 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -22,7 +22,7 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={ + nodes={ "one": chain, }, channels={ @@ -39,12 +39,12 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: async def test_invoke_single_process_in_out_implicit_channels( - mocker: MockerFixture + mocker: MockerFixture, ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel(chains={"one": chain}) + app = Pregel(nodes={"one": chain}) assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == {"title": "PregelOutput"} @@ -59,7 +59,7 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) ) - app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"]) + app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"]) assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == { @@ -75,7 +75,7 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N async def test_invoke_single_process_in_out_reserved_is_last( - mocker: MockerFixture + mocker: MockerFixture, ) -> None: add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1}) @@ -85,7 +85,7 @@ async def test_invoke_single_process_in_out_reserved_is_last( | Channel.write_to("output") ) - app = Pregel(chains={"one": chain}) + app = Pregel(nodes={"one": chain}) assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == {"title": "PregelOutput"} @@ -101,9 +101,7 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={ - "one": chain, - }, + nodes={"one": chain}, output=["output"], ) @@ -121,7 +119,7 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={ + nodes={ "one": chain, }, input=["input"], @@ -143,12 +141,10 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - ) + app = Pregel(nodes={"one": one, "two": two}) assert await app.ainvoke(2) == 4 @@ -188,13 +184,11 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_two = ( - Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") - ) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") pubsub = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, channels={"inbox": Topic(int)}, input=["input", "inbox"], ) @@ -208,15 +202,11 @@ async def test_batch_two_processes_in_out() -> None: await asyncio.sleep(inp / 10) return inp + 1 - chain_one = ( - Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") - ) - chain_two = ( - Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") - ) + one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") + two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, channels={"one": LastValue(int)}, ) @@ -227,14 +217,14 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chains[str(i)] = ( + nodes[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains) + app = Pregel(nodes=nodes) # No state is left over from previous invocations for _ in range(10): @@ -250,14 +240,14 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chains[str(i)] = ( + nodes[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains) + app = Pregel(nodes=nodes) # No state is left over from previous invocations for _ in range(3): @@ -284,10 +274,10 @@ async def test_invoke_two_processes_two_in_two_out_invalid( ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) + app = Pregel(nodes={"one": one, "two": two}) with pytest.raises(InvalidUpdateError): # LastValue channels can only be updated once per iteration @@ -297,11 +287,11 @@ async def test_invoke_two_processes_two_in_two_out_invalid( async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, channels={"output": Topic(int)}, ) @@ -317,7 +307,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: raise ValueError("Input is too large") return input - chain_one = ( + one = ( Channel.subscribe_to(["input"]).join(["total"]) | add_one | Channel.write_to("output", "total") @@ -327,7 +317,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: memory = MemoryCheckpoint() app = Pregel( - chains={"chain_one": chain_one}, + nodes={"one": one}, channels={"total": BinaryOperatorAggregate(int, operator.add)}, saver=memory, ) @@ -363,15 +353,15 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") chain_four = ( Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output") ) app = Pregel( - chains={ - "chain_one": chain_one, + nodes={ + "one": one, "chain_three": chain_three, "chain_four": chain_four, }, @@ -394,17 +384,17 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) inner_app = Pregel( - chains={ + nodes={ "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") } ) - chain_one = ( + one = ( Channel.subscribe_to("input") | add_10_each | Channel.write_to("inbox_one").map() ) - chain_two = ( + two = ( Channel.subscribe_to("inbox_one") | inner_app.map() | sorted @@ -413,9 +403,9 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") app = Pregel( - chains={ - "chain_one": chain_one, - "chain_two": chain_two, + nodes={ + "one": one, + "two": two, "chain_three": chain_three, }, channels={ @@ -436,14 +426,14 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = ( + one = ( Channel.subscribe_to("input") | add_one | Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) - chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") + two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) + app = Pregel(nodes={"one": one, "two": two}) # Then invoke pubsub assert [c async for c in app.astream(2)] == [3, 4] @@ -451,10 +441,10 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") - chain_two = Channel.subscribe_to("between") | add_one + one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") + two = Channel.subscribe_to("between") | add_one - app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) + app = Pregel(nodes={"one": one, "two": two}) # It finishes executing (once no more messages being published) # but returns nothing, as nothing was published to "output" topic @@ -484,13 +474,11 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: cleanup_async() add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_two = ( - Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") - ) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, + nodes={"one": one, "two": two}, channels={ "inbox": Topic(int), "ctx": Context(an_int, an_int_async, typ=int),