Create a Topic channel, Make LastValue the default channel if not specified, Add default input and output keys

- Topic channel combines the features of Inbox, Archive, UniqueInbox, UniqueArchive, which have been removed.
This commit is contained in:
Nuno Campos
2023-11-04 17:40:12 +00:00
parent 124722c0c8
commit 7f76fbb699
19 changed files with 387 additions and 702 deletions
+47 -45
View File
@@ -25,8 +25,8 @@
"from langchain.schema.document import Document\n",
"from langchain.schema import format_document\n",
"\n",
"from permchain import Channel, Pregel, PregelRead\n",
"from permchain.channels import LastValue, Inbox"
"from permchain import Channel, Pregel\n",
"from permchain.channels import LastValue, Topic"
]
},
{
@@ -209,18 +209,15 @@
"source": [
"channels = {\n",
" # input\n",
" \"question\": LastValue(str),\n",
" \"docs\": Inbox(Document),\n",
" \"docs\": Topic(Document),\n",
" # intermediate\n",
" \"docs_to_finalize\": Inbox(Document),\n",
" # output\n",
" \"answer\": LastValue(str),\n",
" \"docs_to_finalize\": Topic(Document),\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 23,
"id": "67370694-86f4-4b64-9d4f-38b2e306abeb",
"metadata": {},
"outputs": [],
@@ -234,12 +231,16 @@
" return Channel.write_to(\"docs_to_finalize\")\n",
"\n",
"\n",
"def split_docs_with_question(input: dict[str, str | list[Document]]) -> list[dict[str, str | list[Document]]]:\n",
" return [\n",
" {\"docs\": docs, \"question\": input[\"question\"]}\n",
" for docs in _split_list_of_docs(input[\"docs\"])\n",
" ]\n",
"\n",
"\n",
"collapse = (\n",
" Channel.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",
" Channel.subscribe_to([\"docs\", \"question\"])\n",
" | split_docs_with_question\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",
@@ -255,7 +256,7 @@
},
{
"cell_type": "code",
"execution_count": 12,
"execution_count": 24,
"id": "3019e7d2-ab7f-4868-b43c-ad898d824a26",
"metadata": {},
"outputs": [],
@@ -274,7 +275,7 @@
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": 25,
"id": "69fcb829-3dae-432a-8db3-11bbb179a7d2",
"metadata": {},
"outputs": [
@@ -283,51 +284,52 @@
"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[0m- collapse({'docs': [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",
" 'question': 'where did harrison work'})\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[0m{'docs': [...], 'docs_to_finalize': [], '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[0m- collapse({'docs': [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",
" 'question': 'where did harrison work'})\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[0m{'docs': [...], 'docs_to_finalize': [], '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[0m- collapse({'docs': [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",
" 'question': 'where did harrison work'})\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[0m{'docs': [], 'docs_to_finalize': [...], '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[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 worked at Kensho.',\n",
" 'docs': (...),\n",
" 'docs_to_finalize': (...),\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 worked at Kensho.'"
"'Harrison used to work at Kensho.'"
]
},
"execution_count": 13,
"execution_count": 25,
"metadata": {},
"output_type": "execute_result"
}
+3 -25
View File
@@ -75,12 +75,6 @@ reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
# application
channels = {
"question": LastValue(str),
"draft": LastValue(str),
"notes": LastValue(str),
}
drafter = (
# subscribe to question channel as a dict with a single key, "question"
Channel.subscribe_to(["question"]) | drafter_chain | Channel.write_to("draft")
@@ -105,7 +99,6 @@ reviser = (
)
draft_revise_loop = Pregel(
channels=channels,
chains={
"drafter": drafter,
"editor": editor,
@@ -113,27 +106,12 @@ draft_revise_loop = Pregel(
},
# input will be a dict with a single key, "question"
input=["question"],
# output will be a dict with keys "draft" and "notes"
output=["draft", "notes"],
# output will be the value of "draft"
output="draft",
# debug logging
debug=True,
)
# run
for draft in draft_revise_loop.stream({"question": "What food do turtles eat?"}):
print(draft)
print("---")
async def main():
async for draft in draft_revise_loop.astream(
{"question": "What food do turtles eat?"}
):
print(draft)
print("---")
# import asyncio
# asyncio.run(main())
print(draft_revise_loop.invoke({"question": "What food do turtles eat?"}))
-2
View File
@@ -1,5 +1,4 @@
from permchain import Channel, Pregel
from permchain.channels import LastValue
grow_value = (
Channel.subscribe_to("value")
@@ -9,7 +8,6 @@ grow_value = (
app = Pregel(
chains={"grow_value": grow_value},
channels={"value": LastValue(str)},
input="value",
output="value",
)
+11 -21
View File
@@ -7,7 +7,8 @@ from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
from langchain.utils.html import extract_sub_links
from permchain import Channel, Pregel
from permchain.channels import Archive, Context, LastValue, UniqueArchive, UniqueInbox
from permchain.channels.context import Context
from permchain.channels.topic import Topic
# Load url with sync httpx client
@@ -85,39 +86,28 @@ def recursive_web_loader(
metadata_extractor = metadata_extractor or _metadata_extractor
# define the channels
channels = {
"base_url": LastValue(str),
"next_urls": UniqueInbox(str),
"documents": Archive(Document),
"visited": UniqueArchive(str),
"next_urls": Topic(str, unique=True),
"documents": Topic(Document, accumulate=True),
"client": Context(httpx_client, httpx_aclient),
}
# the main chain that gets executed recursively
# while there are urls in next_urls
visitor = (
# while there are urls in next_urls
# run the chain below for each url in next_urls
# adding the current values of visited set, base_url and httpx client
Channel.subscribe_to_each("next_urls", key="url").join(
["visited", "client", "base_url"]
)
# adding the current values of base_url and httpx client
Channel.subscribe_to_each("next_urls", key="url").join(["client", "base_url"])
# load the url (with sync and async implementations)
| RunnablePassthrough.assign(body=RunnableLambda(load_url, load_url_async))
| Channel.write_to(
# send this url to the visited set
visited=lambda x: x["url"],
# send a new document to the documents stream
documents=lambda x: Document(
page_content=extractor(x["body"]),
metadata=metadata_extractor(x["body"], x["url"]),
),
# send the next urls to the next_urls set
# only if not visited already
next_urls=lambda x: [
url
for url in extract_sub_links(
x["body"], x["url"], base_url=x["base_url"]
)
if url not in x["visited"] and url != x["url"]
],
# send the next urls to the next_urls topic
next_urls=lambda x: extract_sub_links(
x["body"], x["url"], base_url=x["base_url"]
),
)
)
return Pregel(