examples/notebooks/04-Querying.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d40dfbd7",
   "metadata": {},
   "source": [
    "# Querying"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08d9b60c",
   "metadata": {},
   "source": [
    "## Install toolkit\n",
    "\n",
    "Run the command below to install the graphrag-toolkit. If you've already installed the toolkit, you don't need to install it again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ca98b2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install https://github.com/awslabs/graphrag-toolkit/archive/refs/tags/v1.1.3.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04d08c74",
   "metadata": {},
   "source": [
    "### TraversalBasedRetriever\n",
    "\n",
    "See [TraversalBasedRetriever](https://github.com/awslabs/graphrag-toolkit/blob/main/docs/querying.md#traversalbasedretriever)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ec2a722",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import set_logging_config\n",
    "from graphrag_toolkit import LexicalGraphQueryEngine, format_source\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "set_logging_config('INFO', ['graphrag_toolkit.storage'])\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "query_engine = LexicalGraphQueryEngine.for_traversal_based_search(\n",
    "    graph_store, \n",
    "    vector_store,\n",
    "    post_processors=[\n",
    "        format_source('url')\n",
    "    ]\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What are the differences between Neptune Database and Neptune Analytics?\")\n",
    "\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b5a947e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in response.source_nodes:\n",
    "    print(n.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb78418c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in response.source_nodes:\n",
    "    print(n.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41035822",
   "metadata": {},
   "source": [
    "#### Set subretriever\n",
    "\n",
    "In the example below, the `TraversalBasedRetriever` is configured with a `ChunkBasedSearch` subretriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90d04b97",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphQueryEngine, format_source\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "from graphrag_toolkit.retrieval.retrievers import ChunkBasedSearch\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "query_engine = LexicalGraphQueryEngine.for_traversal_based_search(\n",
    "    graph_store, \n",
    "    vector_store,\n",
    "    retrievers=[ChunkBasedSearch],\n",
    "    post_processors=[\n",
    "        format_source('url')\n",
    "    ]\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What are the differences between Neptune Database and Neptune Analytics?\")\n",
    "\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d174f696",
   "metadata": {},
   "source": [
    "### SemanticGuidedRetriever\n",
    "\n",
    "See [SemanticGuidedRetriever](https://github.com/awslabs/graphrag-toolkit/blob/main/docs/querying.md#semanticguidedretriever)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d1f3184",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphQueryEngine\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "query_engine = LexicalGraphQueryEngine.for_semantic_guided_search(\n",
    "    graph_store, \n",
    "    vector_store\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What are the differences between Neptune Database and Neptune Analytics?\")\n",
    "\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3416845",
   "metadata": {},
   "source": [
    "#### Set subretrievers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1adddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphQueryEngine, format_source\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "from graphrag_toolkit.retrieval.retrievers import StatementCosineSimilaritySearch, KeywordRankingSearch, SemanticBeamGraphSearch\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "query_engine = LexicalGraphQueryEngine.for_semantic_guided_search(\n",
    "    graph_store, \n",
    "    vector_store,\n",
    "    retrievers=[\n",
    "        StatementCosineSimilaritySearch, \n",
    "        KeywordRankingSearch, \n",
    "        SemanticBeamGraphSearch\n",
    "    ]\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What are the differences between Neptune Database and Neptune Analytics?\")\n",
    "\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b234dc61",
   "metadata": {},
   "source": [
    "#### Reranking beam search (CPU)\n",
    "\n",
    "The example below uses a `SentenceReranker` with a `RerankingBeamGraphSearch` to rerank statements while conducting the beam search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "902acaef",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphQueryEngine\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "from graphrag_toolkit.retrieval.retrievers import RerankingBeamGraphSearch, StatementCosineSimilaritySearch, KeywordRankingSearch\n",
    "from graphrag_toolkit.retrieval.post_processors import SentenceReranker\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "cosine_retriever = StatementCosineSimilaritySearch(\n",
    "    vector_store=vector_store,\n",
    "    graph_store=graph_store,\n",
    "    top_k=50\n",
    ")\n",
    "\n",
    "keyword_retriever = KeywordRankingSearch(\n",
    "    vector_store=vector_store,\n",
    "    graph_store=graph_store,\n",
    "    max_keywords=10\n",
    ")\n",
    "\n",
    "reranker = SentenceReranker(\n",
    "    batch_size=128\n",
    ")\n",
    "\n",
    "beam_retriever = RerankingBeamGraphSearch(\n",
    "    vector_store=vector_store,\n",
    "    graph_store=graph_store,\n",
    "    reranker=reranker,\n",
    "    initial_retrievers=[cosine_retriever, keyword_retriever],\n",
    "    max_depth=8,\n",
    "    beam_width=100\n",
    ")\n",
    "\n",
    "query_engine = LexicalGraphQueryEngine.for_semantic_guided_search(\n",
    "    graph_store, \n",
    "    vector_store,\n",
    "    retrievers=[\n",
    "        cosine_retriever,\n",
    "        keyword_retriever,\n",
    "        beam_retriever\n",
    "    ]\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What are the differences between Neptune Database and Neptune Analytics?\")\n",
    "\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "717049be",
   "metadata": {},
   "source": [
    "#### Reranking beam search (GPU)\n",
    "\n",
    "The example below uses a `BGEReranker` with a `RerankingBeamGraphSearch` to rerank statements while conducting the beam search.\n",
    "\n",
    "There will be a delay the first time this runs while the reranker downloads tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06c3b363",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphQueryEngine\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "from graphrag_toolkit.retrieval.retrievers import RerankingBeamGraphSearch, StatementCosineSimilaritySearch, KeywordRankingSearch\n",
    "from graphrag_toolkit.retrieval.post_processors import BGEReranker\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "cosine_retriever = StatementCosineSimilaritySearch(\n",
    "    vector_store=vector_store,\n",
    "    graph_store=graph_store,\n",
    "    top_k=50\n",
    ")\n",
    "\n",
    "keyword_retriever = KeywordRankingSearch(\n",
    "    vector_store=vector_store,\n",
    "    graph_store=graph_store,\n",
    "    max_keywords=10\n",
    ")\n",
    "\n",
    "reranker = BGEReranker(\n",
    "    gpu_id=0,\n",
    "    batch_size=128\n",
    ")\n",
    "\n",
    "beam_retriever = RerankingBeamGraphSearch(\n",
    "    vector_store=vector_store,\n",
    "    graph_store=graph_store,\n",
    "    reranker=reranker,\n",
    "    initial_retrievers=[cosine_retriever, keyword_retriever],\n",
    "    max_depth=8,\n",
    "    beam_width=100\n",
    ")\n",
    "\n",
    "query_engine = LexicalGraphQueryEngine.for_semantic_guided_search(\n",
    "    graph_store, \n",
    "    vector_store,\n",
    "    retrievers=[\n",
    "        cosine_retriever,\n",
    "        keyword_retriever,\n",
    "        beam_retriever\n",
    "    ]\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What are the differences between Neptune Database and Neptune Analytics?\")\n",
    "\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ba2ac35",
   "metadata": {},
   "source": [
    "#### Post-processors \n",
    "\n",
    "The example below uses a `StatementDiversityPostProcessor`, `SentenceReranker` and `StatementEnhancementPostProcessor`.\n",
    "\n",
    "  - `SentenceReranker` - Reranks statements using the `mixedbread-ai/mxbai-rerank-xsmall-v1` model. \n",
    "\n",
    "  - `StatementEnhancementPostProcessor` - Enhances statements by using chunk context and an LLM to improve content while preserving original metadata.\n",
    "\n",
    "  - `StatementDiversityPostProcessor` - Removes similar statements using TF-IDF similarity with a default threshold of 0.975 to ensure diversity in the processed nodes.\n",
    "\n",
    "Before running `StatementDiversityPostProcessor` for the first time, load the following package:\n",
    "\n",
    "```\n",
    "python -m spacy download en_core_web_sm\n",
    "```\n",
    "\n",
    "If you're running on a GPU device, you can replace the `SentenceReranker` with a `BGEReranker`, which reranks statements using the ``BAAI/bge-reranker-v2-minicpm-layerwise`` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f2d46a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy download en_core_web_sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef08819d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphQueryEngine\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "from graphrag_toolkit.retrieval.post_processors import BGEReranker, SentenceReranker, StatementDiversityPostProcessor, StatementEnhancementPostProcessor\n",
    "import os\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "query_engine = LexicalGraphQueryEngine.for_semantic_guided_search(\n",
    "    graph_store, \n",
    "    vector_store,\n",
    "    post_processors=[\n",
    "        SentenceReranker(), \n",
    "        StatementDiversityPostProcessor(), \n",
    "        StatementEnhancementPostProcessor()\n",
    "    ]\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"What are the differences between Neptune Database and Neptune Analytics?\")\n",
    "\n",
    "print(response.response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}