run-llama--llama_index/docs/examples/property_graph/property_graph_basic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Property Graph Index\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/property_graph/property_graph_basic.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "\n",
    "In this notebook, we demonstrate some basic usage of the `PropertyGraphIndex` in LlamaIndex.\n",
    "\n",
    "The property graph index here will take unstructured documents, extract a property graph from it, and provide various methods to query that graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/loganmarkewich/Library/Caches/pypoetry/virtualenvs/llama-index-bXUwlEfH-py3.11/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Parsing nodes: 100%|██████████| 1/1 [00:00<00:00, 25.46it/s]\n",
      "Extracting paths from text: 100%|██████████| 22/22 [00:12<00:00,  1.72it/s]\n",
      "Extracting implicit paths: 100%|██████████| 22/22 [00:00<00:00, 36186.15it/s]\n",
      "Generating embeddings: 100%|██████████| 1/1 [00:00<00:00,  1.14it/s]\n",
      "Generating embeddings: 100%|██████████| 5/5 [00:00<00:00,  5.43it/s]\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core import PropertyGraphIndex\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "index = PropertyGraphIndex.from_documents(\n",
    "    documents,\n",
    "    llm=OpenAI(model=\"gpt-3.5-turbo\", temperature=0.3),\n",
    "    embed_model=OpenAIEmbedding(model_name=\"text-embedding-3-small\"),\n",
    "    show_progress=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So lets recap what exactly just happened\n",
    "- `PropertyGraphIndex.from_documents()` - we loaded documents into an index\n",
    "- `Parsing nodes` - the index parsed the documents into nodes\n",
    "- `Extracting paths from text` - the nodes were passed to an LLM, and the LLM was prompted to generate knowledge graph triples (i.e. paths)\n",
    "- `Extracting implicit paths` - each `node.relationships` property was used to infer implicit paths\n",
    "- `Generating embeddings` - embeddings were generated for each text node and graph node (hence this happens twice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets explore what we created! For debugging purposes, the default `SimplePropertyGraphStore` includes a helper to save a `networkx` representation of the graph to an `html` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.property_graph_store.save_networkx_graph(name=\"./kg.html\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Opening the html in a browser, we can see our graph!\n",
    "\n",
    "If you zoom in, each \"dense\" node with many connections is actually the source chunk, with extracted entities and relations branching off from there."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![example graph](./kg_screenshot.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customizing Low-Level Construction\n",
    "\n",
    "If we wanted, we can do the same ingestion using the low-level API, leverage `kg_extractors`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.indices.property_graph import (\n",
    "    ImplicitPathExtractor,\n",
    "    SimpleLLMPathExtractor,\n",
    ")\n",
    "\n",
    "index = PropertyGraphIndex.from_documents(\n",
    "    documents,\n",
    "    embed_model=OpenAIEmbedding(model_name=\"text-embedding-3-small\"),\n",
    "    kg_extractors=[\n",
    "        ImplicitPathExtractor(),\n",
    "        SimpleLLMPathExtractor(\n",
    "            llm=OpenAI(model=\"gpt-3.5-turbo\", temperature=0.3),\n",
    "            num_workers=4,\n",
    "            max_paths_per_chunk=10,\n",
    "        ),\n",
    "    ],\n",
    "    show_progress=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a full guide on all extractors, see the [detailed usage page](/../../module_guides/indexing/lpg_index_guide#construction)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying\n",
    "\n",
    "Querying a property graph index typically consists of using one or more sub-retrievers and combining results.\n",
    "\n",
    "Graph retrieval can be thought of \n",
    "- selecting node(s)\n",
    "- traversing from those nodes\n",
    "\n",
    "By default, two types of retrieval are used in unison\n",
    "- synoynm/keyword expansion - use the LLM to generate synonyms and keywords from the query\n",
    "- vector retrieval - use embeddings to find nodes in your graph\n",
    "\n",
    "Once nodes are found, you can either\n",
    "- return the paths adjacent to the selected nodes (i.e. triples)\n",
    "- return the paths + the original source text of the chunk (if available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interleaf -> Was -> On the way down\n",
      "Viaweb -> Had -> Code editor\n",
      "Interleaf -> Built -> Impressive technology\n",
      "Interleaf -> Added -> Scripting language\n",
      "Interleaf -> Made -> Scripting language\n",
      "Viaweb -> Suggested -> Take to hospital\n",
      "Interleaf -> Had done -> Something bold\n",
      "Viaweb -> Called -> After\n",
      "Interleaf -> Made -> Dialect of lisp\n",
      "Interleaf -> Got crushed by -> Moore's law\n",
      "Dan giffin -> Worked for -> Viaweb\n",
      "Interleaf -> Had -> Smart people\n",
      "Interleaf -> Had -> Few years to live\n",
      "Interleaf -> Made -> Software\n",
      "Interleaf -> Made -> Software for creating documents\n",
      "Paul graham -> Started -> Viaweb\n",
      "Scripting language -> Was -> Dialect of lisp\n",
      "Scripting language -> Is -> Dialect of lisp\n",
      "Software -> Will be affected by -> Rapid change\n",
      "Code editor -> Was -> In viaweb\n",
      "Software -> Worked via -> Web\n",
      "Programs -> Typed on -> Punch cards\n",
      "Computers -> Skipped -> Step\n",
      "Idea -> Was clear from -> Experience\n",
      "Apartment -> Wasn't -> Rent-controlled\n"
     ]
    }
   ],
   "source": [
    "retriever = index.as_retriever(\n",
    "    include_text=False,  # include source text, default True\n",
    ")\n",
    "\n",
    "nodes = retriever.retrieve(\"What happened at Interleaf and Viaweb?\")\n",
    "\n",
    "for node in nodes:\n",
    "    print(node.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interleaf had smart people and built impressive technology, including adding a scripting language that was a dialect of Lisp. However, despite their efforts, they were eventually impacted by Moore's Law and faced challenges. Viaweb, on the other hand, was started by Paul Graham and had a code editor where users could define their own page styles using Lisp expressions. Viaweb also suggested taking someone to the hospital and called something \"After.\"\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=True,\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What happened at Interleaf and Viaweb?\")\n",
    "\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For full details on customizing retrieval and querying, see [the docs page](/../../module_guides/indexing/lpg_index_guide#retrieval-and-querying)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Storage\n",
    "\n",
    "By default, storage happens using our simple in-memory abstractions - `SimpleVectorStore` for embeddings and `SimplePropertyGraphStore` for the property graph.\n",
    "\n",
    "We can save and load these to/from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.storage_context.persist(persist_dir=\"./storage\")\n",
    "\n",
    "from llama_index.core import StorageContext, load_index_from_storage\n",
    "\n",
    "index = load_index_from_storage(\n",
    "    StorageContext.from_defaults(persist_dir=\"./storage\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vector Stores\n",
    "\n",
    "While some graph databases support vectors (like Neo4j), you can still specify the vector store to use on top of your graph for cases where its not supported, or cases where you want to override.\n",
    "\n",
    "Below we will combine `ChromaVectorStore` with the default `SimplePropertyGraphStore`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-chroma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.graph_stores import SimplePropertyGraphStore\n",
    "from llama_index.vector_stores.chroma import ChromaVectorStore\n",
    "import chromadb\n",
    "\n",
    "client = chromadb.PersistentClient(\"./chroma_db\")\n",
    "collection = client.get_or_create_collection(\"my_graph_vector_db\")\n",
    "\n",
    "index = PropertyGraphIndex.from_documents(\n",
    "    documents,\n",
    "    embed_model=OpenAIEmbedding(model_name=\"text-embedding-3-small\"),\n",
    "    graph_store=SimplePropertyGraphStore(),\n",
    "    vector_store=ChromaVectorStore(collection=collection),\n",
    "    show_progress=True,\n",
    ")\n",
    "\n",
    "index.storage_context.persist(persist_dir=\"./storage\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then to load:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = PropertyGraphIndex.from_existing(\n",
    "    SimplePropertyGraphStore.from_persist_dir(\"./storage\"),\n",
    "    vector_store=ChromaVectorStore(chroma_collection=collection),\n",
    "    llm=OpenAI(model=\"gpt-3.5-turbo\", temperature=0.3),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This looks slightly different than purely using the storage context, but the syntax is more concise now that we've started to mix things together."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-bXUwlEfH-py3.11",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}