From d014b13d6f7f3619e72a1ba097ed390bd000a944 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Fri, 13 Oct 2023 06:44:04 -0700 Subject: [PATCH 1/3] combine docs notebook --- examples/combine_docs.ipynb | 309 ++++++++++++++++++++++++++++++++++++ 1 file changed, 309 insertions(+) create mode 100644 examples/combine_docs.ipynb diff --git a/examples/combine_docs.ipynb b/examples/combine_docs.ipynb new file mode 100644 index 000000000..3e6478f80 --- /dev/null +++ b/examples/combine_docs.ipynb @@ -0,0 +1,309 @@ +{ + "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 operator import itemgetter\n", + "\n", + "from langchain.chat_models.openai import ChatOpenAI\n", + "from langchain.prompts import SystemMessagePromptTemplate, ChatPromptTemplate, PromptTemplate\n", + "from langchain.schema.output_parser import StrOutputParser\n", + "from langchain.runnables.openai_functions import OpenAIFunctionsRouter\n", + "from langchain.schema.runnable import RunnableMap\n", + "from langchain.schema.document import Document\n", + "from langchain.schema import format_document\n", + "\n", + "from permchain import Pregel, channels" + ] + }, + { + "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" + ] + }, + { + "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(lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)).map() | (lambda x: \"\\n\\n\".join(x))" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "29b2668d-e4a6-4876-9b04-bdc841774c62", + "metadata": {}, + "outputs": [], + "source": [ + "docs = [Document(page_content=\"Harrison used to work at Kensho\"), Document(page_content=\"Ankush worked at Facebook\")]" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "17da58b7-8685-4d0a-9a47-c398c085d477", + "metadata": {}, + "outputs": [], + "source": [ + "stuff_chain = {\n", + " \"question\": lambda x: x[\"question\"],\n", + " \"context\": (lambda x: x['docs']) | _combine_documents\n", + "} |ChatPromptTemplate.from_messages([\n", + " (\"system\", \"Answer user questions based on the following documents:\\n\\n{context}\"),\n", + " (\"human\", \"{question}\"),\n", + "]) | ChatOpenAI()" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "87295b71-0afc-4901-b57c-a7b945aa4bd9", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "AIMessage(content='Harrison used to work at Kensho.')" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})" + ] + }, + { + "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": 10, + "id": "b15f5abb-1cfe-4965-a021-c891506c5dd2", + "metadata": {}, + "outputs": [], + "source": [ + "many_docs = docs * 5" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "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" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "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": 12, + "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" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "8d524ba6-0939-4a5d-8db0-4fa1ef06eaeb", + "metadata": {}, + "outputs": [], + "source": [ + "input_inbox = channels.LastValue[str](\"input_inbox\")\n", + "reduce_inbox = channels.LastValue[str](\"reduce_inbox\")\n", + "collapse_inbox = channels.LastValue[str](\"collapse_inbox\")\n", + "output_inbox = channels.LastValue[str](\"output_inbox\")" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "67370694-86f4-4b64-9d4f-38b2e306abeb", + "metadata": {}, + "outputs": [], + "source": [ + "# Decide if should finish or should reduce one more step\n", + "def decide_end(plan):\n", + " if len(plan['docs']) > 1:\n", + " return Pregel.send_to(\"reduce_inbox\")\n", + " else:\n", + " return {\"docs\": lambda x: x[\"docs\"][0], \"question\": lambda x: x[\"question\"]} | stuff_chain | Pregel.send_to(\"output_inbox\")\n", + "\n", + "# Chain that collapses documents then chooses end\n", + "collapse_chain = Pregel.subscribe_to(input=collapse_inbox) | RunnableMap({\n", + " \"docs\": lambda x: _split_list_of_docs(x[\"docs\"]),\n", + " \"question\": lambda x: x[\"question\"]\n", + "}) | decide_end\n", + "\n", + "\n", + "reduce_chain = (\n", + " Pregel.subscribe_to(input=input_inbox)\n", + " | (lambda x: [{\"docs\": d, \"question\": x[\"question\"]} for d in x['docs']])\n", + " | stuff_chain.map() \n", + " | Pregel.send_to({\"collapse_inbox\": {\n", + " \"docs\": lambda x: [Document(page_content=m.content) for m in x],\n", + " }})\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "3019e7d2-ab7f-4868-b43c-ad898d824a26", + "metadata": {}, + "outputs": [ + { + "ename": "ValidationError", + "evalue": "6 validation errors for Pregel\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mValidationError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[23], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m pubsub \u001b[38;5;241m=\u001b[39m \u001b[43mPregel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreduce_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcollapse_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_inbox\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/workplace/permchain/permchain/pregel.py:244\u001b[0m, in \u001b[0;36mPregel.__init__\u001b[0;34m(self, input, output, step_timeout, *processes, **kwargs)\u001b[0m\n\u001b[1;32m 236\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m 237\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 238\u001b[0m \u001b[38;5;241m*\u001b[39mprocesses: PregelInvoke \u001b[38;5;241m|\u001b[39m PregelBatch,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 242\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m 243\u001b[0m ):\n\u001b[0;32m--> 244\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 245\u001b[0m \u001b[43m \u001b[49m\u001b[43mprocesses\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprocesses\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 246\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 247\u001b[0m \u001b[43m \u001b[49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 248\u001b[0m \u001b[43m \u001b[49m\u001b[43mstep_timeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstep_timeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 249\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 250\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/.pyenv/versions/3.10.1/envs/permchain/lib/python3.10/site-packages/langchain/load/serializable.py:90\u001b[0m, in \u001b[0;36mSerializable.__init__\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m 89\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m---> 90\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 91\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_lc_kwargs \u001b[38;5;241m=\u001b[39m kwargs\n", + "File \u001b[0;32m~/.pyenv/versions/3.10.1/envs/permchain/lib/python3.10/site-packages/pydantic/main.py:341\u001b[0m, in \u001b[0;36mpydantic.main.BaseModel.__init__\u001b[0;34m()\u001b[0m\n", + "\u001b[0;31mValidationError\u001b[0m: 6 validation errors for Pregel\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)" + ] + } + ], + "source": [ + "pubsub = Pregel(input_inbox, reduce_inbox, collapse_inbox, input=input_inbox, output=output_inbox)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 101, + "id": "69fcb829-3dae-432a-8db3-11bbb179a7d2", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[AIMessage(content='Harrison used to work at Kensho.', additional_kwargs={}, example=False)]" + ] + }, + "execution_count": 101, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "reduce_agent.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})" + ] + }, + { + "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.10.1" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 1e083958102aaebca9a620b71fa7d69e5dee5e35 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 18 Oct 2023 12:21:57 +0100 Subject: [PATCH 2/3] WIP --- examples/combine_docs.ipynb | 107 ++++++++++++++++++++++++------------ 1 file changed, 73 insertions(+), 34 deletions(-) diff --git a/examples/combine_docs.ipynb b/examples/combine_docs.ipynb index 3e6478f80..7cbf6994b 100644 --- a/examples/combine_docs.ipynb +++ b/examples/combine_docs.ipynb @@ -12,7 +12,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 11, "id": "624c452c-ddd5-4390-9065-7ec55dc64b96", "metadata": {}, "outputs": [], @@ -20,10 +20,14 @@ "from operator import itemgetter\n", "\n", "from langchain.chat_models.openai import ChatOpenAI\n", - "from langchain.prompts import SystemMessagePromptTemplate, ChatPromptTemplate, PromptTemplate\n", + "from langchain.prompts import (\n", + " SystemMessagePromptTemplate,\n", + " ChatPromptTemplate,\n", + " PromptTemplate,\n", + ")\n", "from langchain.schema.output_parser import StrOutputParser\n", "from langchain.runnables.openai_functions import OpenAIFunctionsRouter\n", - "from langchain.schema.runnable import RunnableMap\n", + "from langchain.schema.runnable import RunnableMap, RunnablePassthrough\n", "from langchain.schema.document import Document\n", "from langchain.schema import format_document\n", "\n", @@ -59,7 +63,9 @@ "source": [ "DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template=\"{page_content}\")\n", "\n", - "_combine_documents = RunnableLambda(lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)).map() | (lambda x: \"\\n\\n\".join(x))" + "_combine_documents = RunnableLambda(\n", + " lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n", + ").map() | (lambda x: \"\\n\\n\".join(x))" ] }, { @@ -69,28 +75,40 @@ "metadata": {}, "outputs": [], "source": [ - "docs = [Document(page_content=\"Harrison used to work at Kensho\"), Document(page_content=\"Ankush worked at Facebook\")]" + "docs = [\n", + " Document(page_content=\"Harrison used to work at Kensho\"),\n", + " Document(page_content=\"Ankush worked at Facebook\"),\n", + "]" ] }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 5, "id": "17da58b7-8685-4d0a-9a47-c398c085d477", "metadata": {}, "outputs": [], "source": [ - "stuff_chain = {\n", - " \"question\": lambda x: x[\"question\"],\n", - " \"context\": (lambda x: x['docs']) | _combine_documents\n", - "} |ChatPromptTemplate.from_messages([\n", - " (\"system\", \"Answer user questions based on the following documents:\\n\\n{context}\"),\n", - " (\"human\", \"{question}\"),\n", - "]) | ChatOpenAI()" + "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", + ")" ] }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 6, "id": "87295b71-0afc-4901-b57c-a7b945aa4bd9", "metadata": {}, "outputs": [ @@ -100,7 +118,7 @@ "AIMessage(content='Harrison used to work at Kensho.')" ] }, - "execution_count": 9, + "execution_count": 6, "metadata": {}, "output_type": "execute_result" } @@ -121,7 +139,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 7, "id": "b15f5abb-1cfe-4965-a021-c891506c5dd2", "metadata": {}, "outputs": [], @@ -131,7 +149,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 8, "id": "ccad04a3-fd3f-4e73-b895-29e53535f000", "metadata": {}, "outputs": [], @@ -156,7 +174,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 9, "id": "11cfd337-9f3b-4b26-ba30-251e17b18994", "metadata": {}, "outputs": [ @@ -175,7 +193,7 @@ " Document(page_content='Ankush worked at Facebook')]]" ] }, - "execution_count": 12, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } @@ -188,10 +206,23 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 10, "id": "8d524ba6-0939-4a5d-8db0-4fa1ef06eaeb", "metadata": {}, - "outputs": [], + "outputs": [ + { + "ename": "TypeError", + "evalue": "LastValue() takes no arguments", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[10], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m input_inbox \u001b[38;5;241m=\u001b[39m \u001b[43mchannels\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mLastValue\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43minput_inbox\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m reduce_inbox \u001b[38;5;241m=\u001b[39m channels\u001b[38;5;241m.\u001b[39mLastValue[\u001b[38;5;28mstr\u001b[39m](\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mreduce_inbox\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 3\u001b[0m collapse_inbox \u001b[38;5;241m=\u001b[39m channels\u001b[38;5;241m.\u001b[39mLastValue[\u001b[38;5;28mstr\u001b[39m](\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcollapse_inbox\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "File \u001b[0;32m/opt/homebrew/Cellar/python@3.11/3.11.5/Frameworks/Python.framework/Versions/3.11/lib/python3.11/typing.py:1268\u001b[0m, in \u001b[0;36m_BaseGenericAlias.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1265\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_inst:\n\u001b[1;32m 1266\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mType \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m cannot be instantiated; \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 1267\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124muse \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__origin__\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m() instead\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 1268\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__origin__\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1269\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1270\u001b[0m result\u001b[38;5;241m.\u001b[39m__orig_class__ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\n", + "\u001b[0;31mTypeError\u001b[0m: LastValue() takes no arguments" + ] + } + ], "source": [ "input_inbox = channels.LastValue[str](\"input_inbox\")\n", "reduce_inbox = channels.LastValue[str](\"reduce_inbox\")\n", @@ -208,25 +239,31 @@ "source": [ "# Decide if should finish or should reduce one more step\n", "def decide_end(plan):\n", - " if len(plan['docs']) > 1:\n", + " if len(plan[\"docs\"]) > 1:\n", " return Pregel.send_to(\"reduce_inbox\")\n", " else:\n", - " return {\"docs\": lambda x: x[\"docs\"][0], \"question\": lambda x: x[\"question\"]} | stuff_chain | Pregel.send_to(\"output_inbox\")\n", + " return stuff_chain | Pregel.send_to(\"output_inbox\")\n", + "\n", "\n", "# Chain that collapses documents then chooses end\n", - "collapse_chain = Pregel.subscribe_to(input=collapse_inbox) | RunnableMap({\n", - " \"docs\": lambda x: _split_list_of_docs(x[\"docs\"]),\n", - " \"question\": lambda x: x[\"question\"]\n", - "}) | decide_end\n", + "collapse_chain = (\n", + " Pregel.subscribe_to(docs=collapse_inbox, question=\"question\")\n", + " | RunnablePassthrough.assign(docs=lambda x: _split_list_of_docs(x[\"docs\"]))\n", + " | decide_end\n", + ")\n", "\n", "\n", "reduce_chain = (\n", " Pregel.subscribe_to(input=input_inbox)\n", - " | (lambda x: [{\"docs\": d, \"question\": x[\"question\"]} for d in x['docs']])\n", - " | stuff_chain.map() \n", - " | Pregel.send_to({\"collapse_inbox\": {\n", - " \"docs\": lambda x: [Document(page_content=m.content) for m in x],\n", - " }})\n", + " | (lambda x: [{\"docs\": d, \"question\": x[\"question\"]} for d in x[\"docs\"]])\n", + " | stuff_chain.map()\n", + " | Pregel.send_to(\n", + " {\n", + " \"collapse_inbox\": {\n", + " \"docs\": lambda x: [Document(page_content=m.content) for m in x],\n", + " }\n", + " }\n", + " )\n", ")" ] }, @@ -252,7 +289,9 @@ } ], "source": [ - "pubsub = Pregel(input_inbox, reduce_inbox, collapse_inbox, input=input_inbox, output=output_inbox)\n" + "pubsub = Pregel(\n", + " input_inbox, reduce_inbox, collapse_inbox, input=input_inbox, output=output_inbox\n", + ")" ] }, { @@ -301,7 +340,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.1" + "version": "3.11.5" } }, "nbformat": 4, From 49ff5e2eab3f4bc25c17f0c41f89a6a447d83aa7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 23 Oct 2023 12:01:35 +0100 Subject: [PATCH 3/3] Implement reduce chain --- examples/combine_docs.ipynb | 192 +++++++++++++++++-------------- examples/draft-revise-loop.py | 14 ++- examples/readme.py | 4 +- examples/recursive-web-loader.py | 19 +-- permchain/__init__.py | 5 +- permchain/channels/__init__.py | 4 - permchain/pregel/__init__.py | 32 +++++- permchain/pregel/validate.py | 9 +- tests/test_channels.py | 53 ++++----- tests/test_pregel.py | 103 +++++++++-------- tests/test_pregel_async.py | 97 ++++++++-------- 11 files changed, 290 insertions(+), 242 deletions(-) diff --git a/examples/combine_docs.ipynb b/examples/combine_docs.ipynb index 7cbf6994b..cf4415f73 100644 --- a/examples/combine_docs.ipynb +++ b/examples/combine_docs.ipynb @@ -12,26 +12,20 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 1, "id": "624c452c-ddd5-4390-9065-7ec55dc64b96", "metadata": {}, "outputs": [], "source": [ - "from operator import itemgetter\n", - "\n", "from langchain.chat_models.openai import ChatOpenAI\n", - "from langchain.prompts import (\n", - " SystemMessagePromptTemplate,\n", - " ChatPromptTemplate,\n", - " PromptTemplate,\n", - ")\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.runnables.openai_functions import OpenAIFunctionsRouter\n", - "from langchain.schema.runnable import RunnableMap, RunnablePassthrough\n", "from langchain.schema.document import Document\n", "from langchain.schema import format_document\n", "\n", - "from permchain import Pregel, channels" + "from permchain import Pregel, PregelRead, channels\n" ] }, { @@ -51,7 +45,7 @@ "metadata": {}, "outputs": [], "source": [ - "from langchain.schema.runnable import RunnableLambda" + "from langchain.schema.runnable import RunnableLambda\n" ] }, { @@ -65,7 +59,7 @@ "\n", "_combine_documents = RunnableLambda(\n", " lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n", - ").map() | (lambda x: \"\\n\\n\".join(x))" + ").map() | (lambda x: \"\\n\\n\".join(x))\n" ] }, { @@ -78,7 +72,7 @@ "docs = [\n", " Document(page_content=\"Harrison used to work at Kensho\"),\n", " Document(page_content=\"Ankush worked at Facebook\"),\n", - "]" + "]\n" ] }, { @@ -103,7 +97,8 @@ " ]\n", " )\n", " | ChatOpenAI()\n", - ")" + " | StrOutputParser()\n", + ")\n" ] }, { @@ -115,7 +110,7 @@ { "data": { "text/plain": [ - "AIMessage(content='Harrison used to work at Kensho.')" + "'Harrison used to work at Kensho.'" ] }, "execution_count": 6, @@ -124,7 +119,7 @@ } ], "source": [ - "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})" + "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})\n" ] }, { @@ -144,7 +139,7 @@ "metadata": {}, "outputs": [], "source": [ - "many_docs = docs * 5" + "many_docs = docs * 5\n" ] }, { @@ -169,7 +164,7 @@ " 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" + " return new_result_doc_list\n" ] }, { @@ -201,7 +196,7 @@ "source": [ "# Just to show what its like split\n", "split_docs = _split_list_of_docs(many_docs)\n", - "split_docs" + "split_docs\n" ] }, { @@ -209,110 +204,135 @@ "execution_count": 10, "id": "8d524ba6-0939-4a5d-8db0-4fa1ef06eaeb", "metadata": {}, - "outputs": [ - { - "ename": "TypeError", - "evalue": "LastValue() takes no arguments", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[10], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m input_inbox \u001b[38;5;241m=\u001b[39m \u001b[43mchannels\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mLastValue\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43minput_inbox\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m reduce_inbox \u001b[38;5;241m=\u001b[39m channels\u001b[38;5;241m.\u001b[39mLastValue[\u001b[38;5;28mstr\u001b[39m](\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mreduce_inbox\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 3\u001b[0m collapse_inbox \u001b[38;5;241m=\u001b[39m channels\u001b[38;5;241m.\u001b[39mLastValue[\u001b[38;5;28mstr\u001b[39m](\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcollapse_inbox\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", - "File \u001b[0;32m/opt/homebrew/Cellar/python@3.11/3.11.5/Frameworks/Python.framework/Versions/3.11/lib/python3.11/typing.py:1268\u001b[0m, in \u001b[0;36m_BaseGenericAlias.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1265\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_inst:\n\u001b[1;32m 1266\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mType \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m cannot be instantiated; \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 1267\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124muse \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__origin__\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m() instead\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 1268\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__origin__\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1269\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1270\u001b[0m result\u001b[38;5;241m.\u001b[39m__orig_class__ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\n", - "\u001b[0;31mTypeError\u001b[0m: LastValue() takes no arguments" - ] - } - ], + "outputs": [], "source": [ - "input_inbox = channels.LastValue[str](\"input_inbox\")\n", - "reduce_inbox = channels.LastValue[str](\"reduce_inbox\")\n", - "collapse_inbox = channels.LastValue[str](\"collapse_inbox\")\n", - "output_inbox = channels.LastValue[str](\"output_inbox\")" + "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": 22, + "execution_count": 11, "id": "67370694-86f4-4b64-9d4f-38b2e306abeb", "metadata": {}, "outputs": [], "source": [ - "# Decide if should finish or should reduce one more step\n", - "def decide_end(plan):\n", - " if len(plan[\"docs\"]) > 1:\n", - " return Pregel.send_to(\"reduce_inbox\")\n", + "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", - " return stuff_chain | Pregel.send_to(\"output_inbox\")\n", + " # send to the finalizer if we're ready to produce final answer\n", + " return Pregel.write_to(\"docs_to_finalize\")\n", "\n", "\n", - "# Chain that collapses documents then chooses end\n", - "collapse_chain = (\n", - " Pregel.subscribe_to(docs=collapse_inbox, question=\"question\")\n", - " | RunnablePassthrough.assign(docs=lambda x: _split_list_of_docs(x[\"docs\"]))\n", - " | decide_end\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", - "\n", - "reduce_chain = (\n", - " Pregel.subscribe_to(input=input_inbox)\n", - " | (lambda x: [{\"docs\": d, \"question\": x[\"question\"]} for d in x[\"docs\"]])\n", - " | stuff_chain.map()\n", - " | Pregel.send_to(\n", - " {\n", - " \"collapse_inbox\": {\n", - " \"docs\": lambda x: [Document(page_content=m.content) for m in x],\n", - " }\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": 23, + "execution_count": 12, "id": "3019e7d2-ab7f-4868-b43c-ad898d824a26", "metadata": {}, - "outputs": [ - { - "ename": "ValidationError", - "evalue": "6 validation errors for Pregel\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mValidationError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[23], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m pubsub \u001b[38;5;241m=\u001b[39m \u001b[43mPregel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreduce_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcollapse_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_inbox\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/workplace/permchain/permchain/pregel.py:244\u001b[0m, in \u001b[0;36mPregel.__init__\u001b[0;34m(self, input, output, step_timeout, *processes, **kwargs)\u001b[0m\n\u001b[1;32m 236\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m 237\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 238\u001b[0m \u001b[38;5;241m*\u001b[39mprocesses: PregelInvoke \u001b[38;5;241m|\u001b[39m PregelBatch,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 242\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m 243\u001b[0m ):\n\u001b[0;32m--> 244\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 245\u001b[0m \u001b[43m \u001b[49m\u001b[43mprocesses\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprocesses\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 246\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 247\u001b[0m \u001b[43m \u001b[49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 248\u001b[0m \u001b[43m \u001b[49m\u001b[43mstep_timeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstep_timeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 249\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 250\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/.pyenv/versions/3.10.1/envs/permchain/lib/python3.10/site-packages/langchain/load/serializable.py:90\u001b[0m, in \u001b[0;36mSerializable.__init__\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m 89\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m---> 90\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 91\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_lc_kwargs \u001b[38;5;241m=\u001b[39m kwargs\n", - "File \u001b[0;32m~/.pyenv/versions/3.10.1/envs/permchain/lib/python3.10/site-packages/pydantic/main.py:341\u001b[0m, in \u001b[0;36mpydantic.main.BaseModel.__init__\u001b[0;34m()\u001b[0m\n", - "\u001b[0;31mValidationError\u001b[0m: 6 validation errors for Pregel\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)" - ] - } - ], + "outputs": [], "source": [ - "pubsub = Pregel(\n", - " input_inbox, reduce_inbox, collapse_inbox, input=input_inbox, output=output_inbox\n", - ")" + "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": 101, + "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": [ - "[AIMessage(content='Harrison used to work at Kensho.', additional_kwargs={}, example=False)]" + "'Harrison used to work at Kensho.'" ] }, - "execution_count": 101, + "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "reduce_agent.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})" + "reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})\n" ] }, { diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index ee9970ebe..7384f64f3 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -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, diff --git a/examples/readme.py b/examples/readme.py index 97998b3ce..647396b17 100644 --- a/examples/readme.py +++ b/examples/readme.py @@ -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", ) diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index cbba99a7f..3ca220fbc 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -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 diff --git a/permchain/__init__.py b/permchain/__init__.py index f6e3f8dfb..732952882 100644 --- a/permchain/__init__.py +++ b/permchain/__init__.py @@ -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"] diff --git a/permchain/channels/__init__.py b/permchain/channels/__init__.py index 9fe8df138..a92933e03 100644 --- a/permchain/channels/__init__.py +++ b/permchain/channels/__init__.py @@ -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", diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 06784eb37..b8e49a822 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -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}, ) diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py index aa7287641..bdc201c49 100644 --- a/permchain/pregel/validate.py +++ b/permchain/pregel/validate.py @@ -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: diff --git a/tests/test_channels.py b/tests/test_channels.py index 6345eb6ce..4e27f1a0b 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -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 diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 975a10748..a36b12c2c 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -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"], diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 135192f06..fec8f7892 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -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"],