WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-23 16:12:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Rename chains to nodes

Browse Source
This commit is contained in:
Nuno Campos
2024-01-06 12:43:15 -08:00
parent 17cd659533
commit de599be8a1
8 changed files with 342 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/Untitled.ipynb
+139
View File
@@ -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": [
"<frame at 0x1057c46c0, file '/Users/nuno/dev/permchain/.venv/lib/python3.11/site-packages/langchain/hub.py', line 81, code pull> <module 'langchain.hub' from '/Users/nuno/dev/permchain/.venv/lib/python3.11/site-packages/langchain/hub.py'>\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
}
examples/langgraph.ipynb
+16 -46
View File
@@ -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,
permchain/langgraph/__init__.py
+3 -3
View File
@@ -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,
)
permchain/pregel/__init__.py
+21 -11
View File
@@ -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()
permchain/pregel/model.py
+50
View File
@@ -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):
...
permchain/pregel/validate.py
+10 -10
View File
@@ -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):
tests/test_pregel.py
+52 -66
View File
@@ -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),
tests/test_pregel_async.py
+51 -63
View File
@@ -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),