integrations/llamaindex/single-agent-rag-system/single_agent_rag_system.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a single-agent GraphRAG system with LlamaIndex and Memgraph\n",
    "\n",
    "In this example, we build a single-agent GraphRAG system using LlamaIndex and\n",
    "Memgraph, integrating retrieval-augmented generation (RAG) with graph-based\n",
    "querying and tool-using agents. We'll explore how to:\n",
    "\n",
    "- Set up **Memgraph** as a graph store for structured knowledge retrieval.\n",
    "- Use **LlamaIndex** to create a Property Graph Index and perform Memgraph's\n",
    "  **vector search** on embedded data.\n",
    "- Implement an agent that uses tools for both arithmetic operations and semantic\n",
    "  retrieval.\n",
    "  \n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "1. Make sure you have [Docker](https://www.docker.com/) running in the\n",
    "   background. \n",
    "\n",
    "2. Run Memgraph\n",
    "\n",
    "The easiest way to run Memgraph is by using the following commands:\n",
    "\n",
    "For Linux/macOS: `curl https://install.memgraph.com | sh`\n",
    "\n",
    "For Windows: `iwr https://windows.memgraph.com | iex`\n",
    "\n",
    "3. Install necessary dependencies:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index llama-index-graph-stores-memgraph python-dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "Create `.env` file that contains your OpenAI API key:\n",
    "\n",
    "`OPENAI_API_KEY=sk-proj-...`\n",
    "\n",
    "## Create the script\n",
    "\n",
    "Let's first load our `.env` file and set the LLM model we want to use. In this\n",
    "example, we're using OpenAI's gpt-4 model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "\n",
    "# settings\n",
    "Settings.llm = OpenAI(model=\"gpt-4\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define calculator tools\n",
    "\n",
    "Next, define addition and multiplication tools for calculations and add them to\n",
    "`FunctionTool` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "# function tools\n",
    "def multiply(a: float, b: float) -> float:\n",
    "    \"\"\"Multiply two numbers and return the product\"\"\"\n",
    "    return a * b\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "\n",
    "def add(a: float, b: float) -> float:\n",
    "    \"\"\"Add two numbers and return the sum\"\"\"\n",
    "    return a + b\n",
    "\n",
    "add_tool = FunctionTool.from_defaults(fn=add)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the dataset \n",
    "\n",
    "Besides the basic operations, we also want to create a RAG pipeline and perform\n",
    "retrieval operations on the dataset of our choice. In this example, we're using\n",
    "the PDF file about the Canadian budget for 2023. The file is transformed into PDF\n",
    "and stored in the `data` directory. Let's load that dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./data\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Memgraph graph store\n",
    "\n",
    "We'll now establish a connection to **Memgraph**, using\n",
    "`MemgraphPropertyGraphStore` from LlamaIndex. This allows us to store and\n",
    "retrieve structured data efficiently, enabling **graph-based querying** for\n",
    "retrieval-augmented generation (RAG) pipelines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.graph_stores.memgraph import MemgraphPropertyGraphStore\n",
    "\n",
    "graph_store = MemgraphPropertyGraphStore(\n",
    "    username=\"\",  # Your Memgraph username, default is \"\"\n",
    "    password=\"\",  # Your Memgraph password, default is \"\"\n",
    "    url=\"bolt://localhost:7687\"  # Connection URL for Memgraph\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a knowledge graph in Memgraph\n",
    "\n",
    "This section builds a **Property Graph Index** using `PropertyGraphIndex` from\n",
    "LlamaIndex. This index allows us to store and retrieve structured knowledge in a\n",
    "**graph database (Memgraph)** while leveraging OpenAI embeddings for semantic\n",
    "search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import PropertyGraphIndex\n",
    "from llama_index.core.indices.property_graph import SchemaLLMPathExtractor\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "index = PropertyGraphIndex.from_documents(\n",
    "    documents,\n",
    "    embed_model=OpenAIEmbedding(model_name=\"text-embedding-ada-002\"),\n",
    "    kg_extractors=[\n",
    "        SchemaLLMPathExtractor(\n",
    "            llm=OpenAI(model=\"gpt-4\", temperature=0.0)\n",
    "        )\n",
    "    ],\n",
    "    property_graph_store=graph_store,\n",
    "    show_progress=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RAG Pipeline: query engine and retrieval agent\n",
    "\n",
    "Let's now set up a **Retrieval-Augmented Generation (RAG) pipeline**. The\n",
    "pipeline enables efficient data retrieval from a structured knowledge base\n",
    "(Memgraph) and provides contextual responses using OpenAI's GPT-4.\n",
    "\n",
    "First, we convert the **Property Graph Index** into a **query engine**, allowing\n",
    "structured queries over the indexed data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine()\n",
    "\n",
    "# smoke test\n",
    "response = query_engine.query(\n",
    "    \"What was the total amount of the 2023 Canadian federal budget?\"\n",
    ")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating and running the agent\n",
    "\n",
    "Let's now create a **RAG agent** that can retrieve budget data and perform\n",
    "calculations. First, we define `budget_tool`, which provides facts about the\n",
    "2023 Canadian federal budget. Then, we create a `ReActAgent` that combines this\n",
    "tool with calculation tools, allowing it to both fetch information and handle\n",
    "math operations. Finally, we ask the agent: \"What is the total amount of the\n",
    "2023 Canadian federal budget multiplied by 3?\" and print the response to see it\n",
    "work step by step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent import ReActAgent\n",
    "from llama_index.core.tools import QueryEngineTool\n",
    "\n",
    "# RAG pipeline as a tool\n",
    "budget_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine,\n",
    "    name=\"canadian_budget_2023\",\n",
    "    description=\"A RAG engine with some basic facts about the 2023 Canadian federal budget.\"\n",
    ")\n",
    "\n",
    "# Create the agent with tools\n",
    "agent = ReActAgent.from_tools([multiply_tool, add_tool, budget_tool], verbose=True)\n",
    "\n",
    "# Query the agent\n",
    "response = agent.chat(\"What is the total amount of the 2023 Canadian federal budget multiplied by 3? Go step by step, using a tool to do any math.\")\n",
    "\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 64-bit (microsoft store)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "289d8ae9ac585fcc15d0d9333c941ae27bdf80d3e799883224b20975f2046730"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}