{
 "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 Channel, Pregel\n",
    "from permchain.channels import LastValue, Topic"
   ]
  },
  {
   "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(\n",
    "    lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n",
    ").map() | (lambda x: \"\\n\\n\".join(x))"
   ]
  },
  {
   "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",
    "]"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "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})"
   ]
  },
  {
   "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"
   ]
  },
  {
   "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"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8d524ba6-0939-4a5d-8db0-4fa1ef06eaeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "channels = {\n",
    "    # input\n",
    "    \"docs\": Topic(Document),\n",
    "    # intermediate\n",
    "    \"docs_to_finalize\": Topic(Document),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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 Channel.write_to(\"docs\")\n",
    "    else:\n",
    "        # send to the finalizer if we're ready to produce final answer\n",
    "        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\", \"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",
    ")\n",
    "\n",
    "# Convert final set of docs to an answer\n",
    "finalize = (\n",
    "    Channel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n",
    "    | stuff_chain\n",
    "    | Channel.write_to(\"answer\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "3019e7d2-ab7f-4868-b43c-ad898d824a26",
   "metadata": {},
   "outputs": [],
   "source": [
    "reduce_chain = Pregel(\n",
    "    chains={\n",
    "        \"collapse\": collapse,\n",
    "        \"finalize\": finalize,\n",
    "    },\n",
    "    channels=channels,\n",
    "    input=[\"question\", \"docs\"],\n",
    "    output=\"answer\",\n",
    "    debug=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "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({'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': [...], '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({'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': [...], '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({'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': [], '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[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": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reduce_chain.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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}