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

Merge pull request #9 from langchain-ai/harrison/pregel-combine

Browse Source 
Harrison/pregel combine
This commit is contained in:
Nuno Campos
2023-10-23 12:02:41 +01:00
committed by GitHub
parent 62e4df815a 49ff5e2eab
commit 4f206fcc49
11 changed files with 552 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/combine_docs.ipynb
+368
View File
@@ -0,0 +1,368 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "780c1001-557c-4b03-8ebd-a2a381d5f85d",
"metadata": {},
"source": [
"# Combine Docs\n",
"\n",
"PermChain is a great choice for implementating workflows that involve operating over longer documents because of its recursive nature"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "624c452c-ddd5-4390-9065-7ec55dc64b96",
"metadata": {},
"outputs": [],
"source": [
"from langchain.chat_models.openai import ChatOpenAI\n",
"from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
"from langchain.schema.output_parser import StrOutputParser\n",
"from langchain.schema.runnable import Runnable, RunnablePassthrough\n",
"from langchain.schema.output_parser import StrOutputParser\n",
"from langchain.schema.document import Document\n",
"from langchain.schema import format_document\n",
"\n",
"from permchain import Pregel, PregelRead, channels\n"
]
},
{
"cell_type": "markdown",
"id": "271728d7-b3c8-4ec6-a728-19835e282ec3",
"metadata": {},
"source": [
"## Stuff Documents\n",
"\n",
"Stuff documents is simple - just a chain"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "0462aff0-1b88-49cc-bfe2-3c169d5e1d63",
"metadata": {},
"outputs": [],
"source": [
"from langchain.schema.runnable import RunnableLambda\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "59d6430b-c113-4498-9ffc-f4623f7a0b5c",
"metadata": {},
"outputs": [],
"source": [
"DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template=\"{page_content}\")\n",
"\n",
"_combine_documents = RunnableLambda(\n",
" lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n",
").map() | (lambda x: \"\\n\\n\".join(x))\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "29b2668d-e4a6-4876-9b04-bdc841774c62",
"metadata": {},
"outputs": [],
"source": [
"docs = [\n",
" Document(page_content=\"Harrison used to work at Kensho\"),\n",
" Document(page_content=\"Ankush worked at Facebook\"),\n",
"]\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "17da58b7-8685-4d0a-9a47-c398c085d477",
"metadata": {},
"outputs": [],
"source": [
"stuff_chain = (\n",
" {\n",
" \"question\": lambda x: x[\"question\"],\n",
" \"context\": (lambda x: x[\"docs\"]) | _combine_documents,\n",
" }\n",
" | ChatPromptTemplate.from_messages(\n",
" [\n",
" (\n",
" \"system\",\n",
" \"Answer user questions based on the following documents:\\n\\n{context}\",\n",
" ),\n",
" (\"human\", \"{question}\"),\n",
" ]\n",
" )\n",
" | ChatOpenAI()\n",
" | StrOutputParser()\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "87295b71-0afc-4901-b57c-a7b945aa4bd9",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'Harrison used to work at Kensho.'"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})\n"
]
},
{
"cell_type": "markdown",
"id": "fff324c1-7fbf-41e5-861f-a10ba0112dbd",
"metadata": {},
"source": [
"## Reduce Documents\n",
"\n",
"Reduce documents tries to merge documents recursively."
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "b15f5abb-1cfe-4965-a021-c891506c5dd2",
"metadata": {},
"outputs": [],
"source": [
"many_docs = docs * 5\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "ccad04a3-fd3f-4e73-b895-29e53535f000",
"metadata": {},
"outputs": [],
"source": [
"def _split_list_of_docs(docs, max_length=70):\n",
" new_result_doc_list = []\n",
" _sub_result_docs = []\n",
" for doc in docs:\n",
" _sub_result_docs.append(doc)\n",
" _num_tokens = sum([len(d.page_content) for d in _sub_result_docs])\n",
" if _num_tokens > max_length:\n",
" if len(_sub_result_docs) == 1:\n",
" raise ValueError(\n",
" \"A single document was longer than the context length,\"\n",
" \" we cannot handle this.\"\n",
" )\n",
" new_result_doc_list.append(_sub_result_docs[:-1])\n",
" _sub_result_docs = _sub_result_docs[-1:]\n",
" new_result_doc_list.append(_sub_result_docs)\n",
" return new_result_doc_list\n"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "11cfd337-9f3b-4b26-ba30-251e17b18994",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[[Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook')],\n",
" [Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook')],\n",
" [Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook')],\n",
" [Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook')],\n",
" [Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook')]]"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Just to show what its like split\n",
"split_docs = _split_list_of_docs(many_docs)\n",
"split_docs\n"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "8d524ba6-0939-4a5d-8db0-4fa1ef06eaeb",
"metadata": {},
"outputs": [],
"source": [
"chans = {\n",
" # input\n",
" \"question\": channels.LastValue(str),\n",
" \"docs\": channels.Inbox(Document),\n",
" # intermediate\n",
" \"docs_to_finalize\": channels.Inbox(Document),\n",
" # output\n",
" \"answer\": channels.LastValue(str),\n",
"}\n"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "67370694-86f4-4b64-9d4f-38b2e306abeb",
"metadata": {},
"outputs": [],
"source": [
"def decide(docs: list[Document]) -> Runnable:\n",
" if len(_split_list_of_docs(docs)) > 1:\n",
" # send back to the beginning if we still need to collapse more\n",
" return Pregel.write_to(\"docs\")\n",
" else:\n",
" # send to the finalizer if we're ready to produce final answer\n",
" return Pregel.write_to(\"docs_to_finalize\")\n",
"\n",
"\n",
"collapse = (\n",
" Pregel.subscribe_to(\"docs\")\n",
" | _split_list_of_docs\n",
" | {\"docs_list\": RunnablePassthrough(), \"question\": PregelRead(\"question\")}\n",
" # {docs: list[list[Doc]], question: str} -> list[{docs: list[Doc], question: str}]\n",
" | (lambda x: [{\"docs\": docs, \"question\": x[\"question\"]} for docs in x[\"docs_list\"]])\n",
" | stuff_chain.map() # Collapse each list of docs to a single string\n",
" | (lambda x: [Document(page_content=s) for s in x]) # A new (smaller) list of docs\n",
" | decide\n",
")\n",
"\n",
"# Convert final set of docs to an answer\n",
"finalize = (\n",
" Pregel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n",
" | stuff_chain\n",
" | Pregel.write_to(\"answer\")\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "3019e7d2-ab7f-4868-b43c-ad898d824a26",
"metadata": {},
"outputs": [],
"source": [
"reduce_chain = Pregel(\n",
" chains={\n",
" \"collapse\": collapse,\n",
" \"finalize\": finalize,\n",
" },\n",
" channels=chans,\n",
" input=[\"question\", \"docs\"],\n",
" output=\"answer\",\n",
" debug=True,\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "69fcb829-3dae-432a-8db3-11bbb179a7d2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 0 with 1 task. Next tasks:\n",
"\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook')))\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 0. Channel values:\n",
"\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n",
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 1 with 1 task. Next tasks:\n",
"\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.')))\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 1. Channel values:\n",
"\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n",
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 2 with 1 task. Next tasks:\n",
"\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.')))\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 2. Channel values:\n",
"\u001b[0m{'docs': (...),\n",
" 'docs_to_finalize': (...),\n",
" 'question': 'where did harrison work'}\n",
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 3 with 1 task. Next tasks:\n",
"\u001b[0m- finalize({'docs': (Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'))})\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 3. Channel values:\n",
"\u001b[0m{'answer': 'Harrison used to work at Kensho.',\n",
" 'docs': (...),\n",
" 'docs_to_finalize': (...),\n",
" 'question': 'where did harrison work'}\n"
]
},
{
"data": {
"text/plain": [
"'Harrison used to work at Kensho.'"
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "265b29cd-d4f4-4e48-8d4e-b759e909ac2e",
"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/draft-revise-loop.py
+8 -6
View File
@@ -5,7 +5,7 @@ from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
from langchain.prompts import SystemMessagePromptTemplate
from langchain.schema.output_parser import StrOutputParser
from permchain import Pregel, channels
from permchain import Channels, Pregel
# prompts
@@ -74,6 +74,12 @@ reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
# application
channels = {
"question": Channels.LastValue(str),
"draft": Channels.LastValue(str),
"notes": Channels.LastValue(str),
}
drafter = (
# subscribe to question channel as a dict with a single key, "question"
Pregel.subscribe_to(["question"])
@@ -102,11 +108,7 @@ reviser = (
)
draft_revise_loop = Pregel(
channels={
"question": channels.LastValue(str),
"draft": channels.LastValue(str),
"notes": channels.LastValue(str),
},
channels=channels,
chains={
"drafter": drafter,
"editor": editor,
examples/readme.py
+2 -2
View File
@@ -1,4 +1,4 @@
from permchain import Pregel, channels
from permchain import Channels, Pregel
grow_value = (
Pregel.subscribe_to("value")
@@ -8,7 +8,7 @@ grow_value = (
app = Pregel(
chains={"grow_value": grow_value},
channels={"value": channels.LastValue(str)},
channels={"value": Channels.LastValue(str)},
input="value",
output="value",
)
examples/recursive-web-loader.py
+10 -9
View File
@@ -6,7 +6,7 @@ from langchain.schema import Document
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
from langchain.utils.html import extract_sub_links
from permchain import Pregel, channels
from permchain import Channels, Pregel
# Load url with sync httpx client
@@ -82,6 +82,14 @@ def recursive_web_loader(
# assign default extractors
extractor = extractor or (lambda x: x)
metadata_extractor = metadata_extractor or _metadata_extractor
# define the channels
channels = {
"base_url": Channels.LastValue(str),
"next_urls": Channels.UniqueInbox(str),
"documents": Channels.Stream(Document),
"visited": Channels.Set(str),
"client": Channels.ContextManager(httpx_client, httpx_aclient),
}
# the main chain that gets executed recursively
visitor = (
# while there are urls in next_urls
@@ -112,20 +120,13 @@ def recursive_web_loader(
)
)
return Pregel(
channels=channels,
chains={
# use the base_url as the first url to visit
"input": Pregel.subscribe_to("base_url") | Pregel.write_to("next_urls"),
# add the main chain
"visitor": visitor,
},
# define the channels
channels={
"base_url": channels.LastValue(str),
"next_urls": channels.UniqueInbox(str),
"documents": channels.Stream(Document),
"visited": channels.Set(str),
"client": channels.ContextManager(httpx_client, httpx_aclient),
},
# this will accept a string as input
input="base_url",
# and return a dict with documents and visited set
permchain/__init__.py
+3 -2
View File
@@ -1,4 +1,5 @@
import permchain.channels as channels
import permchain.channels as Channels
from permchain.pregel import Pregel
from permchain.pregel.read import PregelRead
__all__ = ["channels", "Pregel"]
__all__ = ["Channels", "Pregel", "PregelRead"]
permchain/channels/__init__.py
-4
View File
@@ -1,4 +1,3 @@
from permchain.channels.base import Channel, EmptyChannelError, InvalidUpdateError
from permchain.channels.binop import BinaryOperatorAggregate
from permchain.channels.context import ContextManager
from permchain.channels.inbox import Inbox, UniqueInbox
@@ -6,9 +5,6 @@ from permchain.channels.last_value import LastValue
from permchain.channels.stream import Set, Stream
__all__ = [
"Channel",
"EmptyChannelError",
"InvalidUpdateError",
"LastValue",
"Inbox",
"UniqueInbox",
permchain/pregel/__init__.py
+29 -3
View File
@@ -3,7 +3,17 @@ from __future__ import annotations
import asyncio
import concurrent.futures
from collections import defaultdict, deque
from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast
from typing import (
Any,
AsyncIterator,
Iterator,
Mapping,
Optional,
Sequence,
Type,
cast,
overload,
)
from langchain.callbacks.manager import (
AsyncCallbackManagerForChainRun,
@@ -96,14 +106,30 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
**{k: (self.channels[k].ValueType, None) for k in self.output},
)
@overload
@classmethod
def subscribe_to(cls, channels: str | Sequence[str]) -> PregelInvoke:
def subscribe_to(cls, channels: str, key: Optional[str] = None) -> PregelInvoke:
...
@overload
@classmethod
def subscribe_to(cls, channels: Sequence[str], key: None = None) -> PregelInvoke:
...
@classmethod
def subscribe_to(
cls, channels: str | Sequence[str], key: Optional[str] = None
) -> PregelInvoke:
"""Runs process.invoke() each time channels are updated,
with a dict of the channel values as input."""
if not isinstance(channels, str) and key is not None:
raise ValueError(
"Can't specify a key when subscribing to multiple channels"
)
return PregelInvoke(
channels=cast(
Mapping[None, str] | Mapping[str, str],
{None: channels}
{key: channels}
if isinstance(channels, str)
else {chan: chan for chan in channels},
)
permchain/pregel/validate.py
+4 -5
View File
@@ -29,11 +29,10 @@ def validate_chains_channels(
if input not in subscribed_channels:
raise ValueError(f"Input channel {input} is not subscribed to by any chain")
else:
for chan in input:
if chan not in subscribed_channels:
raise ValueError(
f"Input channel {chan} is not subscribed to by any chain"
)
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"
)
if isinstance(output, str):
if output not in channels:
tests/test_channels.py
+27 -26
View File
@@ -6,17 +6,18 @@ import httpx
import pytest
from pytest_mock import MockerFixture
import permchain.channels as channels
import permchain.channels as Channels
from permchain.channels.base import EmptyChannelError, InvalidUpdateError
def test_last_value() -> None:
with channels.LastValue(int).empty() as channel:
with Channels.LastValue(int).empty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
with pytest.raises(channels.EmptyChannelError):
with pytest.raises(EmptyChannelError):
channel.get()
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
channel.update([5, 6])
channel.update([3])
@@ -26,13 +27,13 @@ def test_last_value() -> None:
async def test_last_value_async() -> None:
async with channels.LastValue(int).aempty() as channel:
async with Channels.LastValue(int).aempty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
with pytest.raises(channels.EmptyChannelError):
with pytest.raises(EmptyChannelError):
channel.get()
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
channel.update([5, 6])
channel.update([3])
@@ -42,11 +43,11 @@ async def test_last_value_async() -> None:
def test_inbox() -> None:
with channels.Inbox(str).empty() as channel:
with Channels.Inbox(str).empty() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, Sequence[str]]
with pytest.raises(channels.EmptyChannelError):
with pytest.raises(EmptyChannelError):
channel.get()
channel.update(["a", "b"])
@@ -56,11 +57,11 @@ def test_inbox() -> None:
async def test_inbox_async() -> None:
async with channels.Inbox(str).aempty() as channel:
async with Channels.Inbox(str).aempty() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, Sequence[str]]
with pytest.raises(channels.EmptyChannelError):
with pytest.raises(EmptyChannelError):
channel.get()
channel.update(["a", "b"])
@@ -71,7 +72,7 @@ async def test_inbox_async() -> None:
def test_set() -> None:
with channels.Set(str).empty() as channel:
with Channels.Set(str).empty() as channel:
assert channel.ValueType is FrozenSet[str]
assert channel.UpdateType is str
@@ -83,7 +84,7 @@ def test_set() -> None:
async def test_set_async() -> None:
async with channels.Set(str).aempty() as channel:
async with Channels.Set(str).aempty() as channel:
assert channel.ValueType is FrozenSet[str]
assert channel.UpdateType is str
@@ -95,11 +96,11 @@ async def test_set_async() -> None:
def test_binop() -> None:
with channels.BinaryOperatorAggregate(int, operator.add).empty() as channel:
with Channels.BinaryOperatorAggregate(int, operator.add).empty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
with pytest.raises(channels.EmptyChannelError):
with pytest.raises(EmptyChannelError):
channel.get()
channel.update([1, 2, 3])
@@ -109,11 +110,11 @@ def test_binop() -> None:
async def test_binop_async() -> None:
async with channels.BinaryOperatorAggregate(int, operator.add).aempty() as channel:
async with Channels.BinaryOperatorAggregate(int, operator.add).aempty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
with pytest.raises(channels.EmptyChannelError):
with pytest.raises(EmptyChannelError):
channel.get()
channel.update([1, 2, 3])
@@ -134,17 +135,17 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
finally:
cleanup()
with channels.ContextManager(an_int, None, int).empty() as channel:
with Channels.ContextManager(an_int, None, int).empty() as channel:
assert setup.call_count == 1
assert cleanup.call_count == 0
assert channel.ValueType is int
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is None
assert channel.get() == 5
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
channel.update([5]) # type: ignore
assert setup.call_count == 1
@@ -152,14 +153,14 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
with channels.ContextManager(httpx.Client).empty() as channel:
with Channels.ContextManager(httpx.Client).empty() as channel:
assert channel.ValueType is httpx.Client
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is None
assert isinstance(channel.get(), httpx.Client)
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
channel.update([5]) # type: ignore
@@ -182,17 +183,17 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None:
finally:
cleanup()
async with channels.ContextManager(an_int_sync, an_int, int).aempty() as channel:
async with Channels.ContextManager(an_int_sync, an_int, int).aempty() as channel:
assert setup.call_count == 1
assert cleanup.call_count == 0
assert channel.ValueType is int
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is None
assert channel.get() == 5
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
channel.update([5]) # type: ignore
assert setup.call_count == 1
tests/test_pregel.py
+52 -51
View File
@@ -7,7 +7,8 @@ import pytest
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Pregel, channels
from permchain import Channels, Pregel
from permchain.channels.base import InvalidUpdateError
def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -19,8 +20,8 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
"one": chain,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output="output",
@@ -40,8 +41,8 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
"one": chain,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output=["output"],
@@ -65,8 +66,8 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
"one": chain,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input=["input"],
output=["output"],
@@ -93,9 +94,9 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
},
input="input",
output="output",
@@ -112,9 +113,9 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
},
input=["input", "inbox"],
output="output",
@@ -138,9 +139,9 @@ def test_batch_two_processes_in_out() -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"one": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"one": Channels.LastValue(int),
},
input="input",
output="output",
@@ -154,13 +155,13 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"-1": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = channels.LastValue(int)
chans[str(i)] = Channels.LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
)
@@ -182,13 +183,13 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"-1": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = channels.LastValue(int)
chans[str(i)] = Channels.LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
)
@@ -224,14 +225,14 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output="output",
)
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
# LastValue channels can only be updated once per iteration
app.invoke(2)
@@ -245,8 +246,8 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.Inbox(int),
},
input="input",
output="output",
@@ -271,9 +272,9 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
"chain_four": chain_four,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
},
input="input",
output="output",
@@ -298,8 +299,8 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
"one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output="output",
@@ -323,10 +324,10 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
"chain_three": chain_three,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox_one": channels.Inbox(int),
"outbox_one": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox_one": Channels.Inbox(int),
"outbox_one": Channels.LastValue(int),
},
input="input",
output="output",
@@ -352,9 +353,9 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"between": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
},
input="input",
output="output",
@@ -371,9 +372,9 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"between": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
},
input="input",
output="output",
@@ -394,9 +395,9 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"between": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
},
input="input",
output="output",
@@ -422,10 +423,10 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"ctx": channels.ContextManager(an_int, typ=int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"ctx": Channels.ContextManager(an_int, typ=int),
},
input="input",
output=["inbox", "output"],
tests/test_pregel_async.py
+49 -48
View File
@@ -6,7 +6,8 @@ import pytest
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Pregel, channels
from permchain import Channels, Pregel
from permchain.channels.base import InvalidUpdateError
async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -18,8 +19,8 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
"one": chain,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output="output",
@@ -37,8 +38,8 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
"one": chain,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output=["output"],
@@ -62,8 +63,8 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
"one": chain,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input=["input"],
output=["output"],
@@ -90,9 +91,9 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
},
input="input",
output="output",
@@ -109,9 +110,9 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
pubsub = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
},
input=["input", "inbox"],
output="output",
@@ -136,9 +137,9 @@ async def test_batch_two_processes_in_out() -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"one": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"one": Channels.LastValue(int),
},
input="input",
output="output",
@@ -152,13 +153,13 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"-1": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = channels.LastValue(int)
chans[str(i)] = Channels.LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
)
@@ -181,13 +182,13 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"-1": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = channels.LastValue(int)
chans[str(i)] = Channels.LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
)
@@ -229,14 +230,14 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output="output",
)
with pytest.raises(channels.InvalidUpdateError):
with pytest.raises(InvalidUpdateError):
# LastValue channels can only be updated once per iteration
await app.ainvoke(2)
@@ -250,8 +251,8 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.Inbox(int),
},
input="input",
output="output",
@@ -276,9 +277,9 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
"chain_four": chain_four,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
},
input="input",
output="output",
@@ -304,8 +305,8 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
"one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
},
input="input",
output="output",
@@ -329,10 +330,10 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
"chain_three": chain_three,
},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox_one": channels.Inbox(int),
"outbox_one": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox_one": Channels.Inbox(int),
"outbox_one": Channels.LastValue(int),
},
input="input",
output="output",
@@ -360,9 +361,9 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"between": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
},
input="input",
output="output",
@@ -380,9 +381,9 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"between": channels.LastValue(int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
},
input="input",
output="output",
@@ -423,10 +424,10 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": channels.LastValue(int),
"output": channels.LastValue(int),
"inbox": channels.Inbox(int),
"ctx": channels.ContextManager(an_int, an_int_async, typ=int),
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"ctx": Channels.ContextManager(an_int, an_int_async, typ=int),
},
input="input",
output=["inbox", "output"],