S8.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv(dotenv_path='.env')\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import PyPDFLoader\n",
    "\n",
    "file_path = \"./other_data/_the_Dementia_Wars__Santa's_Salvation_.pdf\"\n",
    "\n",
    "book_loader = PyPDFLoader(file_path)\n",
    "book_data = book_loader.load_and_split()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create the Graph Representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('Chapter 1', \"Santa Claus's Call to Arms\", 'is called'),\n",
       " ('Santa Claus', 'heavily', 'sighs'),\n",
       " ('Santa Claus', 'Rudolph', 'tells'),\n",
       " ('Santa Claus', 'what does this mean for Christmas?', 'responds'),\n",
       " ('Santa Claus', 'gratefully', 'smiles'),\n",
       " ('Santa Claus', 'his allies', 'turns to'),\n",
       " ('Santa Claus', 'we will fight this war together', 'declares'),\n",
       " ('Santa Claus', 'our mission remains unchanged', 'says'),\n",
       " ('Rudolph', 'with determination', 'steps forward'),\n",
       " ('Reindeer', 'one by one', 'nod in agreement'),\n",
       " ('Mrs. Claus', 'her voice is filled with compassion', 'steps forward')]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.indexes import GraphIndexCreator\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "llm = OpenAI()\n",
    "index_creator = GraphIndexCreator(llm=llm)\n",
    "graph = index_creator.from_text(book_data[20].page_content)\n",
    "\n",
    "graph.get_triples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"1917pt\" height=\"188pt\" viewBox=\"0.00 0.00 1917.13 188.00\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 184)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-184 1913.13,-184 1913.13,4 -4,4\"/>\n",
       "<!-- Chapter 1 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>Chapter 1</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"110.72\" cy=\"-162\" rx=\"47.26\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"110.72\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">Chapter 1</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#39;s Call to Arms -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>Santa Claus's Call to Arms</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"110.72\" cy=\"-90\" rx=\"110.72\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"110.72\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">Santa Claus's Call to Arms</text>\n",
       "</g>\n",
       "<!-- Chapter 1&#45;&gt;Santa Claus&#39;s Call to Arms -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>Chapter 1-&gt;Santa Claus's Call to Arms</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M110.72,-143.7C110.72,-136.41 110.72,-127.73 110.72,-119.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"114.22,-119.61 110.72,-109.61 107.22,-119.61 114.22,-119.61\"/>\n",
       "</g>\n",
       "<!-- Santa Claus -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>Santa Claus</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"792.72\" cy=\"-162\" rx=\"54.93\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"792.72\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">Santa Claus</text>\n",
       "</g>\n",
       "<!-- heavily -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>heavily</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"277.72\" cy=\"-90\" rx=\"38.04\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"277.72\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">heavily</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#45;&gt;heavily -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>Santa Claus-&gt;heavily</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M738.33,-158.97C650.34,-154.77 471.81,-142.43 324.72,-108 321.99,-107.36 319.19,-106.61 316.39,-105.79\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"317.76,-102.56 307.17,-102.81 315.61,-109.22 317.76,-102.56\"/>\n",
       "</g>\n",
       "<!-- Rudolph -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>Rudolph</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"376.72\" cy=\"-90\" rx=\"42.65\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"376.72\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">Rudolph</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#45;&gt;Rudolph -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>Santa Claus-&gt;Rudolph</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M740.36,-156.22C669.16,-149.08 538.15,-133.65 428.72,-108 425.75,-107.3 422.7,-106.51 419.64,-105.66\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"420.69,-102.32 410.1,-102.82 418.69,-109.03 420.69,-102.32\"/>\n",
       "</g>\n",
       "<!-- what does this mean for Christmas? -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>what does this mean for Christmas?</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"582.72\" cy=\"-90\" rx=\"145.01\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"582.72\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">what does this mean for Christmas?</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#45;&gt;what does this mean for Christmas? -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>Santa Claus-&gt;what does this mean for Christmas?</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M755.18,-148.49C723.82,-138.03 678.55,-122.94 642.26,-110.85\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"643.83,-107.68 633.24,-107.84 641.62,-114.32 643.83,-107.68\"/>\n",
       "</g>\n",
       "<!-- gratefully -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>gratefully</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"792.72\" cy=\"-90\" rx=\"46.75\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"792.72\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">gratefully</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#45;&gt;gratefully -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>Santa Claus-&gt;gratefully</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M792.72,-143.7C792.72,-136.41 792.72,-127.73 792.72,-119.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"796.22,-119.61 792.72,-109.61 789.22,-119.61 796.22,-119.61\"/>\n",
       "</g>\n",
       "<!-- his allies -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>his allies</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"900.72\" cy=\"-90\" rx=\"43.68\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"900.72\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">his allies</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#45;&gt;his allies -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>Santa Claus-&gt;his allies</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M816.69,-145.46C831.96,-135.57 851.93,-122.62 868.58,-111.83\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"870.17,-114.97 876.66,-106.59 866.37,-109.1 870.17,-114.97\"/>\n",
       "</g>\n",
       "<!-- we will fight this war together -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>we will fight this war together</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1085.7\" cy=\"-90\" rx=\"123.52\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1085.7\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">we will fight this war together</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#45;&gt;we will fight this war together -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>Santa Claus-&gt;we will fight this war together</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M836.27,-150.6C882.58,-139.53 956.64,-121.84 1012.1,-108.6\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1012.83,-112.03 1021.73,-106.28 1011.19,-105.22 1012.83,-112.03\"/>\n",
       "</g>\n",
       "<!-- our mission remains unchanged -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>our mission remains unchanged</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1356.7\" cy=\"-90\" rx=\"129.66\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1356.7\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">our mission remains unchanged</text>\n",
       "</g>\n",
       "<!-- Santa Claus&#45;&gt;our mission remains unchanged -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>Santa Claus-&gt;our mission remains unchanged</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M843.27,-154.73C934.77,-143.37 1129.7,-119.18 1251.3,-104.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1251.41,-107.6 1260.9,-102.89 1250.54,-100.66 1251.41,-107.6\"/>\n",
       "</g>\n",
       "<!-- with determination -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>with determination</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"376.72\" cy=\"-18\" rx=\"81.55\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"376.72\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">with determination</text>\n",
       "</g>\n",
       "<!-- Rudolph&#45;&gt;with determination -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>Rudolph-&gt;with determination</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M376.72,-71.7C376.72,-64.41 376.72,-55.73 376.72,-47.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"380.22,-47.62 376.72,-37.62 373.22,-47.62 380.22,-47.62\"/>\n",
       "</g>\n",
       "<!-- Reindeer -->\n",
       "<g id=\"node12\" class=\"node\">\n",
       "<title>Reindeer</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1556.7\" cy=\"-162\" rx=\"44.19\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1556.7\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">Reindeer</text>\n",
       "</g>\n",
       "<!-- one by one -->\n",
       "<g id=\"node13\" class=\"node\">\n",
       "<title>one by one</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1556.7\" cy=\"-90\" rx=\"51.86\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1556.7\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">one by one</text>\n",
       "</g>\n",
       "<!-- Reindeer&#45;&gt;one by one -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>Reindeer-&gt;one by one</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1556.7,-143.7C1556.7,-136.41 1556.7,-127.73 1556.7,-119.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1560.2,-119.61 1556.7,-109.61 1553.2,-119.61 1560.2,-119.61\"/>\n",
       "</g>\n",
       "<!-- Mrs. Claus -->\n",
       "<g id=\"node14\" class=\"node\">\n",
       "<title>Mrs. Claus</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1767.7\" cy=\"-162\" rx=\"52.38\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1767.7\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">Mrs. Claus</text>\n",
       "</g>\n",
       "<!-- her voice is filled with compassion -->\n",
       "<g id=\"node15\" class=\"node\">\n",
       "<title>her voice is filled with compassion</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1767.7\" cy=\"-90\" rx=\"141.43\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1767.7\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">her voice is filled with compassion</text>\n",
       "</g>\n",
       "<!-- Mrs. Claus&#45;&gt;her voice is filled with compassion -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>Mrs. Claus-&gt;her voice is filled with compassion</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1767.7,-143.7C1767.7,-136.41 1767.7,-127.73 1767.7,-119.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1771.2,-119.61 1767.7,-109.61 1764.2,-119.61 1771.2,-119.61\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import SVG\n",
    "graph.draw_graphviz(path=\"graph.svg\")\n",
    "SVG(\"graph.svg\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "graphs = [\n",
    "    index_creator.from_text(doc.page_content)\n",
    "    for doc in book_data\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1941"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import networkx as nx\n",
    "\n",
    "graph_nx = graphs[0]._graph\n",
    "for g in graphs[1:]:\n",
    "    graph_nx = nx.compose(graph_nx, g._graph)\n",
    "\n",
    "graph_nx\n",
    "graph_nx.number_of_nodes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<langchain_community.graphs.networkx_graph.NetworkxEntityGraph at 0x1308de830>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.graphs.networkx_graph import NetworkxEntityGraph\n",
    "\n",
    "graph = NetworkxEntityGraph(graph_nx)\n",
    "graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "#graph.draw_graphviz(path=\"graph.pdf\", prog=\"fdp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Augmenting LLMs with a Knowledge Base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphQAChain chain...\u001b[0m\n",
      "Entities Extracted:\n",
      "\u001b[32;1m\u001b[1;3mNorth Pole\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3mNorth Pole crumbled under President Harrison's delusions\n",
      "North Pole had been defended spirit of Christmas\n",
      "North Pole becomes battleground\n",
      "North Pole is secure once again\n",
      "North Pole is in danger of the war\n",
      "North Pole is the fate of Christmas\n",
      "North Pole begins to regain former glory\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The North Pole is experiencing a war or conflict that is endangering the fate of Christmas. However, it is also being defended and has regained some of its former glory.'"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chains import GraphQAChain\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "\n",
    "chain = GraphQAChain.from_llm(\n",
    "    llm=llm,\n",
    "    graph=graph,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "chain.run(\"What is happening in the North Pole?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphQAChain chain...\u001b[0m\n",
      "Entities Extracted:\n",
      "\u001b[32;1m\u001b[1;3mTimmy\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3mTimmy watched in awe\n",
      "Timmy watched in fear\n",
      "Timmy saw hope\n",
      "Timmy stands alongside Santa Claus\n",
      "Timmy stumbles through the desolate and snow-covered landscape\n",
      "Timmy spots a group of figures\n",
      "Timmy runs towards them\n",
      "Timmy nods a mixture of relief and hope\n",
      "Timmy watched with wide eyes\n",
      "Timmy felt surge of courage\n",
      "Timmy vowed to help\n",
      "Timmy was caught in the crossfire\n",
      "Timmy felt a glimmer of hope\n",
      "Timmy remembered the stories his parents had told him about Santa Claus and the magic of Christmas\n",
      "Timmy believed that there was still a chance for peace\n",
      "Timmy was searching for allies of Santa Claus\n",
      "Timmy had found the allies he had been searching for\n",
      "Timmy caught sight of a familiar figure\n",
      "Timmy had been searching for the allies he had been searching for\n",
      "Timmy trudged through the snow\n",
      "Timmy called out for Santa Claus to help him\n",
      "Timmy looked up at his newfoun\n",
      "Timmy encountered the war's devastation\n",
      "Timmy saw the place he had called home\n",
      "Timmy reduced to rubble\n",
      "Timmy had to get out of there\n",
      "Timmy stumbled out of the wreckage\n",
      "Timmy ran through the chaos\n",
      "Timmy focus on reaching Santa Claus and his allies\n",
      "Timmy caught glimpses of the devastation around him\n",
      "Timmy surveys scene\n",
      "Timmy makes his way through wreckage\n",
      "Timmy is drawn to small family\n",
      "Timmy realizes magnitude of tragedy\n",
      "Timmy asks Are you okay?\n",
      "Timmy offers small smile of reassurance\n",
      "Timmy believes in power of Santa Claus and his allies\n",
      "Timmy nods We will protect them\n",
      "Timmy feels sense of responsibility\n",
      "Timmy is in group of orphaned children\n",
      "Timmy became their protector\n",
      "Timmy continued his voice steady\n",
      "Timmy blushed at the praise\n",
      "Timmy leads the way\n",
      "Timmy scans the area\n",
      "Timmy feels the weight of responsibility\n",
      "Timmy refuses to let it deter him\n",
      "Timmy motions for the children to follow\n",
      "Timmy cautiously approaches a partially collapsed building\n",
      "Timmy pushes the door open\n",
      "Timmy whispers stay close\n",
      "Timmy trusts to lead them safely through the darkness\n",
      "Timmy relies on his bravery and resourcefulness\n",
      "Timmy moves deeper into the building\n",
      "Timmy comes across a row of shelves\n",
      "Timmy realizes they have stumbled upon a hidden stash of supplies\n",
      "Timmy exclaims we've hit the jackpot\n",
      "Timmy looked around his gaze falling on faces of children\n",
      "Timmy spoke in midst of war\n",
      "Timmy closed his eyes his heart filled with gratitude\n",
      "Timmy showed incredible strength\n",
      "Timmy nodded a renewed sense of hope\n",
      "Timmy feels determination\n",
      "Timmy has seen devastation\n",
      "Timmy watched Elf Lieutenant Sparkle\n",
      "Timmy kneels down beside injured creature\n",
      "Timmy says We have to help him\n",
      "Timmy says We'll nurse him back to health\n",
      "Timmy whispers You're free now\n",
      "Timmy feels a profound sense of accomplishment\n",
      "Timmy continues their journey\n",
      "Timmy fights to protect the spirit of Christmas\n",
      "Timmy brings peace to the war-torn North Pole)\n",
      "END OF EXAMPL\n",
      "Timmy admires handiwork\n",
      "Timmy felt surge of hope\n",
      "Timmy witnessed acts of selflessness\n",
      "Timmy witnessed compassion\n",
      "Timmy took it upon himself to assist\n",
      "Timmy helped carry supplies\n",
      "Timmy comforted frightened children\n",
      "Timmy shared stories\n",
      "Timmy admitted used to think\n",
      "Timmy nodded his heart filled with a renewed sense of purpose\n",
      "Timmy feels determination in his own heart\n",
      "Timmy declares can't let President Harrison's war destroy the North Pole\n",
      "Timmy says have to do whatever it takes to protect the spirit of Christmas\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Timmy is a young boy who is determined to protect the spirit of Christmas and help Santa Claus and his allies.'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"Who is Timmy?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the Diffbot Graph Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_experimental.graph_transformers.diffbot import DiffbotGraphTransformer\n",
    "\n",
    "diffbot_nlp = DiffbotGraphTransformer(\n",
    "    diffbot_api_key=os.getenv(\"DIFFBOT_API_KEY\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import WikipediaLoader\n",
    "\n",
    "query = \"Warren Buffett\"\n",
    "\n",
    "raw_docs = WikipediaLoader(query=query).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_docs = diffbot_nlp.convert_to_graph_documents(raw_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Relationship(source=Node(id='http://www.wikidata.org/entity/Q217583', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='CHIEF_EXECUTIVE_OFFICER', properties={'evidence': 'Warren Edward Buffett ( BUF-it; born August 30, 1930) is an American businessman, investor, and philanthropist who currently serves as the co-founder, chairman and CEO of Berkshire Hathaway.', 'isCurrent': 'true'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q217583', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='FOUNDED_BY', properties={'evidence': 'Warren Edward Buffett ( BUF-it; born August 30, 1930) is an American businessman, investor, and philanthropist who currently serves as the co-founder, chairman and CEO of Berkshire Hathaway.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), target=Node(id='Leila', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), target=Node(id='Leila', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Leila', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='HAS_CHILD', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Leila', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Leila', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q5284', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q203807', type='Organization'), type='EMPLOYEE_OR_MEMBER_OF', properties={'evidence': 'He founded the Giving Pledge in 2010 with Bill Gates, whereby billionaires pledge to give away at least half of their fortunes.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), target=Node(id='Stahl', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), target=Node(id='Stahl', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q16254414', type='Organization'), type='EDUCATED_AT', properties={'evidence': 'He attended New York Institute of Finance to focus on his economics background and soon pursued a business career.', 'isNotCurrent': 'true'}),\n",
       " Relationship(source=Node(id='Buffett Partnership Ltd.', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q217583', type='Organization'), type='ACQUIRED_BY', properties={'evidence': 'He created Buffett Partnership Ltd. in 1956 and his investment firm eventually acquired a textile manufacturing firm, Berkshire Hathaway, assuming its name to create a diversified holding company.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q203807', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q5284', type='Person'), type='FOUNDED_BY', properties={'evidence': 'He founded the Giving Pledge in 2010 with Bill Gates, whereby billionaires pledge to give away at least half of their fortunes.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q8033204', type='Organization'), type='EDUCATED_AT', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".', 'isNotCurrent': 'true', 'endTime': '1947'}),\n",
       " Relationship(source=Node(id='Leila', type='Person'), target=Node(id='Stahl', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Alice Deal Junior High School', type='Organization'), type='EDUCATED_AT', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".', 'isNotCurrent': 'true'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q1553', type='Location'), type='PERSON_LOCATION', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q1553', type='Location'), type='PLACE_OF_BIRTH', properties={'evidence': 'As of January 2024, he had a net worth of $122 billion, making him the tenth-richest person in the world.Buffett was born in Omaha, Nebraska.'}),\n",
       " Relationship(source=Node(id='Leila', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Leila', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q61', type='Location'), type='PERSON_LOCATION', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q61', type='Location'), type='PERSON_LOCATION', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q30', type='Location'), type='NATIONALITY', properties={'evidence': 'Warren Edward Buffett ( BUF-it; born August 30, 1930) is an American businessman, investor, and philanthropist who currently serves as the co-founder, chairman and CEO of Berkshire Hathaway.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q30', type='Location'), type='PERSON_LOCATION', properties={'evidence': 'Warren Edward Buffett ( BUF-it; born August 30, 1930) is an American businessman, investor, and philanthropist who currently serves as the co-founder, chairman and CEO of Berkshire Hathaway.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q217583', type='Organization'), type='EMPLOYEE_OR_MEMBER_OF', properties={'evidence': \"Since 1970, Buffett has presided as the chairman and largest shareholder of Berkshire Hathaway, one of America's foremost holding companies and world's leading corporate conglomerates.\", 'isCurrent': 'true', 'positionHeld': 'CEO'}),\n",
       " Relationship(source=Node(id='Stahl', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Stahl', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Leila', type='Person'), type='HAS_PARENT', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Leila', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Leila', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Alice Deal Junior High School', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q61', type='Location'), type='ORGANIZATION_LOCATIONS', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q918630', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='WORK_RELATIONSHIP', properties={'evidence': 'In 1978, fellow investor and long-time business associate Charlie Munger joined Buffett as vice-chairman.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q5284', type='Person'), type='WORK_RELATIONSHIP', properties={'evidence': 'He founded the Giving Pledge in 2010 with Bill Gates, whereby billionaires pledge to give away at least half of their fortunes.'}),\n",
       " Relationship(source=Node(id='Warren', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q8033204', type='Organization'), type='EDUCATED_AT', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".', 'isNotCurrent': 'true', 'endTime': '1947'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q203807', type='Organization'), type='EMPLOYEE_OR_MEMBER_OF', properties={'evidence': 'He founded the Giving Pledge in 2010 with Bill Gates, whereby billionaires pledge to give away at least half of their fortunes.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q43199', type='Location'), type='PERSON_LOCATION', properties={'evidence': 'As of January 2024, he had a net worth of $122 billion, making him the tenth-richest person in the world.Buffett was born in Omaha, Nebraska.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q43199', type='Location'), type='PLACE_OF_BIRTH', properties={'evidence': 'As of January 2024, he had a net worth of $122 billion, making him the tenth-richest person in the world.Buffett was born in Omaha, Nebraska.'}),\n",
       " Relationship(source=Node(id='Warren', type='Person'), target=Node(id='Alice Deal Junior High School', type='Organization'), type='EDUCATED_AT', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".', 'isNotCurrent': 'true'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q8033204', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q61', type='Location'), type='ORGANIZATION_LOCATIONS', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".', 'isCurrent': 'true'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), type='HAS_PARENT', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Buffett Partnership Ltd.', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q1505660', type='Skill'), type='INDUSTRY', properties={'evidence': 'He created Buffett Partnership Ltd. in 1956 and his investment firm eventually acquired a textile manufacturing firm, Berkshire Hathaway, assuming its name to create a diversified holding company.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q5284', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='WORK_RELATIONSHIP', properties={'evidence': 'He founded the Giving Pledge in 2010 with Bill Gates, whereby billionaires pledge to give away at least half of their fortunes.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Stahl', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Stahl', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Stahl', type='Person'), type='HAS_PARENT', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q918630', type='Person'), type='WORK_RELATIONSHIP', properties={'evidence': 'In 1978, fellow investor and long-time business associate Charlie Munger joined Buffett as vice-chairman.'}),\n",
       " Relationship(source=Node(id='Warren', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q61', type='Location'), type='PERSON_LOCATION', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".'}),\n",
       " Relationship(source=Node(id='Warren', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q61', type='Location'), type='PERSON_LOCATION', properties={'evidence': 'In 1942, his father was elected to the first of four terms in the United States Congress, and after moving with his family to Washington, D.C., Warren finished elementary school, attended Alice Deal Junior High School and graduated from what was then Woodrow Wilson High School in 1947, where his senior yearbook picture reads: \"likes math; a future stockbroker\".'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='Buffett Partnership Ltd.', type='Organization'), type='EMPLOYEE_OR_MEMBER_OF', properties={'evidence': 'He created Buffett Partnership Ltd. in 1956 and his investment firm eventually acquired a textile manufacturing firm, Berkshire Hathaway, assuming its name to create a diversified holding company.', 'isNotCurrent': 'true', 'startTime': '1956'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q217583', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q1505660', type='Skill'), type='INDUSTRY', properties={'evidence': 'He created Buffett Partnership Ltd. in 1956 and his investment firm eventually acquired a textile manufacturing firm, Berkshire Hathaway, assuming its name to create a diversified holding company.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q723488', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='HAS_CHILD', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q8134', type='Skill'), type='INTERESTED_IN', properties={'evidence': 'He attended New York Institute of Finance to focus on his economics background and soon pursued a business career.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q8134', type='Skill'), type='INTERESTED_IN', properties={'evidence': 'He attended New York Institute of Finance to focus on his economics background and soon pursued a business career.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q8134', type='Skill'), type='FIELD_OF_WORK', properties={'evidence': 'He attended New York Institute of Finance to focus on his economics background and soon pursued a business career.'}),\n",
       " Relationship(source=Node(id='http://www.wikidata.org/entity/Q203807', type='Organization'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='FOUNDED_BY', properties={'evidence': 'He founded the Giving Pledge in 2010 with Bill Gates, whereby billionaires pledge to give away at least half of their fortunes.'}),\n",
       " Relationship(source=Node(id='Stahl', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='FAMILY_MEMBER', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Stahl', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='SOCIAL_RELATIONSHIP', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'}),\n",
       " Relationship(source=Node(id='Stahl', type='Person'), target=Node(id='http://www.wikidata.org/entity/Q47213', type='Person'), type='HAS_CHILD', properties={'evidence': 'Warren Edward Buffett was born on August 30, 1930 in Omaha, Nebraska, as the second of three children and the only son of Leila (née Stahl) and Congressman Howard Buffett.'})]"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph_docs[0].relationships"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "book_graph = diffbot_nlp.convert_to_graph_documents(book_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a Local Graph Database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.graphs import Neo4jGraph\n",
    "\n",
    "graph_db = Neo4jGraph(\n",
    "    url=\"bolt://localhost:7687\",\n",
    "    username=\"neo4j\",\n",
    "    password=\"pleaseletmein\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_db.add_graph_documents(book_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_db.refresh_schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Augmenting an LLM with the Graph Database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import GraphCypherQAChain\n",
    "\n",
    "cypher_llm = ChatOpenAI(temperature=0, model_name=\"gpt-4\")\n",
    "qa_llm = ChatOpenAI(temperature=0, model_name=\"gpt-3.5-turbo\")\n",
    "\n",
    "chain = GraphCypherQAChain.from_llm(\n",
    "    cypher_llm=cypher_llm,\n",
    "    qa_llm=qa_llm,\n",
    "    graph=graph_db,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mTask:Generate Cypher statement to query a graph database.\n",
      "Instructions:\n",
      "Use only the provided relationship types and properties in the schema.\n",
      "Do not use any other relationship types or properties that are not provided.\n",
      "Schema:\n",
      "Node properties are the following:\n",
      "Person {id: STRING, name: STRING},Location {id: STRING, name: STRING},Skill {id: STRING, name: STRING},Organization {name: STRING, id: STRING}\n",
      "Relationship properties are the following:\n",
      "PERSON_LOCATION {evidence: STRING},INTERESTED_IN {evidence: STRING},SOCIAL_RELATIONSHIP {evidence: STRING},ROMANTIC_RELATIONSHIP {evidence: STRING},EMPLOYEE_OR_MEMBER_OF {evidence: STRING, isCurrent: STRING}\n",
      "The relationships are the following:\n",
      "(:Person)-[:PERSON_LOCATION]->(:Location),(:Person)-[:INTERESTED_IN]->(:Skill),(:Person)-[:ROMANTIC_RELATIONSHIP]->(:Person),(:Person)-[:SOCIAL_RELATIONSHIP]->(:Person),(:Person)-[:EMPLOYEE_OR_MEMBER_OF]->(:Organization)\n",
      "Note: Do not include any explanations or apologies in your responses.\n",
      "Do not respond to any questions that might ask anything else than for you to construct a Cypher statement.\n",
      "Do not include any text except the generated Cypher statement.\n",
      "\n",
      "The question is:\n",
      "What is happening in the North Pole?\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (p:Person)-[r:PERSON_LOCATION]->(l:Location) WHERE l.name = \"North Pole\" RETURN p,r,l;\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[]\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant that helps to form nice and human understandable answers.\n",
      "The information part contains the provided information that you must use to construct an answer.\n",
      "The provided information is authoritative, you must never doubt it or try to use your internal knowledge to correct it.\n",
      "Make the answer sound as a response to the question. Do not mention that you based the result on the given information.\n",
      "If the provided information is empty, say that you don't know the answer.\n",
      "Information:\n",
      "[]\n",
      "\n",
      "Question: What is happening in the North Pole?\n",
      "Helpful Answer:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"I'm sorry, but I don't have any information on what is currently happening in the North Pole.\""
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.callbacks import StdOutCallbackHandler\n",
    "\n",
    "\n",
    "chain.run(\"What is happening in the North Pole?\", callbacks=[StdOutCallbackHandler()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mTask:Generate Cypher statement to query a graph database.\n",
      "Instructions:\n",
      "Use only the provided relationship types and properties in the schema.\n",
      "Do not use any other relationship types or properties that are not provided.\n",
      "Schema:\n",
      "Node properties are the following:\n",
      "Person {id: STRING, name: STRING},Location {id: STRING, name: STRING},Skill {id: STRING, name: STRING},Organization {name: STRING, id: STRING}\n",
      "Relationship properties are the following:\n",
      "PERSON_LOCATION {evidence: STRING},INTERESTED_IN {evidence: STRING},SOCIAL_RELATIONSHIP {evidence: STRING},ROMANTIC_RELATIONSHIP {evidence: STRING},EMPLOYEE_OR_MEMBER_OF {evidence: STRING, isCurrent: STRING}\n",
      "The relationships are the following:\n",
      "(:Person)-[:PERSON_LOCATION]->(:Location),(:Person)-[:INTERESTED_IN]->(:Skill),(:Person)-[:ROMANTIC_RELATIONSHIP]->(:Person),(:Person)-[:SOCIAL_RELATIONSHIP]->(:Person),(:Person)-[:EMPLOYEE_OR_MEMBER_OF]->(:Organization)\n",
      "Note: Do not include any explanations or apologies in your responses.\n",
      "Do not respond to any questions that might ask anything else than for you to construct a Cypher statement.\n",
      "Do not include any text except the generated Cypher statement.\n",
      "\n",
      "The question is:\n",
      "Who is Timmy?\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (p:Person {name: 'Timmy'}) RETURN p\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'p': {'name': 'Timmy', 'id': 'Timmy'}}]\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant that helps to form nice and human understandable answers.\n",
      "The information part contains the provided information that you must use to construct an answer.\n",
      "The provided information is authoritative, you must never doubt it or try to use your internal knowledge to correct it.\n",
      "Make the answer sound as a response to the question. Do not mention that you based the result on the given information.\n",
      "If the provided information is empty, say that you don't know the answer.\n",
      "Information:\n",
      "[{'p': {'name': 'Timmy', 'id': 'Timmy'}}]\n",
      "\n",
      "Question: Who is Timmy?\n",
      "Helpful Answer:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Timmy is a person whose name is Timmy.'"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"Who is Timmy?\", callbacks=[StdOutCallbackHandler()])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}