legal_recursive_retrieval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recursive Retrieval on a Lexical Graph "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "- Create a venv if desired\n",
    "  - `python3 -m venv venv`\n",
    "  - `source venv/bin/activate`\n",
    "- Run `pip install -r requirements.txt` to install the necessary packages.\n",
    "\n",
    "## Important\n",
    "- If you are providing a workspace, it should be empty and should only contain documents used build the lexical graph.\n",
    "\n",
    "## Integration with Reducto\n",
    "\n",
    "- As of the notebook's creation date, any developers interested in the Reducto API must reach out to founders to get a key for the demo. \n",
    "- However, the [free demo on the website](https://app.reducto.ai/) allows up to 10 pages of text to be analyzed\n",
    "- The API response JSONs can be copied and saved as `<pdf filename without extension>.json` under `docs_structure/reducto/`. \n",
    "\n",
    "\n",
    "\n",
    "![Reducto Page](./img/reducto_page.png)\n",
    "![Reducto Save JSON](./img/reducto_save_json.png)\n",
    "\n",
    "- In our example, `bnm_compliance.pdf` in `/docs` was saved as `bnm_compliance.json` in `/docs_structure/reducto/`.\n",
    "- This already has been done for you so you can run the notebook as-is – follow the steps above if you want to try with your own documents."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Link Detection\n",
    "<!-- <img src=\"./img/diagram.png\" width=\"700px\"> -->\n",
    "![Diagram](./img/diagram.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "- We use pymupdf to extract links from the documents.\n",
    "- Pymupdf lets us access the locations of the links and the destination of the links (this applies to internal links that link to other parts of the document)\n",
    "- We locate one Reducto element that is closest to the link's location, and another element closest to the link's destination, and establish a link between the two Reducto elements.\n",
    "- This is what happens in the `extract_links_and_sections_to_elements` function.\n",
    "- The output is a dataframe to visualize the links in the pdf\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Links tuples:  [('f5f9f96d5d805f4e7d1ce832cc0e4937', '5feb986fe20fb5de66a9d110a602f280'), ('2cd1bcffc9bbdd0e7e590bfb22c52c1f', '1ce9f6267c7404bc6f29825b40c66c77'), ('c21822c0338a1563d5e8100467df1bc2', '1cc747588ddb4be5060534dd2d53eaba'), ('fabbe7be4a8a55626bc1d90782d54852', 'a878a69b7a0e10e0b5dd1f3c02c82049'), ('552896ac6f1046fe4f601afcdf5c6325', 'c097483e4d870f1651a3fbbf31e1f380'), ('42c2c25fccac1b063f84600d5807f3ff', 'da2a46ebf01d7a06d709c6c2fb8a0bda'), ('b432259b7584d0d856d4e444985de1c2', '499dcdb601259bb3cbce85ce364ed632'), ('2c517489527da348b5fc153b17b87057', '91f5ccdce7837fc5c8ab0b3ed7e14e54'), ('003d2a95aa3e44c56b28ca467f614446', '6cdf24899c11cedfa54a1f48658033a1'), ('39415dc21008fef7877cde99174cae62', '9898aabafc24b8c858516d18f8ca38fc'), ('baa81dc508fb8531ce35ede2c2246829', '0c30c80e8079513bcfb42217d016f720'), ('1cc747588ddb4be5060534dd2d53eaba', 'e4308cee3126ddd2a949650aeb69a368'), ('bbba75c05798bba059a8269b624befcb', '63f9a7375e9087b9974612325f288586'), ('bbba75c05798bba059a8269b624befcb', 'fc2e0f615469f18d6f4a14975f4d2aea'), ('cd9d52958cf53b9295b73a9434cb2ba9', '6cdf24899c11cedfa54a1f48658033a1'), ('cd9d52958cf53b9295b73a9434cb2ba9', '9898aabafc24b8c858516d18f8ca38fc'), ('c6086e80eb5855c34404bd1e6e7eb06c', 'd9432f9931c5bf206db8317a67e69052'), ('04e391fb570851132e2aba3610cd3f86', '4cd12c4a3ddeb6778b65630b4e624564'), ('4dc2f2eb485c0582ed2fc60c0f291557', '334fdd9aecc048c657f8ad1d9c0d8a24'), ('bc08a3ba18aa7a08d732dbe4c97dacd4', '570902ce5ab35809597f3727115bcd64'), ('bc08a3ba18aa7a08d732dbe4c97dacd4', '0535171ad1204d0c05d28b16a437a17c'), ('4dbea03fb547a6c5b4e4e8160f226fb3', 'f9fb2b33eabb778bf7f7073b6acffb66'), ('4dbea03fb547a6c5b4e4e8160f226fb3', '4ffb5b3a6776c69a888ebce7a3630bcd'), ('bd4533882412131b8d98e84499fe7d9c', '98d7434318b67e1256f563bf7fd13472'), ('b29501497f35ec60a1e37a6e5200a75b', '9898aabafc24b8c858516d18f8ca38fc'), ('9aee9c464ce513503b26252dc35e3ad1', '081d34781b0f7ab41f94dbe52896a23e'), ('6cf04fc7a955f2d7df8ea100ec7af539', '0535171ad1204d0c05d28b16a437a17c')]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>page_number</th>\n",
       "      <th>link_text</th>\n",
       "      <th>guess_link_reducto_element</th>\n",
       "      <th>link_type</th>\n",
       "      <th>link_destination</th>\n",
       "      <th>named_destination</th>\n",
       "      <th>guess_destination_reducto_element</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>PART A      OVERVIEW \\n..........................</td>\n",
       "      <td>PART A OVERVIEW</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 3</td>\n",
       "      <td>Coordinates: Point(68.0, 83.91998)</td>\n",
       "      <td>PART A OVERVIEW</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>1 \\nIntroduction ................................</td>\n",
       "      <td>1 Introduction</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 3</td>\n",
       "      <td>Coordinates: Point(68.0, 112.91998)</td>\n",
       "      <td>1 Introduction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>2 \\nApplicability ...............................</td>\n",
       "      <td>2 Applicability</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 3</td>\n",
       "      <td>Coordinates: Point(68.0, 671.92)</td>\n",
       "      <td>2.1 This policy document is applicable to all ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2</td>\n",
       "      <td>3 \\nLegal provisions ............................</td>\n",
       "      <td>3 Legal provisions</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 4</td>\n",
       "      <td>Coordinates: Point(68.0, 83.91998)</td>\n",
       "      <td>3 Legal provisions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2</td>\n",
       "      <td>4 \\nEffective date ..............................</td>\n",
       "      <td>4 Effective date</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 4</td>\n",
       "      <td>Coordinates: Point(68.0, 229.91999)</td>\n",
       "      <td>4 Effective date</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>2</td>\n",
       "      <td>5 \\nInterpretation ..............................</td>\n",
       "      <td>5 Interpretation</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 4</td>\n",
       "      <td>Coordinates: Point(68.0, 291.91999)</td>\n",
       "      <td>5 Interpretation</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2</td>\n",
       "      <td>PART B      POLICY REQUIREMENTS \\n...............</td>\n",
       "      <td>PART B POLICY REQUIREMENTS</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 6</td>\n",
       "      <td>Coordinates: Point(68.0, 83.91998)</td>\n",
       "      <td>PART B POLICY REQUIREMENTS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>2</td>\n",
       "      <td>6 \\nResponsibilities of the board and senior m...</td>\n",
       "      <td>6 Responsibilities of the board and senior man...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 6</td>\n",
       "      <td>Coordinates: Point(68.0, 112.91998)</td>\n",
       "      <td>6 Responsibilities of the board and senior man...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>2</td>\n",
       "      <td>7 \\nOrganisation of the compliance function .....</td>\n",
       "      <td>7 Organisation of the compliance function</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 8</td>\n",
       "      <td>Coordinates: Point(68.0, 345.91999)</td>\n",
       "      <td>\\t- 7 Organisation of the compliance function</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>2</td>\n",
       "      <td>8 \\nResponsibilities of the compliance functio...</td>\n",
       "      <td>8 Responsibilities of the compliance function</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 10</td>\n",
       "      <td>Coordinates: Point(68.0, 441.91999)</td>\n",
       "      <td>\\t- 8 Responsibilities of the compliance function</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>2</td>\n",
       "      <td>9 \\nResponsibilities of the internal audit fun...</td>\n",
       "      <td>9 Responsibilities of the internal audit function</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 12</td>\n",
       "      <td>Coordinates: Point(68.0, 524.92)</td>\n",
       "      <td>\\t- 9 Responsibilities of the internal audit f...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>3</td>\n",
       "      <td>5.2.</td>\n",
       "      <td>2.1 This policy document is applicable to all ...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 4</td>\n",
       "      <td>Coordinates: Point(68.0, 381.91999)</td>\n",
       "      <td>- 5.2 For the purpose of this policy document-</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>4</td>\n",
       "      <td>8.4</td>\n",
       "      <td>\\t- \"compliance function\" refers to officers c...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 11</td>\n",
       "      <td>Coordinates: Point(67.0, 179.91999)</td>\n",
       "      <td>- S 8.4 The compliance function must identify ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>4</td>\n",
       "      <td>8.11,</td>\n",
       "      <td>\\t- \"compliance function\" refers to officers c...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 12</td>\n",
       "      <td>Coordinates: Point(67.0, 414.91999)</td>\n",
       "      <td>- S 8.11 The compliance function is responsibl...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>7</td>\n",
       "      <td>7</td>\n",
       "      <td>\\t(c) establish a compliance function commensu...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 8</td>\n",
       "      <td>Coordinates: Point(68.0, 345.91999)</td>\n",
       "      <td>\\t- 7 Organisation of the compliance function</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>7</td>\n",
       "      <td>8;</td>\n",
       "      <td>\\t(c) establish a compliance function commensu...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 10</td>\n",
       "      <td>Coordinates: Point(68.0, 441.91999)</td>\n",
       "      <td>\\t- 8 Responsibilities of the compliance function</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>7</td>\n",
       "      <td>6.2(d);</td>\n",
       "      <td>\\t(j) inform the board of the CCO's cessation ...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 6</td>\n",
       "      <td>Coordinates: Point(67.0, 423.91999)</td>\n",
       "      <td>- S 6.3 In relation to the position of the CCO...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>7</td>\n",
       "      <td>6.4(a),</td>\n",
       "      <td>- S 6.5 In relation to paragraph 6.4(a), senio...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 7</td>\n",
       "      <td>Coordinates: Point(67.0, 138.91999)</td>\n",
       "      <td>- S 6.4 Senior management is collectively resp...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>8</td>\n",
       "      <td>6.4(h),</td>\n",
       "      <td>- S 6.6 For the purpose of paragraph 6.4(h), s...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 7</td>\n",
       "      <td>Coordinates: Point(67.0, 400.91999)</td>\n",
       "      <td>\\t(h) report to the board regularly on complia...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>8</td>\n",
       "      <td>6.4(i),</td>\n",
       "      <td>- S 6.7 In relation to paragraph 6.4(i), senio...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 7</td>\n",
       "      <td>Coordinates: Point(67.0, 441.91999)</td>\n",
       "      <td>\\t(i) report to the board at least annually on...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>8</td>\n",
       "      <td>8.8.</td>\n",
       "      <td>- S 6.7 In relation to paragraph 6.4(i), senio...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 11</td>\n",
       "      <td>Coordinates: Point(67.0, 552.92)</td>\n",
       "      <td>- S 8.8 The CCO must report to senior manageme...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>9</td>\n",
       "      <td>7.2</td>\n",
       "      <td>- S 7.4 In relation to paragraphs 7.2 and 7.3-</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 8</td>\n",
       "      <td>Coordinates: Point(67.0, 462.91999)</td>\n",
       "      <td>- S 7.2 Where compliance function responsibili...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>9</td>\n",
       "      <td>7.3–</td>\n",
       "      <td>- S 7.4 In relation to paragraphs 7.2 and 7.3-</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 9</td>\n",
       "      <td>Coordinates: Point(67.0, 138.91999)</td>\n",
       "      <td>- S 7.3 Where the CCO also assumes responsibil...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>9</td>\n",
       "      <td>9.1</td>\n",
       "      <td>\\t\\t(b) compliance function responsibilities c...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 12</td>\n",
       "      <td>Coordinates: Point(66.0, 559.92)</td>\n",
       "      <td>- S 9.1 A financial institution must ensure th...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>9</td>\n",
       "      <td>8.</td>\n",
       "      <td>- S 7.5 A large financial institution is requi...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 10</td>\n",
       "      <td>Coordinates: Point(68.0, 441.91999)</td>\n",
       "      <td>\\t- 8 Responsibilities of the compliance function</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>10</td>\n",
       "      <td>7.9</td>\n",
       "      <td>- S 7.10 Where such arrangements referred to i...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 9</td>\n",
       "      <td>Coordinates: Point(67.0, 648.92)</td>\n",
       "      <td>- G 7.9 A constructive and cooperative working...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>12</td>\n",
       "      <td>8.8</td>\n",
       "      <td>- S 8.9 The CCO must ensure that the reports r...</td>\n",
       "      <td>internal</td>\n",
       "      <td>page 11</td>\n",
       "      <td>Coordinates: Point(67.0, 552.92)</td>\n",
       "      <td>- S 8.8 The CCO must report to senior manageme...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    page_number                                          link_text  \\\n",
       "0             2  PART A      OVERVIEW \\n..........................   \n",
       "1             2  1 \\nIntroduction ................................   \n",
       "2             2  2 \\nApplicability ...............................   \n",
       "3             2  3 \\nLegal provisions ............................   \n",
       "4             2  4 \\nEffective date ..............................   \n",
       "5             2  5 \\nInterpretation ..............................   \n",
       "6             2  PART B      POLICY REQUIREMENTS \\n...............   \n",
       "7             2  6 \\nResponsibilities of the board and senior m...   \n",
       "8             2  7 \\nOrganisation of the compliance function .....   \n",
       "9             2  8 \\nResponsibilities of the compliance functio...   \n",
       "10            2  9 \\nResponsibilities of the internal audit fun...   \n",
       "11            3                                               5.2.   \n",
       "12            4                                                8.4   \n",
       "13            4                                              8.11,   \n",
       "14            7                                                  7   \n",
       "15            7                                                 8;   \n",
       "16            7                                            6.2(d);   \n",
       "17            7                                            6.4(a),   \n",
       "18            8                                            6.4(h),   \n",
       "19            8                                            6.4(i),   \n",
       "20            8                                               8.8.   \n",
       "21            9                                                7.2   \n",
       "22            9                                               7.3–   \n",
       "23            9                                                9.1   \n",
       "24            9                                                 8.   \n",
       "25           10                                                7.9   \n",
       "26           12                                                8.8   \n",
       "\n",
       "                           guess_link_reducto_element link_type  \\\n",
       "0                                     PART A OVERVIEW  internal   \n",
       "1                                      1 Introduction  internal   \n",
       "2                                     2 Applicability  internal   \n",
       "3                                  3 Legal provisions  internal   \n",
       "4                                    4 Effective date  internal   \n",
       "5                                    5 Interpretation  internal   \n",
       "6                          PART B POLICY REQUIREMENTS  internal   \n",
       "7   6 Responsibilities of the board and senior man...  internal   \n",
       "8           7 Organisation of the compliance function  internal   \n",
       "9       8 Responsibilities of the compliance function  internal   \n",
       "10  9 Responsibilities of the internal audit function  internal   \n",
       "11  2.1 This policy document is applicable to all ...  internal   \n",
       "12  \\t- \"compliance function\" refers to officers c...  internal   \n",
       "13  \\t- \"compliance function\" refers to officers c...  internal   \n",
       "14  \\t(c) establish a compliance function commensu...  internal   \n",
       "15  \\t(c) establish a compliance function commensu...  internal   \n",
       "16  \\t(j) inform the board of the CCO's cessation ...  internal   \n",
       "17  - S 6.5 In relation to paragraph 6.4(a), senio...  internal   \n",
       "18  - S 6.6 For the purpose of paragraph 6.4(h), s...  internal   \n",
       "19  - S 6.7 In relation to paragraph 6.4(i), senio...  internal   \n",
       "20  - S 6.7 In relation to paragraph 6.4(i), senio...  internal   \n",
       "21     - S 7.4 In relation to paragraphs 7.2 and 7.3-  internal   \n",
       "22     - S 7.4 In relation to paragraphs 7.2 and 7.3-  internal   \n",
       "23  \\t\\t(b) compliance function responsibilities c...  internal   \n",
       "24  - S 7.5 A large financial institution is requi...  internal   \n",
       "25  - S 7.10 Where such arrangements referred to i...  internal   \n",
       "26  - S 8.9 The CCO must ensure that the reports r...  internal   \n",
       "\n",
       "   link_destination                    named_destination  \\\n",
       "0            page 3   Coordinates: Point(68.0, 83.91998)   \n",
       "1            page 3  Coordinates: Point(68.0, 112.91998)   \n",
       "2            page 3     Coordinates: Point(68.0, 671.92)   \n",
       "3            page 4   Coordinates: Point(68.0, 83.91998)   \n",
       "4            page 4  Coordinates: Point(68.0, 229.91999)   \n",
       "5            page 4  Coordinates: Point(68.0, 291.91999)   \n",
       "6            page 6   Coordinates: Point(68.0, 83.91998)   \n",
       "7            page 6  Coordinates: Point(68.0, 112.91998)   \n",
       "8            page 8  Coordinates: Point(68.0, 345.91999)   \n",
       "9           page 10  Coordinates: Point(68.0, 441.91999)   \n",
       "10          page 12     Coordinates: Point(68.0, 524.92)   \n",
       "11           page 4  Coordinates: Point(68.0, 381.91999)   \n",
       "12          page 11  Coordinates: Point(67.0, 179.91999)   \n",
       "13          page 12  Coordinates: Point(67.0, 414.91999)   \n",
       "14           page 8  Coordinates: Point(68.0, 345.91999)   \n",
       "15          page 10  Coordinates: Point(68.0, 441.91999)   \n",
       "16           page 6  Coordinates: Point(67.0, 423.91999)   \n",
       "17           page 7  Coordinates: Point(67.0, 138.91999)   \n",
       "18           page 7  Coordinates: Point(67.0, 400.91999)   \n",
       "19           page 7  Coordinates: Point(67.0, 441.91999)   \n",
       "20          page 11     Coordinates: Point(67.0, 552.92)   \n",
       "21           page 8  Coordinates: Point(67.0, 462.91999)   \n",
       "22           page 9  Coordinates: Point(67.0, 138.91999)   \n",
       "23          page 12     Coordinates: Point(66.0, 559.92)   \n",
       "24          page 10  Coordinates: Point(68.0, 441.91999)   \n",
       "25           page 9     Coordinates: Point(67.0, 648.92)   \n",
       "26          page 11     Coordinates: Point(67.0, 552.92)   \n",
       "\n",
       "                    guess_destination_reducto_element  \n",
       "0                                     PART A OVERVIEW  \n",
       "1                                      1 Introduction  \n",
       "2   2.1 This policy document is applicable to all ...  \n",
       "3                                  3 Legal provisions  \n",
       "4                                    4 Effective date  \n",
       "5                                    5 Interpretation  \n",
       "6                          PART B POLICY REQUIREMENTS  \n",
       "7   6 Responsibilities of the board and senior man...  \n",
       "8       \\t- 7 Organisation of the compliance function  \n",
       "9   \\t- 8 Responsibilities of the compliance function  \n",
       "10  \\t- 9 Responsibilities of the internal audit f...  \n",
       "11     - 5.2 For the purpose of this policy document-  \n",
       "12  - S 8.4 The compliance function must identify ...  \n",
       "13  - S 8.11 The compliance function is responsibl...  \n",
       "14      \\t- 7 Organisation of the compliance function  \n",
       "15  \\t- 8 Responsibilities of the compliance function  \n",
       "16  - S 6.3 In relation to the position of the CCO...  \n",
       "17  - S 6.4 Senior management is collectively resp...  \n",
       "18  \\t(h) report to the board regularly on complia...  \n",
       "19  \\t(i) report to the board at least annually on...  \n",
       "20  - S 8.8 The CCO must report to senior manageme...  \n",
       "21  - S 7.2 Where compliance function responsibili...  \n",
       "22  - S 7.3 Where the CCO also assumes responsibil...  \n",
       "23  - S 9.1 A financial institution must ensure th...  \n",
       "24  \\t- 8 Responsibilities of the compliance function  \n",
       "25  - G 7.9 A constructive and cooperative working...  \n",
       "26  - S 8.8 The CCO must report to senior manageme...  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import fitz\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import math\n",
    "import json\n",
    "from collections import defaultdict\n",
    "import hashlib\n",
    "from typing import Dict, Any\n",
    "FILENAME = \"bnm_compliance\"\n",
    "\n",
    "# Unstructured didn't need this, but Reducto needed a vertical stretch to get accurate results\n",
    "HEIGHT_ADJUSTMENT_FACTOR = 0.96\n",
    "\n",
    "\n",
    "def calculate_distance(\n",
    "    point1: tuple[float, float], point2: tuple[float, float]\n",
    ") -> float:\n",
    "    \"\"\"Calculate the Euclidean distance between two points where a point is a two-item tuple for xy coords.\"\"\"\n",
    "    return math.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2)\n",
    "\n",
    "\n",
    "def point_to_segment_distance(\n",
    "    px: float, py: float, ax: float, ay: float, bx: float, by: float\n",
    ") -> float:\n",
    "    \"\"\"Calculate the minimum distance from point (px, py) to the segment (ax, ay) - (bx, by).\"\"\"\n",
    "    dx, dy = bx - ax, by - ay\n",
    "    if dx == dy == 0:  # the segment is a point\n",
    "        return calculate_distance((px, py), (ax, ay))\n",
    "\n",
    "    t = max(0, min(1, ((px - ax) * dx + (py - ay) * dy) / (dx * dx + dy * dy)))\n",
    "    nearest = (ax + t * dx, ay + t * dy)\n",
    "    return calculate_distance((px, py), nearest)\n",
    "\n",
    "\n",
    "def min_distance_to_bbox(\n",
    "    point: tuple[float, float], element_bbox_data: dict, pdf_path: str\n",
    ") -> float:\n",
    "    \"\"\"Calculate the minimum distance from a point to a 4-point bounding box.\"\"\"\n",
    "\n",
    "    # Normalise pymupdf point to pymupdf's page resolution\n",
    "    pymupdf_width, pymupdf_height = get_page_resolution_pymupdf(\n",
    "        pdf_path, element_bbox_data[\"page\"]\n",
    "    )\n",
    "    px, py = point\n",
    "    px /= pymupdf_width\n",
    "    py /= pymupdf_height\n",
    "\n",
    "    # Normalise bbox points to Reducto's page resolution\n",
    "    left, top, width, height = (\n",
    "        element_bbox_data[\"left\"],\n",
    "        element_bbox_data[\"top\"],\n",
    "        element_bbox_data[\"width\"],\n",
    "        element_bbox_data[\"height\"],\n",
    "    )\n",
    "    points = [\n",
    "        (left, top),\n",
    "        (left + width, top),\n",
    "        (left + width, top + height),\n",
    "        (left, top + height),\n",
    "    ]\n",
    "    bbox = [(x, y * HEIGHT_ADJUSTMENT_FACTOR) for x, y in points]\n",
    "\n",
    "    distances = []\n",
    "\n",
    "    # Calculate all combinations of point-to-segment distances\n",
    "    for i in range(len(bbox)):\n",
    "        ax, ay = bbox[i]\n",
    "        bx, by = bbox[(i + 1) % len(bbox)]\n",
    "        distances.append(point_to_segment_distance(px, py, ax, ay, bx, by))\n",
    "    return min(distances)\n",
    "\n",
    "\n",
    "def return_closest_elements(\n",
    "    page_number: int,\n",
    "    point: tuple[float, float],\n",
    "    elements_by_pages: defaultdict,\n",
    "    pdf_path: str,\n",
    ") -> list:\n",
    "    \"\"\"\n",
    "    Given pymupdf point coordinates and a page number, return a list of Reducto elements on the same\n",
    "    page sorted by distance to the point. Used to find the corresponding Reducto element to a PDF link\n",
    "    destination.\n",
    "    \"\"\"\n",
    "    elements = elements_by_pages[page_number]\n",
    "    sorted_elements = sorted(\n",
    "        elements,\n",
    "        key=lambda element: min_distance_to_bbox(point, element[\"bbox\"], pdf_path),\n",
    "    )\n",
    "    return sorted_elements\n",
    "\n",
    "\n",
    "def get_elements_sort_by_pages(file_name: str) -> defaultdict:\n",
    "    \"\"\"Retrieve and organise elements by pages from a Reducto output JSON file.\"\"\"\n",
    "    elements_file = os.path.join(\"./docs_structure\", \"reducto\", f\"{file_name}.json\")\n",
    "\n",
    "    with open(elements_file, \"r\") as f:\n",
    "        data = json.load(f)\n",
    "\n",
    "    elements_by_pages = defaultdict(list)\n",
    "\n",
    "    for chunk in data[\"result\"][\"chunks\"]:\n",
    "        for element in chunk[\"blocks\"]:\n",
    "            elements_by_pages[element[\"bbox\"][\"page\"]].append(element)\n",
    "\n",
    "    return elements_by_pages\n",
    "\n",
    "\n",
    "def get_elements(file_name: str) -> Dict[str, Any]:\n",
    "    \"\"\"Retrieve elements from a Reducto output JSON file.\"\"\"\n",
    "    elements_file = os.path.join(\"./docs_structure\", \"reducto\", f\"{file_name}.json\")\n",
    "    elements = {}\n",
    "\n",
    "    with open(elements_file, \"r\") as f:\n",
    "        data = json.load(f)\n",
    "\n",
    "        for chunk in data[\"result\"][\"chunks\"]:\n",
    "            for element in chunk[\"blocks\"]:\n",
    "                page_number = int(element[\"bbox\"][\"page\"])\n",
    "                element_id = hash_element(element)\n",
    "                elements[element_id] = element\n",
    "\n",
    "    return elements\n",
    "\n",
    "\n",
    "def get_page_resolution_pymupdf(pdf_path: str, page_number: int) -> tuple[float, float]:\n",
    "    \"\"\"Get the width and height of a specific page in a PDF document using pymupdf.\"\"\"\n",
    "    pdf_document = fitz.open(pdf_path)\n",
    "\n",
    "    # Ensure the page number is valid\n",
    "    if page_number < 1 or page_number > pdf_document.page_count:\n",
    "        raise ValueError(\"Page number is out of range.\")\n",
    "\n",
    "    # Load the specified page (page_number is 1-based index)\n",
    "    page = pdf_document.load_page(page_number - 1)\n",
    "\n",
    "    # Get the page dimensions in points\n",
    "    width = page.rect.width\n",
    "    height = page.rect.height\n",
    "\n",
    "    return width, height\n",
    "\n",
    "\n",
    "def hash_element(element: dict) -> str:\n",
    "    \"\"\"Generate a hash for an element using its content and bounding box.\"\"\"\n",
    "    content = element[\"content\"]\n",
    "    bbox = str(element[\"bbox\"]).encode(\"utf-8\")\n",
    "    return hashlib.md5(content.encode(\"utf-8\") + bbox).hexdigest()\n",
    "\n",
    "def whyhow_unique_node_name(content, current_element_hash_id):\n",
    "    \"\"\"Generate a unique node name for WhyHow based on the content and element hash.\"\"\"\n",
    "    clean_content = ' '.join(content.split())\n",
    "    node_name_content_preview = clean_content[:50] + \"...\" if len(clean_content) > 50 else clean_content\n",
    "    # node name has to be unique on whyhow systems, so add element hash/id to the end to make it unique\n",
    "    node_name_content_preview += \"_\"\n",
    "    node_name_content_preview += current_element_hash_id\n",
    "    return node_name_content_preview\n",
    "\n",
    "def extract_links_and_sections_to_elements(\n",
    "    pdf_path: str, elements_by_pages: defaultdict\n",
    ") -> tuple[pd.DataFrame, list[tuple]]:\n",
    "    \"\"\"\n",
    "    Extract links from a PDF document and associate them with Reducto elements based on proximity.\n",
    "\n",
    "    Returns a DataFrame containing the link information and a list of tuples representing links\n",
    "    between elements.\n",
    "    \"\"\"\n",
    "    document = fitz.open(pdf_path)\n",
    "    links_data = []\n",
    "    links_tuples = []\n",
    "\n",
    "    for page_num in range(len(document)):\n",
    "        page = document.load_page(page_num)\n",
    "        links = page.get_links()\n",
    "\n",
    "        for link in links:\n",
    "            link_data = {}\n",
    "            link_data[\"page_number\"] = page_num + 1\n",
    "\n",
    "            # Extract link text\n",
    "            rect = fitz.Rect(link[\"from\"])\n",
    "            text = page.get_text(\"text\", clip=rect)\n",
    "            link_data[\"link_text\"] = text.strip() if text else \"No text\"\n",
    "\n",
    "            x, y, _, _ = link[\"from\"]  # Get the coordinates of the link\n",
    "            closest_elems = return_closest_elements(\n",
    "                page_num + 1, (x, y), elements_by_pages, pdf_path\n",
    "            )\n",
    "            closest_elem = closest_elems[0] if closest_elems else None\n",
    "            if closest_elem:\n",
    "                link_data[\"guess_link_reducto_element\"] = closest_elem[\"content\"]\n",
    "                link_element_id = hash_element(closest_elem)\n",
    "            else:\n",
    "                link_data[\"guess_link_reducto_element\"] = \"N/A\"\n",
    "                link_element_id = None\n",
    "\n",
    "            # Determine link type and location\n",
    "            if link[\"kind\"] == fitz.LINK_GOTO:\n",
    "                link_data[\"link_type\"] = \"internal\"\n",
    "                link_data[\"link_destination\"] = \"page \" + str(\n",
    "                    link[\"page\"] + 1 if link[\"page\"] >= 0 else \"Unknown\"\n",
    "                )\n",
    "\n",
    "                # Extract named destination if available\n",
    "                dest = link.get(\"to\", None)\n",
    "                if dest:\n",
    "                    link_data[\"named_destination\"] = f\"Coordinates: {dest}\"\n",
    "                    x, y = dest\n",
    "                    target_page_num = link[\"page\"] + 1\n",
    "                    closest_elems = return_closest_elements(\n",
    "                        target_page_num, (x, y), elements_by_pages, pdf_path\n",
    "                    )\n",
    "                    closest_elem = closest_elems[0] if closest_elems else None\n",
    "                    if closest_elem:\n",
    "                        link_data[\"guess_destination_reducto_element\"] = closest_elem[\n",
    "                            \"content\"\n",
    "                        ]\n",
    "                        link_destination_element_id = hash_element(closest_elem)\n",
    "                    else:\n",
    "                        link_data[\"guess_destination_reducto_element\"] = \"N/A\"\n",
    "                        link_destination_element_id = None\n",
    "                else:\n",
    "                    link_data[\"named_destination\"] = \"N/A\"\n",
    "                    link_data[\"guess_destination_reducto_element\"] = \"N/A\"\n",
    "                    link_destination_element_id = None\n",
    "\n",
    "            elif link[\"kind\"] == fitz.LINK_URI:\n",
    "                link_data[\"link_type\"] = \"external\"\n",
    "                link_data[\"link_destination\"] = link[\"uri\"]\n",
    "                link_data[\"named_destination\"] = \"N/A\"\n",
    "                link_data[\"guess_destination_reducto_element\"] = \"N/A\"\n",
    "                link_destination_element_id = link[\"uri\"]\n",
    "            else:\n",
    "                link_data[\"link_type\"] = \"unknown\"\n",
    "                link_data[\"link_destination\"] = \"Unknown location\"\n",
    "                link_data[\"named_destination\"] = \"N/A\"\n",
    "                link_data[\"guess_destination_reducto_element\"] = \"N/A\"\n",
    "                link_destination_element_id = None\n",
    "\n",
    "            # Keep track of tuples to create links later in Neo4j\n",
    "            links_tuples.append((link_element_id, link_destination_element_id))\n",
    "            links_data.append(link_data)\n",
    "\n",
    "    df = pd.DataFrame(links_data)\n",
    "\n",
    "    # Replace NaN and infinite values\n",
    "    df.replace([np.inf, -np.inf, np.nan], \"N/A\", inplace=True)\n",
    "\n",
    "    return df, links_tuples\n",
    "\n",
    "\n",
    "def save_to_csv(df: pd.DataFrame, output_path: str) -> None:\n",
    "    \"\"\"Save a DataFrame to a CSV file at the specified path.\"\"\"\n",
    "    # Create directory if it doesn't exist\n",
    "    os.makedirs(os.path.dirname(output_path), exist_ok=True)\n",
    "    df.to_csv(output_path, index=False)\n",
    "\n",
    "\n",
    "def analyse_links(\n",
    "    filename: str, elements_by_pages: defaultdict\n",
    ") -> tuple[pd.DataFrame, list[tuple]]:\n",
    "    \"\"\"\n",
    "    Analyse links in a PDF document, matching them to Reducto elements, and saving the results to a CSV file.\n",
    "\n",
    "    Returns a DataFrame with the link data and a list of tuples representing links between elements.\n",
    "    \"\"\"\n",
    "    # Path to your PDF file\n",
    "    pdf_path = f\"docs/{filename}.pdf\"\n",
    "\n",
    "    # Extract links to df\n",
    "    links_df, links_tuples = extract_links_and_sections_to_elements(\n",
    "        pdf_path, elements_by_pages\n",
    "    )\n",
    "    print(\"Links tuples: \", links_tuples)\n",
    "\n",
    "    # Save DataFrame to CSV\n",
    "    output_csv_path = f\"docs_links/{filename}-with-reducto.csv\"\n",
    "    save_to_csv(links_df, output_csv_path)\n",
    "\n",
    "    return links_df, links_tuples\n",
    "\n",
    "\n",
    "# Example usage\n",
    "links_df, links_tuples = analyse_links(\n",
    "    FILENAME, get_elements_sort_by_pages(FILENAME)\n",
    ")\n",
    "links_df.head(200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the Definitions Graph with WhyHow\n",
    "\n",
    "## Setup WhyHow Workspace\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FETCH/CREATE WORKSPACE: Looking for workspaces:\n",
      "FETCH/CREATE WORKSPACE: Found workspace with matching name BNM Compliance Test id=66d7fd24779a727e458a5945.\n",
      "UPDATE SCHEMA: Checking for existing schemas with same name...\n",
      "UPDATE SCHEMA: Updating schema from schema/lexical_schema.json...\n",
      "UPDATE SCHEMA: Uploaded schema lexical_schema.json with ID: 66d822bceda622fc957cc7fe\n",
      "\n",
      "UPDATE SCHEMA: Checking for existing schemas with same name...\n",
      "UPDATE SCHEMA: Updating schema from schema/definitions_schema.json...\n",
      "UPDATE SCHEMA: Uploaded schema definitions_schema.json with ID: 66d822bd779a727e458a5aba\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from typing import List, Tuple, Dict, Any, Optional\n",
    "import csv\n",
    "from whyhow.schemas import Triple, Node, Relation, Chunk, Document, Workspace, Schema\n",
    "from whyhow import WhyHow\n",
    "import sys\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "WHYHOW_API_KEY = os.getenv(\"WHYHOW_API_KEY\")\n",
    "WHYHOW_API_URL = os.getenv(\"WHYHOW_API_URL\")\n",
    "\n",
    "\n",
    "def fetch_or_create_new_workspace(\n",
    "    workspace_name: str, client: Optional[WhyHow] = None\n",
    ") -> Workspace:\n",
    "    \"\"\"\n",
    "    Fetches an existing workspace by name or creates a new one if none exists.\n",
    "    \"\"\"\n",
    "    # SDK only supports getting workspace by ID, since one can create workspaces with duplicate names\n",
    "    # We assume for convenience that the workspace name is unique so we can fetch workspace by name over and over again\n",
    "\n",
    "    # Don't have to keep creating new workspaces every time cells are run\n",
    "    if client is None:\n",
    "        client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "\n",
    "    print(\"FETCH/CREATE WORKSPACE: Looking for workspaces:\")\n",
    "    workspaces = client.workspaces.get_all(limit=20, name=workspace_name)\n",
    "\n",
    "    # return the first instance of the workspace with the same name\n",
    "    for workspace in workspaces:\n",
    "        if workspace.name == workspace_name:\n",
    "            print(\n",
    "                f\"FETCH/CREATE WORKSPACE: Found workspace with matching name {workspace_name} id={workspace.workspace_id}.\"\n",
    "            )\n",
    "            return workspace\n",
    "    print(\n",
    "        f\"FETCH/CREATE WORKSPACE: No workspace found with name {workspace_name}. Creating new workspace...\"\n",
    "    )\n",
    "    return client.workspaces.create(workspace_name)\n",
    "\n",
    "\n",
    "WORKSPACE_ID = fetch_or_create_new_workspace(\"Central Bank Compliance\").workspace_id\n",
    "FILENAME = \"bnm_compliance\"\n",
    "\n",
    "\n",
    "def update_schema(schema_dir: str, client: Optional[WhyHow] = None) -> Schema:\n",
    "    \"\"\"\n",
    "    Updates a schema on WhyHow by uploading a new schema file. If a schema with the same name exists, it is deleted.\n",
    "    \"\"\"\n",
    "    # Convention: only one schema should exist on WhyHow for every schema file here.\n",
    "    # Otherwise, constantly running this cell will create duplicate, messy schemas.\n",
    "\n",
    "    if client is None:\n",
    "        client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "    # check that schema_dir points to a file that exists\n",
    "    if not os.path.isfile(schema_dir):\n",
    "        raise FileNotFoundError(f\"Schema file {schema_dir} not found.\")\n",
    "\n",
    "    print(\"UPDATE SCHEMA: Checking for existing schemas with same name...\")\n",
    "    schema_name = os.path.basename(schema_dir)\n",
    "\n",
    "    # Delete all schemas of the same name.\n",
    "    workspace_schemas = client.schemas.get_all(workspace_id=WORKSPACE_ID)\n",
    "    if not workspace_schemas:\n",
    "        matched_schemas = [\n",
    "            schema for schema in workspace_schemas if schema.name == schema_name\n",
    "        ]\n",
    "\n",
    "        if matched_schemas:\n",
    "            print(\n",
    "                f\"UPDATE SCHEMA: Found {len(matched_schemas)} existing schema(s) with same name. Deleting...\"\n",
    "            )\n",
    "            for schema in matched_schemas:\n",
    "                client.schemas.delete(schema.schema_id)\n",
    "\n",
    "    print(f\"UPDATE SCHEMA: Updating schema from {schema_dir}...\")\n",
    "    entities, relations, patterns = client.schemas.load_json(schema_dir)\n",
    "    new_schema = client.schemas.create(\n",
    "        workspace_id=WORKSPACE_ID,\n",
    "        name=schema_name,\n",
    "        entities=entities,\n",
    "        relations=relations,\n",
    "        patterns=patterns,\n",
    "    )\n",
    "    print(\n",
    "        f\"UPDATE SCHEMA: Uploaded schema {schema_name} with ID: {new_schema.schema_id}\\n\"\n",
    "    )\n",
    "    return new_schema\n",
    "\n",
    "\n",
    "import shutil\n",
    "\n",
    "\n",
    "def update_document(document_dir: str, client: Optional[WhyHow] = None) -> Document:\n",
    "    \"\"\"\n",
    "    Updates a document on WhyHow by uploading a new document file. If a document with the same name exists, it is deleted.\n",
    "    \"\"\"\n",
    "    if client is None:\n",
    "        client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "    # Convention: only one document should exist on WhyHow for every document file here.\n",
    "    # Otherwise, constantly running this cell will create duplicate, messy documents.\n",
    "\n",
    "    # check that document_dir points to a file that exists\n",
    "    if not os.path.isfile(document_dir):\n",
    "        raise FileNotFoundError(f\"Document file {document_dir} not found.\")\n",
    "\n",
    "    document_name = os.path.basename(document_dir)\n",
    "\n",
    "    existing_docs = list(\n",
    "        client.documents.get_all(workspace_id=WORKSPACE_ID, filename=document_name)\n",
    "    )\n",
    "    if existing_docs:\n",
    "        print(\"UPDATE DOC: Found old existing document in workspace. Deleting...\")\n",
    "        _ = client.documents.delete(existing_docs[0].document_id)\n",
    "    else:\n",
    "        print(f\"UPDATE DOC: No old existing document found in workspace.\")\n",
    "\n",
    "    print(f\"UPDATE DOC: Creating new document {document_name} in workspace...\")\n",
    "    try:\n",
    "        document = client.documents.upload(path=document_dir, workspace_id=WORKSPACE_ID)\n",
    "        return document\n",
    "    except Exception as e:\n",
    "        print(\n",
    "            f\"UPDATE DOC: Failed to upload document: renaming... {document_name}. {e}\"\n",
    "        )\n",
    "        # get local doc and make a copy with a number at the end:\n",
    "        base_name, ext = os.path.splitext(document_name)\n",
    "        new_document_name = f\"{base_name}_1{ext}\"\n",
    "        new_document_dir = os.path.join(\n",
    "            os.path.dirname(document_dir), new_document_name\n",
    "        )\n",
    "\n",
    "        # copy the local doc to the new name\n",
    "        shutil.copy(document_dir, new_document_dir)\n",
    "\n",
    "        print(f\"UPDATE DOC: Retrying with new document name {new_document_name}...\")\n",
    "        try:\n",
    "            document = client.documents.upload(\n",
    "                path=new_document_dir, workspace_id=WORKSPACE_ID\n",
    "            )\n",
    "            return document\n",
    "        except Exception as e:\n",
    "            print(\n",
    "                f\"UPDATE DOC: Failed to upload document again: {new_document_name}. {e}\"\n",
    "            )\n",
    "            raise\n",
    "\n",
    "\n",
    "def retrieve_document(document_name, client=None):\n",
    "    \"\"\"\n",
    "    Retrieves a document by name from the WhyHow workspace.\n",
    "    \"\"\"\n",
    "    if client is None:\n",
    "        client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "\n",
    "    print(f\"RETRIEVE DOC: Fetching document {document_name} from workspace...\")\n",
    "    documents = list(\n",
    "        client.documents.get_all(\n",
    "            limit=20, filename=document_name, workspace_id=WORKSPACE_ID\n",
    "        )\n",
    "    )\n",
    "    if not documents:\n",
    "        raise FileNotFoundError(\n",
    "            f\"RETRIEVE DOC:  Failed to find existing document by filename.\"\n",
    "        )\n",
    "    else:\n",
    "        print(\n",
    "            f\"RETRIEVE DOC: Found document in workspace. Fetching first instance with id {documents[0].document_id}\"\n",
    "        )\n",
    "        return documents[0]\n",
    "\n",
    "\n",
    "LEXICAL_SCHEMA_NAME = \"Lexical Graph Schema\"\n",
    "LEXICAL_SCHEMA_DIR = os.path.join(\"schema\", \"lexical_schema.json\")\n",
    "DEFINITIONS_SCHEMA_NAME = \"Definition Graph Schema\"\n",
    "DEFINITIONS_SCHEMA_DIR = os.path.join(\"schema\", \"definitions_schema.json\")\n",
    "LEXICAL_SCHEMA_ID = update_schema(LEXICAL_SCHEMA_DIR).schema_id\n",
    "\n",
    "DEFINITIONS_SCHEMA_ID = update_schema(DEFINITIONS_SCHEMA_DIR).schema_id\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the definition Graph on WhyHow\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "page 4: [TermDefinition(term='S', definition='denotes a standard, an obligation, a requirement, specification, direction, condition and any interpretative, supplemental and transitional provisions that must be complied with. Non-compliance may result in enforcement action;'), TermDefinition(term='G', definition='denotes guidance which may consist of statements or information intended to promote common understanding and advice or recommendations that are encouraged to be adopted;'), TermDefinition(term='board', definition='means the board of directors of a financial institution, including a committee of the board where the responsibilities of the board set out in this policy document have been delegated to such a committee;'), TermDefinition(term='senior management', definition='refers to the chief executive officer and senior officers of a financial institution;'), TermDefinition(term='chief compliance officer', definition='means the senior officer of a financial institution, however styled, who is the central point of authority for a financial institution’s compliance matters and is responsible for providing an institution-wide view on the management of compliance risk;'), TermDefinition(term='compliance function', definition='refers to officers carrying out compliance function responsibilities of a financial institution, as described in paragraphs 8.4 to 8.11, including the CCO;')]\n",
      "page 5: [TermDefinition(term='compliance risk', definition='the risk of legal or regulatory sanctions, financial loss or reputational damage which a financial institution may suffer as a result of its failure to comply with legal and regulatory requirements applicable to its activities'), TermDefinition(term='control function', definition='a function that has a responsibility independent from business lines to provide objective assessment, reporting or assurance. This includes the risk management function, the compliance function and the internal audit function'), TermDefinition(term='financial institution', definition='a licensed person under the FSA; (b) a licensed person under the IFSA; (c) a development financial institution prescribed under the DFIA; and (d) a financial holding company'), TermDefinition(term='large financial institution', definition='a financial institution with one or more business lines that are significant in terms of market share in the relevant industry; or (b) a financial institution with a large network of offices within or outside the country through operations of branches and subsidiaries'), TermDefinition(term='legal and regulatory requirements', definition='all laws, rules, standards and regulatory requirements (including any ruling of the Shariah Advisory Council) relevant to a financial institution’s activities in all jurisdictions in which the financial institution, or any of its branches or subsidiaries, conducts activities')]\n",
      "Created new CSV file, saved to ./whyhow_csv_chunks/bnm_compliance_definitions.csv\n",
      "UPDATE DOC: Found old existing document in workspace. Deleting...\n",
      "UPDATE DOC: Creating new document bnm_compliance_definitions.csv in workspace...\n",
      "UPDATE SCHEMA: Checking for existing schemas with same name...\n",
      "UPDATE SCHEMA: Updating schema from schema/definitions_schema.json...\n",
      "UPDATE SCHEMA: Uploaded schema definitions_schema.json with ID: 66d822d0eda622fc957cc80b\n",
      "\n",
      "Created graph with ID: 66d822d1eda622fc957cc80c\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import fitz  # PyMuPDF\n",
    "import csv\n",
    "from pydantic import BaseModel\n",
    "from openai import OpenAI\n",
    "from dotenv import load_dotenv\n",
    "from typing import Tuple, List\n",
    "\n",
    "from whyhow import WhyHow\n",
    "from whyhow.schemas import Triple, Node, Relation, Chunk, Graph\n",
    "\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Get OpenAI API key from environment variables\n",
    "openai_api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "if openai_api_key is None:\n",
    "    raise ValueError(\"OPENAI_API_KEY not found in environment variables.\")\n",
    "\n",
    "# Initialize OpenAI client\n",
    "openai_client = OpenAI()\n",
    "openai_client.api_key = openai_api_key\n",
    "\n",
    "\n",
    "class TermDefinition(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe the structured output of legal term-definition pairs for the GPT-4o model.\n",
    "    \"\"\"\n",
    "\n",
    "    term: str\n",
    "    definition: str\n",
    "\n",
    "\n",
    "class TermDefinitionCSV(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe the structured output of legal term-definition pairs CSv for GPT-4o\n",
    "    \"\"\"\n",
    "\n",
    "    content: List[TermDefinition]\n",
    "\n",
    "\n",
    "class DefinitionGraphBuilder(BaseModel):\n",
    "    \"\"\"\n",
    "    Given a range of pages, the builder extracts legal terms and their definitions from a PDF.\n",
    "    This PDF is found using the global FILENAME.\n",
    "    The builder can save extracted terms as a CSV file, and upload to WhyHow.\n",
    "\n",
    "    If build_csv is False, no new CSV is created, only the existing one is uploaded.\n",
    "    \"\"\"\n",
    "\n",
    "    # page range starts count from 1\n",
    "    definitions_page_ranges: List[Tuple[int, int]]\n",
    "    file_path: str = f\"./docs/{FILENAME}.pdf\"\n",
    "    definitions_csv_path: str = f\"./whyhow_csv_chunks/{FILENAME}_definitions.csv\"\n",
    "    definitions_csv_filename: str = os.path.basename(definitions_csv_path)\n",
    "    schema_dir: str = DEFINITIONS_SCHEMA_DIR\n",
    "\n",
    "    # if build_csv is False, the builder will not create a new local CSV, it will just upload the existing one\n",
    "    build_csv: bool = True\n",
    "    graph_name: str = \"Legal Definitions Graph\"\n",
    "    graph: Optional[Graph] = None\n",
    "\n",
    "    def extract_and_save_to_csv(self):\n",
    "        doc = fitz.open(self.file_path)\n",
    "\n",
    "        all_terms_and_definitions = []\n",
    "\n",
    "        # Flatten the list of ranges into individual pages\n",
    "        page_numbers = []\n",
    "        for start, end in self.definitions_page_ranges:\n",
    "            page_numbers.extend(range(start, end + 1))\n",
    "\n",
    "        for page_num in page_numbers:\n",
    "            page = doc.load_page(page_num - 1)  # subtract 1 for 0-based index\n",
    "            text = page.get_text(\"text\")\n",
    "\n",
    "            prompt = f\"\"\"\n",
    "You will be shown a text from a page of a document containing legal terms and their definitions. Your task is to extract the legal terms and their corresponding definitions (verbatim) from the text and return them as the schema JSON.\n",
    "\n",
    "Text:\n",
    "{text}\n",
    "\n",
    "Provide the extracted terms and definitions in the specified CSV format.\n",
    "\"\"\"\n",
    "            completion = openai_client.beta.chat.completions.parse(\n",
    "                model=\"gpt-4o-2024-08-06\",\n",
    "                messages=[\n",
    "                    {\"role\": \"system\", \"content\": \"Extract the terms and definitions.\"},\n",
    "                    {\"role\": \"user\", \"content\": prompt},\n",
    "                ],\n",
    "                response_format=TermDefinitionCSV,\n",
    "            )\n",
    "\n",
    "            terms_and_definitions = completion.choices[0].message.parsed.content\n",
    "            print(f\"page {page_num}: {terms_and_definitions}\")\n",
    "\n",
    "            # Ensure that terms_and_definitions is a list of TermDefinition objects\n",
    "            if isinstance(terms_and_definitions, list):\n",
    "                for item in terms_and_definitions:\n",
    "                    if isinstance(item, TermDefinition):\n",
    "                        all_terms_and_definitions.append((item.term, item.definition))\n",
    "                    else:\n",
    "                        print(\"Unexpected item format:\", item)\n",
    "\n",
    "        with open(self.definitions_csv_path, \"w\", newline=\"\") as csvfile:\n",
    "            writer = csv.writer(csvfile)\n",
    "            writer.writerow([\"Legal Term\", \"Definition\"])\n",
    "            for term, definition in all_terms_and_definitions:\n",
    "                writer.writerow([term, definition])\n",
    "\n",
    "        print(f\"Created new CSV file, saved to {self.definitions_csv_path}\")\n",
    "\n",
    "    def build_definitions_graph(self):\n",
    "        client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "        if self.build_csv:\n",
    "            self.extract_and_save_to_csv()\n",
    "            document = update_document(self.definitions_csv_path)\n",
    "        else:\n",
    "            document = retrieve_document(self.definitions_csv_filename)\n",
    "\n",
    "        schema = update_schema(self.schema_dir)\n",
    "\n",
    "        self.graph = client.graphs.create(\n",
    "            name=\"Legal Definitions Graph\",\n",
    "            workspace_id=WORKSPACE_ID,\n",
    "            document_ids=[document.document_id],\n",
    "            schema_id=schema.schema_id,\n",
    "            mode=\"structured\",\n",
    "        )\n",
    "        print(f\"Created graph with ID: {self.graph.graph_id}\")\n",
    "\n",
    "        return self.graph\n",
    "\n",
    "\n",
    "definitions_graph = DefinitionGraphBuilder(\n",
    "    definitions_page_ranges=[(4, 5)],\n",
    "    file_path=f\"./docs/{FILENAME}.pdf\",\n",
    "    definitions_csv_path=f\"./whyhow_csv_chunks/{FILENAME}_definitions.csv\",\n",
    "    build_csv=True,\n",
    "    graph_name=\"Legal Definitions Graph\",\n",
    ").build_definitions_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the Lexical Graph with WhyHow\n",
    "\n",
    "- Note: WhyHow only has one identifying field for the node, node_name, and it has to be unique. \n",
    "  - Due to this, Node name and id mean the same thing. On the UI, it shows node_name to the user\n",
    "  - The way we structured this notebook is to use a preview (first 50 characters of element content) + a hash of the reducto element to use as node_name/node_id\n",
    "  - This makes things a lot more readable on the UI\n",
    "  - This also has a nice side effect of node_ids also containing a little bit of content/context for LLMs. See more in the \"Updating Node Context with Link Info\":\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "from typing import Tuple, List\n",
    "\n",
    "from whyhow import WhyHow\n",
    "from whyhow.schemas import Triple, Node, Relation, Chunk, Graph, Workspace\n",
    "\n",
    "from enum import Enum\n",
    "\n",
    "\n",
    "class ElementType(str, Enum):  # subclass of str so JSON serializable\n",
    "    PAGE = \"Page\"\n",
    "    TITLE = \"Title\"\n",
    "    FIGURE = \"Figure\"\n",
    "    FOOTER = \"Footer\"\n",
    "    LIST_ITEM = \"List Item\"\n",
    "    HEADER = \"Header\"\n",
    "    DOCUMENT = \"Document\"\n",
    "    SECTION_HEADER = \"Section Header\"\n",
    "    TEXT = \"Text\"\n",
    "    TABLE = \"Table\"\n",
    "\n",
    "\n",
    "ELEMENTS = [ElementType.value for ElementType in ElementType]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "from whyhow.schemas import Triple, Node, Relation, Chunk, Graph\n",
    "\n",
    "\n",
    "class WhyHowLexicalGraphBuilder(BaseModel):\n",
    "    class Config:\n",
    "        # WhyHow SDK doesn't have validators for WhyHow Client\n",
    "        arbitrary_types_allowed = True\n",
    "\n",
    "    file_name: str\n",
    "    client: Optional[WhyHow] = None\n",
    "    workspace: Workspace = fetch_or_create_new_workspace(\"Central Bank Compliance\")\n",
    "    build_and_upload_csvs: bool = True\n",
    "\n",
    "    json_data: Optional[Dict[str, Any]] = None\n",
    "    graph: Optional[Graph] = None\n",
    "    schema_dir: str = LEXICAL_SCHEMA_DIR\n",
    "\n",
    "    triples: List[Triple] = []\n",
    "\n",
    "    # Data structures to store Elements\n",
    "    elements_by_page_number: Optional[Dict[str, List[Dict[str, Any]]]] = None\n",
    "    elements: Optional[List[Dict[str, Any]]] = None\n",
    "    # Other data structures to store links and footers\n",
    "    links_tuples: Optional[List[Tuple[str, str]]] = None\n",
    "    footers_by_page_number: Optional[Dict[str, List[str]]] = None\n",
    "    # Init a lookup map that links an element hash id to its corresponding chunk on WhyHow\n",
    "    hash_lookup_chunk: Dict[str, Chunk] = {}\n",
    "    # Init a lookup map that links an element hash id to its corresponding node on WhyHow\n",
    "    hash_lookup_node: Dict[str, Node] = {}\n",
    "    # Init a map that links a chunk_id to its corresponding chunk\n",
    "    chunk_map: Dict[str, Chunk] = {}\n",
    "\n",
    "    def __init__(self, **data):\n",
    "        super().__init__(**data)\n",
    "        # Load the JSON data for elements\n",
    "        print(\"Loading JSON data...\")\n",
    "        try:\n",
    "            with open(\n",
    "                os.path.join(\"./docs_structure\", \"reducto\", f\"{self.file_name}.json\"),\n",
    "                \"r\",\n",
    "            ) as f:\n",
    "                self.json_data = json.load(f)\n",
    "                print(\"JSON data loaded.\")\n",
    "            if not self.json_data:\n",
    "                raise ValueError(\"No data found in JSON file.\")\n",
    "        except FileNotFoundError:\n",
    "            raise FileNotFoundError(\"JSON file not found.\")\n",
    "\n",
    "        # Data structures to store Elements\n",
    "        self.elements_by_page_number = get_elements_sort_by_pages(self.file_name)\n",
    "        self.elements = get_elements(self.file_name)\n",
    "        # Other data structures to store links and footers\n",
    "        self.links_tuples = analyse_links(self.file_name, self.elements_by_page_number)[\n",
    "            1\n",
    "        ]\n",
    "        self.footers_by_page_number = defaultdict(list)\n",
    "\n",
    "        if self.client is None:\n",
    "            self.client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "\n",
    "    def build_csvs(self):\n",
    "        print(\"Initialising CSV files...\")\n",
    "\n",
    "        for element_type in ELEMENTS:\n",
    "            file_name_element = element_type.lower().replace(\" \", \"_\")\n",
    "            csv_path = os.path.join(\n",
    "                \"./whyhow_csv_chunks\", f\"{self.file_name}_{file_name_element}.csv\"\n",
    "            )\n",
    "\n",
    "            with open(csv_path, \"w\", newline=\"\") as f:\n",
    "                writer = csv.writer(f)\n",
    "                if element_type == \"Page\":\n",
    "                    writer.writerow(\n",
    "                        [element_type, \"Type\", \"Hash\", \"Node Name\", \"Page Number\"]\n",
    "                    )\n",
    "                elif element_type == \"Document\":\n",
    "                    writer.writerow([element_type, \"Type\"])\n",
    "                else:\n",
    "                    writer.writerow(\n",
    "                        [\n",
    "                            element_type,\n",
    "                            \"Type\",\n",
    "                            \"Hash\",\n",
    "                            \"Node Name\",\n",
    "                            \"Content\",\n",
    "                            \"Page Number\",\n",
    "                            \"Location\",\n",
    "                        ]\n",
    "                    )\n",
    "\n",
    "        # write CSVs consisting of elements\n",
    "        print(\"Writing elements to CSV files...\")\n",
    "        for reducto_chunk in self.json_data[\"result\"][\"chunks\"]:\n",
    "            for element in reducto_chunk[\"blocks\"]:\n",
    "                type = element[\"type\"]\n",
    "                page_number = element[\"bbox\"][\"page\"]\n",
    "                element_id_hash = hash_element(element)\n",
    "\n",
    "                # Escape double quotes in content\n",
    "                content = element[\"content\"].replace('\"', '\"\"')\n",
    "                # for pydantic issues where str cannot be \"\"\n",
    "                if content == \"\":\n",
    "                    content = \"(empty content)\"\n",
    "\n",
    "                node_name = whyhow_unique_node_name(content, element_id_hash)\n",
    "\n",
    "                if type in ELEMENTS and type not in [\"Page\", \"Document\"]:\n",
    "                    file_name_element = type.lower().replace(\" \", \"_\")\n",
    "                    csv_path = os.path.join(\n",
    "                        \"./whyhow_csv_chunks\",\n",
    "                        f\"{self.file_name}_{file_name_element}.csv\",\n",
    "                    )\n",
    "                    # Header Row for Figure/Footer/Header/List Item/Section Header/Table/Text/Title CSVs\n",
    "                    # [element_type, \"Type\", \"Hash\", \"Node Name\", \"Content\", \"Page Number\", \"Location\"]\n",
    "                    data_row = [\n",
    "                        element_id_hash,\n",
    "                        node_name,\n",
    "                        f'\"{content}\"',\n",
    "                        page_number,\n",
    "                        f'\"{element[\"bbox\"]}\"',\n",
    "                    ]\n",
    "                    with open(csv_path, \"a\", newline=\"\") as f:\n",
    "                        writer = csv.writer(f)\n",
    "                        writer.writerow([f'\"{content}\"', type] + data_row)\n",
    "\n",
    "        # Write a single row to represent the document\n",
    "        document_csv_path = os.path.join(\n",
    "            \"./whyhow_csv_chunks\", f\"{self.file_name}_document.csv\"\n",
    "        )\n",
    "        with open(document_csv_path, \"a\", newline=\"\") as f:\n",
    "            writer = csv.writer(f)\n",
    "            # Header row for Document CSV is [\"Document\", \"Type\"]\n",
    "            writer.writerow([self.file_name, \"Document\"])\n",
    "\n",
    "        # Write a single row to represent the cover page. Can add on more for other pages e.g. Table of Contents, Appendix.\n",
    "        # Need to come up with a programmatic way for write_csvs to handle custom pages\n",
    "        pages_csv_path = os.path.join(\n",
    "            \"./whyhow_csv_chunks\", f\"{self.file_name}_page.csv\"\n",
    "        )\n",
    "        # Header row for Pages CSV is [\"Page\", \"Type\", \"Hash\", \"Node Name\", \"Page Number\"]\n",
    "        with open(pages_csv_path, \"a\", newline=\"\") as f:\n",
    "            writer = csv.writer(f)\n",
    "            writer.writerow([\"Cover Page\", \"Page\", \"cover_page\", \"Cover Page\", 1])\n",
    "\n",
    "    def upload_csvs_to_whyhow(self):\n",
    "        # Add CSV documents to workspace:\n",
    "        print(\"Uploading CSV files to WhyHow workspace...\")\n",
    "        for element_type in ELEMENTS:\n",
    "            element_type = element_type.lower()\n",
    "            file_name_element = element_type.replace(\" \", \"_\")\n",
    "            csv_path = os.path.join(\n",
    "                \"./whyhow_csv_chunks\", f\"{self.file_name}_{file_name_element}.csv\"\n",
    "            )\n",
    "            doc = update_document(csv_path)\n",
    "\n",
    "    def update_hash_lookups(self):\n",
    "        print(\"Updating hash_chunkid_map...\")\n",
    "        docs = self.client.documents.get_all(workspace_id=self.workspace.workspace_id)\n",
    "        for doc in docs:\n",
    "            print(f\"Processing document {doc.metadata.filename}\")\n",
    "            chunks = self.client.chunks.get_all(document_filename=doc.metadata.filename)\n",
    "            skip_doc = False  # Flag to determine if the document should be skipped\n",
    "            # we retrieve all docs from WhyHow, but only want to process the csvs for document elements\n",
    "\n",
    "            for chunk in chunks:\n",
    "                if \"Type\" not in chunk.content:\n",
    "                    skip_doc = True  # Set the flag to skip the document\n",
    "                    print(\n",
    "                        f\"Skipping document {doc.metadata.filename}, it is not an elements csv file. Probably a definitions csv file.\"\n",
    "                    )\n",
    "                    break  # Exit the chunk loop and skip the document\n",
    "\n",
    "                chunk_type = chunk.content[\"Type\"]\n",
    "                if chunk_type == \"Page\":\n",
    "                    self.hash_lookup_chunk[chunk.content[\"Hash\"]] = chunk\n",
    "                elif chunk_type == \"Document\":\n",
    "                    self.hash_lookup_chunk[self.file_name] = chunk\n",
    "                else:\n",
    "                    self.hash_lookup_chunk[chunk.content[\"Hash\"]] = chunk\n",
    "\n",
    "                # Add chunk to chunk_map\n",
    "                self.chunk_map[chunk.chunk_id] = chunk\n",
    "\n",
    "            if skip_doc:\n",
    "                continue  # Skip processing the rest of this document\n",
    "\n",
    "    def create_new_node(\n",
    "        self, element_hash_id, type, node_name_content_preview, properties, chunk_ids=[]\n",
    "    ):\n",
    "        # create a new node with SDK\n",
    "        self.hash_lookup_node[element_hash_id] = Node(\n",
    "            label=type,\n",
    "            name=node_name_content_preview,\n",
    "            chunk_ids=chunk_ids,\n",
    "            properties=properties,\n",
    "        )\n",
    "        print(f\"Created new node: {element_hash_id} with type {type}\")\n",
    "\n",
    "    def link_nodes(self, head_hash, tail_hash, relation):\n",
    "        # link any two nodes, given their hash ids and a relation\n",
    "        chunk_ids = [\n",
    "            self.hash_lookup_chunk[head_hash].chunk_id,\n",
    "            self.hash_lookup_chunk[tail_hash].chunk_id,\n",
    "        ]\n",
    "\n",
    "        # for pydantic reasons, we cannot have empty strings in chunk_ids\n",
    "        if \"\" in chunk_ids:\n",
    "            chunk_ids = [id for id in chunk_ids if id != \"\"]\n",
    "\n",
    "        self.triples.append(\n",
    "            Triple(\n",
    "                head=self.hash_lookup_node[head_hash],\n",
    "                tail=self.hash_lookup_node[tail_hash],\n",
    "                relation=relation,\n",
    "                chunk_ids=chunk_ids,\n",
    "            )\n",
    "        )\n",
    "        print(f\"Linked {head_hash} to {tail_hash} with relation {relation.name}\")\n",
    "\n",
    "    def build_graph_whyhow(self):\n",
    "        # If we already have uploaded the CSVs, we can skip the CSV creation and uploading steps\n",
    "        # update_hash_lookups updates local map with data from whyhow\n",
    "        if not self.build_and_upload_csvs:\n",
    "            self.update_hash_lookups()\n",
    "        else:\n",
    "            self.build_csvs()\n",
    "            self.upload_csvs_to_whyhow()\n",
    "            self.update_hash_lookups()\n",
    "\n",
    "        last_section_id = None\n",
    "        last_list_item_id = None\n",
    "\n",
    "        # Denotes an element containing another element. Used for collecting indented list items under a single list item\n",
    "        contains_relation = Relation(name=\"contains\")\n",
    "\n",
    "        # Denotes a section being a parent of elements .\n",
    "        is_parent_relation = Relation(name=\"is_parent\")\n",
    "\n",
    "        # Denotes a link in the pdf, linking from one section/element to another\n",
    "        links_to_relation = Relation(name=\"links_to\")\n",
    "\n",
    "        # Used to keep track of section header flows in a document\n",
    "        follows_relation = Relation(name=\"follows\")\n",
    "\n",
    "        print(\"Initialising Triples and Nodes...\")\n",
    "        # create new document\n",
    "        self.hash_lookup_node[self.file_name] = Node(\n",
    "            label=\"Document\",\n",
    "            name=self.file_name,\n",
    "            properties=self.hash_lookup_chunk[self.file_name].content,\n",
    "            chunk_ids=[self.hash_lookup_chunk[self.file_name].chunk_id],\n",
    "        )\n",
    "\n",
    "        # create cover page\n",
    "        self.hash_lookup_node[\"cover_page\"] = Node(\n",
    "            label=\"Page\",\n",
    "            name=\"Cover Page\",\n",
    "            properties=self.hash_lookup_chunk[\"cover_page\"].content,\n",
    "            chunk_ids=[self.hash_lookup_chunk[\"cover_page\"].chunk_id],\n",
    "        )\n",
    "\n",
    "        # link cover page to document\n",
    "        self.link_nodes(self.file_name, \"cover_page\", contains_relation)\n",
    "\n",
    "        for reducto_chunk in self.json_data[\"result\"][\"chunks\"]:\n",
    "            for element in reducto_chunk[\"blocks\"]:\n",
    "                page_number = int(element[\"bbox\"][\"page\"])\n",
    "                type = element[\"type\"]\n",
    "                content = element[\"content\"]\n",
    "\n",
    "                current_element_hash_id = hash_element(element)\n",
    "                # print(\"Current element hash id: \", current_element_hash_id)\n",
    "\n",
    "                # come up with the name for the node on WhyHow\n",
    "                node_name_content_preview = whyhow_unique_node_name(\n",
    "                    content, current_element_hash_id\n",
    "                )\n",
    "                # print(\n",
    "                #     f\"Processing new {type} element on page {page_number} with content: {node_name_content_preview}\"\n",
    "                # )\n",
    "\n",
    "                # Assign footers\n",
    "                if type == \"Footer\":\n",
    "                    self.footers_by_page_number[page_number].append(element)\n",
    "\n",
    "                if type == \"Page Number\":\n",
    "                    continue\n",
    "\n",
    "                if page_number == 1:\n",
    "                    # create new id for element\n",
    "                    # self.create_new_element(session, element_id, page_number, content, type)\n",
    "                    # self.link_element_to_page(session, cover_page_id, element_id, type)\n",
    "                    try:\n",
    "                        self.create_new_node(\n",
    "                            current_element_hash_id,\n",
    "                            type,\n",
    "                            node_name_content_preview,\n",
    "                            self.hash_lookup_chunk[current_element_hash_id].content,\n",
    "                            chunk_ids=[\n",
    "                                self.hash_lookup_chunk[current_element_hash_id].chunk_id\n",
    "                            ],\n",
    "                        )\n",
    "                    except:\n",
    "                        print(\n",
    "                            f\"Error creating new node: current_element_hash_id: {current_element_hash_id}, type: {type}, node name preview: {node_name_content_preview}\"\n",
    "                        )\n",
    "\n",
    "                    try:\n",
    "                        self.link_nodes(\n",
    "                            \"cover_page\", current_element_hash_id, contains_relation\n",
    "                        )\n",
    "                    except:\n",
    "                        print(\n",
    "                            f\"Error linking nodes: current_element_hash_id: cover_page,{current_element_hash_id}\"\n",
    "                        )\n",
    "\n",
    "                else:\n",
    "                    try:\n",
    "                        self.create_new_node(\n",
    "                            current_element_hash_id,\n",
    "                            type,\n",
    "                            node_name_content_preview,\n",
    "                            self.hash_lookup_chunk[current_element_hash_id].content,\n",
    "                            chunk_ids=[\n",
    "                                self.hash_lookup_chunk[current_element_hash_id].chunk_id\n",
    "                            ],\n",
    "                        )\n",
    "\n",
    "                    except:\n",
    "                        print(\n",
    "                            f\"Error creating new node: current_element_hash_id: {current_element_hash_id}, type: {type}, node name preview: {node_name_content_preview}\"\n",
    "                        )\n",
    "                    if last_section_id is None:\n",
    "                        # link element to doc if there is no section heading above it\n",
    "                        try:\n",
    "                            self.link_nodes(\n",
    "                                self.file_name,\n",
    "                                current_element_hash_id,\n",
    "                                contains_relation,\n",
    "                            )\n",
    "                        except:\n",
    "                            print(\n",
    "                                f\"Error linking nodes: file_name: {self.file_name}, current_element_hash_id: {current_element_hash_id}\"\n",
    "                            )\n",
    "\n",
    "                    if type == \"Section Header\":\n",
    "                        try:\n",
    "                            self.link_nodes(\n",
    "                                last_section_id,\n",
    "                                current_element_hash_id,\n",
    "                                follows_relation,\n",
    "                            )\n",
    "                        except:\n",
    "                            print(\n",
    "                                f\"Error linking nodes:  last_section_id: {last_section_id}, current_element_hash_id: {current_element_hash_id}\"\n",
    "                            )\n",
    "\n",
    "                        # update last section id since we have a new section\n",
    "                        last_section_id = current_element_hash_id\n",
    "                        last_list_item_id = (\n",
    "                            None  # reset last list item id since start of new section\n",
    "                        )\n",
    "                        try:\n",
    "                            self.link_nodes(\n",
    "                                self.file_name,\n",
    "                                current_element_hash_id,\n",
    "                                contains_relation,\n",
    "                            )\n",
    "                        except:\n",
    "                            print(\n",
    "                                f\"Error linking nodes: file_name: {self.file_name}, current_element_hash_id: {current_element_hash_id}\"\n",
    "                            )\n",
    "\n",
    "                    elif type == \"List Item\":\n",
    "                        if (\n",
    "                            last_list_item_id is not None\n",
    "                        ):  # this isn't the first list item in the section\n",
    "                            last_tabs_count = self.elements[last_list_item_id][\n",
    "                                \"content\"\n",
    "                            ].count(\"\\t\")\n",
    "                            current_tabs_count = content.count(\"\\t\")\n",
    "\n",
    "                            if (\n",
    "                                current_tabs_count > last_tabs_count\n",
    "                            ):  # we have reached an indented list\n",
    "                                parent_id = last_list_item_id\n",
    "                                # self.link_listitem_to_listitem(\n",
    "                                # session, parent_id, element_hash_id)\n",
    "                                try:\n",
    "                                    self.link_nodes(\n",
    "                                        parent_id,\n",
    "                                        current_element_hash_id,\n",
    "                                        contains_relation,\n",
    "                                    )\n",
    "                                except:\n",
    "                                    print(\n",
    "                                        f\"Error linking nodes: parent_id: {parent_id}, current_element_hash_id: {current_element_hash_id}\"\n",
    "                                    )\n",
    "                            else:\n",
    "                                last_list_item_id = current_element_hash_id\n",
    "                                try:\n",
    "                                    self.link_nodes(\n",
    "                                        last_section_id,\n",
    "                                        current_element_hash_id,\n",
    "                                        is_parent_relation,\n",
    "                                    )\n",
    "                                except:\n",
    "                                    print(\n",
    "                                        f\"Error linking nodes: last_section_id: {last_section_id}, current_element_hash_id: {current_element_hash_id}\"\n",
    "                                    )\n",
    "                            # new list item\n",
    "                        else:  # this is the first list item in the section\n",
    "                            last_list_item_id = current_element_hash_id\n",
    "                            try:\n",
    "                                self.link_nodes(\n",
    "                                    last_section_id,\n",
    "                                    current_element_hash_id,\n",
    "                                    is_parent_relation,\n",
    "                                )\n",
    "                            except:\n",
    "                                print(\n",
    "                                    f\"Error linking nodes: last_section_id: {last_section_id}, current_element_hash_id: {current_element_hash_id}\"\n",
    "                                )\n",
    "                    else:\n",
    "                        try:\n",
    "                            self.link_nodes(\n",
    "                                last_section_id,\n",
    "                                current_element_hash_id,\n",
    "                                is_parent_relation,\n",
    "                            )\n",
    "                        except:\n",
    "                            print(\n",
    "                                f\"Error linking nodes: last_section_id: {last_section_id}, current_element_hash_id: {current_element_hash_id}\"\n",
    "                            )\n",
    "        # link two elements to each other.\n",
    "        # a tuple is a pair of elements hash ids\n",
    "        for tuple in links_tuples:\n",
    "            if tuple[0] is not None and tuple[1] is not None:\n",
    "                try:\n",
    "                    self.link_nodes(tuple[0], tuple[1], links_to_relation)\n",
    "                except:\n",
    "                    print(\n",
    "                        f\"Error linking nodes: tuple[0]: {tuple[0]}, tuple[1]: {tuple[1]}\"\n",
    "                    )\n",
    "\n",
    "        print(\"Building graph...\")\n",
    "        try:\n",
    "            self.graph = self.client.graphs.create_graph_from_triples(\n",
    "                name=f\"{self.file_name} Lexical Graph\",\n",
    "                workspace_id=self.workspace.workspace_id,\n",
    "                triples=self.triples,\n",
    "                schema_id=LEXICAL_SCHEMA_ID,  # self.schema.schema_id\n",
    "            )\n",
    "            print(\"Done with building graph!\")\n",
    "            print(\"Graph: \", self.graph)\n",
    "        except Exception as e:\n",
    "            print(f\"Error building graph. {e}\")\n",
    "        print(\"Number of elements: \", len(self.elements))\n",
    "        return self.graph\n",
    "\n",
    "\n",
    "central_bank_doc = WhyHowLexicalGraphBuilder(\n",
    "    file_name=FILENAME,\n",
    "    # clears up workspace on init. Alternatively can call clear_workspace() method\n",
    "    # clear_workspace_on_init=False,\n",
    "    # if true, rebuilds a fresh csv and uploads to whyhow when graph builds\n",
    "    # if false, downloads CSVs from whyhow and uses them to build graph\n",
    "    upload_new_csvs_on_build=True,\n",
    "    client=WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL),\n",
    ")\n",
    "LEXICAL_GRAPH_ID = central_bank_doc.build_graph_whyhow()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download Data to Notebook\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Grab all chunks\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing document bnm_compliance_table.csv\n",
      "Processing document bnm_compliance_text.csv\n",
      "Processing document bnm_compliance_section_header.csv\n",
      "Processing document bnm_compliance_document.csv\n",
      "Processing document bnm_compliance_header.csv\n",
      "Processing document bnm_compliance_list_item.csv\n",
      "Processing document bnm_compliance_footer.csv\n",
      "Processing document bnm_compliance_figure.csv\n",
      "Processing document bnm_compliance_title.csv\n",
      "Processing document bnm_compliance_page.csv\n",
      "Processing document bnm_compliance_definitions.csv\n",
      "Skipping document bnm_compliance_definitions.csv, it is not an elements csv file. Probably a definitions csv file.\n"
     ]
    }
   ],
   "source": [
    "from whyhow.schemas import Chunk, Node\n",
    "\n",
    "# Grab all chunks\n",
    "docs = central_bank_doc.client.documents.get_all(\n",
    "    workspace_id=WORKSPACE_ID\n",
    ")\n",
    "client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "chunks_dict: Dict[str, Chunk] = {}\n",
    "\n",
    "\n",
    "for doc in docs:\n",
    "    print(f\"Processing document {doc.metadata.filename}\")\n",
    "    current_doc_chunks = list(\n",
    "        client.chunks.get_all(\n",
    "            document_filename=doc.metadata.filename, include_embeddings=True\n",
    "        )\n",
    "    )\n",
    "\n",
    "    skip_doc = False  # Flag to determine if the document should be skipped\n",
    "\n",
    "    for chunk in current_doc_chunks:\n",
    "        if \"Type\" not in chunk.content:\n",
    "            skip_doc = True  # Set the flag to skip the document\n",
    "            print(\n",
    "                f\"Skipping document {doc.metadata.filename}, it is not an elements csv file. Probably a definitions csv file.\"\n",
    "            )\n",
    "            break  # Exit the chunk loop and skip the document\n",
    "\n",
    "    if skip_doc:\n",
    "        continue  # Skip processing the rest of this document\n",
    "    else:\n",
    "        for chunk in current_doc_chunks:\n",
    "            chunks_dict[chunk.chunk_id] = (\n",
    "                chunk  # Add chunk to the dictionary with chunk ID as the key\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of chunks: 195\n"
     ]
    }
   ],
   "source": [
    "print(f\"number of chunks: {len(chunks_dict)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get all WhyHow Nodes \n",
    "(Use Backend without SDK since not supported yet)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Lexical Graph ID: 66d82333eda622fc957cc89e\n",
      "Number of nodes fetched: 195\n"
     ]
    }
   ],
   "source": [
    "import requests\n",
    "from typing import List, Dict, Any\n",
    "from datetime import datetime\n",
    "from whyhow.schemas import Node\n",
    "import time\n",
    "\n",
    "# wait 10 seconds for the graph to be processed\n",
    "time.sleep(10)\n",
    "\n",
    "print(f\"Lexical Graph ID: {LEXICAL_GRAPH_ID.graph_id}\")\n",
    "def fetch_and_convert_nodes() -> Dict[str, Node]:\n",
    "    \"\"\"\n",
    "    Fetch nodes from API and convert to Node objects.\n",
    "    \"\"\"\n",
    "    api_url = (\n",
    "        f\"{WHYHOW_API_URL}/graphs/{LEXICAL_GRAPH_ID.graph_id}/nodes?limit=-1\"\n",
    "    )\n",
    "\n",
    "    headers = {\"x-api-key\": WHYHOW_API_KEY}\n",
    "    response = requests.get(api_url, headers=headers)\n",
    "    response.raise_for_status()\n",
    "\n",
    "    node_dict = {}\n",
    "    response_json = response.json()\n",
    "\n",
    "    # Access the 'nodes' key within the response\n",
    "    nodes_list = response_json.get(\"nodes\", [])\n",
    "\n",
    "    for item in nodes_list:\n",
    "        node_id = str(item[\"_id\"])\n",
    "        node = Node(\n",
    "            node_id=node_id,\n",
    "            label=item[\"label\"],\n",
    "            name=item[\"name\"],\n",
    "            chunk_ids=item.get(\"chunks\", []),\n",
    "            properties=item.get(\"properties\", {}),\n",
    "            created_at=(\n",
    "                datetime.fromisoformat(\n",
    "                    item[\"created_at\"][\"$date\"].replace(\"Z\", \"+00:00\")\n",
    "                )\n",
    "                if \"created_at\" in item\n",
    "                else None\n",
    "            ),\n",
    "            updated_at=(\n",
    "                datetime.fromisoformat(\n",
    "                    item[\"updated_at\"][\"$date\"].replace(\"Z\", \"+00:00\")\n",
    "                )\n",
    "                if \"updated_at\" in item\n",
    "                else None\n",
    "            ),\n",
    "        )\n",
    "        node_dict[node_id] = node\n",
    "\n",
    "    return node_dict\n",
    "\n",
    "\n",
    "whyhow_nodes_dict = fetch_and_convert_nodes()\n",
    "print(f\"Number of nodes fetched: {len(whyhow_nodes_dict)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Index with LlamaIndex\n",
    "### Convert WhyHow Nodes to Local Nodes, indexable with LlamaIndex\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, ConfigDict\n",
    "from typing import List, Dict, Any, Optional, Literal\n",
    "from whyhow.schemas import Triple, Node\n",
    "\n",
    "from llama_index.core import SimpleDirectoryReader\n",
    "from llama_index.core import StorageContext\n",
    "from llama_index.core import Settings\n",
    "from llama_index.core.schema import TextNode\n",
    "from llama_index.llms.openai import OpenAI as LLamaOpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "\n",
    "Settings.llm = LLamaOpenAI(model=\"gpt-4o\")\n",
    "Settings.embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-small\")\n",
    "storage_context = StorageContext.from_defaults()\n",
    "\n",
    "client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "\n",
    "\n",
    "class MultiAgentSearchLocalNode(TextNode):\n",
    "    \"\"\"\n",
    "    An extended version of LlamaIndex's TextNode, with support for Element Types from Reducto\n",
    "    Also support for chunk IDs from WhyHow\n",
    "    Supports a tree-recursive print function for outputting structure as YAML for LLM calss\n",
    "    \"\"\"\n",
    "\n",
    "    # A TextNode class was used because it supports indexing with LlamaIndex.\n",
    "    class Config:\n",
    "        arbitrary_types_allowed = True\n",
    "        json_encoders = {\n",
    "            ElementType: lambda v: str(\n",
    "                v\n",
    "            )  # Ensure ElementType is serialized as a string\n",
    "        }\n",
    "\n",
    "    type: ElementType\n",
    "    chunk_id: Optional[str] = None\n",
    "    embedding: Optional[List[float]] = None\n",
    "    metadata: Optional[Dict[str, Any]] = None\n",
    "    children: Dict[str, \"MultiAgentSearchLocalNode\"] = {}\n",
    "    context: Optional[Dict[str, Any]] = {}\n",
    "\n",
    "    def __init__(self, *, whyhow_node: Node, **data: Any):\n",
    "        super().__init__(**data)\n",
    "        try:\n",
    "            chunk = chunks_dict[whyhow_node.chunk_ids[0]]\n",
    "        except:\n",
    "            print(\n",
    "                f\"Error getting chunk with chunk_id: {whyhow_node.chunk_ids[0]} for node {whyhow_node.name}\"\n",
    "            )\n",
    "            raise ValueError(\"Error getting chunk with chunk_id\")\n",
    "\n",
    "        self.id_ = whyhow_node.name\n",
    "        self.chunk_id = chunk.chunk_id\n",
    "        self.embedding = chunk.embedding\n",
    "        self.metadata = whyhow_node.properties\n",
    "        self.text = chunk.content.get(\n",
    "            \"Content\", chunk.content.get(\"Type\", \"(empty content)\")\n",
    "        )\n",
    "\n",
    "    # Implement hash and eq for use in dictionaries and sets\n",
    "    def __hash__(self) -> int:\n",
    "        return hash(\n",
    "            (\n",
    "                self.node_id,\n",
    "                self.chunk_id,\n",
    "                tuple(self.embedding) if self.embedding else None,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # Implement hash and eq for use in dictionaries and sets\n",
    "    def __eq__(self, other: object) -> bool:\n",
    "        if isinstance(other, MultiAgentSearchLocalNode):\n",
    "            return self.node_id == other.node_id\n",
    "        return False\n",
    "\n",
    "    # print node for prompt. use yaml to save tokens + better performance than json\n",
    "    # recursive tree printing function\n",
    "    # GPT will return with JSON\n",
    "    def print_node_prompt(self, indent_level: int = 0, show_type=False) -> str:\n",
    "        cleaned_text = \" \".join(self.text.split())\n",
    "        indent_unit = \"  \"\n",
    "        indent = indent_unit * indent_level\n",
    "\n",
    "        prompt = f\"\"\"{indent}- node_id: `{self.node_id}`\"\"\"\n",
    "        \n",
    "        if show_type:\n",
    "            prompt += f\"{indent}{indent_unit*2}type: {self.type}\\n\"\n",
    "\n",
    "        prompt += f\"{indent}{indent_unit*2}content: {cleaned_text}\"\n",
    "        \n",
    "        if self.context.items():\n",
    "            context = \"\"\n",
    "            for key, value in self.context.items():\n",
    "                if value not in [None, \"\", \" \"]:\n",
    "                    context += f\"{indent}{indent_unit*3}- {key}: {value}\\n\"\n",
    "            prompt += f\"{indent}{indent_unit*2}context:\\n\"\n",
    "            prompt += f\"{context}\\n\"\n",
    "        children = list(self.children.values())\n",
    "        if len(children) > 0:\n",
    "            prompt += f\"{indent}{indent_unit*2}children: \\n\"\n",
    "            for child in children:\n",
    "                prompt += f\"{child.print_node_prompt(indent_level = indent_level + 5)}\"\n",
    "        return prompt\n",
    "\n",
    "\n",
    "def whyhow_node_to_local_node(whyhow_node: Node) -> MultiAgentSearchLocalNode:\n",
    "    return MultiAgentSearchLocalNode(\n",
    "        whyhow_node=whyhow_node, type=ElementType(whyhow_node.properties[\"Type\"])\n",
    "    )\n",
    "\n",
    "\n",
    "def unique_whyhow_nodes_from_triples(\n",
    "    triples: List[Triple], part_to_return: Literal[\"Heads\", \"Tails\", \"Both\"] = \"Both\"\n",
    ") -> List[Node]:\n",
    "    unique_nodes = {}\n",
    "    for triple in triples:\n",
    "        if part_to_return == \"Heads\":\n",
    "            unique_nodes[triple.head.name] = triple.head\n",
    "        elif part_to_return == \"Tails\":\n",
    "            unique_nodes[triple.tail.name] = triple.tail\n",
    "        else:\n",
    "            unique_nodes[triple.head.name] = triple.head\n",
    "            unique_nodes[triple.tail.name] = triple.tail\n",
    "\n",
    "    return list(unique_nodes.values())\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Nodes for multi-agent search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "local_nodes_map: Dict[str, MultiAgentSearchLocalNode] = {}\n",
    "for whyhow_node in whyhow_nodes_dict.values():\n",
    "    # print(f\"Processing node: {whyhow_node}\")\n",
    "    local_nodes_map[whyhow_node.name] = MultiAgentSearchLocalNode(\n",
    "        whyhow_node=whyhow_node, type=ElementType(whyhow_node.properties[\"Type\"])\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Making sure all nodes have the same embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Minimum embedding length for nodes: 1536\n",
      "Minimum embedding length for chunks: 1536\n",
      "Maximum embedding length for nodes: 1536\n",
      "Maximum embedding length for chunks: 1536\n"
     ]
    }
   ],
   "source": [
    "min_node_embedding_length = min(\n",
    "    len(node.embedding) for node in local_nodes_map.values()\n",
    ")\n",
    "min_chunk_embedding_length = min(len(chunk.embedding) for chunk in chunks_dict.values())\n",
    "\n",
    "print(\"Minimum embedding length for nodes:\", min_node_embedding_length)\n",
    "print(\"Minimum embedding length for chunks:\", min_chunk_embedding_length)\n",
    "\n",
    "max_node_embedding_length = max(\n",
    "    len(node.embedding) for node in local_nodes_map.values()\n",
    ")\n",
    "max_chunk_embedding_length = max(len(chunk.embedding) for chunk in chunks_dict.values())\n",
    "\n",
    "print(\"Maximum embedding length for nodes:\", max_node_embedding_length)\n",
    "print(\"Maximum embedding length for chunks:\", max_chunk_embedding_length)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hybrid Retriever (BM25 + Vector Search) with LlamaIndex\n",
    "- A combination of BM25 and Vector Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import QueryBundle\n",
    "\n",
    "# import NodeWithScore\n",
    "from llama_index.core.schema import NodeWithScore\n",
    "\n",
    "from llama_index.core import SimpleKeywordTableIndex\n",
    "\n",
    "# Retrievers\n",
    "from llama_index.core.retrievers import (\n",
    "    BaseRetriever,\n",
    "    VectorIndexRetriever,\n",
    "    KeywordTableSimpleRetriever,\n",
    ")\n",
    "from llama_index.retrievers.bm25 import BM25Retriever\n",
    "from typing import List\n",
    "import Stemmer\n",
    "\n",
    "# Create a vector store using local nodes\n",
    "vector_index = VectorStoreIndex(\n",
    "    nodes=list(local_nodes_map.values()), storage_context=storage_context\n",
    ")\n",
    "\n",
    "keyword_index = SimpleKeywordTableIndex(\n",
    "    nodes=list(local_nodes_map.values()), storage_context=storage_context\n",
    ")\n",
    "\n",
    "\n",
    "class BM25Keyword(BaseRetriever):\n",
    "    \"\"\"Custom retriever that performs both keyword search and BM25 search.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        mode: str = \"OR\",\n",
    "        keyword_top_k=5,\n",
    "        bm25_top_k=3,\n",
    "    ) -> None:\n",
    "        \"\"\"Init params.\"\"\"\n",
    "\n",
    "        self._keyword_retriever = KeywordTableSimpleRetriever(\n",
    "            index=keyword_index,\n",
    "            max_keywords_per_query=keyword_top_k,\n",
    "        )\n",
    "        self._bm25_retriever = BM25Retriever.from_defaults(\n",
    "            nodes=list(local_nodes_map.values()),\n",
    "            similarity_top_k=bm25_top_k,\n",
    "            # Optional: We can pass in the stemmer and set the language for stopwords\n",
    "            # This is important for removing stopwords and stemming the query + text\n",
    "            # The default is english for both\n",
    "            stemmer=Stemmer.Stemmer(\"english\"),\n",
    "            language=\"english\",\n",
    "        )\n",
    "\n",
    "        if mode not in (\"AND\", \"OR\"):\n",
    "            raise ValueError(\"Invalid mode.\")\n",
    "        self._mode = mode\n",
    "        super().__init__()\n",
    "\n",
    "    def _retrieve(self, query_bundle: QueryBundle) -> List[MultiAgentSearchLocalNode]:\n",
    "        \"\"\"Retrieve nodes given query.\"\"\"\n",
    "\n",
    "        keyword_nodes = self._keyword_retriever.retrieve(query_bundle)\n",
    "        bm25_nodes = self._bm25_retriever.retrieve(query_bundle)\n",
    "\n",
    "        keyword_nodes_ids = {n.node.node_id for n in keyword_nodes}\n",
    "        bm25_nodes_ids = {n.node.node_id for n in bm25_nodes}\n",
    "\n",
    "        combined_dict = {n.node.node_id: n for n in keyword_nodes}\n",
    "        combined_dict.update({n.node.node_id: n for n in bm25_nodes})\n",
    "\n",
    "        if self._mode == \"AND\":\n",
    "            retrieve_ids = keyword_nodes_ids.intersection(bm25_nodes_ids)\n",
    "        else:\n",
    "            retrieve_ids = keyword_nodes_ids.union(bm25_nodes_ids)\n",
    "\n",
    "        retrieve_nodes = [combined_dict[rid] for rid in retrieve_ids]\n",
    "\n",
    "        # print(f\"Retrieved {len(retrieve_nodes)} nodes.\")\n",
    "\n",
    "        return retrieve_nodes\n",
    "\n",
    "\n",
    "class VectorBM25(BaseRetriever):\n",
    "    \"\"\"Custom retriever that performs both semantic search and BM25 search.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        mode: str = \"OR\",\n",
    "        vector_top_k=2,\n",
    "        bm25_top_k=6,\n",
    "    ) -> None:\n",
    "        \"\"\"Init params.\"\"\"\n",
    "\n",
    "        self._vector_retriever = VectorIndexRetriever(\n",
    "            index=vector_index, similarity_top_k=vector_top_k, verbose=True\n",
    "        )\n",
    "        self._bm25_retriever = BM25Retriever.from_defaults(\n",
    "            nodes=list(local_nodes_map.values()),\n",
    "            similarity_top_k=bm25_top_k,\n",
    "            # Optional: We can pass in the stemmer and set the language for stopwords\n",
    "            # This is important for removing stopwords and stemming the query + text\n",
    "            # The default is english for both\n",
    "            stemmer=Stemmer.Stemmer(\"english\"),\n",
    "            language=\"english\",\n",
    "        )\n",
    "        self._vector_top_k = vector_top_k\n",
    "        self._bm25_top_k = bm25_top_k\n",
    "        if mode not in (\"AND\", \"OR\"):\n",
    "            raise ValueError(\"Invalid mode.\")\n",
    "        self._mode = mode\n",
    "        super().__init__()\n",
    "\n",
    "    def _retrieve(self, query_bundle: QueryBundle) -> List[MultiAgentSearchLocalNode]:\n",
    "        \"\"\"Retrieve nodes given query.\"\"\"\n",
    "\n",
    "        vector_nodes = self._vector_retriever.retrieve(query_bundle)\n",
    "        bm25_nodes = self._bm25_retriever.retrieve(query_bundle)\n",
    "\n",
    "        vector_ids = {n.node.node_id for n in vector_nodes}\n",
    "        keyword_ids = {n.node.node_id for n in bm25_nodes}\n",
    "\n",
    "        combined_dict = {n.node.node_id: n for n in vector_nodes}\n",
    "        combined_dict.update({n.node.node_id: n for n in bm25_nodes})\n",
    "\n",
    "        if self._mode == \"AND\":\n",
    "            retrieve_ids = vector_ids.intersection(keyword_ids)\n",
    "        else:\n",
    "            retrieve_ids = vector_ids.union(keyword_ids)\n",
    "\n",
    "        retrieve_nodes = [combined_dict[rid] for rid in retrieve_ids]\n",
    "\n",
    "        # print(f\"Retrieved {len(retrieve_nodes)} nodes.\")\n",
    "\n",
    "        return retrieve_nodes\n",
    "\n",
    "\n",
    "from llama_index.core import get_response_synthesizer\n",
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "\n",
    "# Note: \"AND\" mode has qualitatively much better retriever performance,\n",
    "# However, it only works best with embeddings of dim 3072. If you use \"AND\" mode on dim=1536, you retrieve zero nodes.\n",
    "# If you want to use dim=3072, you cannot use the embeddings obtained from whyhow Nodes\n",
    "# Use text-embedding-3-large. In init function of MultiAgentSearchLocalNode,\n",
    "# Change this line:\n",
    "# self.embedding = chunk.embedding\n",
    "# To:\n",
    "# self.embedding = get_text_embedding(chunk.content.get(\"Content\", \"\")\n",
    "\n",
    "\n",
    "# def get_text_embedding(text: str) -> List[float]:\n",
    "#     response = openai_client.embeddings.create(\n",
    "#         input=text,\n",
    "#         model=\"text-embedding-3-large\",\n",
    "#     )\n",
    "#     return response.data[0].embedding\n",
    "\n",
    "# define custom retriever\n",
    "vector_bm25_retriever = VectorBM25(vector_top_k=2, bm25_top_k=6, mode=\"OR\")\n",
    "# bm25_retriever = BM25Keyword(keyword_top_k=5, bm25_top_k=3, mode=\"OR\")\n",
    "# keyword_retriever = BM25Keyword(keyword_top_k=5, bm25_top_k=3, mode=\"OR\")\n",
    "\n",
    "keyword_retriever = KeywordTableSimpleRetriever(\n",
    "    index=keyword_index,\n",
    "    max_keywords_per_query=6,\n",
    ")\n",
    "bm25_retriever = BM25Retriever.from_defaults(\n",
    "    nodes=list(local_nodes_map.values()),\n",
    "    similarity_top_k=6,\n",
    "    # Optional: We can pass in the stemmer and set the language for stopwords\n",
    "    # This is important for removing stopwords and stemming the query + text\n",
    "    # The default is english for both\n",
    "    stemmer=Stemmer.Stemmer(\"english\"),\n",
    "    language=\"english\",\n",
    ")\n",
    "\n",
    "keywords_bm25_retriever = BM25Keyword(keyword_top_k=5, bm25_top_k=3, mode=\"OR\")\n",
    "\n",
    "\n",
    "response_synthesizer = get_response_synthesizer()\n",
    "\n",
    "\n",
    "def retrieve_with_vector_bm25(query, verbose = False):\n",
    "    retrieved_nodes: NodeWithScore = vector_bm25_retriever.retrieve(\n",
    "        QueryBundle(query_str=query)\n",
    "    )\n",
    "    retrieved_nodes_local_form: MultiAgentSearchLocalNode = [\n",
    "        local_nodes_map[node.node.node_id] for node in retrieved_nodes\n",
    "    ]\n",
    "    if verbose:\n",
    "        for node in retrieved_nodes:\n",
    "            print(node)\n",
    "    return retrieved_nodes_local_form\n",
    "\n",
    "\n",
    "def retrieve_with_bm25(query, verbose = False):\n",
    "    retrieved_nodes: NodeWithScore = bm25_retriever.retrieve(\n",
    "        QueryBundle(query_str=query)\n",
    "    )\n",
    "    retrieved_nodes_local_form: MultiAgentSearchLocalNode = [\n",
    "        local_nodes_map[node.node.node_id] for node in retrieved_nodes\n",
    "    ]\n",
    "    if verbose:\n",
    "        for node in retrieved_nodes:\n",
    "            print(node)\n",
    "    return retrieved_nodes_local_form\n",
    "\n",
    "\n",
    "def retrieve_with_keywords(query, verbose = False):\n",
    "    retrieved_nodes: NodeWithScore = keyword_retriever.retrieve(\n",
    "        QueryBundle(query_str=query)\n",
    "    )\n",
    "    retrieved_nodes_local_form: MultiAgentSearchLocalNode = [\n",
    "        local_nodes_map[node.node.node_id] for node in retrieved_nodes\n",
    "    ]\n",
    "    if verbose:\n",
    "        for node in retrieved_nodes:\n",
    "            print(node)\n",
    "    return retrieved_nodes_local_form\n",
    "\n",
    "\n",
    "def retrieve_with_keywords_bm25(query, verbose = False):\n",
    "    retrieved_nodes: NodeWithScore = keywords_bm25_retriever.retrieve(\n",
    "        QueryBundle(query_str=query)\n",
    "    )\n",
    "    retrieved_nodes_local_form: MultiAgentSearchLocalNode = [\n",
    "        local_nodes_map[node.node.node_id] for node in retrieved_nodes\n",
    "    ]\n",
    "    if verbose:\n",
    "        for node in retrieved_nodes:\n",
    "            print(node)\n",
    "    return retrieved_nodes_local_form\n",
    "\n",
    "\n",
    "# Run this code to test the retrieval of documents\n",
    "# vector_bm25_retriever(\"Return paragraph 7.3\")\n",
    "retrieve_with_keywords(\n",
    "    \"Return paragraph 7.3 and 7.4 to answer the query: How can the Board and the CCO manage control functions?\", verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hybrid Retriever (BM25 + Vector Search) test with LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this code to test the retriever + LLM answering\n",
    "# print(response)\n",
    "# retrieved_nodes = custom_retriever.retrieve(\n",
    "#     QueryBundle(query_str=\"How can the Board and the CCO manage control functions?\")\n",
    "# )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recursive tree functions on subgraphs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from typing import Dict\n",
    "from whyhow.schemas import Node\n",
    "from typing import List, Any, Optional, Dict\n",
    "\n",
    "\n",
    "def prune_node_from_nodes(\n",
    "    nodes: List[MultiAgentSearchLocalNode],\n",
    "    node_id: str,\n",
    ") -> None:\n",
    "    \"\"\"Prune a node tree by node_id.\"\"\"\n",
    "    \n",
    "    def prune(node: MultiAgentSearchLocalNode) -> bool:\n",
    "        if node.node_id == node_id:\n",
    "            return True\n",
    "        to_delete = []\n",
    "        for child_id, child_node in node.children.items():\n",
    "            if prune(child_node):\n",
    "                to_delete.append(child_id)\n",
    "        for child_id in to_delete:\n",
    "            del node.children[child_id]\n",
    "        return False\n",
    "\n",
    "    nodes[:] = [root_node for root_node in nodes if not prune(root_node)]\n",
    "\n",
    "\n",
    "def recursive_update(\n",
    "    nodes: List[MultiAgentSearchLocalNode],\n",
    "    node_id: str,\n",
    "    property: str,\n",
    "    property_value: Any,\n",
    ") -> None:\n",
    "    \"\"\"Recursively search a node tree by node_id and update the specified property with the given value.\"\"\"\n",
    "\n",
    "    def update_node(node: MultiAgentSearchLocalNode) -> bool:\n",
    "        if node.node_id == node_id:\n",
    "            setattr(node, property, property_value)\n",
    "            return True\n",
    "        for child in node.children.values():\n",
    "            if update_node(child):\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "    for root_node in nodes:\n",
    "        if update_node(root_node):\n",
    "            break\n",
    "\n",
    "\n",
    "def recursive_add_context(\n",
    "    nodes: List[MultiAgentSearchLocalNode],\n",
    "    node_id: str,\n",
    "    context_key: str,\n",
    "    context_value: str,\n",
    ") -> None:\n",
    "    \"\"\"Recursively search a node tree by node_id and add a key-value pair to the context dictionary of the matching node.\"\"\"\n",
    "\n",
    "    def add_context_to_node(node: MultiAgentSearchLocalNode) -> bool:\n",
    "        if node.node_id == node_id:\n",
    "            node.context[context_key] = context_value\n",
    "            return True\n",
    "        for child in node.children.values():\n",
    "            if add_context_to_node(child):\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "    for root_node in nodes:\n",
    "        if add_context_to_node(root_node):\n",
    "            break\n",
    "\n",
    "    return nodes\n",
    "\n",
    "def recursive_append_context_items(\n",
    "    nodes: List[MultiAgentSearchLocalNode],\n",
    "    node_id: str,\n",
    "    context_key: str,\n",
    "    new_items: List[str],\n",
    ") -> None:\n",
    "    \"\"\"Recursively search a node tree by node_id and add a key-value pair to the context dictionary of the matching node.\"\"\"\n",
    "\n",
    "    def add_context_link_to_node(node: MultiAgentSearchLocalNode) -> bool:\n",
    "        if node.node_id == node_id:\n",
    "            node_links = node.context.get(context_key, [])\n",
    "            node_links.append(new_items)\n",
    "            node.context[context_key] = node_links\n",
    "\n",
    "            return True\n",
    "        for child in node.children.values():\n",
    "            if add_context_link_to_node(child):\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "    for root_node in nodes:\n",
    "        if add_context_link_to_node(root_node):\n",
    "            break\n",
    "\n",
    "    return nodes\n",
    "\n",
    "def recursive_check_if_exists(\n",
    "    nodes: List[MultiAgentSearchLocalNode],\n",
    "    node_id: str,\n",
    ") -> bool:\n",
    "    \"\"\" Recursively search a node tree by node_id and check if the node exists.\"\"\"\n",
    "    \n",
    "    def check_if_exists(node: MultiAgentSearchLocalNode) -> bool:\n",
    "        if node.node_id == node_id:\n",
    "            return True\n",
    "        for child in node.children.values():\n",
    "            if check_if_exists(child):\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "    for root_node in nodes:\n",
    "        if check_if_exists(root_node):\n",
    "            return True\n",
    "\n",
    "    return False\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tools for Agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filtering out nodes that don't match a query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "starting:\n",
      "- node_id: `(b) ensure that the CCO has sufficient stature to ..._ade5bcb25050d514e6d10e163c1581ec`    content: \" (b) ensure that the CCO has sufficient stature to allow for effective engagement with the CEO and other members of senior management;\"\n",
      "- node_id: `(c) engage with the CCO on a regular basis2 to pro..._75d9df850bf8422349822accdcb9338e`    content: \" (c) engage with the CCO on a regular basis2 to provide the opportunity for the CCO to discuss issues faced by the compliance function;\"\n",
      "- node_id: `(d) provide the CCO with direct and unimpeded acce..._753739a7f2c2cf4c206b6e102cf20aca`    content: \" (d) provide the CCO with direct and unimpeded access to the board;\"\n",
      "- node_id: `(e) ensure that the CCO is supported with sufficie..._b5cb66fa679baafeacc244e3dbe13938`    content: \" (e) ensure that the CCO is supported with sufficient resources, including competent officers, to perform his duties effectively; and\"\n",
      "- node_id: `(f) where the CCO also carries out responsibilitie..._684a6af3d1c0785ec56eeb2cbe6c1644`    content: \" (f) where the CCO also carries out responsibilities in respect of other control functions3, be satisfied that a sound overall control environment will not be compromised by the combination of responsibilities performed by the CCO.\"\n",
      "- node_id: `1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8`    content: \"1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011. 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time. 3 Refer to paragraphs 7.3 and 7.4.\"\n",
      "- node_id: `Issued on: 10 May 2016_b350089add36cc840ad1fb778a27cd62`    content: \"Issued on: 10 May 2016\"\n",
      "- node_id: `Compliance_747e07c6b6fb06898766b61480775e55`    content: \"Compliance\"\n",
      "- node_id: `Responsibilities of senior management_26342c4f3533e56f55513945dccc9f9c`    content: \"Responsibilities of senior management\"\n",
      "- node_id: `- S 6.4 Senior management is collectively responsi..._4cd12c4a3ddeb6778b65630b4e624564`    content: \"- S 6.4 Senior management is collectively responsible for the effective management of a financial institution's compliance risk. In discharging this responsibility, senior management must-\"\n",
      "==================================================\n",
      "\n",
      "- node_id: `(b) ensure that the CCO has sufficient stature to ..._ade5bcb25050d514e6d10e163c1581ec`    content: \" (b) ensure that the CCO has sufficient stature to allow for effective engagement with the CEO and other members of senior management;\"\n",
      "- node_id: `(c) engage with the CCO on a regular basis2 to pro..._75d9df850bf8422349822accdcb9338e`    content: \" (c) engage with the CCO on a regular basis2 to provide the opportunity for the CCO to discuss issues faced by the compliance function;\"\n",
      "- node_id: `(d) provide the CCO with direct and unimpeded acce..._753739a7f2c2cf4c206b6e102cf20aca`    content: \" (d) provide the CCO with direct and unimpeded access to the board;\"\n",
      "- node_id: `(e) ensure that the CCO is supported with sufficie..._b5cb66fa679baafeacc244e3dbe13938`    content: \" (e) ensure that the CCO is supported with sufficient resources, including competent officers, to perform his duties effectively; and\"\n",
      "- node_id: `(f) where the CCO also carries out responsibilitie..._684a6af3d1c0785ec56eeb2cbe6c1644`    content: \" (f) where the CCO also carries out responsibilities in respect of other control functions3, be satisfied that a sound overall control environment will not be compromised by the combination of responsibilities performed by the CCO.\"\n",
      "- node_id: `- S 6.4 Senior management is collectively responsi..._4cd12c4a3ddeb6778b65630b4e624564`    content: \"- S 6.4 Senior management is collectively responsible for the effective management of a financial institution's compliance risk. In discharging this responsibility, senior management must-\"\n"
     ]
    }
   ],
   "source": [
    "from typing import Dict\n",
    "from whyhow.schemas import Node\n",
    "from typing import List, Any, Optional, Dict\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Get OpenAI API key from environment variables\n",
    "openai_api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "if openai_api_key is None:\n",
    "    raise ValueError(\"OPENAI_API_KEY not found in environment variables.\")\n",
    "\n",
    "# Initialize OpenAI client\n",
    "openai_client = OpenAI()\n",
    "openai_client.api_key = openai_api_key\n",
    "\n",
    "\n",
    "class PruneNodeIdResponse(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe a node via node_id being returned by the LLM.\n",
    "    \"\"\"\n",
    "\n",
    "    node_id: str\n",
    "    reason: str\n",
    "\n",
    "\n",
    "class PruneNodeIdsResponse(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe nodes via node_ids being returned by the LLM.\n",
    "    \"\"\"\n",
    "\n",
    "    content: List[PruneNodeIdResponse]\n",
    "\n",
    "\n",
    "def prune_nodes(\n",
    "    query: str, current_nodes: List[MultiAgentSearchLocalNode]\n",
    ") -> List[MultiAgentSearchLocalNode]:\n",
    "    \"\"\"\n",
    "    LLM call to remove any nodes that don't contain relevant information to the query.\n",
    "    \"\"\"\n",
    "\n",
    "    starter_prompt = f\"\"\"\n",
    "    You are an intelligent agent responsible for post-retrieval filtration of retrieved nodes from a document. The nodes are to answer the query: \n",
    "    ```{query}```\n",
    "    \n",
    "    Below is a list of nodes that were automatically retrieved. Your task is to identify the nodes that do not contain relevant information to the query, and a short reason as to why they should be pruned.\n",
    "    \n",
    "    Be warned that deleting nodes will delete any sub-nodes listed within them.\n",
    "\n",
    "    Nodes are identified by node_id. Return a list of (node_id, reason) where reason is a short explanation as to why the node should be pruned. node_ids must be quoted in backticks.\n",
    "    \n",
    "    \"\"\"\n",
    "\n",
    "    all_flattened_nodes : Dict[str, MultiAgentSearchLocalNode] = {}\n",
    "\n",
    "    # node_id, (footer_text, further_explore_footer)\n",
    "\n",
    "    def recurse_collect_nodes(current_nodes: List[MultiAgentSearchLocalNode]):\n",
    "        \"\"\"\n",
    "        Recursively iterate via DFS on nodes, to index them by page number.\n",
    "        \"\"\"\n",
    "        for node in current_nodes:\n",
    "            all_flattened_nodes[node.id_] = node\n",
    "\n",
    "            if len(node.children) > 0:\n",
    "                recurse_collect_nodes(list(node.children.values()))\n",
    "\n",
    "    recurse_collect_nodes(current_nodes)\n",
    "\n",
    "    printout = \"\"\n",
    "    for node in all_flattened_nodes.values():\n",
    "        printout += node.print_node_prompt()\n",
    "\n",
    "    completion = openai_client.beta.chat.completions.parse(\n",
    "        model=\"gpt-4o-2024-08-06\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": starter_prompt},\n",
    "            {\"role\": \"user\", \"content\": printout},\n",
    "        ],\n",
    "        response_format=PruneNodeIdsResponse,\n",
    "    )\n",
    "\n",
    "    node_id_response = completion.choices[0].message.parsed.content\n",
    "\n",
    "    prune_dictionary_results: Dict[str, str] = {}\n",
    "\n",
    "    if isinstance(node_id_response, list):\n",
    "        for item in node_id_response:\n",
    "            if isinstance(item, PruneNodeIdResponse):\n",
    "                if item.node_id[0] == \"`\":\n",
    "                    node_id = item.node_id[1:-1]\n",
    "                else:\n",
    "                    node_id = item.node_id\n",
    "                # print(f\"Pruning node with node_id: {node_id} due to reason: {item.reason}\")\n",
    "                prune_dictionary_results[node_id] = item.reason\n",
    "                prune_node_from_nodes(current_nodes, node_id)\n",
    "\n",
    "\n",
    "\n",
    "    return current_nodes, prune_dictionary_results\n",
    "\n",
    "\n",
    "\n",
    "local_nodes = list(local_nodes_map.values())\n",
    "starting_nodes = local_nodes[100:110]\n",
    "\n",
    "\n",
    "print(\"starting:\")\n",
    "for node in starting_nodes:\n",
    "    print(node.print_node_prompt())\n",
    "    \n",
    "print(\"=\"*50)\n",
    "print()\n",
    "\n",
    "nodes_test1 , _= prune_nodes(\"How can CCO manage control functions?\", starting_nodes)\n",
    "\n",
    "\n",
    "for node in nodes_test1:\n",
    "    print(node.print_node_prompt())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Collecting List Item Nodes for a clause\n",
    "\n",
    "A clause be represented by a single List Item Node, or with a combination of List Item Nodes:\n",
    "\n",
    "```\n",
    "    CLAUSE 1                   CLAUSE 2                   \n",
    "    +------------------------+   +--------------------+   \n",
    "    | - Clause Title 1       |   | - Clause Title 2   |   \n",
    "    |    - Subclause 1.1     |   +--------------------+   \n",
    "    |    - Subclause 1.2     |                            \n",
    "    +------------------------+                            \n",
    "```\n",
    "\n",
    "- There is a one to many relationship between Clause Title and Subclauses\n",
    "  - Clause Title 1 -> Subclause 1.1\n",
    "  - Clause Title 1 -> Subclause 2.1\n",
    "- Vector/keyword search might only fetch a title or the subclause.\n",
    "- We would have to traverse the nodes to fetch the whole clause.\n",
    "  - Example (Assume there are only two levels: clause and subclause):\n",
    "  ```\n",
    "    # Example\n",
    "    # Input:  [Section_Heading, Text, List_Item(Subclause 1.2), Image]\n",
    "\n",
    "    # Output: [Section_Heading, Text, List_Item(Clause 1 Title), Image]\n",
    "    #                                   └── [List_Item(Subclause 1.1), List_Item(Subclause 1.2)]\n",
    "  ```\n",
    "\n",
    "  Since we have a depth of 2 max, we will have to run the query twice to guarantee we will fetch the whole clause\n",
    "    - Step 1: fetch any titles of subclauses. this step will have fetched us all clause titles\n",
    "      - We will get Triples of X contains List Item Y\n",
    "        - If a non List Item X contains List Item Y, by deduction, we can say that List Item Y is a clause title since there is nothing above it\n",
    "        - If a List Item X contains List Item Y, by deduction, we can say that List Item Y is a subclause since there is a List Item above containing it. And by definition, List Item X is a clause title\n",
    "    - Step 2: based on all clause titles fetched in Step 1, we can fetch all subclauses.\n",
    "      - We look for Triples of List Item X contains List Item Y\n",
    "      - List Item X is clause title from Step 1, and thus List Item Y is a subclause\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROMPT:\n",
      "- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"\n",
      "- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    children: \n",
      "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"\n"
     ]
    }
   ],
   "source": [
    "from collections import defaultdict\n",
    "from typing import Set, Dict, List\n",
    "from whyhow.schemas import Triple, Node, Query\n",
    "\n",
    "def print_triple(triple: Triple):\n",
    "    print(\"----------------------------------------------------\")\n",
    "    print(f\"Head:{triple.head.name} of type {triple.head.label}\")\n",
    "    print(f\"Relation: {triple.relation.name}\")\n",
    "    print(f\"Tail: {triple.tail.name} of type {triple.tail.label}\")\n",
    "    \n",
    "    \n",
    "\n",
    "def fetch_whole_clause_lists(\n",
    "    current_nodes: List[MultiAgentSearchLocalNode], query_client=None\n",
    ") -> List[MultiAgentSearchLocalNode]:\n",
    "    \"\"\"\n",
    "    Given a list of nodes will retrieve and replace subclauses with full clauses.\n",
    "    Assumes that we have a flat sequence of nodes, not a tree of nodes.\n",
    "    \"\"\"\n",
    "    # (Assume there are only two levels: clause and subclause)\n",
    "    # Example\n",
    "    # Input:  [Section_Heading, Text, List_Item(Subclause 1.2), Image]\n",
    "\n",
    "    # Output: [Section_Heading, Text, List_Item(Clause 1 Title), Image]\n",
    "    #                                   └── [List_Item(Subclause 1.1), List_Item(Subclause 1.2)]\n",
    "    if query_client is None:\n",
    "        query_client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "\n",
    "    # this query is to fetch all clause titles\n",
    "    first_pass_query: Query = query_client.graphs.query_structured(\n",
    "        graph_id=LEXICAL_GRAPH_ID.graph_id,\n",
    "        entities=[ElementType.LIST_ITEM],\n",
    "        relations=[\"contains\", \"is_parent\"],\n",
    "        values=[\n",
    "            node.id_ for node in current_nodes if node.type == ElementType.LIST_ITEM\n",
    "        ],\n",
    "    )\n",
    "    \n",
    "    # keep track of all clause title ids and subclause ids\n",
    "    clause_title_ids : Set[str] = set()\n",
    "    subclause_ids : Set[str] = set()\n",
    "    \n",
    "    # keep track of list_item to nodes\n",
    "    id_to_node_dict : Dict[str, MultiAgentSearchLocalNode] = {}\n",
    "    \n",
    "    # keep track of which clauses map to which subclauses and vice versa\n",
    "    clause_to_subclauses_map: Dict[str, List[str]] = defaultdict(list)\n",
    "    subclause_to_clause_map: Dict[str, str] = {}\n",
    "    \n",
    "\n",
    "    for triple in first_pass_query.triples:\n",
    "        # print_triple(triple)\n",
    "        # if a list contains a list item, head is clause title, tail is subclause\n",
    "        if triple.head.label == ElementType.LIST_ITEM and triple.tail.label == ElementType.LIST_ITEM:\n",
    "            clause_title_ids.add(triple.head.name) \n",
    "            subclause_ids.add(triple.tail.name)\n",
    "            \n",
    "            clause_to_subclauses_map[triple.head.name].append(triple.tail.name)\n",
    "            subclause_to_clause_map[triple.tail.name] = triple.head.name\n",
    "            \n",
    "            # subclause node\n",
    "            tail_node = whyhow_node_to_local_node(triple.tail)\n",
    "            # if clause title node exists, add subclause to its children\n",
    "            if triple.head.name in id_to_node_dict:\n",
    "                id_to_node_dict[triple.head.name].children[triple.tail.name]=tail_node\n",
    "            else:\n",
    "                # else just create a new clause title node with no children\n",
    "                id_to_node_dict[triple.head.name] = whyhow_node_to_local_node(triple.head)\n",
    "            # add subclause node to the map\n",
    "            id_to_node_dict[triple.tail.name] = tail_node\n",
    "            \n",
    "        # if a list contains a non-list item, tail is clause title (we don't know if they might have subclauses below. that's for the 2nd pass)\n",
    "        elif triple.head.label != ElementType.LIST_ITEM and triple.tail.label == ElementType.LIST_ITEM:\n",
    "            clause_title_ids.add(triple.tail.name)\n",
    "            id_to_node_dict[triple.tail.name] = whyhow_node_to_local_node(triple.tail)\n",
    "        \n",
    "            \n",
    "\n",
    "    # we have determined all title clauses from the first pass.\n",
    "    # second pass of query is to fetch all subclauses\n",
    "    second_pass_query = query_client.graphs.query_structured(\n",
    "        graph_id=LEXICAL_GRAPH_ID.graph_id,\n",
    "        entities=[ElementType.LIST_ITEM],\n",
    "        relations=[\"contains\"],\n",
    "        values=list(clause_title_ids)\n",
    "    )\n",
    "\n",
    "    for triple in second_pass_query.triples:\n",
    "        # print_triple(triple)\n",
    "        # if a list contains a list item, head is clause title, tail is subclause\n",
    "        if triple.head.label == ElementType.LIST_ITEM and triple.tail.label == ElementType.LIST_ITEM:\n",
    "            clause_title_ids.add(triple.head.name)\n",
    "            subclause_ids.add(triple.tail.name)\n",
    "            \n",
    "            clause_to_subclauses_map[triple.head.name].append(triple.tail.name)\n",
    "            subclause_to_clause_map[triple.tail.name] = triple.head.name\n",
    "\n",
    "            # subclause node\n",
    "            tail_node = whyhow_node_to_local_node(triple.tail)\n",
    "            # if clause title node exists, add subclause to its children\n",
    "            if triple.head.name in id_to_node_dict:\n",
    "                id_to_node_dict[triple.head.name].children[tail_node.id_]=tail_node\n",
    "            else:\n",
    "                # else just create a new clause title node with no children\n",
    "                id_to_node_dict[triple.head.name] = whyhow_node_to_local_node(triple.head)\n",
    "            # add subclause node to the map\n",
    "            id_to_node_dict[triple.tail.name] = tail_node\n",
    "\n",
    "    # create a set of unseen clause titles\n",
    "    clause_titles_ids_set: Set[str] = set(clause_title_ids)\n",
    "\n",
    "    final_nodes = []\n",
    "    for current_node in current_nodes:\n",
    "        # print(f\"Current node: {current_node.type} {current_node.id_}\")\n",
    "        # if not a list item, ignore\n",
    "        if current_node.type != ElementType.LIST_ITEM:\n",
    "            final_nodes.append(current_node)\n",
    "            continue\n",
    "        # elif we see a clause item, link any subclauses to it\n",
    "        elif current_node.id_ in clause_titles_ids_set:\n",
    "            # mark clause as seen by removing it from the set of unseen clauses\n",
    "            clause_titles_ids_set.remove(current_node.id_)\n",
    "            # replace clause node with one that contains all subclauses\n",
    "            final_node_to_append = id_to_node_dict[current_node.id_]\n",
    "            final_nodes.append(final_node_to_append)\n",
    "        # elif we see a subclause item\n",
    "        elif current_node.id_ in subclause_ids:\n",
    "            # look up the subclause's parent clause id\n",
    "            clause = subclause_to_clause_map[current_node.id_]\n",
    "            # if its clause id is unseen, replace the subclause with the full clause that contains all subclauses\n",
    "            if clause in clause_titles_ids_set:\n",
    "                clause_titles_ids_set.remove(clause)\n",
    "                final_node_to_append = id_to_node_dict[clause]\n",
    "                final_nodes.append(final_node_to_append)\n",
    "            # otherwise, if this subclause is part of a clause we have seen before, don't add it in\n",
    "\n",
    "    return final_nodes\n",
    "\n",
    "\n",
    "local_nodes = list(local_nodes_map.values())\n",
    "starting_nodes = local_nodes[139:141]\n",
    "\n",
    "nodes_test = fetch_whole_clause_lists(starting_nodes)\n",
    "\n",
    "print(\"PROMPT:\")\n",
    "for node in nodes_test:\n",
    "    print(node.print_node_prompt())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition Agent Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relevant Definitions:\n",
      "Board: The board of directors of a financial institution, including a committee of the board where the responsibilities of the board set out in this policy document have been delegated to such a committee.\n",
      "Control Function: A function that has a responsibility independent from business lines to provide objective assessment, reporting or assurance. This includes the risk management function, the compliance function and the internal audit function.\n",
      "Chief Compliance Officer: The senior officer of a financial institution, however styled, who is the central point of authority for a financial institution's compliance matters and is responsible for providing an institution-wide view on the management of compliance risk.\n",
      "\n",
      "{'Board': 'The board of directors of a financial institution, including a committee of the board where the responsibilities of the board set out in this policy document have been delegated to such a committee.', 'Control Function': 'A function that has a responsibility independent from business lines to provide objective assessment, reporting or assurance. This includes the risk management function, the compliance function and the internal audit function.', 'Chief Compliance Officer': \"The senior officer of a financial institution, however styled, who is the central point of authority for a financial institution's compliance matters and is responsible for providing an institution-wide view on the management of compliance risk.\"}\n"
     ]
    }
   ],
   "source": [
    "from whyhow import WhyHow\n",
    "\n",
    "def definitions_search(query: str, client: Optional[WhyHow]=None) -> Dict[str, str]:\n",
    "    \"\"\"\n",
    "    Search for definitions of terms in a question prompt and return them as a dictionary.\n",
    "    \"\"\"\n",
    "    if client is None:\n",
    "        client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "        \n",
    "    query_prompt = f\"\"\"Return me definitions for the terms in this query: \"{query}\" Return only relevant terms that exist in the query. Ensure the term-definition pairs are separated by newlines, properly capitalised. \"\"\"\n",
    "\n",
    "    definitions_response = client.graphs.query_unstructured(\n",
    "        graph_id=definitions_graph.graph_id,\n",
    "        query=query_prompt,\n",
    "    )\n",
    "    \n",
    "    response_text = definitions_response.answer\n",
    "    term_def_pairs = response_text.split('\\n')\n",
    "    definitions_dict = {}\n",
    "    \n",
    "    for pair in term_def_pairs:\n",
    "        if ':' in pair:\n",
    "            term, definition = pair.split(':', 1)\n",
    "            definitions_dict[term.strip()] = definition.strip()\n",
    "    \n",
    "    return definitions_dict\n",
    "\n",
    "\n",
    "\n",
    "definitions_dict = definitions_search(\"How can the Board and the CCO manage control functions?\")\n",
    "\n",
    "\n",
    "def print_prompt_definitions_dict(definitions_dict):\n",
    "    prompt = \"Relevant Definitions:\\n\"\n",
    "    for term, definition in definitions_dict.items():\n",
    "        prompt += f\"{term}: {definition}\\n\"\n",
    "    return prompt\n",
    "\n",
    "print(print_prompt_definitions_dict(definitions_dict))\n",
    "print(definitions_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Updating Node Context with Footer Info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- node_id: `- S 6.3 In relation to the position of the CCO, th..._d9432f9931c5bf206db8317a67e69052`    content: \"- S 6.3 In relation to the position of the CCO, the board must-\"    children: \n",
      "          - node_id: `(c) engage with the CCO on a regular basis2 to pro..._75d9df850bf8422349822accdcb9338e`              content: \" (c) engage with the CCO on a regular basis2 to provide the opportunity for the CCO to discuss issues faced by the compliance function;\"              context:\n",
      "                - footer info: The board should also consider engaging the CCO without the presence of other members of senior management from time to time.\n",
      "                - further explore footer?: False\n",
      "\n",
      "          - node_id: `(d) provide the CCO with direct and unimpeded acce..._753739a7f2c2cf4c206b6e102cf20aca`              content: \" (d) provide the CCO with direct and unimpeded access to the board;\"          - node_id: `(e) ensure that the CCO is supported with sufficie..._b5cb66fa679baafeacc244e3dbe13938`              content: \" (e) ensure that the CCO is supported with sufficient resources, including competent officers, to perform his duties effectively; and\"          - node_id: `(f) where the CCO also carries out responsibilitie..._684a6af3d1c0785ec56eeb2cbe6c1644`              content: \" (f) where the CCO also carries out responsibilities in respect of other control functions3, be satisfied that a sound overall control environment will not be compromised by the combination of responsibilities performed by the CCO.\"              context:\n",
      "                - footer info: Refer to paragraphs 7.3 and 7.4.\n",
      "                - further explore footer?: True\n",
      "\n",
      "          - node_id: `(a) approve the appointment, remuneration and term..._f5abf05a077c9b80dd1a4a0a793064c8`              content: \" (a) approve the appointment, remuneration and termination of the CCO¹;\"              context:\n",
      "                - footer info: Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011.\n",
      "                - further explore footer?: False\n",
      "\n",
      "          - node_id: `(b) ensure that the CCO has sufficient stature to ..._ade5bcb25050d514e6d10e163c1581ec`              content: \" (b) ensure that the CCO has sufficient stature to allow for effective engagement with the CEO and other members of senior management;\"\n",
      "- node_id: `1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8`    content: \"1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011. 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time. 3 Refer to paragraphs 7.3 and 7.4.\"\n",
      "- node_id: `Issued on: 10 May 2016_b350089add36cc840ad1fb778a27cd62`    content: \"Issued on: 10 May 2016\"\n",
      "- node_id: `Compliance_747e07c6b6fb06898766b61480775e55`    content: \"Compliance\"\n",
      "- node_id: `Responsibilities of senior management_26342c4f3533e56f55513945dccc9f9c`    content: \"Responsibilities of senior management\"\n",
      "- node_id: `- S 6.4 Senior management is collectively responsi..._4cd12c4a3ddeb6778b65630b4e624564`    content: \"- S 6.4 Senior management is collectively responsible for the effective management of a financial institution's compliance risk. In discharging this responsibility, senior management must-\"    children: \n",
      "          - node_id: `(a) establish a written compliance policy and ensu..._4b734f83a89aeb246dc8bef15b47b6f3`              content: \" (a) establish a written compliance policy and ensure that it is kept up to date;\"          - node_id: `(b) communicate the policy to all officers and ens..._4606b79fb7b6a0c04ff96837b8140fd2`              content: \" (b) communicate the policy to all officers and ensure that appropriate remedial or disciplinary actions are taken if the compliance policy is breached;\"          - node_id: `(c) establish a compliance function commensurate w..._cd9d52958cf53b9295b73a9434cb2ba9`              content: \" (c) establish a compliance function commensurate with the size, nature of operations and complexity of the financial institution, having regard to the requirements in paragraphs 7 and 8;\"          - node_id: `(d) provide sufficient resources for the complianc..._bc8b34fe51d9616562cdc72c0ca06cd7`              content: \" (d) provide sufficient resources for the compliance function, including officers with the appropriate competencies and experience;\"          - node_id: `(e) ensure that the compliance function is able to..._10dd1dfdaa45d1fca7a368814c9ebaac`              content: \" (e) ensure that the compliance function is able to secure assistance from other functions with specific expertise (for example, legal, Shariah review or risk management) and has clear authority to engage with any officers and obtain access to relevant information for purposes of discharging its responsibilities;\"          - node_id: `(f) ensure that the compliance function is able to..._ec892fe9679abba80831f186d7bd8c85`              content: \" (f) ensure that the compliance function is able to engage relevant external expertise to undertake compliance assessments in specific areas4 where necessary;\"              context:\n",
      "                - footer info: 4 For example, to conduct investigations of possible incidents of non-compliance.\n",
      "                - further explore footer?: False\n",
      "\n",
      "          - node_id: `(g) ensure that the compliance function is kept in..._4bffa3bce40cc3cff49e1b2a7ca41860`              content: \" (g) ensure that the compliance function is kept informed of any organisational developments to facilitate the timely identification of compliance risk;\"          - node_id: `(h) report to the board regularly on compliance is..._334fdd9aecc048c657f8ad1d9c0d8a24`              content: \" (h) report to the board regularly on compliance issues and promptly on any material incidents of non-compliance (for example, failures that may attract a significant risk of legal or regulatory sanction);\"          - node_id: `(i) report to the board at least annually on the e..._570902ce5ab35809597f3727115bcd64`              content: \" (i) report to the board at least annually on the effectiveness of the financial institution's overall management of compliance risk in such a manner as to assist the board in carrying out its responsibilities as set out in paragraph 6.2(d); and\"          - node_id: `(j) inform the board of the CCO's cessation from o..._c6086e80eb5855c34404bd1e6e7eb06c`              content: \" (j) inform the board of the CCO's cessation from office and the reasons leading to the cessation.\"\n",
      "Number of nodes changed: 4\n"
     ]
    }
   ],
   "source": [
    "from typing import Dict\n",
    "from whyhow.schemas import Node\n",
    "from typing import List, Any, Optional, Dict\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Get OpenAI API key from environment variables\n",
    "openai_api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "if openai_api_key is None:\n",
    "    raise ValueError(\"OPENAI_API_KEY not found in environment variables.\")\n",
    "\n",
    "# Initialize OpenAI client\n",
    "openai_client = OpenAI()\n",
    "openai_client.api_key = openai_api_key\n",
    "\n",
    "\n",
    "footer_nodes_by_page_number: Dict[int, List[MultiAgentSearchLocalNode]] = defaultdict(\n",
    "    list\n",
    ")\n",
    "\n",
    "# build a map of footer nodes by page number\n",
    "for node_id, node in local_nodes_map.items():\n",
    "    if node.type == ElementType.FOOTER:\n",
    "        page_number = int(node.metadata[\"Page Number\"])\n",
    "        footer_nodes_by_page_number[page_number].append(node)\n",
    "\n",
    "\n",
    "class NodeIdFooterResponse(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe the structured output of node-id to footer info pair\n",
    "    \"\"\"\n",
    "\n",
    "    node_id: str\n",
    "    footer_text: str\n",
    "    further_explore_footer: bool\n",
    "\n",
    "\n",
    "class NodeIdFooterResponses(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe the structured output of node-id to footer info pairs.\n",
    "    \"\"\"\n",
    "\n",
    "    content: List[NodeIdFooterResponse]\n",
    "\n",
    "\n",
    "def fill_nodes_with_footer_context(\n",
    "    query: str, current_nodes: List[MultiAgentSearchLocalNode]\n",
    ") -> Tuple[List[MultiAgentSearchLocalNode], int]:\n",
    "    \"\"\"\n",
    "    Takes nodes and returns a list of footer contexts to fill them with.\n",
    "    \"\"\"\n",
    "\n",
    "    starter_prompt = f\"\"\"\n",
    "    You are an agent responsible for the structured extraction of a document page to answer the question:\n",
    "    ```{query}```\n",
    "    Below are two lists of OCR'd text, one for nodes and one for footers. All the text was taken from one page. Some of these nodes make references to the footers below, possibly through footnote numbers.\n",
    "\n",
    "    Your task is, for each node_id, to identify if it contains (1) any footer information, and (2) if the footer requires further reading ON THE SAME DOCUMENT. Return only node_ids that contain footer info.\n",
    "    \n",
    "    For each node:\n",
    "    1. deduce from the context and footnode numbering if any footer content might be relevant to the node given the query. If so, extract that footer content (it can be a partial or whole extract) and assign it to the node. If there is no footer info, SKIP the node. Only return nodes that have footer content.\n",
    "    \n",
    "    2. Some of the footer content may require the reader to consult other parts of the document (it must be an internal link pointing to other paragraphs/sections, not external reading). If so, set the further_explore_footer flag to true, else false.\n",
    "\n",
    "    Then return a list of (node_id, relevant_footer_content, further_explore_footer). node_ids must be quoted in backticks.\n",
    "    \n",
    "    \"\"\"\n",
    "\n",
    "    nodes_by_page_number: Dict[int, List[MultiAgentSearchLocalNode]] = defaultdict(list)\n",
    "    \n",
    "    # collection of which nodes to mark with footer context\n",
    "                            #node_id, (footer_text, further_explore_footer) \n",
    "    node_id_footer_context: Dict[str, Tuple[str, bool]] = {}\n",
    "    \n",
    "    \n",
    "    def recurse_index_nodes(current_nodes: List[MultiAgentSearchLocalNode]):\n",
    "        \"\"\"\n",
    "        Recursively iterate via DFS on nodes, to index them by page number.\n",
    "        \"\"\"\n",
    "        for node in current_nodes:\n",
    "            page_number = int(node.metadata[\"Page Number\"])\n",
    "            nodes_by_page_number[page_number].append(node)\n",
    "\n",
    "            if len(node.children) > 0:\n",
    "                recurse_index_nodes(list(node.children.values()))\n",
    "\n",
    "    recurse_index_nodes(current_nodes)\n",
    "\n",
    "    for page_number, nodes in nodes_by_page_number.items():\n",
    "\n",
    "        printout_of_curent_nodes = f\"Nodes on page {page_number}\"\n",
    "        for node in nodes:\n",
    "            printout_of_curent_nodes += node.print_node_prompt()\n",
    "\n",
    "        printout_of_footers = f\"Footers on page {page_number}:\"\n",
    "        for footer in footer_nodes_by_page_number[page_number]:\n",
    "            printout_of_footers += footer.print_node_prompt()\n",
    "\n",
    "        completion = openai_client.beta.chat.completions.parse(\n",
    "            model=\"gpt-4o-2024-08-06\",\n",
    "            messages=[\n",
    "                {\"role\": \"system\", \"content\": starter_prompt},\n",
    "                {\"role\": \"user\", \"content\": printout_of_curent_nodes},\n",
    "                {\"role\": \"user\", \"content\": printout_of_footers},\n",
    "            ],\n",
    "            response_format=NodeIdFooterResponses,\n",
    "        )\n",
    "\n",
    "        node_id_footer_response = completion.choices[0].message.parsed.content\n",
    "\n",
    "        if isinstance(node_id_footer_response, list):\n",
    "            for item in node_id_footer_response:\n",
    "                if isinstance(item, NodeIdFooterResponse):\n",
    "                    # sometimes GPT returns node ids with or without backticks\n",
    "                    if item.node_id[0] == \"`\":\n",
    "                        node_id = item.node_id[1:-1]\n",
    "                    else:\n",
    "                        node_id = item.node_id\n",
    "                    node_id_footer_context[node_id] = (item.footer_text, item.further_explore_footer)\n",
    "\n",
    "    change_count = 0\n",
    "    # Based on the dict of node_id - footer context\n",
    "    for node_id, (footer_text, further_explore_footer) in node_id_footer_context.items():\n",
    "        # some LLM calls \n",
    "        if footer_text not in [\"\", \" \", None]:\n",
    "            # print(f\"node_id: {node_id}, footer_text: '{footer_text}', further_explore_footer: {further_explore_footer}\")\n",
    "            recursive_add_context(current_nodes, node_id, \"footer info\", footer_text)\n",
    "            recursive_add_context(current_nodes, node_id, \"further explore footer?\", str(further_explore_footer))\n",
    "            change_count += 1\n",
    "\n",
    "    return current_nodes, change_count\n",
    "\n",
    "\n",
    "local_nodes = list(local_nodes_map.values())\n",
    "starting_nodes = local_nodes[100:110]\n",
    "\n",
    "nodes_test1 = fetch_whole_clause_lists(starting_nodes)\n",
    "nodes_test2, change_count = fill_nodes_with_footer_context(\"What general information is in my footer?\", nodes_test1)\n",
    "\n",
    "\n",
    "for node in nodes_test2:\n",
    "    print(node.print_node_prompt())\n",
    "    \n",
    "print(f\"Number of nodes changed: {change_count}\")\n",
    "    \n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Updating Node Context with Link Info\n",
    "- For a given node, add hints to the nodes that say it is linked to other nodes\n",
    "\n",
    "\n",
    "- Note: WhyHow only has one identifying field for a node, called 'node_name', and it has to be unique. \n",
    "  - Due to this, Node name and id mean the same thing. On the UI, it shows node_name to the user\n",
    "  - The way we structured this notebook is to use a preview (first 50 characters of element content) + a hash of the reducto element to use as node_name/node_id\n",
    "- This comes with some benefits. For every node with links, we make it have a hint that says (to the LLM reading it), \"hey, I have a link to these node_ids: [list of linked nodes by ID]\"\n",
    "- These node_ids also contain a preview (first 50 characters), so the LLM can decide if it needs to explore/follow the links or not"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- node_id: `- S 6.3 In relation to the position of the CCO, th..._d9432f9931c5bf206db8317a67e69052`    content: \"- S 6.3 In relation to the position of the CCO, the board must-\"    children: \n",
      "          - node_id: `(c) engage with the CCO on a regular basis2 to pro..._75d9df850bf8422349822accdcb9338e`              content: \" (c) engage with the CCO on a regular basis2 to provide the opportunity for the CCO to discuss issues faced by the compliance function;\"              context:\n",
      "                - footer info: The board should also consider engaging the CCO without the presence of other members of senior management from time to time.\n",
      "                - further explore footer?: False\n",
      "\n",
      "          - node_id: `(d) provide the CCO with direct and unimpeded acce..._753739a7f2c2cf4c206b6e102cf20aca`              content: \" (d) provide the CCO with direct and unimpeded access to the board;\"          - node_id: `(e) ensure that the CCO is supported with sufficie..._b5cb66fa679baafeacc244e3dbe13938`              content: \" (e) ensure that the CCO is supported with sufficient resources, including competent officers, to perform his duties effectively; and\"          - node_id: `(f) where the CCO also carries out responsibilitie..._684a6af3d1c0785ec56eeb2cbe6c1644`              content: \" (f) where the CCO also carries out responsibilities in respect of other control functions3, be satisfied that a sound overall control environment will not be compromised by the combination of responsibilities performed by the CCO.\"              context:\n",
      "                - footer info: Refer to paragraphs 7.3 and 7.4.\n",
      "                - further explore footer?: True\n",
      "\n",
      "          - node_id: `(a) approve the appointment, remuneration and term..._f5abf05a077c9b80dd1a4a0a793064c8`              content: \" (a) approve the appointment, remuneration and termination of the CCO¹;\"              context:\n",
      "                - footer info: Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011.\n",
      "                - further explore footer?: False\n",
      "\n",
      "          - node_id: `(b) ensure that the CCO has sufficient stature to ..._ade5bcb25050d514e6d10e163c1581ec`              content: \" (b) ensure that the CCO has sufficient stature to allow for effective engagement with the CEO and other members of senior management;\"\n",
      "- node_id: `1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8`    content: \"1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011. 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time. 3 Refer to paragraphs 7.3 and 7.4.\"\n",
      "- node_id: `Issued on: 10 May 2016_b350089add36cc840ad1fb778a27cd62`    content: \"Issued on: 10 May 2016\"\n",
      "- node_id: `Compliance_747e07c6b6fb06898766b61480775e55`    content: \"Compliance\"\n",
      "- node_id: `Responsibilities of senior management_26342c4f3533e56f55513945dccc9f9c`    content: \"Responsibilities of senior management\"\n",
      "- node_id: `- S 6.4 Senior management is collectively responsi..._4cd12c4a3ddeb6778b65630b4e624564`    content: \"- S 6.4 Senior management is collectively responsible for the effective management of a financial institution's compliance risk. In discharging this responsibility, senior management must-\"    children: \n",
      "          - node_id: `(a) establish a written compliance policy and ensu..._4b734f83a89aeb246dc8bef15b47b6f3`              content: \" (a) establish a written compliance policy and ensure that it is kept up to date;\"          - node_id: `(b) communicate the policy to all officers and ens..._4606b79fb7b6a0c04ff96837b8140fd2`              content: \" (b) communicate the policy to all officers and ensure that appropriate remedial or disciplinary actions are taken if the compliance policy is breached;\"          - node_id: `(c) establish a compliance function commensurate w..._cd9d52958cf53b9295b73a9434cb2ba9`              content: \" (c) establish a compliance function commensurate with the size, nature of operations and complexity of the financial institution, having regard to the requirements in paragraphs 7 and 8;\"              context:\n",
      "                - link: [\"['- 7 Organisation of the compliance function_6cdf24899c11cedfa54a1f48658033a1', '- 8 Responsibilities of the compliance function_9898aabafc24b8c858516d18f8ca38fc']\"]\n",
      "\n",
      "          - node_id: `(d) provide sufficient resources for the complianc..._bc8b34fe51d9616562cdc72c0ca06cd7`              content: \" (d) provide sufficient resources for the compliance function, including officers with the appropriate competencies and experience;\"          - node_id: `(e) ensure that the compliance function is able to..._10dd1dfdaa45d1fca7a368814c9ebaac`              content: \" (e) ensure that the compliance function is able to secure assistance from other functions with specific expertise (for example, legal, Shariah review or risk management) and has clear authority to engage with any officers and obtain access to relevant information for purposes of discharging its responsibilities;\"          - node_id: `(f) ensure that the compliance function is able to..._ec892fe9679abba80831f186d7bd8c85`              content: \" (f) ensure that the compliance function is able to engage relevant external expertise to undertake compliance assessments in specific areas4 where necessary;\"              context:\n",
      "                - footer info: 4 For example, to conduct investigations of possible incidents of non-compliance.\n",
      "                - further explore footer?: False\n",
      "\n",
      "          - node_id: `(g) ensure that the compliance function is kept in..._4bffa3bce40cc3cff49e1b2a7ca41860`              content: \" (g) ensure that the compliance function is kept informed of any organisational developments to facilitate the timely identification of compliance risk;\"          - node_id: `(h) report to the board regularly on compliance is..._334fdd9aecc048c657f8ad1d9c0d8a24`              content: \" (h) report to the board regularly on compliance issues and promptly on any material incidents of non-compliance (for example, failures that may attract a significant risk of legal or regulatory sanction);\"          - node_id: `(i) report to the board at least annually on the e..._570902ce5ab35809597f3727115bcd64`              content: \" (i) report to the board at least annually on the effectiveness of the financial institution's overall management of compliance risk in such a manner as to assist the board in carrying out its responsibilities as set out in paragraph 6.2(d); and\"          - node_id: `(j) inform the board of the CCO's cessation from o..._c6086e80eb5855c34404bd1e6e7eb06c`              content: \" (j) inform the board of the CCO's cessation from office and the reasons leading to the cessation.\"              context:\n",
      "                - link: [\"['- S 6.3 In relation to the position of the CCO, th..._d9432f9931c5bf206db8317a67e69052']\"]\n",
      "\n",
      "\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "def fill_nodes_with_link_hints(\n",
    "    current_nodes: List[MultiAgentSearchLocalNode], query_client=None\n",
    ") -> Tuple[List[MultiAgentSearchLocalNode], int]:\n",
    "    \"\"\"\n",
    "    Given a tree/list of nodes, update the nodes with hints to the LLM if any of the nodes contain links.\n",
    "    \"\"\"\n",
    "\n",
    "    if query_client is None:\n",
    "        query_client = WhyHow(api_key=WHYHOW_API_KEY, base_url=WHYHOW_API_URL)\n",
    "\n",
    "    # some nodes are nested, so recurse to get all nodes\n",
    "    flattened_nodes: Dict[str, MultiAgentSearchLocalNode] = {}\n",
    "\n",
    "    link_dict : Dict[str, List[str]] = defaultdict(list)\n",
    "\n",
    "    def recurse_index_nodes(current_nodes: List[MultiAgentSearchLocalNode]):\n",
    "        for node in current_nodes:\n",
    "            flattened_nodes[node.id_] = node\n",
    "\n",
    "\n",
    "            if len(node.children) > 0:\n",
    "                recurse_index_nodes(list(node.children.values()))\n",
    "\n",
    "    recurse_index_nodes(current_nodes)\n",
    "\n",
    "    # fetch all elements are in a link_to triple containing the current_nodes\n",
    "    links_query: Query = query_client.graphs.query_structured(\n",
    "        graph_id=LEXICAL_GRAPH_ID.graph_id,\n",
    "        entities=ELEMENTS,\n",
    "        relations=[\"links_to\"],\n",
    "        values= list(flattened_nodes.keys()),\n",
    "    )\n",
    "\n",
    "    # head links_to tail\n",
    "    # if the head is in our flattened nodes, then it links to tail\n",
    "    # we could repeat in reverse to get backlinks (like Obsidian), but out of scope.\n",
    "    for triple in links_query.triples:\n",
    "        if triple.head.name in flattened_nodes.keys():\n",
    "            link_dict[triple.head.name].append(triple.tail.name)\n",
    "            \n",
    "    # recurse through the whole graph and add context markers inside any node with a link\n",
    "    # show the llm hints that some clauses/text/etc link to external parts\n",
    "    change_count = 0\n",
    "    for link_node_id, destinations in link_dict.items():\n",
    "        recursive_append_context_items(current_nodes, link_node_id,\"link\", str(destinations))\n",
    "        change_count +=1\n",
    "        \n",
    "    return current_nodes, change_count\n",
    "\n",
    "\n",
    "local_nodes = list(local_nodes_map.values())\n",
    "starting_nodes = local_nodes[100:120]\n",
    "nodes_test2 = fetch_whole_clause_lists(starting_nodes)\n",
    "nodes_test3, change_count = fill_nodes_with_link_hints(nodes_test2)\n",
    "nodes_test4, change_count = fill_nodes_with_footer_context(\"What general information is in my footer?\",nodes_test3)\n",
    "\n",
    "\n",
    "for node in nodes_test4:\n",
    "    print(node.print_node_prompt())\n",
    "\n",
    "print(change_count)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MY Central Bank Question:\n",
    "\n",
    "- How can the Board and the CCO manage control functions?\n",
    "\n",
    "## Ideal flow:\n",
    "\n",
    "- Gives me a definition of CCO from definition page\n",
    "- Retrieval of clauses from 6.3, 7.2.\n",
    "- Detection of “Refer to Para 7.3 and 7.4”.\n",
    "- Pulls in Para 7.3 and 7.4, which mentions Para 9.1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Langgraph\n",
    "\n",
    "![](./langgraph.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "# Multi-Agent Log\n",
       "## Pass 1\n",
       "### Initial Search Agent\n",
       " - retrieved 5 nodes, pruned to 5 nodes\n",
       "### Definitions Agent\n",
       "Retrieved 3 definitions for: \n",
       "  - **Board** : The board of directors of a financial institution, including a committee of the board where the responsibilities of the board set out in this policy document have been delegated to such a committee.\n",
       "  - **Chief Compliance Officer** : The senior officer of a financial institution, however styled, who is the central point of authority for a financial institution's compliance matters and is responsible for providing an institution-wide view on the management of compliance risk.\n",
       "  - **Control Function** : A function that has a responsibility independent from business lines to provide objective assessment, reporting or assurance. This includes the risk management function, the compliance function and the internal audit function.\n",
       "### Context Fetched for nodes:\n",
       "- fetched full clauses, added to 3 nodes. 2 new links\n",
       "- filled nodes with footer context\n",
       "\n",
       "```\n",
       "- node_id: `- S 6.3 In relation to the position of the CCO, th..._d9432f9931c5bf206db8317a67e69052`    content: \"- S 6.3 In relation to the position of the CCO, the board must-\"    children: \n",
       "          - node_id: `(a) approve the appointment, remuneration and term..._f5abf05a077c9b80dd1a4a0a793064c8`              content: \" (a) approve the appointment, remuneration and termination of the CCO¹;\"              context:\n",
       "                - footer info: 1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(b) ensure that the CCO has sufficient stature to ..._ade5bcb25050d514e6d10e163c1581ec`              content: \" (b) ensure that the CCO has sufficient stature to allow for effective engagement with the CEO and other members of senior management;\"          - node_id: `(c) engage with the CCO on a regular basis2 to pro..._75d9df850bf8422349822accdcb9338e`              content: \" (c) engage with the CCO on a regular basis2 to provide the opportunity for the CCO to discuss issues faced by the compliance function;\"              context:\n",
       "                - footer info: 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(d) provide the CCO with direct and unimpeded acce..._753739a7f2c2cf4c206b6e102cf20aca`              content: \" (d) provide the CCO with direct and unimpeded access to the board;\"          - node_id: `(e) ensure that the CCO is supported with sufficie..._b5cb66fa679baafeacc244e3dbe13938`              content: \" (e) ensure that the CCO is supported with sufficient resources, including competent officers, to perform his duties effectively; and\"          - node_id: `(f) where the CCO also carries out responsibilitie..._684a6af3d1c0785ec56eeb2cbe6c1644`              content: \" (f) where the CCO also carries out responsibilities in respect of other control functions3, be satisfied that a sound overall control environment will not be compromised by the combination of responsibilities performed by the CCO.\"              context:\n",
       "                - footer info: 3 Refer to paragraphs 7.3 and 7.4.\n",
       "                - further explore footer?: True\n",
       "\n",
       "- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "\n",
       "```\n",
       "### Supervisor Agent\n",
       "- **supervisor**: continue\n",
       "### Router Agent\n",
       "- found 2 nodes to fetch links for\n",
       "    - '- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66'\n",
       "    - '- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472'\n",
       " - found 1 nodes to fetch footers for\n",
       "    - `- S 6.3 In relation to the position of the CCO, th..._d9432f9931c5bf206db8317a67e69052` \n",
       "      - **Search query:** Return paragraphs 7.3 and 7.4 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "### Recursive Retrieval\n",
       "2 nodes with link(s)\n",
       "1 nodes with footer info\n",
       " - - S 6.3 In relation to the position of the CCO, th..._d9432f9931c5bf206db8317a67e69052\n",
       "  - **Search query:** Return paragraphs 7.3 and 7.4 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "  - **Returned:** \n",
       "    - - S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3\n",
       "    - - S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd\n",
       "# PASS 2\n",
       "### Context Fetched for nodes:\n",
       "- fetched full clauses, added to 2 nodes. 2 new links\n",
       "- filled nodes with footer context\n",
       "\n",
       "```\n",
       "- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"\n",
       "```\n",
       "### Supervisor Agent\n",
       "- **supervisor**: continue\n",
       "### Router Agent\n",
       "- found 1 nodes to fetch links for\n",
       "    - - S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472\n",
       " - found 2 nodes to fetch footers for\n",
       "    - `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3` \n",
       "      - **Search query:** Return paragraphs 7.2 and 7.3 to answer the query 'How can the Board and the CCO manage control functions?'.\n",
       "    - `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c` \n",
       "      - **Search query:** Return paragraph 9.1 to answer the query 'How can the Board and the CCO manage control functions?'.\n",
       "### Recursive Retrieval\n",
       "1 nodes with link(s)\n",
       "  - - S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472\n",
       "2 nodes with footer info\n",
       " - - S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3\n",
       "  - **Search query:** Return paragraphs 7.2 and 7.3 to answer the query 'How can the Board and the CCO manage control functions?'.\n",
       "  - **Returned:** \n",
       "    - - S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd\n",
       "    - 1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8\n",
       "    - - S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66\n",
       "    - (b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c\n",
       "    - (c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864\n",
       " - (b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c\n",
       "  - **Search query:** Return paragraph 9.1 to answer the query 'How can the Board and the CCO manage control functions?'.\n",
       "  - **Returned:** \n",
       "    - - S 7.6 Notwithstanding paragraph 7.3, the CCO of ..._614f4361a3bd64ad022d9d007b16da05\n",
       "    - - 1.1 A strong compliance culture reflects a corpo..._1fe85abd327658de88e7eed861b59a93\n",
       "    - - S 7.10 Where such arrangements referred to in pa..._9aee9c464ce513503b26252dc35e3ad1\n",
       "    - Compliance_95521f3c1af881ea7cf1d71ec7c00c9a\n",
       "    - (b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c\n",
       "    - - S 8.9 The CCO must ensure that the reports refer..._6cf04fc7a955f2d7df8ea100ec7af539\n",
       "    - (c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864\n",
       "# PASS 3\n",
       "### Context Fetched for nodes:\n",
       "- fetched full clauses, added to 10 nodes. 4 new links\n",
       "- filled nodes with footer context\n",
       "\n",
       "```\n",
       "- node_id: `- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472`    content: \"- S 9.1 A financial institution must ensure that there is a clear separation of the internal audit function and other functions carrying out compliance function responsibilities. Compliance risk must be included in the risk assessment methodology of the internal audit function, and an audit programme that covers the adequacy and effectiveness of the other functions carrying out compliance function responsibilities should be established, including testing of controls commensurate with the perceived level of risk.\"    context:\n",
       "      - footer info: Issued on: 10 May 2016\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8`    content: \"1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011. 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time. 3 Refer to paragraphs 7.3 and 7.4.\"    context:\n",
       "      - footer info: 1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011. 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time. 3 Refer to paragraphs 7.3 and 7.4.\n",
       "      - further explore footer?: True\n",
       "\n",
       "- node_id: `- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66`    content: \"- S 7.2 Where compliance function responsibilities are shared between a dedicated compliance unit and other control functions-\"    children: \n",
       "          - node_id: `(a) the allocation of the compliance function resp..._32a2cb84e0fdce2e66fc429c12918716`              content: \" (a) the allocation of the compliance function responsibilities, including that for timely communication and escalation of compliance issues to senior management and the board, must be clearly defined and documented in the compliance policy;\"          - node_id: `(b) the CCO must have overall responsibility for c..._d1a0a67335e5aa245b6eeeaa082e845d`              content: \" (b) the CCO must have overall responsibility for coordinating the identification and management of compliance risk at the institution-wide level, and ensuring that compliance monitoring and testing are carried out consistently across the institution;\"          - node_id: `(c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864`              content: \" (c) for the purpose of paragraph 7.2(b), the CCO must have a sound understanding of compliance risks which are under the purview of other control functions, including an understanding of controls applied to manage these risks;\"          - node_id: `(d) arrangements for coordination among the contro..._2c7647be1e92b542fc76be5e0a9e14a4`              content: \" (d) arrangements for coordination among the control functions with the CCO and the compliance unit must be in place to ensure that the CCO is able to perform his responsibilities effectively. Arrangements should promote a view and approach to the management of compliance risk that is consistent across the organisation, including through adequate information flows and avenues to seek advice on compliance issues; and\"          - node_id: `(e) senior management must ensure that officers in..._be843c24dda3f19f1ae16dcb3ee0acd0`              content: \" (e) senior management must ensure that officers in other control functions have the capacity and expertise to deliver their broader mandates while providing adequate focus to their compliance function responsibilities.\"              context:\n",
       "                - footer info: 5 For example, through representation on a new products committee.\n",
       "                - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.6 Notwithstanding paragraph 7.3, the CCO of ..._614f4361a3bd64ad022d9d007b16da05`    content: \"- S 7.6 Notwithstanding paragraph 7.3, the CCO of a large financial institution is not allowed to assume the responsibilities of other control functions.\"    children: \n",
       "          - node_id: `- Independence_7e6261e9f5d7eb958e43140add281d2e`              content: \" - Independence\"- node_id: `- 1.1 A strong compliance culture reflects a corpo..._1fe85abd327658de88e7eed861b59a93`    content: \"- 1.1 A strong compliance culture reflects a corporate culture of high ethical standards and integrity in which the board and senior management lead by example. A financial institution should hold itself to high standards in carrying on business, and at all times observe both the spirit and the letter of the law and regulations. Failure to effectively manage compliance risk may result in adverse consequences for the financial institution's customers, shareholders, officers and the financial institution itself.\"- node_id: `- S 7.10 Where such arrangements referred to in pa..._9aee9c464ce513503b26252dc35e3ad1`    content: \"- S 7.10 Where such arrangements referred to in paragraph 7.9 exist, a financial institution must ensure that-\"    context:\n",
       "      - link: [\"['- G 7.9 A constructive and cooperative working rel..._081d34781b0f7ab41f94dbe52896a23e']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the compliance function is not placed in a pos..._03f9f972a4a09290c72409b0826aa893`              content: \" (a) the compliance function is not placed in a position of conflict;\"          - node_id: `(b) the accountability of the compliance function ..._e301200856ef131fd6a08194a2532e5e`              content: \" (b) the accountability of the compliance function is properly documented6 and\"              context:\n",
       "                - footer info: 6 In the case of a committee, the role of the compliance function may be outlined in the terms of reference or charter.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(c) the compliance function is not prevented from ..._8f4f0fef22f99ea2f589709d5b54496e`              content: \" (c) the compliance function is not prevented from highlighting compliance issues relating to any business decisions to the board or senior management, where necessary.\"          - node_id: `- Resources_0df524eaa89ae4d09e3b49998eb3b771`              content: \" - Resources\"- node_id: `Compliance_95521f3c1af881ea7cf1d71ec7c00c9a`    content: \"Compliance\"- node_id: `- S 8.9 The CCO must ensure that the reports refer..._6cf04fc7a955f2d7df8ea100ec7af539`    content: \"- S 8.9 The CCO must ensure that the reports referred to in paragraph 8.8 are readily available to the internal audit function of the financial institution, the Bank and other regulatory authorities upon request.\"    context:\n",
       "      - link: [\"['- S 8.8 The CCO must report to senior management o..._0535171ad1204d0c05d28b16a437a17c']\"]\n",
       "      - footer info: Issued on: 10 May 2016\n",
       "      - further explore footer?: False\n",
       "\n",
       "    children: \n",
       "          - node_id: `- Advisory_57633bcdded147683c36017c16c03fc8`              content: \" - Advisory\"\n",
       "```\n",
       "### Supervisor Agent\n",
       "- **supervisor**: continue\n",
       "### Router Agent\n",
       "- found 1 nodes to fetch links for\n",
       "    - - S 8.8 The CCO must report to senior management o..._0535171ad1204d0c05d28b16a437a17c\n",
       " - found 1 nodes to fetch footers for\n",
       "    - `1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8` \n",
       "      - **Search query:** Return paragraphs 7.3 and 7.4 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "### Recursive Retrieval\n",
       "1 nodes with link(s)\n",
       "  - - S 8.8 The CCO must report to senior management o..._0535171ad1204d0c05d28b16a437a17c\n",
       "1 nodes with footer info\n",
       " - 1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8\n",
       "  - **Search query:** Return paragraphs 7.3 and 7.4 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "  - **Returned:** \n",
       "    - - S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3\n",
       "    - - S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd\n",
       "# PASS 4\n",
       "### Context Fetched for nodes:\n",
       "- fetched full clauses, added to 3 nodes. 2 new links\n",
       "- filled nodes with footer context\n",
       "\n",
       "```\n",
       "- node_id: `- S 8.8 The CCO must report to senior management o..._0535171ad1204d0c05d28b16a437a17c`    content: \"- S 8.8 The CCO must report to senior management on a regular basis the findings and analyses of compliance risk. The report must include at minimum-\"    context:\n",
       "      - footer info: 8 This may include validating risk control self-assessment reports submitted to the compliance function.\n",
       "      - further explore footer?: False\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the results of the compliance risk assessment ..._26ef19b6fb4a8ebfebff05e1287c3767`              content: \" (a) the results of the compliance risk assessment undertaken during the assessment period, highlighting key changes in the compliance risk profile of a financial institution as well as areas where greater attention by senior management would be needed;\"          - node_id: `(b) a summary of incidents of non-compliance and d..._201744c9299b9836705b537f22245003`              content: \" (b) a summary of incidents of non-compliance and deficiencies in the management of compliance risk in various parts of the financial institution;\"          - node_id: `(c) an assessment of the impact (both financial an..._195aa161c122d502ac0dea7aa47cf042`              content: \" (c) an assessment of the impact (both financial and non-financial) of such incidents on the financial institution (for example, fines, administrative or\"          - node_id: `- otherwise, or other disciplinary actions taken b..._7f1ae1dd2688b1c823daeb82c389dcd6`              content: \" - otherwise, or other disciplinary actions taken by any regulatory authority in respect of any officers of the financial institution);\"          - node_id: `(d) recommendations of corrective measures to addr..._6b56503ce4ba8da3d9ba5d92c8e4aed9`              content: \" (d) recommendations of corrective measures to address incidents of non- compliance and deficiencies in the management of compliance risk, including disciplinary actions;\"          - node_id: `(e) a record of corrective measures already taken ..._2003efeabd745460f4969ec624fd665d`              content: \" (e) a record of corrective measures already taken and an assessment of the adequacy and effectiveness of such measures; and\"          - node_id: `(f) insights and observations regarding the compli..._dd079c2998c82fa7814742328bbba975`              content: \" (f) insights and observations regarding the compliance culture that exists in the organisation or in specific parts of the organisation that may give rise to compliance concerns.\"- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "\n",
       "```\n",
       "### Supervisor Agent\n",
       "- **supervisor**: continue\n",
       "### Router Agent\n",
       "- found 1 nodes to fetch links for\n",
       "    - - S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472\n",
       " - found 2 nodes to fetch footers for\n",
       "    - `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3` \n",
       "      - **Search query:** Return paragraphs 7.2 and 7.3 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "    - `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c` \n",
       "      - **Search query:** Return paragraph 9.1 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "### Recursive Retrieval\n",
       "1 nodes with link(s)\n",
       "  - - S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472\n",
       "2 nodes with footer info\n",
       " - - S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3\n",
       "  - **Search query:** Return paragraphs 7.2 and 7.3 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "  - **Returned:** \n",
       "    - - S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd\n",
       "    - - S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66\n",
       "    - (b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c\n",
       "    - - S 7.6 Notwithstanding paragraph 7.3, the CCO of ..._614f4361a3bd64ad022d9d007b16da05\n",
       "    - (c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864\n",
       " - (b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c\n",
       "  - **Search query:** Return paragraph 9.1 to answer the query 'How can the Board and the CCO manage control functions?'\n",
       "  - **Returned:** \n",
       "    - - S 7.6 Notwithstanding paragraph 7.3, the CCO of ..._614f4361a3bd64ad022d9d007b16da05\n",
       "    - - 1.1 A strong compliance culture reflects a corpo..._1fe85abd327658de88e7eed861b59a93\n",
       "    - - S 7.10 Where such arrangements referred to in pa..._9aee9c464ce513503b26252dc35e3ad1\n",
       "    - (b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c\n",
       "    - (c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864\n",
       "# PASS 5\n",
       "### Context Fetched for nodes:\n",
       "- fetched full clauses, added to 7 nodes. 3 new links\n",
       "- filled nodes with footer context\n",
       "\n",
       "```\n",
       "- node_id: `- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472`    content: \"- S 9.1 A financial institution must ensure that there is a clear separation of the internal audit function and other functions carrying out compliance function responsibilities. Compliance risk must be included in the risk assessment methodology of the internal audit function, and an audit programme that covers the adequacy and effectiveness of the other functions carrying out compliance function responsibilities should be established, including testing of controls commensurate with the perceived level of risk.\"- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66`    content: \"- S 7.2 Where compliance function responsibilities are shared between a dedicated compliance unit and other control functions-\"    children: \n",
       "          - node_id: `(a) the allocation of the compliance function resp..._32a2cb84e0fdce2e66fc429c12918716`              content: \" (a) the allocation of the compliance function responsibilities, including that for timely communication and escalation of compliance issues to senior management and the board, must be clearly defined and documented in the compliance policy;\"          - node_id: `(b) the CCO must have overall responsibility for c..._d1a0a67335e5aa245b6eeeaa082e845d`              content: \" (b) the CCO must have overall responsibility for coordinating the identification and management of compliance risk at the institution-wide level, and ensuring that compliance monitoring and testing are carried out consistently across the institution;\"          - node_id: `(c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864`              content: \" (c) for the purpose of paragraph 7.2(b), the CCO must have a sound understanding of compliance risks which are under the purview of other control functions, including an understanding of controls applied to manage these risks;\"          - node_id: `(d) arrangements for coordination among the contro..._2c7647be1e92b542fc76be5e0a9e14a4`              content: \" (d) arrangements for coordination among the control functions with the CCO and the compliance unit must be in place to ensure that the CCO is able to perform his responsibilities effectively. Arrangements should promote a view and approach to the management of compliance risk that is consistent across the organisation, including through adequate information flows and avenues to seek advice on compliance issues; and\"          - node_id: `(e) senior management must ensure that officers in..._be843c24dda3f19f1ae16dcb3ee0acd0`              content: \" (e) senior management must ensure that officers in other control functions have the capacity and expertise to deliver their broader mandates while providing adequate focus to their compliance function responsibilities.\"- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.6 Notwithstanding paragraph 7.3, the CCO of ..._614f4361a3bd64ad022d9d007b16da05`    content: \"- S 7.6 Notwithstanding paragraph 7.3, the CCO of a large financial institution is not allowed to assume the responsibilities of other control functions.\"    children: \n",
       "          - node_id: `- Independence_7e6261e9f5d7eb958e43140add281d2e`              content: \" - Independence\"- node_id: `- 1.1 A strong compliance culture reflects a corpo..._1fe85abd327658de88e7eed861b59a93`    content: \"- 1.1 A strong compliance culture reflects a corporate culture of high ethical standards and integrity in which the board and senior management lead by example. A financial institution should hold itself to high standards in carrying on business, and at all times observe both the spirit and the letter of the law and regulations. Failure to effectively manage compliance risk may result in adverse consequences for the financial institution's customers, shareholders, officers and the financial institution itself.\"- node_id: `- S 7.10 Where such arrangements referred to in pa..._9aee9c464ce513503b26252dc35e3ad1`    content: \"- S 7.10 Where such arrangements referred to in paragraph 7.9 exist, a financial institution must ensure that-\"    context:\n",
       "      - link: [\"['- G 7.9 A constructive and cooperative working rel..._081d34781b0f7ab41f94dbe52896a23e']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the compliance function is not placed in a pos..._03f9f972a4a09290c72409b0826aa893`              content: \" (a) the compliance function is not placed in a position of conflict;\"          - node_id: `(b) the accountability of the compliance function ..._e301200856ef131fd6a08194a2532e5e`              content: \" (b) the accountability of the compliance function is properly documented6 and\"              context:\n",
       "                - footer info: 6 In the case of a committee, the role of the compliance function may be outlined in the terms of reference or charter.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(c) the compliance function is not prevented from ..._8f4f0fef22f99ea2f589709d5b54496e`              content: \" (c) the compliance function is not prevented from highlighting compliance issues relating to any business decisions to the board or senior management, where necessary.\"          - node_id: `- Resources_0df524eaa89ae4d09e3b49998eb3b771`              content: \" - Resources\"\n",
       "```\n",
       "### Supervisor Agent\n",
       "### Answering Agent\n",
       "## Answer to Query: \n",
       "The Board and the Chief Compliance Officer (CCO) can manage control functions by following several guidelines:\n",
       "\n",
       "1. **Board Responsibilities**:\n",
       "   - The board must approve the key aspects of the CCO’s role, including appointment, remuneration, and termination ([S 6.3(a)](- S 6.3)).\n",
       "   - They should ensure that the CCO has the stature for effective engagement with the CEO and senior management ([S 6.3(b)](- S 6.3)).\n",
       "   - The board should engage with the CCO regularly and occasionally without other senior management present to discuss compliance issues ([S 6.3(c)](- S 6.3)).\n",
       "\n",
       "2. **CCO Resources and Independence**:\n",
       "   - The CCO must be provided with direct access to the board and sufficient resources to perform duties effectively ([S 6.3(d), (e)](- S 6.3)).\n",
       "   - If the CCO holds responsibilities in other control areas, the board should ensure the overall control environment remains sound ([S 6.3(f)](- S 6.3)).\n",
       "\n",
       "3. **Control Function Sharing**:\n",
       "   - The board must approve the sharing of compliance function responsibilities between a dedicated unit and other control functions ([S 7.4(a)](- S 7.4)).\n",
       "   - The CCO's other responsibilities should not compromise his independence or ability to focus adequately on compliance ([S 7.3](- S 7.3)).\n",
       "\n",
       "4. **Coordination and Documentation**:\n",
       "   - There needs to be clear documentation of compliance responsibilities and maintenance of a consistent compliance risk management approach across the organization ([S 7.2(a)](- S 7.2)).\n",
       "\n",
       "5. **Independence from Internal Audit**:\n",
       "   - Separation must be maintained between the compliance function and internal audit to preserve the integrity of independent reviews ([S 7.4(b)](- S 7.4), [S 9.1](- S 9.1)).\n",
       "\n",
       "These mechanisms ensure that both the board and the CCO effectively manage and oversee compliance and other control functions within an organization.\n",
       "## Source Nodes: \n",
       "```\n",
       "- node_id: `- S 6.3 In relation to the position of the CCO, th..._d9432f9931c5bf206db8317a67e69052`    content: \"- S 6.3 In relation to the position of the CCO, the board must-\"    children: \n",
       "          - node_id: `(a) approve the appointment, remuneration and term..._f5abf05a077c9b80dd1a4a0a793064c8`              content: \" (a) approve the appointment, remuneration and termination of the CCO¹;\"              context:\n",
       "                - footer info: 1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(b) ensure that the CCO has sufficient stature to ..._ade5bcb25050d514e6d10e163c1581ec`              content: \" (b) ensure that the CCO has sufficient stature to allow for effective engagement with the CEO and other members of senior management;\"          - node_id: `(c) engage with the CCO on a regular basis2 to pro..._75d9df850bf8422349822accdcb9338e`              content: \" (c) engage with the CCO on a regular basis2 to provide the opportunity for the CCO to discuss issues faced by the compliance function;\"              context:\n",
       "                - footer info: 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(d) provide the CCO with direct and unimpeded acce..._753739a7f2c2cf4c206b6e102cf20aca`              content: \" (d) provide the CCO with direct and unimpeded access to the board;\"          - node_id: `(e) ensure that the CCO is supported with sufficie..._b5cb66fa679baafeacc244e3dbe13938`              content: \" (e) ensure that the CCO is supported with sufficient resources, including competent officers, to perform his duties effectively; and\"          - node_id: `(f) where the CCO also carries out responsibilitie..._684a6af3d1c0785ec56eeb2cbe6c1644`              content: \" (f) where the CCO also carries out responsibilities in respect of other control functions3, be satisfied that a sound overall control environment will not be compromised by the combination of responsibilities performed by the CCO.\"              context:\n",
       "                - footer info: 3 Refer to paragraphs 7.3 and 7.4.\n",
       "                - further explore footer?: True\n",
       "\n",
       "- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"- node_id: `- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472`    content: \"- S 9.1 A financial institution must ensure that there is a clear separation of the internal audit function and other functions carrying out compliance function responsibilities. Compliance risk must be included in the risk assessment methodology of the internal audit function, and an audit programme that covers the adequacy and effectiveness of the other functions carrying out compliance function responsibilities should be established, including testing of controls commensurate with the perceived level of risk.\"    context:\n",
       "      - footer info: Issued on: 10 May 2016\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `1 Refer to the policy documents on Fit and Proper ..._f4de0fe932a8ef8aabbd57b80f39f1b8`    content: \"1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011. 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time. 3 Refer to paragraphs 7.3 and 7.4.\"    context:\n",
       "      - footer info: 1 Refer to the policy documents on Fit and Proper Criteria issued on 28 June 2013 and Guidelines on Fit and Proper for Key Responsible Persons for Development Financial Institutions issued on 15 September 2011. 2 The board should also consider engaging the CCO without the presence of other members of senior management from time to time. 3 Refer to paragraphs 7.3 and 7.4.\n",
       "      - further explore footer?: True\n",
       "\n",
       "- node_id: `- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66`    content: \"- S 7.2 Where compliance function responsibilities are shared between a dedicated compliance unit and other control functions-\"    children: \n",
       "          - node_id: `(a) the allocation of the compliance function resp..._32a2cb84e0fdce2e66fc429c12918716`              content: \" (a) the allocation of the compliance function responsibilities, including that for timely communication and escalation of compliance issues to senior management and the board, must be clearly defined and documented in the compliance policy;\"          - node_id: `(b) the CCO must have overall responsibility for c..._d1a0a67335e5aa245b6eeeaa082e845d`              content: \" (b) the CCO must have overall responsibility for coordinating the identification and management of compliance risk at the institution-wide level, and ensuring that compliance monitoring and testing are carried out consistently across the institution;\"          - node_id: `(c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864`              content: \" (c) for the purpose of paragraph 7.2(b), the CCO must have a sound understanding of compliance risks which are under the purview of other control functions, including an understanding of controls applied to manage these risks;\"          - node_id: `(d) arrangements for coordination among the contro..._2c7647be1e92b542fc76be5e0a9e14a4`              content: \" (d) arrangements for coordination among the control functions with the CCO and the compliance unit must be in place to ensure that the CCO is able to perform his responsibilities effectively. Arrangements should promote a view and approach to the management of compliance risk that is consistent across the organisation, including through adequate information flows and avenues to seek advice on compliance issues; and\"          - node_id: `(e) senior management must ensure that officers in..._be843c24dda3f19f1ae16dcb3ee0acd0`              content: \" (e) senior management must ensure that officers in other control functions have the capacity and expertise to deliver their broader mandates while providing adequate focus to their compliance function responsibilities.\"              context:\n",
       "                - footer info: 5 For example, through representation on a new products committee.\n",
       "                - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.6 Notwithstanding paragraph 7.3, the CCO of ..._614f4361a3bd64ad022d9d007b16da05`    content: \"- S 7.6 Notwithstanding paragraph 7.3, the CCO of a large financial institution is not allowed to assume the responsibilities of other control functions.\"    children: \n",
       "          - node_id: `- Independence_7e6261e9f5d7eb958e43140add281d2e`              content: \" - Independence\"- node_id: `- 1.1 A strong compliance culture reflects a corpo..._1fe85abd327658de88e7eed861b59a93`    content: \"- 1.1 A strong compliance culture reflects a corporate culture of high ethical standards and integrity in which the board and senior management lead by example. A financial institution should hold itself to high standards in carrying on business, and at all times observe both the spirit and the letter of the law and regulations. Failure to effectively manage compliance risk may result in adverse consequences for the financial institution's customers, shareholders, officers and the financial institution itself.\"- node_id: `- S 7.10 Where such arrangements referred to in pa..._9aee9c464ce513503b26252dc35e3ad1`    content: \"- S 7.10 Where such arrangements referred to in paragraph 7.9 exist, a financial institution must ensure that-\"    context:\n",
       "      - link: [\"['- G 7.9 A constructive and cooperative working rel..._081d34781b0f7ab41f94dbe52896a23e']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the compliance function is not placed in a pos..._03f9f972a4a09290c72409b0826aa893`              content: \" (a) the compliance function is not placed in a position of conflict;\"          - node_id: `(b) the accountability of the compliance function ..._e301200856ef131fd6a08194a2532e5e`              content: \" (b) the accountability of the compliance function is properly documented6 and\"              context:\n",
       "                - footer info: 6 In the case of a committee, the role of the compliance function may be outlined in the terms of reference or charter.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(c) the compliance function is not prevented from ..._8f4f0fef22f99ea2f589709d5b54496e`              content: \" (c) the compliance function is not prevented from highlighting compliance issues relating to any business decisions to the board or senior management, where necessary.\"          - node_id: `- Resources_0df524eaa89ae4d09e3b49998eb3b771`              content: \" - Resources\"- node_id: `Compliance_95521f3c1af881ea7cf1d71ec7c00c9a`    content: \"Compliance\"- node_id: `- S 8.9 The CCO must ensure that the reports refer..._6cf04fc7a955f2d7df8ea100ec7af539`    content: \"- S 8.9 The CCO must ensure that the reports referred to in paragraph 8.8 are readily available to the internal audit function of the financial institution, the Bank and other regulatory authorities upon request.\"    context:\n",
       "      - link: [\"['- S 8.8 The CCO must report to senior management o..._0535171ad1204d0c05d28b16a437a17c']\"]\n",
       "      - footer info: Issued on: 10 May 2016\n",
       "      - further explore footer?: False\n",
       "\n",
       "    children: \n",
       "          - node_id: `- Advisory_57633bcdded147683c36017c16c03fc8`              content: \" - Advisory\"- node_id: `- S 8.8 The CCO must report to senior management o..._0535171ad1204d0c05d28b16a437a17c`    content: \"- S 8.8 The CCO must report to senior management on a regular basis the findings and analyses of compliance risk. The report must include at minimum-\"    context:\n",
       "      - footer info: 8 This may include validating risk control self-assessment reports submitted to the compliance function.\n",
       "      - further explore footer?: False\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the results of the compliance risk assessment ..._26ef19b6fb4a8ebfebff05e1287c3767`              content: \" (a) the results of the compliance risk assessment undertaken during the assessment period, highlighting key changes in the compliance risk profile of a financial institution as well as areas where greater attention by senior management would be needed;\"          - node_id: `(b) a summary of incidents of non-compliance and d..._201744c9299b9836705b537f22245003`              content: \" (b) a summary of incidents of non-compliance and deficiencies in the management of compliance risk in various parts of the financial institution;\"          - node_id: `(c) an assessment of the impact (both financial an..._195aa161c122d502ac0dea7aa47cf042`              content: \" (c) an assessment of the impact (both financial and non-financial) of such incidents on the financial institution (for example, fines, administrative or\"          - node_id: `- otherwise, or other disciplinary actions taken b..._7f1ae1dd2688b1c823daeb82c389dcd6`              content: \" - otherwise, or other disciplinary actions taken by any regulatory authority in respect of any officers of the financial institution);\"          - node_id: `(d) recommendations of corrective measures to addr..._6b56503ce4ba8da3d9ba5d92c8e4aed9`              content: \" (d) recommendations of corrective measures to address incidents of non- compliance and deficiencies in the management of compliance risk, including disciplinary actions;\"          - node_id: `(e) a record of corrective measures already taken ..._2003efeabd745460f4969ec624fd665d`              content: \" (e) a record of corrective measures already taken and an assessment of the adequacy and effectiveness of such measures; and\"          - node_id: `(f) insights and observations regarding the compli..._dd079c2998c82fa7814742328bbba975`              content: \" (f) insights and observations regarding the compliance culture that exists in the organisation or in specific parts of the organisation that may give rise to compliance concerns.\"- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472`    content: \"- S 9.1 A financial institution must ensure that there is a clear separation of the internal audit function and other functions carrying out compliance function responsibilities. Compliance risk must be included in the risk assessment methodology of the internal audit function, and an audit programme that covers the adequacy and effectiveness of the other functions carrying out compliance function responsibilities should be established, including testing of controls commensurate with the perceived level of risk.\"- node_id: `- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd`    content: \"- S 7.3 Where the CCO also assumes responsibilities for other control functions, the CCO must ensure that his independence and ability to provide sufficient time, focus and commitment to his responsibilities in respect of the compliance function is not impaired.\"    context:\n",
       "      - footer info: 5 For example, through representation on a new products committee.\n",
       "      - further explore footer?: False\n",
       "\n",
       "- node_id: `- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66`    content: \"- S 7.2 Where compliance function responsibilities are shared between a dedicated compliance unit and other control functions-\"    children: \n",
       "          - node_id: `(a) the allocation of the compliance function resp..._32a2cb84e0fdce2e66fc429c12918716`              content: \" (a) the allocation of the compliance function responsibilities, including that for timely communication and escalation of compliance issues to senior management and the board, must be clearly defined and documented in the compliance policy;\"          - node_id: `(b) the CCO must have overall responsibility for c..._d1a0a67335e5aa245b6eeeaa082e845d`              content: \" (b) the CCO must have overall responsibility for coordinating the identification and management of compliance risk at the institution-wide level, and ensuring that compliance monitoring and testing are carried out consistently across the institution;\"          - node_id: `(c) for the purpose of paragraph 7.2(b), the CCO m..._ef2ecf18947fbf24d2f59a7ad3967864`              content: \" (c) for the purpose of paragraph 7.2(b), the CCO must have a sound understanding of compliance risks which are under the purview of other control functions, including an understanding of controls applied to manage these risks;\"          - node_id: `(d) arrangements for coordination among the contro..._2c7647be1e92b542fc76be5e0a9e14a4`              content: \" (d) arrangements for coordination among the control functions with the CCO and the compliance unit must be in place to ensure that the CCO is able to perform his responsibilities effectively. Arrangements should promote a view and approach to the management of compliance risk that is consistent across the organisation, including through adequate information flows and avenues to seek advice on compliance issues; and\"          - node_id: `(e) senior management must ensure that officers in..._be843c24dda3f19f1ae16dcb3ee0acd0`              content: \" (e) senior management must ensure that officers in other control functions have the capacity and expertise to deliver their broader mandates while providing adequate focus to their compliance function responsibilities.\"- node_id: `- S 7.4 In relation to paragraphs 7.2 and 7.3-_4dbea03fb547a6c5b4e4e8160f226fb3`    content: \"- S 7.4 In relation to paragraphs 7.2 and 7.3-\"    context:\n",
       "      - link: [\"['- S 7.2 Where compliance function responsibilities..._f9fb2b33eabb778bf7f7073b6acffb66', '- S 7.3 Where the CCO also assumes responsibilitie..._4ffb5b3a6776c69a888ebce7a3630bcd']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the board must approve the sharing of complian..._c98e9a54d02d6b13b61c5d1ed01ff30f`              content: \" (a) the board must approve the sharing of compliance function responsibilities between a dedicated compliance unit and other control functions; and\"          - node_id: `(b) compliance function responsibilities cannot be..._bd4533882412131b8d98e84499fe7d9c`              content: \" (b) compliance function responsibilities cannot be shared with, nor can the CCO assume responsibilities for, internal audit as such practices would render the independent review process described in paragraph 9.1 ineffective.\"              context:\n",
       "                - link: [\"['- S 9.1 A financial institution must ensure that t..._98d7434318b67e1256f563bf7fd13472']\"]\n",
       "\n",
       "- node_id: `- S 7.6 Notwithstanding paragraph 7.3, the CCO of ..._614f4361a3bd64ad022d9d007b16da05`    content: \"- S 7.6 Notwithstanding paragraph 7.3, the CCO of a large financial institution is not allowed to assume the responsibilities of other control functions.\"    children: \n",
       "          - node_id: `- Independence_7e6261e9f5d7eb958e43140add281d2e`              content: \" - Independence\"- node_id: `- 1.1 A strong compliance culture reflects a corpo..._1fe85abd327658de88e7eed861b59a93`    content: \"- 1.1 A strong compliance culture reflects a corporate culture of high ethical standards and integrity in which the board and senior management lead by example. A financial institution should hold itself to high standards in carrying on business, and at all times observe both the spirit and the letter of the law and regulations. Failure to effectively manage compliance risk may result in adverse consequences for the financial institution's customers, shareholders, officers and the financial institution itself.\"- node_id: `- S 7.10 Where such arrangements referred to in pa..._9aee9c464ce513503b26252dc35e3ad1`    content: \"- S 7.10 Where such arrangements referred to in paragraph 7.9 exist, a financial institution must ensure that-\"    context:\n",
       "      - link: [\"['- G 7.9 A constructive and cooperative working rel..._081d34781b0f7ab41f94dbe52896a23e']\"]\n",
       "\n",
       "    children: \n",
       "          - node_id: `(a) the compliance function is not placed in a pos..._03f9f972a4a09290c72409b0826aa893`              content: \" (a) the compliance function is not placed in a position of conflict;\"          - node_id: `(b) the accountability of the compliance function ..._e301200856ef131fd6a08194a2532e5e`              content: \" (b) the accountability of the compliance function is properly documented6 and\"              context:\n",
       "                - footer info: 6 In the case of a committee, the role of the compliance function may be outlined in the terms of reference or charter.\n",
       "                - further explore footer?: False\n",
       "\n",
       "          - node_id: `(c) the compliance function is not prevented from ..._8f4f0fef22f99ea2f589709d5b54496e`              content: \" (c) the compliance function is not prevented from highlighting compliance issues relating to any business decisions to the board or senior management, where necessary.\"          - node_id: `- Resources_0df524eaa89ae4d09e3b49998eb3b771`              content: \" - Resources\"\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "InvalidUpdateError",
     "evalue": "At key 'query': Can receive only one value per step. Use an Annotated key to handle multiple values.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mInvalidUpdateError\u001b[0m                        Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[23], line 536\u001b[0m\n\u001b[1;32m    532\u001b[0m         \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    533\u001b[0m             \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[0;32m--> 536\u001b[0m \u001b[43mmulti_agent_query\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[23], line 527\u001b[0m, in \u001b[0;36mmulti_agent_query\u001b[0;34m(query)\u001b[0m\n\u001b[1;32m    512\u001b[0m initial_state \u001b[38;5;241m=\u001b[39m AgentState(\n\u001b[1;32m    513\u001b[0m     query\u001b[38;5;241m=\u001b[39mquery,\n\u001b[1;32m    514\u001b[0m     definitions\u001b[38;5;241m=\u001b[39m{},\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    522\u001b[0m     pass_count\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m    523\u001b[0m )\n\u001b[1;32m    525\u001b[0m events \u001b[38;5;241m=\u001b[39m graph\u001b[38;5;241m.\u001b[39mstream(initial_state, {\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrecursion_limit\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;241m150\u001b[39m})\n\u001b[0;32m--> 527\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43ms\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mevents\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m    528\u001b[0m \u001b[43m    \u001b[49m\u001b[43mclear_output\u001b[49m\u001b[43m(\u001b[49m\u001b[43mwait\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    529\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfirst_value\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mnext\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43miter\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43ms\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalues\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Documents/GitHub/github-app/recursive-retrieval/venv/lib/python3.11/site-packages/langgraph/pregel/__init__.py:935\u001b[0m, in \u001b[0;36mPregel.stream\u001b[0;34m(self, input, config, stream_mode, output_keys, interrupt_before, interrupt_after, debug)\u001b[0m\n\u001b[1;32m    904\u001b[0m (\n\u001b[1;32m    905\u001b[0m     debug,\n\u001b[1;32m    906\u001b[0m     stream_modes,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    917\u001b[0m     debug\u001b[38;5;241m=\u001b[39mdebug,\n\u001b[1;32m    918\u001b[0m )\n\u001b[1;32m    920\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m SyncPregelLoop(\n\u001b[1;32m    921\u001b[0m     \u001b[38;5;28minput\u001b[39m,\n\u001b[1;32m    922\u001b[0m     config\u001b[38;5;241m=\u001b[39mconfig,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    933\u001b[0m     \u001b[38;5;66;03m# channels are guaranteed to be immutable for the duration of the step,\u001b[39;00m\n\u001b[1;32m    934\u001b[0m     \u001b[38;5;66;03m# with channel updates applied only at the transition between steps\u001b[39;00m\n\u001b[0;32m--> 935\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[43mloop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtick\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    936\u001b[0m \u001b[43m        \u001b[49m\u001b[43minput_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minput_channels\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    937\u001b[0m \u001b[43m        \u001b[49m\u001b[43minterrupt_before\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minterrupt_before\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    938\u001b[0m \u001b[43m        \u001b[49m\u001b[43minterrupt_after\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minterrupt_after\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    939\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmanager\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    940\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m:\n\u001b[1;32m    941\u001b[0m         \u001b[38;5;66;03m# debug flag\u001b[39;00m\n\u001b[1;32m    942\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m debug:\n\u001b[1;32m    943\u001b[0m             print_step_checkpoint(\n\u001b[1;32m    944\u001b[0m                 loop\u001b[38;5;241m.\u001b[39mcheckpoint_metadata,\n\u001b[1;32m    945\u001b[0m                 loop\u001b[38;5;241m.\u001b[39mchannels,\n\u001b[1;32m    946\u001b[0m                 \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstream_channels_list,\n\u001b[1;32m    947\u001b[0m             )\n",
      "File \u001b[0;32m~/Documents/GitHub/github-app/recursive-retrieval/venv/lib/python3.11/site-packages/langgraph/pregel/loop.py:210\u001b[0m, in \u001b[0;36mPregelLoop.tick\u001b[0;34m(self, input_keys, interrupt_after, interrupt_before, manager)\u001b[0m\n\u001b[1;32m    208\u001b[0m writes \u001b[38;5;241m=\u001b[39m [w \u001b[38;5;28;01mfor\u001b[39;00m t \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtasks \u001b[38;5;28;01mfor\u001b[39;00m w \u001b[38;5;129;01min\u001b[39;00m t\u001b[38;5;241m.\u001b[39mwrites]\n\u001b[1;32m    209\u001b[0m \u001b[38;5;66;03m# all tasks have finished\u001b[39;00m\n\u001b[0;32m--> 210\u001b[0m mv_writes \u001b[38;5;241m=\u001b[39m \u001b[43mapply_writes\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    211\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheckpoint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    212\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mchannels\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    213\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtasks\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    214\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheckpointer_get_next_version\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    215\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    216\u001b[0m \u001b[38;5;66;03m# apply writes to managed values\u001b[39;00m\n\u001b[1;32m    217\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, values \u001b[38;5;129;01min\u001b[39;00m mv_writes\u001b[38;5;241m.\u001b[39mitems():\n",
      "File \u001b[0;32m~/Documents/GitHub/github-app/recursive-retrieval/venv/lib/python3.11/site-packages/langgraph/pregel/algo.py:211\u001b[0m, in \u001b[0;36mapply_writes\u001b[0;34m(checkpoint, channels, tasks, get_next_version)\u001b[0m\n\u001b[1;32m    209\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m chan, vals \u001b[38;5;129;01min\u001b[39;00m pending_writes_by_channel\u001b[38;5;241m.\u001b[39mitems():\n\u001b[1;32m    210\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m chan \u001b[38;5;129;01min\u001b[39;00m channels:\n\u001b[0;32m--> 211\u001b[0m         updated \u001b[38;5;241m=\u001b[39m \u001b[43mchannels\u001b[49m\u001b[43m[\u001b[49m\u001b[43mchan\u001b[49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mupdate\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    212\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m updated \u001b[38;5;129;01mand\u001b[39;00m get_next_version \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    213\u001b[0m             checkpoint[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mchannel_versions\u001b[39m\u001b[38;5;124m\"\u001b[39m][chan] \u001b[38;5;241m=\u001b[39m get_next_version(\n\u001b[1;32m    214\u001b[0m                 max_version, channels[chan]\n\u001b[1;32m    215\u001b[0m             )\n",
      "File \u001b[0;32m~/Documents/GitHub/github-app/recursive-retrieval/venv/lib/python3.11/site-packages/langgraph/channels/last_value.py:55\u001b[0m, in \u001b[0;36mLastValue.update\u001b[0;34m(self, values)\u001b[0m\n\u001b[1;32m     53\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(values) \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m---> 55\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m InvalidUpdateError(\n\u001b[1;32m     56\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAt key \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkey\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m: Can receive only one value per step. Use an Annotated key to handle multiple values.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     57\u001b[0m     )\n\u001b[1;32m     59\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvalue \u001b[38;5;241m=\u001b[39m values[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\n\u001b[1;32m     60\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "\u001b[0;31mInvalidUpdateError\u001b[0m: At key 'query': Can receive only one value per step. Use an Annotated key to handle multiple values."
     ]
    }
   ],
   "source": [
    "# Define the State\n",
    "import operator\n",
    "from typing import Annotated, Sequence, TypedDict\n",
    "from langchain_core.messages import BaseMessage\n",
    "from langgraph.graph import END, StateGraph, START\n",
    "\n",
    "import tiktoken\n",
    "\n",
    "tokenizer = tiktoken.get_encoding(\"o200k_base\")\n",
    "\n",
    "TOKEN_LIMIT = 32000\n",
    "\n",
    "\n",
    "def count_tokens(text: str) -> int:\n",
    "    return len(tokenizer.encode(text))\n",
    "\n",
    "\n",
    "def count_tokens_for_nodes(nodes: List[MultiAgentSearchLocalNode]) -> int:\n",
    "    return sum(count_tokens(node.print_node_prompt()) for node in nodes)\n",
    "\n",
    "\n",
    "debug_markdown = \"# Multi-Agent Log\\n\"\n",
    "debug_markdown += \"## Pass 1\\n\"\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    # we keep track of existing messages\n",
    "    query: str\n",
    "    definitions: Dict[str, str]\n",
    "    previous_nodes: List[MultiAgentSearchLocalNode]\n",
    "    last_fetched_context_nodes: List[MultiAgentSearchLocalNode]\n",
    "    node_links_to_fetch: List[str]\n",
    "    node_footers_to_fetch: Dict[str, str]\n",
    "    allow_continue: bool\n",
    "    search_failures: List[str]\n",
    "    markdown_debug: str\n",
    "    pass_count: int\n",
    "\n",
    "\n",
    "def inital_search_agent(state: AgentState) -> AgentState:\n",
    "    state[\"markdown_debug\"] += f\"### Initial Search Agent\\n\"\n",
    "    \n",
    "    vector_bm25_retriever = VectorBM25(vector_top_k=1, bm25_top_k=7, mode=\"OR\")\n",
    "    def retrieve_with_vector_bm25(query, verbose = False):\n",
    "        retrieved_nodes: NodeWithScore = vector_bm25_retriever.retrieve(\n",
    "            QueryBundle(query_str=query)\n",
    "        )\n",
    "        retrieved_nodes_local_form: MultiAgentSearchLocalNode = [\n",
    "            local_nodes_map[node.node.node_id] for node in retrieved_nodes\n",
    "        ]\n",
    "        if verbose:\n",
    "            for node in retrieved_nodes:\n",
    "                print(node)\n",
    "        return retrieved_nodes_local_form\n",
    "\n",
    "    # fetch nodes by vector and bm25 search\n",
    "    initial_retrieved_nodes = retrieve_with_vector_bm25(state[\"query\"])\n",
    "\n",
    "    # LLM prunes nodes that are not relevant\n",
    "    initial_pruned_nodes, _ = prune_nodes(state[\"query\"], initial_retrieved_nodes)\n",
    "\n",
    "    state[\n",
    "        \"markdown_debug\"\n",
    "    ] += f\" - retrieved {len(initial_retrieved_nodes)} nodes, pruned to {len(initial_pruned_nodes)} nodes\\n\"\n",
    "\n",
    "    state[\"last_fetched_context_nodes\"] = initial_pruned_nodes\n",
    "    return state\n",
    "\n",
    "\n",
    "def definition_agent(state: AgentState) -> AgentState:\n",
    "    # fetch definitions\n",
    "    state[\"markdown_debug\"] += f\"### Definitions Agent\\n\"\n",
    "    print(f\"query: {state['query']}\")\n",
    "    definitions_dict = definitions_search(state[\"query\"])\n",
    "    print(f\"dfns: {definitions_dict}\")\n",
    "    state[\"definitions\"] = definitions_dict\n",
    "\n",
    "    state[\"markdown_debug\"] += f\"Retrieved {len(definitions_dict)} definitions for: \\n\"\n",
    "    for term, definition in definitions_dict.items():\n",
    "        state[\"markdown_debug\"] += f\"  - **{term}** : {definition}\\n\"\n",
    "    return state\n",
    "\n",
    "\n",
    "def context_fetch_tool(state: AgentState) -> AgentState:\n",
    "    state[\"markdown_debug\"] += f\"### Context Fetched for nodes:\\n\"\n",
    "    current_nodes = state[\"last_fetched_context_nodes\"]\n",
    "\n",
    "    # expand search to get full clauses\n",
    "    full_clause_nodes = fetch_whole_clause_lists(current_nodes)\n",
    "\n",
    "    # mark links with nodes\n",
    "    full_clause_nodes_with_link_hints, change_count = fill_nodes_with_link_hints(\n",
    "        full_clause_nodes\n",
    "    )\n",
    "\n",
    "    state[\"last_fetched_context_nodes\"] = full_clause_nodes_with_link_hints\n",
    "    state[\n",
    "        \"markdown_debug\"\n",
    "    ] += f\"- fetched full clauses, added to {len(full_clause_nodes_with_link_hints)} nodes. {change_count} new links\\n\"\n",
    "\n",
    "    return state\n",
    "\n",
    "\n",
    "def footer_parsing_agent(state: AgentState) -> AgentState:\n",
    "    # fill nodes with footer context\n",
    "    try:\n",
    "        current_nodes, changes_count = fill_nodes_with_footer_context(\n",
    "            state[\"query\"], state[\"last_fetched_context_nodes\"]\n",
    "        )\n",
    "\n",
    "        state[\"last_fetched_context_nodes\"] = current_nodes\n",
    "        state[\"markdown_debug\"] += f\"- filled nodes with footer context\\n\"\n",
    "\n",
    "        state[\"markdown_debug\"] += f\"\\n```\\n\"\n",
    "        for node in current_nodes:\n",
    "            state[\"markdown_debug\"] += node.print_node_prompt()\n",
    "\n",
    "        state[\"markdown_debug\"] += f\"\\n```\\n\"\n",
    "        \n",
    "    except:\n",
    "        state[\"search_failures\"].append(\"footer parsing failed.\")\n",
    "        state[\"markdown_debug\"] += f\"- footer parsing failed\\n\"\n",
    "\n",
    "    return state\n",
    "\n",
    "\n",
    "def supervisor_agent(state: AgentState) -> AgentState:\n",
    "    state[\"markdown_debug\"] += f\"### Supervisor Agent\\n\"\n",
    "    # Look for search failures. This might be an instance where multiple searches were made for certain parts of the document, but no relevant information was found.\n",
    "    # This means that the search has to be ended prematurely to prevent infinite loops.\n",
    "    printout = \"\"\n",
    "    for node in state[\"previous_nodes\"]:\n",
    "        printout += node.print_node_prompt()\n",
    "    for node in state[\"last_fetched_context_nodes\"]:\n",
    "        printout += node.print_node_prompt()\n",
    "\n",
    "    prompt = f\"\"\"\n",
    "You are a supervisor agent overseeing the multi-agent retrieval process of graph nodes from a document. The nodes are to answer the query:\n",
    "```{state['query']}```\n",
    "\n",
    "\n",
    "Below is a list of nodes that were automatically retrieved, followed by a list of errors. If there are many similar, repeated errors in the retrieval process , where no further linked or relevant nodes could be retrieved, return END to end the process. Else return CONTINUE. \n",
    "Return only a single word, either END or CONTINUE.\n",
    "\"\"\"\n",
    "\n",
    "    completion = openai_client.beta.chat.completions.parse(\n",
    "        model=\"gpt-4o-2024-08-06\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": prompt},\n",
    "            {\"role\": \"user\", \"content\": printout},\n",
    "            {\"role\": \"user\", \"content\": str(state[\"search_failures\"])},\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    response: str = completion.choices[0].message.content\n",
    "\n",
    "    if \"end\" in response.lower():\n",
    "        state[\"allow_continue\"] = False\n",
    "        return state\n",
    "    state[\"markdown_debug\"] += f\"- **supervisor**: continue\\n\"\n",
    "\n",
    "    # check if requests will be in token limit\n",
    "    # chatgpt is 32k, use 32k as baseline\n",
    "    tokens = count_tokens_for_nodes(state[\"previous_nodes\"]) + count_tokens_for_nodes(\n",
    "        state[\"last_fetched_context_nodes\"]\n",
    "    )\n",
    "    if tokens > TOKEN_LIMIT:\n",
    "        state[\"markdown_debug\"] += f\"- warning: tokens over recommended limit\\n\"\n",
    "        previous_nodes, changes_1 = prune_nodes(state[\"query\"], state[\"previous_nodes\"])\n",
    "        last_fetched_context_nodes, changes_2 = prune_nodes(\n",
    "            state[\"query\"], state[\"last_fetched_context_nodes\"]\n",
    "        )\n",
    "        new_tokens = count_tokens_for_nodes(state[\"previous_nodes\"]) + count_tokens(\n",
    "            state[\"last_fetched_context_nodes\"]\n",
    "        )\n",
    "\n",
    "        if new_tokens > TOKEN_LIMIT:\n",
    "            # raise ValueError(\"Too many tokens\")\n",
    "            print(\"Warning: Tried pruning, but still too many tokens\")\n",
    "            state[\n",
    "                \"markdown_debug\"\n",
    "            ] += f\"- Warning: Tried pruning, but still too many tokens\\n\"\n",
    "        if new_tokens < tokens:\n",
    "            state[\"previous_nodes\"] = previous_nodes\n",
    "            state[\"last_fetched_context_nodes\"] = last_fetched_context_nodes\n",
    "\n",
    "            for node_id, reason in changes_1.items():\n",
    "                state[\"markdown_debug\"] += f\"- pruned node {node_id} due to {reason}\\n\"\n",
    "\n",
    "            for node_id, reason in changes_2.items():\n",
    "                state[\"markdown_debug\"] += f\"- pruned node {node_id} due to {reason}\\n\"\n",
    "\n",
    "    return state\n",
    "\n",
    "\n",
    "class FooterSearchResponse(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to contain the response for a node that needs a footer search, and the query involved.\n",
    "    \"\"\"\n",
    "\n",
    "    node_id: str\n",
    "    search_query: str\n",
    "\n",
    "\n",
    "class RouterAgentFooterResponse(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe the structured output of the router agent.\n",
    "    \"\"\"\n",
    "\n",
    "    node_ids_for_footer_search: List[FooterSearchResponse]\n",
    "    \n",
    "class RouterAgentLinkResponse(BaseModel):\n",
    "    \"\"\"\n",
    "    Pydantic schema used to describe the structured output of the router agent.\n",
    "    \"\"\"\n",
    "\n",
    "    node_ids_for_link_retrieval: List[str]\n",
    "\n",
    "\n",
    "def router_agent(state: AgentState) -> AgentState:\n",
    "    # decide if process should should stop or continue\n",
    "    state[\"markdown_debug\"] += f\"### Router Agent\\n\"\n",
    "\n",
    "    starter_prompt_footer = f\"\"\"\n",
    "        You are an intelligent agent overseeing a multi-agent retrieval process of graph nodes from a document. These nodes are to answer the query: \n",
    "        ```{state['query']}```\n",
    "        \n",
    "        Below this request is a list of nodes that were automatically retrieved. \n",
    "        \n",
    "        You must determine if the list of nodes is enough to answer the query. If there isn't enough information, you must identify any relevant footer information in the nodes.\n",
    "        \n",
    "        A node can footer information asking to look in another section/part of the document, which will require a separate natural language search. \n",
    "        Example: If the footer says \"see paragraph x\", a search query e.g. \"Return paragraph x to answer the query '{state['query']}'\" should be made. \n",
    "    \n",
    "        If there are no further nodes worth analysing, return an empty response. ONLY RETURN QUERIES FOR FOOTERS THAT ARE RELEVANT TO ANSWERING THE QUERY\n",
    "        \n",
    "        Else, if any relevant nodes require a footer search, specify the node_id and the search query.\n",
    "        Nodes are identified by node_id and must be quoted in backticks.\n",
    "\n",
    "        \n",
    "    \"\"\"\n",
    "    \n",
    "    starter_prompt_link = f\"\"\"\n",
    "        You are an intelligent agent overseeing a multi-agent retrieval process of graph nodes from a document. These nodes are to answer the query: \n",
    "        ```{state['query']}```\n",
    "        \n",
    "        Below this request is a list of nodes that were automatically retrieved. \n",
    "        \n",
    "        You must determine if the list of nodes is enough to answer the query. If there isn't enough information, you must identify any linked nodes that could be worth exploring.\n",
    "        \n",
    "        If there are no further nodes worth analysing, return an empty response.\n",
    "        \n",
    "        Return a list of node_ids. ONLY RETURN NODE_IDS for NODES THAT ARE RELEVANT TO ANSWERING THE QUERY. Nodes are identified by node_id and must be quoted in backticks.\n",
    "    \"\"\"\n",
    "    \n",
    "    # collect latest nodes, and all nodes\n",
    "    last_fetched_nodes_flattened: Dict[str, MultiAgentSearchLocalNode] = {}\n",
    "    all_nodes_flattened: Dict[str, MultiAgentSearchLocalNode] = {}\n",
    "\n",
    "    def recurse_collect_last_fetched_nodes(\n",
    "        current_nodes: List[MultiAgentSearchLocalNode],\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Recursively iterate via DFS on nodes, to index them by page number.\n",
    "        \"\"\"\n",
    "        for node in current_nodes:\n",
    "            last_fetched_nodes_flattened[node.id_] = node\n",
    "            all_nodes_flattened[node.id_] = node\n",
    "\n",
    "            if len(node.children) > 0:\n",
    "                recurse_collect_last_fetched_nodes(list(node.children.values()))\n",
    "\n",
    "    recurse_collect_last_fetched_nodes(state[\"last_fetched_context_nodes\"])\n",
    "\n",
    "    def recurse_collect_all_nodes(current_nodes: List[MultiAgentSearchLocalNode]):\n",
    "        \"\"\"\n",
    "        Recursively iterate via DFS on nodes, to index them by page number.\n",
    "        \"\"\"\n",
    "        for node in current_nodes:\n",
    "            all_nodes_flattened[node.id_] = node\n",
    "\n",
    "            if len(node.children) > 0:\n",
    "                recurse_collect_all_nodes(list(node.children.values()))\n",
    "\n",
    "    recurse_collect_all_nodes(state[\"last_fetched_context_nodes\"])\n",
    "\n",
    "    printout = \"\"\n",
    "    for node in last_fetched_nodes_flattened.values():\n",
    "        printout += node.print_node_prompt()\n",
    "\n",
    "    completion = openai_client.beta.chat.completions.parse(\n",
    "        model=\"gpt-4o-2024-08-06\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": starter_prompt_link},\n",
    "            {\"role\": \"user\", \"content\": printout},\n",
    "        ],\n",
    "        response_format=RouterAgentLinkResponse,\n",
    "    )\n",
    "\n",
    "    # Parse the response content \n",
    "    explore_response = json.loads(completion.choices[0].message.content)\n",
    "    print(f\"explore_response: {explore_response}\")\n",
    "\n",
    "    # Extract node_ids_for_link_retrieval\n",
    "    # Ensure that the node_ids are not already in the all_nodes_flattened\n",
    "    node_ids_for_link_retrieval = []\n",
    "    for node_id in explore_response[\"node_ids_for_link_retrieval\"]:\n",
    "        # sometimes GPT returns node ids with or without backticks\n",
    "        if node_id[0] == \"`\":\n",
    "            node_id = node_id[1:-1]\n",
    "        if all_nodes_flattened.get(node_id) is None:\n",
    "            node_ids_for_link_retrieval.append(node_id)\n",
    "\n",
    "    print(f\"node_ids_for_link_retrieval: {node_ids_for_link_retrieval}\")\n",
    "\n",
    "\n",
    "    completion = openai_client.beta.chat.completions.parse(\n",
    "        model=\"gpt-4o-2024-08-06\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": starter_prompt_footer},\n",
    "            {\"role\": \"user\", \"content\": printout},\n",
    "        ],\n",
    "        response_format=RouterAgentFooterResponse,\n",
    "    )\n",
    "\n",
    "    # Parse the response content \n",
    "    explore_response = json.loads(completion.choices[0].message.content)\n",
    "    print(f\"explore_response: {explore_response}\")\n",
    "\n",
    "    # Initialize the dictionary for node_ids_for_footer_search\n",
    "    node_ids_for_footer_search = {}\n",
    "\n",
    "    # Populate the dictionary with node_id and search_query pairs\n",
    "    # Must be new pairs, so not existing in all_nodes_flattened\n",
    "    for footer_response in explore_response[\"node_ids_for_footer_search\"]:\n",
    "        node_id = footer_response[\"node_id\"]\n",
    "        # sometimes GPT returns node ids with or without backticks\n",
    "        if node_id[0] == \"`\":\n",
    "            node_id = node_id[1:-1]\n",
    "\n",
    "        node_ids_for_footer_search[node_id] = footer_response[\"search_query\"]\n",
    "\n",
    "    # Update the state with the node_ids_for_link_retrieval and node_ids_for_footer_search\n",
    "    state[\"node_links_to_fetch\"] = node_ids_for_link_retrieval\n",
    "    state[\"node_footers_to_fetch\"] = node_ids_for_footer_search\n",
    "\n",
    "    # debug prints\n",
    "    state[\n",
    "        \"markdown_debug\"\n",
    "    ] += f\"- found {len(node_ids_for_link_retrieval)} nodes to fetch links for\\n\"\n",
    "    for node_id in node_ids_for_link_retrieval:\n",
    "        state[\"markdown_debug\"] += f\"    - {node_id}\\n\"\n",
    "    state[\n",
    "        \"markdown_debug\"\n",
    "    ] += f\" - found {len(node_ids_for_footer_search)} nodes to fetch footers for\\n\"\n",
    "    for node_id, search_query in node_ids_for_footer_search.items():\n",
    "        state[\"markdown_debug\"] += f\"    - `{node_id}` \\n\"\n",
    "        state[\"markdown_debug\"] += f\"      - **Search query:** {search_query}\\n\"\n",
    "\n",
    "    return state\n",
    "\n",
    "\n",
    "def answering_agent(state: AgentState) -> AgentState:\n",
    "    state[\"markdown_debug\"] += f\"### Answering Agent\\n\"\n",
    "    # answer the query\n",
    "    prompt = f\"\"\"\n",
    "You are an answering agent. You will be given a list of document nodes that were automatically retrieved by the system. These nodes are to answer the query:\n",
    "```{state['query']}```\n",
    "\n",
    "Give references to sections/paragraphs if possible, but do not output full node ids with backticks and the hash. \n",
    "    \"\"\"\n",
    "\n",
    "    printout = \"\"\n",
    "    for node in state[\"previous_nodes\"]:\n",
    "        printout += node.print_node_prompt()\n",
    "\n",
    "    for node in state[\"last_fetched_context_nodes\"]:\n",
    "        printout += node.print_node_prompt()\n",
    "\n",
    "    completion = openai_client.chat.completions.create(\n",
    "        model=\"gpt-4o-2024-08-06\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": prompt},\n",
    "            {\"role\": \"user\", \"content\": printout},\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    message = completion.choices[0].message.content\n",
    "    print(f\"response: {message}\")\n",
    "\n",
    "    state[\"markdown_debug\"] += f\"## Answer to Query: \\n{message}\\n\"\n",
    "    state[\"markdown_debug\"] += f\"## Source Nodes: \\n```\\n{printout}\\n```\"\n",
    "\n",
    "    return state\n",
    "\n",
    "\n",
    "def recursive_retrieval(state: AgentState) -> AgentState:\n",
    "    state[\"markdown_debug\"] += f\"### Recursive Retrieval\\n\"\n",
    "\n",
    "    current_nodes = state[\"last_fetched_context_nodes\"]\n",
    "    \n",
    "    for current_node in current_nodes:\n",
    "        state[\"previous_nodes\"].append(current_node)\n",
    "        \n",
    "\n",
    "    new_current_nodes = []\n",
    "\n",
    "    # look up the nodes to fetch by id\n",
    "    state[\n",
    "        \"markdown_debug\"\n",
    "    ] += f\"{len(state['node_links_to_fetch'])} nodes with link(s)\\n\"\n",
    "    \n",
    "    \n",
    "    for node_id in state[\"node_links_to_fetch\"]:\n",
    "        # sometimes GPT returns node ids with or without backticks\n",
    "        if node_id[0] == \"`\":\n",
    "            node_id = node_id[1:-1]\n",
    "        if node_id in local_nodes_map:\n",
    "            new_current_nodes.append(local_nodes_map[node_id])\n",
    "            state[\"markdown_debug\"] += f\"  - {node_id}\\n\"\n",
    "        else:\n",
    "            state[\"search_failures\"].append(f\"Failed to fetch node with id: {node_id}\")\n",
    "\n",
    "    state[\n",
    "        \"markdown_debug\"\n",
    "    ] += f\"{len(state['node_footers_to_fetch'])} nodes with footer info\\n\"\n",
    "\n",
    "    for node_id, search_query in state[\"node_footers_to_fetch\"].items():\n",
    "        # fetch nodes by keyword and bm25 search\n",
    "        footer_retrieved_nodes = retrieve_with_keywords_bm25(search_query)\n",
    "        # LLM prunes nodes that are not relevant\n",
    "        footer_retrieved_nodes, _ = prune_nodes(search_query, footer_retrieved_nodes)\n",
    "        state[\"markdown_debug\"] += f\" - {node_id}\\n\"\n",
    "        state[\"markdown_debug\"] += f\"  - **Search query:** {search_query}\\n\"\n",
    "        state[\"markdown_debug\"] += f\"  - **Returned:** \\n\"\n",
    "\n",
    "        for node in footer_retrieved_nodes:\n",
    "            state[\"markdown_debug\"] += f\"    - {node.id_}\\n\"\n",
    "            new_current_nodes.append(node)\n",
    "\n",
    "        # if no nodes fetched, log failure\n",
    "        if len(footer_retrieved_nodes) == 0:\n",
    "            state[\"search_failures\"].append(\n",
    "                f\"Failed to fetch nodes for query: {search_query}\"\n",
    "            )\n",
    "\n",
    "    state[\"last_fetched_context_nodes\"] = new_current_nodes\n",
    "    state[\"pass_count\"] += 1\n",
    "    state[\"markdown_debug\"] += f\"# PASS {state['pass_count']}\\n\"\n",
    "    state[\"node_footers_to_fetch\"] = {}\n",
    "    state[\"node_links_to_fetch\"] = []\n",
    "\n",
    "    return state\n",
    "\n",
    "\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "\n",
    "def should_continue_to_router(state: AgentState) -> str:\n",
    "    return state[\"allow_continue\"]\n",
    "\n",
    "\n",
    "def should_continue_to_recursive_retrieval(state: AgentState) -> str:\n",
    "    if len(state[\"node_links_to_fetch\"]) + len(state[\"node_footers_to_fetch\"]) > 0:\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "workflow.add_node(\"Initial Search Agent\", inital_search_agent)\n",
    "workflow.add_node(\"Definition Agent\", definition_agent)\n",
    "workflow.add_node(\"Context Fetch Tool\", context_fetch_tool)\n",
    "workflow.add_node(\"Footer Parsing Agent\", footer_parsing_agent)\n",
    "workflow.add_node(\"Supervisor Agent\", supervisor_agent)\n",
    "workflow.add_node(\"Router Agent\", router_agent)\n",
    "workflow.add_node(\"Answering Agent\", answering_agent)\n",
    "workflow.add_node(\"Recursive Retrieval\", recursive_retrieval)\n",
    "\n",
    "workflow.add_edge(START, \"Initial Search Agent\")\n",
    "workflow.add_edge(\"Initial Search Agent\", \"Definition Agent\")\n",
    "workflow.add_edge(\"Definition Agent\", \"Context Fetch Tool\")\n",
    "workflow.add_edge(\"Context Fetch Tool\", \"Footer Parsing Agent\")\n",
    "workflow.add_edge(\"Footer Parsing Agent\", \"Supervisor Agent\")\n",
    "workflow.add_edge(\"Supervisor Agent\", \"Router Agent\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"Supervisor Agent\",\n",
    "    should_continue_to_router,\n",
    "    {\n",
    "        True: \"Router Agent\",\n",
    "        False: \"Answering Agent\",\n",
    "    },\n",
    ")\n",
    "workflow.add_conditional_edges(\n",
    "    \"Router Agent\",\n",
    "    should_continue_to_recursive_retrieval,\n",
    "    {\n",
    "        True: \"Recursive Retrieval\",\n",
    "        False: \"Answering Agent\",\n",
    "    },\n",
    ")\n",
    "workflow.add_edge(\"Recursive Retrieval\", \"Context Fetch Tool\")\n",
    "workflow.add_edge(\"Answering Agent\", END)\n",
    "\n",
    "\n",
    "graph = workflow.compile()\n",
    "\n",
    "\n",
    "from IPython.display import display, Markdown, clear_output\n",
    "\n",
    "\n",
    "def multi_agent_query(query=\"How can the Board and the CCO manage control functions?\"):\n",
    "    initial_state = AgentState(\n",
    "        query=query,\n",
    "        definitions={},\n",
    "        previous_nodes=[],\n",
    "        last_fetched_context_nodes=[],\n",
    "        node_links_to_fetch=[],\n",
    "        node_footers_to_fetch={},\n",
    "        allow_continue=True,\n",
    "        search_failures=[],\n",
    "        markdown_debug=debug_markdown,\n",
    "        pass_count=1,\n",
    "    )\n",
    "\n",
    "    events = graph.stream(initial_state, {\"recursion_limit\": 150})\n",
    "\n",
    "    for s in events:\n",
    "        clear_output(wait=True)\n",
    "        first_value = next(iter(s.values()))\n",
    "        try:\n",
    "            display(Markdown(first_value[\"markdown_debug\"]))\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "\n",
    "multi_agent_query()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'venv (Python 3.11.7)' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: '/Users/timothychung/Documents/GitHub/github-app/recursive-retrieval/venv/bin/python -m pip install ipykernel -U --force-reinstall'"
     ]
    }
   ],
   "source": [
    "# if I want to use 3072-dim embedding instead\n",
    "\n",
    "# Settings.llm = LLamaOpenAI(model=\"gpt-4o\")\n",
    "# Settings.embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-large\")\n",
    "# storage_context = StorageContext.from_defaults()\n",
    "\n",
    "# from openai import OpenAI\n",
    "# openai_client = OpenAI(api_key=os.getenv(\"OPENAI_API_KEY\"))\n",
    "\n",
    "# might want to make storage_context persist in dir\n",
    "\n",
    "\n",
    "# def get_text_embedding(text: str) -> List[float]:\n",
    "#     response = openai_client.embeddings.create(\n",
    "#         input=text,\n",
    "#         model=\"text-embedding-3-large\",\n",
    "#     )\n",
    "#     return response.data[0].embedding\n",
    "\n",
    "# from tqdm import tqdm\n",
    "\n",
    "# nodes = []\n",
    "# # every chunk represents a node\n",
    "\n",
    "# for chunk in tqdm(chunks):\n",
    "#     text = chunk.content.get(\"Content\", chunk.content.get(\"Type\", \"No content\"))\n",
    "#     nodes.append(TextNode(\n",
    "#         id_=chunk.chunk_id,\n",
    "#         embedding = get_text_embedding(text),\n",
    "#         metadata = chunk.content,\n",
    "#         text=text,\n",
    "#         excluded_embed_metadata_keys=list(chunk.content.keys())\n",
    "#     ))\n",
    "\n",
    "# storage_context.docstore.add_documents(nodes)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base_env2",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}