end-to-end-lupus.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# GraphRAG Python package End-to-End Example\n",
    "\n",
    "This notebook contains an end-to-end worked example using the [GraphRAG Python package](https://neo4j.com/docs/neo4j-graphrag-python/current/index.html) for Neo4j. It starts with unstructured documents (in this case pdfs), and progresses through knowledge graph construction, knowledge graph retriever design, and complete GraphRAG pipelines. \n",
    "\n",
    "Research papers on Lupus are used as the data source. We design a couple of different retrievers based on different knowledge graph retrieval patterns. \n",
    "\n",
    "For more details and explanations around each of the below steps, see the [corresponding blog post](https://neo4j.com/blog/graphrag-python-package/) which contains a full write-up, in-depth comparison of the retrieval patterns, and additional learning resources. \n",
    "\n",
    "## Pre-Requisites\n",
    "\n",
    "1. __Create a Neo4j Database__: To work through this RAG example, you need a database for storing and retrieving data. There are many options for this. You can quickly start a free Neo4j Graph Database using [Neo4j AuraDB](https://neo4j.com/product/auradb/?ref=neo4j-home-hero). You can use __AuraDB Free__ or start an __AuraDB Professional (Pro) free trial__ for higher ingestion and retrieval performance. The Pro instances have a bit more RAM; we recommend them for the best user experience.\n",
    "2. __Obtain an OpenAI Key__: This example requires an OpenAI API key to use language models and embedders. The cost should be very minimal. If you do not yet have an OpenAI API key you can [create an OpenAI account](https://platform.openai.com/signup) or [sign in](https://platform.openai.com/login). Next, navigate to the [API key page](https://platform.openai.com/account/api-keys) and click \"Create new secret key\". Optionally naming the key. \n",
    "3. __Fill in Credentials__: Either by copying the [`.env.template`](.env.template) file, naming it `.env`, and filling in the appropriate credentials, or by manually putting the credentials into the second code cell below. You will need:\n",
    "    1. The Neo4j URI, username, and password variables from when you created the database. If you created your database on AuraDB, they are in the file you downloaded.\n",
    "    2. Your OpenAI API key.\n",
    "\n"
   ],
   "id": "cf0be1ddc2bd2434"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Setup",
   "id": "799a7b09475ed2a4"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "%%capture\n",
    "%pip install fsspec langchain-text-splitters tiktoken openai python-dotenv numpy torch neo4j-graphrag"
   ],
   "id": "9820f541adf30bfd",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "\n",
    "# load neo4j credentials (and openai api key in background).\n",
    "load_dotenv('.env', override=True)\n",
    "NEO4J_URI = os.getenv('NEO4J_URI')\n",
    "NEO4J_USERNAME = os.getenv('NEO4J_USERNAME')\n",
    "NEO4J_PASSWORD = os.getenv('NEO4J_PASSWORD')\n",
    "\n",
    "#uncomment this line if you aren't using a .env file\n",
    "# os.environ['OPENAI_API_KEY'] = 'copy_paste_the_openai_key_here'"
   ],
   "id": "a023c71324bf6e7f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Knowledge Graph Building\n",
    "\n",
    "The `SimpleKGPipeline` class allows you to automatically build a knowledge graph with a few key inputs, including\n",
    "- a driver to connect to Neo4j,\n",
    "- an LLM for entity extraction, and\n",
    "- an embedding model to create vectors on text chunks for similarity search.\n",
    "\n",
    "There are also some optional inputs, such as node labels, relationship types, and a custom prompt template, which we will use to improve the quality of the knowledge graph. For full details on this, see [the blog](https://neo4j.com/blog/graphrag-python-package/).\n"
   ],
   "id": "418bd212bae7f492"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import neo4j\n",
    "from neo4j_graphrag.llm import OpenAILLM\n",
    "from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings\n",
    "\n",
    "driver = neo4j.GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))\n",
    "\n",
    "ex_llm=OpenAILLM(\n",
    "    model_name=\"gpt-4o-mini\",\n",
    "    model_params={\n",
    "        \"response_format\": {\"type\": \"json_object\"}, # use json_object formatting for best results\n",
    "        \"temperature\": 0 # turning temperature down for more deterministic results\n",
    "    }\n",
    ")\n",
    "\n",
    "#create text embedder\n",
    "embedder = OpenAIEmbeddings()"
   ],
   "id": "fd60a2f512b9701d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "#define node labels\n",
    "basic_node_labels = [\"Object\", \"Entity\", \"Group\", \"Person\", \"Organization\", \"Place\"]\n",
    "\n",
    "academic_node_labels = [\"ArticleOrPaper\", \"PublicationOrJournal\"]\n",
    "\n",
    "medical_node_labels = [\"Anatomy\", \"BiologicalProcess\", \"Cell\", \"CellularComponent\", \n",
    "                       \"CellType\", \"Condition\", \"Disease\", \"Drug\",\n",
    "                       \"EffectOrPhenotype\", \"Exposure\", \"GeneOrProtein\", \"Molecule\",\n",
    "                       \"MolecularFunction\", \"Pathway\"]\n",
    "\n",
    "node_labels = basic_node_labels + academic_node_labels + medical_node_labels\n",
    "\n",
    "# define relationship types\n",
    "rel_types = [\"ACTIVATES\", \"AFFECTS\", \"ASSESSES\", \"ASSOCIATED_WITH\", \"AUTHORED\",\n",
    "    \"BIOMARKER_FOR\", \"CAUSES\", \"CITES\", \"CONTRIBUTES_TO\", \"DESCRIBES\", \"EXPRESSES\",\n",
    "    \"HAS_REACTION\", \"HAS_SYMPTOM\", \"INCLUDES\", \"INTERACTS_WITH\", \"PRESCRIBED\",\n",
    "    \"PRODUCES\", \"RECEIVED\", \"RESULTS_IN\", \"TREATS\", \"USED_FOR\"]\n"
   ],
   "id": "6cba83fa4638e21d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "prompt_template = '''\n",
    "You are a medical researcher tasks with extracting information from papers \n",
    "and structuring it in a property graph to inform further medical and research Q&A.\n",
    "\n",
    "Extract the entities (nodes) and specify their type from the following Input text.\n",
    "Also extract the relationships between these nodes. the relationship direction goes from the start node to the end node. \n",
    "\n",
    "\n",
    "Return result as JSON using the following format:\n",
    "{{\"nodes\": [ {{\"id\": \"0\", \"label\": \"the type of entity\", \"properties\": {{\"name\": \"name of entity\" }} }}],\n",
    "  \"relationships\": [{{\"type\": \"TYPE_OF_RELATIONSHIP\", \"start_node_id\": \"0\", \"end_node_id\": \"1\", \"properties\": {{\"details\": \"Description of the relationship\"}} }}] }}\n",
    "\n",
    "- Use only the information from the Input text.  Do not add any additional information.  \n",
    "- If the input text is empty, return empty Json. \n",
    "- Make sure to create as many nodes and relationships as needed to offer rich medical context for further research.\n",
    "- An AI knowledge assistant must be able to read this graph and immediately understand the context to inform detailed research questions. \n",
    "- Multiple documents will be ingested from different sources and we are using this property graph to connect information, so make sure entity types are fairly general. \n",
    "\n",
    "Use only fhe following nodes and relationships (if provided):\n",
    "{schema}\n",
    "\n",
    "Assign a unique ID (string) to each node, and reuse it to define relationships.\n",
    "Do respect the source and target node types for relationship and\n",
    "the relationship direction.\n",
    "\n",
    "Do not return any additional information other than the JSON in it.\n",
    "\n",
    "Examples:\n",
    "{examples}\n",
    "\n",
    "Input text:\n",
    "\n",
    "{text}\n",
    "'''"
   ],
   "id": "8cbcdedcd1757b1d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import FixedSizeSplitter\n",
    "from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline\n",
    "\n",
    "kg_builder_pdf = SimpleKGPipeline(\n",
    "    llm=ex_llm,\n",
    "    driver=driver,\n",
    "    text_splitter=FixedSizeSplitter(chunk_size=500, chunk_overlap=100),\n",
    "    embedder=embedder,\n",
    "    entities=node_labels,\n",
    "    relations=rel_types,\n",
    "    prompt_template=prompt_template,\n",
    "    from_pdf=True\n",
    ")"
   ],
   "id": "94e1dfac980e6527",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Below, we run the `SimpleKGPipeline` to construct our knowledge graph from 3 pdf documents and store in Neo4j.",
   "id": "4d8ef81d0b5ac70c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "pdf_file_paths = ['truncated-pdfs/biomolecules-11-00928-v2-trunc.pdf', \n",
    "             'truncated-pdfs/GAP-between-patients-and-clinicians_2023_Best-Practice-trunc.pdf', \n",
    "             'truncated-pdfs/pgpm-13-39-trunc.pdf']\n",
    "\n",
    "for path in pdf_file_paths:\n",
    "    print(f\"Processing : {path}\")\n",
    "    pdf_result = await kg_builder_pdf.run_async(file_path=path)\n",
    "    print(f\"Result: {pdf_result}\")"
   ],
   "id": "edeee98826a970a8",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Knowledge Graph Retrieval",
   "id": "a9c0fd965b15b143"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "We will leverage Neo4j's vector search capabilities here. To do this, we need to begin by creating a vector index on the text chunks from the PDFs, which are stored on `Chunk` nodes in our knowledge graph.",
   "id": "4d7b98e310104246"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from neo4j_graphrag.indexes import create_vector_index\n",
    "\n",
    "create_vector_index(driver, name=\"text_embeddings\", label=\"Chunk\",\n",
    "                    embedding_property=\"embedding\", dimensions=1536, similarity_fn=\"cosine\")"
   ],
   "id": "940b051107b89204",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Now that the index is set up, we will start simple with a __VectorRetriever__.  The __VectorRetriever__ just queries `Chunk` nodes via vector search, bringing back the text and some metadata.",
   "id": "ec95391c989ee694"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from neo4j_graphrag.retrievers import VectorRetriever\n",
    "\n",
    "vector_retriever = VectorRetriever(\n",
    "    driver,\n",
    "    index_name=\"text_embeddings\",\n",
    "    embedder=embedder,\n",
    "    return_properties=[\"text\"],\n",
    ")"
   ],
   "id": "eeda7e519c60a02b",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Below we visualize the context we get back when submitting a search prompt. ",
   "id": "982112cc273a472e"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import json\n",
    "\n",
    "vector_res = vector_retriever.get_search_results(query_text = \"How is precision medicine applied to Lupus?\", \n",
    "                                                 top_k=3)\n",
    "for i in vector_res.records: print(\"====\\n\" + json.dumps(i.data(), indent=4))"
   ],
   "id": "ca27731a1f99d71d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The GraphRAG Python Package offers [a wide range of useful retrievers](https://neo4j.com/docs/neo4j-graphrag-python/current/user_guide_rag.html#retriever-configuration), each covering different knowledge graph retrieval patterns.\n",
    "\n",
    "Below we will use the __`VectorCypherRetriever`__, which allows you to run a graph traversal after finding nodes with vector search.  This uses Cypher, Neo4j's graph query language, to define the logic for traversing the graph. \n",
    "\n",
    "As a simple starting point, we'll traverse up to 3 hops out from each Chunk, capture the relationships encountered, and include them in the response alongside our text chunks.\n"
   ],
   "id": "2b2ee3da2339365f"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from neo4j_graphrag.retrievers import VectorCypherRetriever\n",
    "\n",
    "vc_retriever = VectorCypherRetriever(\n",
    "    driver,\n",
    "    index_name=\"text_embeddings\",\n",
    "    embedder=embedder,\n",
    "    retrieval_query=\"\"\"\n",
    "//1) Go out 2-3 hops in the entity graph and get relationships\n",
    "WITH node AS chunk\n",
    "MATCH (chunk)<-[:FROM_CHUNK]-()-[relList:!FROM_CHUNK]-{1,2}()\n",
    "UNWIND relList AS rel\n",
    "\n",
    "//2) collect relationships and text chunks\n",
    "WITH collect(DISTINCT chunk) AS chunks, \n",
    "  collect(DISTINCT rel) AS rels\n",
    "\n",
    "//3) format and return context\n",
    "RETURN '=== text ===\\n' + apoc.text.join([c in chunks | c.text], '\\n---\\n') + '\\n\\n=== kg_rels ===\\n' +\n",
    "  apoc.text.join([r in rels | startNode(r).name + ' - ' + type(r) + '(' + coalesce(r.details, '') + ')' +  ' -> ' + endNode(r).name ], '\\n---\\n') AS info\n",
    "\"\"\"\n",
    ")"
   ],
   "id": "2c81ad815a7b1bf9",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Below we visualize the context we get back when submitting a search prompt. ",
   "id": "ef5183123679ca2d"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "vc_res = vc_retriever.get_search_results(query_text = \"How is precision medicine applied to Lupus?\", top_k=3)\n",
    "\n",
    "# print output\n",
    "kg_rel_pos = vc_res.records[0]['info'].find('\\n\\n=== kg_rels ===\\n')\n",
    "print(\"# Text Chunk Context:\")\n",
    "print(vc_res.records[0]['info'][:kg_rel_pos])\n",
    "print(\"# KG Context From Relationships:\")\n",
    "print(vc_res.records[0]['info'][kg_rel_pos:])"
   ],
   "id": "4a1df583f2450f35",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## GraphRAG\n",
    " \n",
    " You can construct GraphRAG pipelines with the `GraphRAG` class.  At a minimum, you will need to pass the constructor an LLM and a retriever. You can optionally pass a custom prompt template. We will do so here just to provide a bit more guidance for the LLM to stick to information from our data source.\n",
    " \n",
    "Below we create `GraphRAG` objects for both the vector and vector-cypher retrievers. "
   ],
   "id": "f2e55b8b3511cf1"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from neo4j_graphrag.llm import OpenAILLM as LLM\n",
    "from neo4j_graphrag.generation import RagTemplate\n",
    "from neo4j_graphrag.generation.graphrag import GraphRAG\n",
    "\n",
    "llm = LLM(model_name=\"gpt-4o\",  model_params={\"temperature\": 0.0})\n",
    "\n",
    "rag_template = RagTemplate(template='''Answer the Question using the following Context. Only respond with information mentioned in the Context. Do not inject any speculative information not mentioned. \n",
    "\n",
    "# Question:\n",
    "{query_text}\n",
    " \n",
    "# Context:\n",
    "{context}\n",
    "\n",
    "# Answer:\n",
    "''', expected_inputs=['query_text', 'context'])\n",
    "\n",
    "v_rag  = GraphRAG(llm=llm, retriever=vector_retriever, prompt_template=rag_template)\n",
    "vc_rag = GraphRAG(llm=llm, retriever=vc_retriever, prompt_template=rag_template)"
   ],
   "id": "8e2cf317a83d59ae",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Now we can run GraphRAG and examine the outputs. ",
   "id": "fc8e13302fdeadd3"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "q = \"How is precision medicine applied to Lupus? provide in list format.\"\n",
    "print(f\"Vector Response: \\n{v_rag.search(q, retriever_config={'top_k':5}).answer}\")\n",
    "print(\"\\n===========================\\n\")\n",
    "print(f\"Vector + Cypher Response: \\n{vc_rag.search(q, retriever_config={'top_k':5}).answer}\")"
   ],
   "id": "4d0aff643e689611",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "q = \"Can you summarize systemic lupus erythematosus (SLE)? including common effects, biomarkers, and treatments? Provide in detailed list format.\"\n",
    "\n",
    "v_rag_result = v_rag.search(q, retriever_config={'top_k': 5}, return_context=True)\n",
    "vc_rag_result = vc_rag.search(q, retriever_config={'top_k': 5}, return_context=True)\n",
    "\n",
    "print(f\"Vector Response: \\n{v_rag_result.answer}\")\n",
    "print(\"\\n===========================\\n\")\n",
    "print(f\"Vector + Cypher Response: \\n{vc_rag_result.answer}\")"
   ],
   "id": "a2c44f834bdb13ad",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "for i in v_rag_result.retriever_result.items: print(json.dumps(eval(i.content), indent=1))",
   "id": "1088766f7109b6b3",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "vc_ls = vc_rag_result.retriever_result.items[0].content.split('\\\\n---\\\\n')\n",
    "for i in vc_ls:\n",
    "    if \"biomarker\" in i: print(i)"
   ],
   "id": "8b10fdab707a0d4a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "vc_ls = vc_rag_result.retriever_result.items[0].content.split('\\\\n---\\\\n')\n",
    "for i in vc_ls:\n",
    "    if \"treat\" in i: print(i)"
   ],
   "id": "7f1510ef53a9d253",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "q = \"Can you summarize systemic lupus erythematosus (SLE)? including common effects, biomarkers, treatments, and current challenges faced by Physicians and patients? provide in list format with details for each item.\"\n",
    "print(f\"Vector Response: \\n{v_rag.search(q, retriever_config={'top_k': 5}).answer}\")\n",
    "print(\"\\n===========================\\n\")\n",
    "print(f\"Vector + Cypher Response: \\n{vc_rag.search(q, retriever_config={'top_k': 5}).answer}\")"
   ],
   "id": "d03ec43f6f02f608",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "c7d66300af9eac12",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}