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 +}