Update.

Author: antejavor

memgraph-graphRAG/graphRAG.ipynb

@@ -338,6 +338,30 @@
     "    return response.choices[0].message.content\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Running the graph RAG\n",
+    "\n",
+    "Now, it all comes together in the `main` function: \n",
+    "\n",
+    "1. Connect to the database \n",
+    "2. Load the .env file with the `OPENAI_API_KEY=` defined\n",
+    "3. Compute and store the node embeddings \n",
+    "4. Compute the question embedding based on key information \n",
+    "5. Perform the vector search to find the most semantically similar node\n",
+    "6. Get the relevant data that is a few hops away from the pivot node\n",
+    "7. Ask LLM the question with the relevant data "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -402,6 +426,45 @@
     "    main()\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Here are a few examples of questions and answers: \n",
+    "\n",
+    "`To whom was Viserys Targaryen loyal to?`\n",
+    "\n",
+    "The response is:\n",
+    "**Based on the provided data, Viserys Targaryen was loyal to House Targaryen. This information is derived from the relationships indicating that Viserys Targaryen was loyal to House Targaryen and the connections between them**\n",
+    "\n",
+    "`Who killed Viserys Targaryen in Game of thrones?`\n",
+    "\n",
+    "The response is:\n",
+    "\n",
+    "**Based on the provided relevance expansion data, Khal Drogo killed Viserys Targaryen in \"Game of Thrones.\" This information is inferred from the relationship where Khal Drogo is linked to Viserys Targaryen with the action of being \"KILLED\" by Khal Drogo. The data does not show any other character directly killing Viserys Targaryen.\n",
+    "**\n",
+    "\n",
+    "`\"What was the weapon used to kill Viserys Targaryen in Game of Thrones?\"`\n",
+    "\n",
+    "The response is: \n",
+    "**\n",
+    "Based on the provided data, the weapon used to kill Viserys Targaryen in Game of Thrones was not explicitly mentioned. The data only shows that Khal Drogo was involved in the killing of Viserys Targaryen. There is no specific mention of the weapon used in the relevance expansion data. Therefore, I do not have enough information to answer the question about the weapon used to kill Viserys Targaryen.**\n",
+    "\n",
+    "This response is wrong, there is a method mentioned, not weapon, but LLM didn't catch the context due to different naming. \n",
+    "\n",
+    "`\"Who betrayed Viserys Targaryen in Game of Thrones?\"`\n",
+    "\n",
+    "Based on the provided data, Khal Drogo betrayed Viserys Targaryen in Game of Thrones by killing him. This conclusion is drawn from the relationship between Khal Drogo and Viserys Targaryen where it is stated that Khal Drogo killed Viserys Targaryen.\n",
+    "\n",
+    "This response is based on the killing relationship, but betrayal could have different consequences. \n",
+    "\n",
+    "Let's expand this knowledge. \n",
+    "\n",
+    "`\"Who betrayed Viserys Targaryen in Game of Thrones?\"`\n",
+    "Based on the provided data, Khal Drogo betrayed Viserys Targaryen in Game of Thrones. This conclusion is derived from the relationship between Viserys Targaryen and Khal Drogo, where Khal Drogo is connected to Viserys Targaryen through the 'BETRAYED_BY' relationship, indicating that Khal Drogo betrayed Viserys Targaryen."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -419,12 +482,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Sample text summary for processing\n",
-    "summary=\"Viserys Targaryen is the last living son of the former king, Aerys II Targaryen (the 'Mad King'). As one of the last known Targaryen heirs, Viserys Targaryen is obsessed with reclaiming the Iron Throne and restoring his family’s rule over Westeros. Ambitious and arrogant, he often treats his younger sister, Daenerys Targaryen, as a pawn, seeing her only as a means to gain power. His ruthless ambition leads him to make a marriage alliance with Khal Drogo, a powerful Dothraki warlord, hoping Khal Drogo will give him the army he needs. However, Viserys Targaryen’s impatience and disrespect toward the Dothraki culture lead to his downfall; he is ultimately killed by Khal Drogo in a brutal display of 'a crown for a king' – having molten gold poured over his head. Khal Drogo is a prominent warlord and leader of the Dothraki people, known for his fearsome reputation and formidable combat skills. He enters into a marriage with Princess Daenerys Targaryen as part of an alliance orchestrated by her brother, Prince Viserys Targaryen. Though initially aloof and intimidating to Daenerys Targaryen, Khal Drogo grows to care deeply for her, and their relationship evolves into one of mutual respect and love. Khal Drogo becomes devoted to Daenerys Targaryen and her dreams, including her goal of reclaiming the Iron Throne for House Targaryen. However, Khal Drogo’s fate takes a tragic turn after he sustains a serious wound in a skirmish, which becomes infected. A healer named Mirri Maz Duur performs a ritual that leaves Khal Drogo in a vegetative state, robbing him of his strength and dignity. Heartbroken, Daenerys Targaryen ends Khal Drogo’s life mercifully, signaling both the end of her first love and a significant turning point in her journey. King Joffrey Baratheon is the eldest son of Queen Cersei Lannister and, officially, King Robert Baratheon, though he is actually the result of an incestuous relationship between Queen Cersei and her twin brother, Ser Jaime Lannister. Joffrey Baratheon's personality is marked by sadism, cruelty, and impulsive behavior, traits that make him a despised ruler and widely loathed by the people of Westeros. Following King Robert Baratheon’s death, Joffrey Baratheon takes the throne and quickly reveals himself as a tyrannical ruler, prone to rash decisions and heedless of the advice given by those around him. His cruelty is particularly directed at Lady Sansa Stark, his former fiancée, whom he torments and humiliates on multiple occasions. As king, he alienates his allies and fosters unrest with his reckless brutality. Joffrey Baratheon’s reign ends abruptly when he is poisoned at his wedding feast to Margaery Tyrell, an event known as the 'Purple Wedding' which brings relief to many who suffered under his rule. His death marks a turning point in the power struggles of King’s Landing.\""
+    "summary = \"\"\"\n",
+    "    Viserys Targaryen is the last living son of the former king, Aerys II Targaryen (the 'Mad King').\n",
+    "    As one of the last known Targaryen heirs, Viserys Targaryen is obsessed with reclaiming the Iron Throne and \n",
+    "    restoring his family’s rule over Westeros. Ambitious and arrogant, he often treats his younger sister, Daenerys Targaryen, \n",
+    "    as a pawn, seeing her only as a means to gain power. His ruthless ambition leads him to make a marriage alliance with \n",
+    "    Khal Drogo, a powerful Dothraki warlord, hoping Khal Drogo will give him the army he needs. \n",
+    "    However, Viserys Targaryen’s impatience and disrespect toward the Dothraki culture lead to his downfall;\n",
+    "    he is ultimately killed by Khal Drogo in a brutal display of 'a crown for a king' – having molten gold poured over his head. \n",
+    "    \"\"\""
    ]
   },
   {
@@ -436,10 +507,7 @@
     "TODO: add links \n",
     "\n",
     "The first step in the process is to extract entities from the summary using\n",
-    "SpaCy’s large language model. SpaCy is an advanced NLP (natural language\n",
-    "processing) library in Python, designed for tasks such as entity recognition,\n",
-    "part-of-speech tagging, and dependency parsing. It’s widely used for its speed\n",
-    "and accuracy in processing text.\n",
+    "SpaCy’s LLM.\n",
     "\n",
     "To begin, we need to install SpaCy and the specific model we wll be using."
    ]
@@ -455,29 +523,6 @@
     "%python -m spacy download en_core_web_md"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Next, set up your OpenAI API key."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "from wasabi import msg\n",
-    "\n",
-    "os.environ[\"OPENAI_API_KEY\"] = \"<YOUR_OPENAI_KEY>\"\n",
-    "\n",
-    "# Check for OpenAI API key\n",
-    "if not os.getenv(\"OPENAI_API_KEY\"):\n",
-    "    msg.fail(\"OPENAI_API_KEY environment variable not set. Please set it to proceed.\", exits=1)"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -501,54 +546,47 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
+    "import os\n",
+    "import spacy\n",
+    "from spacy_llm.util import assemble\n",
     "import json\n",
     "from collections import Counter\n",
     "from pathlib import Path\n",
     "\n",
-    "import spacy\n",
-    "from spacy_llm.util import assemble\n",
-    "\n",
-    "# Load the spaCy model\n",
-    "nlp = spacy.load(\"en_core_web_md\")\n",
-    "\n",
     "# Split document into sentences\n",
-    "def split_document_sent(text):\n",
+    "def split_document_sent(text, nlp):\n",
     "    doc = nlp(text)\n",
     "    return [sent.text.strip() for sent in doc.sents]\n",
     "\n",
-    "def process_text(text, verbose=False):\n",
+    "\n",
+    "def process_text(text, nlp, verbose=False):\n",
     "    doc = nlp(text)\n",
     "    if verbose:\n",
-    "        msg.text(f\"Text: {doc.text}\")\n",
-    "        msg.text(f\"Entities: {[(ent.text, ent.label_) for ent in doc.ents]}\")\n",
+    "        print(f\"Text: {doc.text}\")\n",
+    "        print(f\"Entities: {[(ent.text, ent.label_) for ent in doc.ents]}\")\n",
     "    return doc\n",
     "\n",
+    "\n",
     "# Pipeline to run entity extraction\n",
-    "def extract_entities(text, verbose=False):\n",
+    "def extract_entities(text, nlp, verbose=False):\n",
     "    processed_data = []\n",
     "    entity_counts = Counter()\n",
     "\n",
-    "    sentences = split_document_sent(text)\n",
+    "    sentences = split_document_sent(text, nlp)\n",
     "    for sent in sentences:\n",
-    "        doc = process_text(sent, verbose)\n",
+    "        doc = process_text(sent, nlp, verbose)\n",
     "        entities = [(ent.text, ent.label_) for ent in doc.ents]\n",
     "\n",
     "        # Store processed data for each sentence\n",
-    "        processed_data.append({'text': doc.text, 'entities': entities})\n",
+    "        processed_data.append({\"text\": doc.text, \"entities\": entities})\n",
     "\n",
     "        # Update counters\n",
     "        entity_counts.update([ent[1] for ent in entities])\n",
     "\n",
     "    # Export to JSON\n",
-    "    with open('processed_data.json', 'w') as f:\n",
+    "    with open(\"processed_data.json\", \"w\") as f:\n",
     "        json.dump(processed_data, f)\n",
-    "\n",
-    "    msg.text(f\"Entity counts: {entity_counts}\")\n",
-    "\n",
-    "# Run the pipeline on the summary text\n",
-    "verbose = True\n",
-    "extract_entities(summary, verbose)\n"
+    "\n"
    ]
   },
   {
@@ -572,45 +610,39 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import json\n",
-    "import openai\n",
-    "from pathlib import Path\n",
+    "    \n",
+    "    # Load the spaCy model\n",
+    "\n",
+    "    nlp = spacy.load(\"en_core_web_md\")\n",
+    "    extract_entities(summary, nlp)\n",
+    "\n",
+    "\n",
+    "    # Load processed data from JSON\n",
+    "    json_path = Path(\"processed_data.json\")\n",
+    "    with open(json_path, \"r\") as f:\n",
+    "        processed_data = json.load(f)\n",
+    "\n",
+    "    # Prepare nodes and relationships\n",
+    "    nodes = []\n",
+    "    relationships = []\n",
+    "\n",
+    "    # Formulate a prompt for GPT-4\n",
+    "    prompt = (\n",
+    "        \"Extract entities and relationships from the following JSON data. For each entry in data['entities'], \"\n",
+    "        \"create a 'node' dictionary with fields 'id' (unique identifier), 'name' (entity text), and 'type' (entity label). \"\n",
+    "        \"For entities that have meaningful connections, define 'relationships' as dictionaries with 'source' (source node id), \"\n",
+    "        \"'target' (target node id), and 'relationship' (type of connection). Create max 30 nodes, format relationships in the format of capital letters and _ inbetween words and format the entire response in the JSON output containing only variables nodes and relationships without any text inbetween. Use following labels for nodes: Character, Title, Location, House, Death, Event, Allegiance and following relationship types: HAPPENED_IN, SIBLING_OF, PARENT_OF, MARRIED_TO, HEALED_BY, RULES, KILLED, LOYAL_TO, BETRAYED_BY. Make sure the entire JSON file fits in the output\"\n",
+    "        \"JSON data:\\n\"\n",
+    "        f\"{json.dumps(processed_data)}\"\n",
+    "    )\n",
     "\n",
-    "# Load processed data from JSON\n",
-    "json_path = Path(\"processed_data.json\")\n",
-    "with open(json_path, \"r\") as f:\n",
-    "    processed_data = json.load(f)\n",
-    "\n",
-    "# Prepare nodes and relationships\n",
-    "nodes = []\n",
-    "relationships = []\n",
-    "\n",
-    "# Formulate a prompt for GPT-4\n",
-    "prompt = (\n",
-    "    \"Extract entities and relationships from the following JSON data. For each entry in data['entities'], \"\n",
-    "    \"create a 'node' dictionary with fields 'id' (unique identifier), 'name' (entity text), and 'type' (entity label). \"\n",
-    "    \"For entities that have meaningful connections, define 'relationships' as dictionaries with 'source' (source node id), \"\n",
-    "    \"'target' (target node id), and 'relationship' (type of connection). Create max 30 nodes, format relationships in the format of capital letters and _ inbetween words and format the entire response in the JSON output containing only variables nodes and relationships without any text inbetween. Use following labels for nodes: Character, Title, Location, House, Death, Event, Allegiance and following relationship types: HAPPENED_IN, SIBLING_OF, PARENT_OF, MARRIED_TO, HEALED_BY, RULES, KILLED, LOYAL_TO, BETRAYED_BY. Make sure the entire JSON file fits in the output\" \n",
-    "    \"JSON data:\\n\"\n",
-    "    f\"{json.dumps(processed_data)}\"\n",
-    ")\n",
-    "\n",
-    "# Call GPT-4 to analyze the JSON and extract structured nodes and relationships\n",
-    "response = openai.ChatCompletion.create(\n",
-    "    model=\"gpt-4\",\n",
-    "    messages=[{\"role\": \"system\", \"content\": \"You are a helpful assistant that structures data into nodes and relationships.\"},\n",
-    "              {\"role\": \"user\", \"content\": prompt}],\n",
-    "    max_tokens=2000\n",
-    ")\n",
-    "\n",
-    "# Parse GPT-4 response and add to nodes and relationships lists\n",
-    "output = response['choices'][0]['message']['content']\n",
-    "print(output)\n",
-    "structured_data = json.loads(output)  # Assuming GPT-4 outputs structured JSON\n",
-    "\n",
-    "# Populate nodes and relationships lists\n",
-    "nodes.extend(structured_data.get(\"nodes\", []))\n",
-    "relationships.extend(structured_data.get(\"relationships\", []))\n"
+    "    response = asyncio.run(get_response(client, prompt))\n",
+    "    structured_data = json.loads(response)  # Assuming GPT-4 outputs structured JSON\n",
+    "\n",
+    "    # Populate nodes and relationships lists\n",
+    "    nodes.extend(structured_data.get(\"nodes\", []))\n",
+    "    relationships.extend(structured_data.get(\"relationships\", []))\n",
+    "\n"
    ]
   },
   {
@@ -669,26 +701,36 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from neo4j import GraphDatabase\n",
-    "\n",
-    "# Initialize the Neo4j driver for Memgraph (modify the URI if necessary)\n",
-    "uri = \"bolt://localhost:7687\"\n",
-    "user = \"\"\n",
-    "password = \"\"\n",
-    "driver = GraphDatabase.driver(uri, auth=(user, password))\n",
-    "\n",
-    "# Function to execute Cypher queries in Memgraph\n",
-    "def execute_cypher_queries(queries):\n",
-    "    with driver.session() as session:\n",
-    "        for query in queries:\n",
-    "            try:\n",
-    "                session.run(query)\n",
-    "                msg.good(f\"Executed query: {query}\")\n",
-    "            except Exception as e:\n",
-    "                msg.fail(f\"Error executing query: {query}. Error: {e}\")\n",
-    "\n",
-    "# Execute the generated Cypher queries\n",
-    "execute_cypher_queries(cypher_queries)\n",
+    "with driver.session() as session:\n",
+    "    for query in cypher_queries:\n",
+    "        try:\n",
+    "            session.run(query)\n",
+    "            print(f\"Executed query: {query}\")\n",
+    "        except Exception as e:\n",
+    "            print(f\"Error executing query: {query}. Error: {e}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The dataset now has additional knowledge: \n",
+    "`\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 2}) CREATE (a)-[:PARENT_OF]->(b)\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 3}) CREATE (a)-[:SIBLING_OF]->(b)\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 4}) CREATE (a)-[:LOYAL_TO]->(b)\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 5}) CREATE (a)-[:HAPPENED_IN]->(b)\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 6}) CREATE (a)-[:HAPPENED_IN]->(b)\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 7}) CREATE (a)-[:SIBLING_OF]->(b)\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 8}) CREATE (a)-[:MARRIED_TO]->(b)\n",
+    "Executed query: MATCH (a {id: 1}), (b {id: 9}) CREATE (a)-[:HEALED_BY]->(b)\n",
+    "Executed query: MATCH (a {id: 8}), (b {id: 9}) CREATE (a)-[:RULES]->(b)\n",
+    "Executed query: MATCH (a {id: 8}), (b {id: 1}) CREATE (a)-[:BETRAYED_BY]->(b)\n",
+    "` \n",
+    "\n",
+    "Now again asking the same question yields a correct answer: `\"Who betrayed Viserys Targaryen in Game of Thrones?\"`\n",
+    "\n",
+    "Based on the provided data, Khal Drogo betrayed Viserys Targaryen in Game of Thrones. This conclusion is derived from the relationship between Viserys Targaryen and Khal Drogo, where Khal Drogo is connected to Viserys Targaryen through the 'BETRAYED_BY' relationship, indicating that Khal Drogo betrayed Viserys Targaryen.\n",
     "\n"
    ]
   }