integrations/llamaindex/multi-agent-rag-system/multi_agent_rag_system.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building multi-agent GraphRAG system with LlamaIndex and Memgraph\n",
    "\n",
    "In this example, we build a multi-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 function agents for both arithmetic operations and semantic\n",
    "  retrieval.\n",
    "- Design an **AgentWorkflow** that combines retrieval and computation to answer\n",
    "  complex queries.\n",
    "\n",
    "By the end, we'll have a fully functional GraphRAG pipeline capable of answering\n",
    "structured queries while performing calculations on retrieved data.\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:"
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv()\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, addition and subtraction tools for calculations and a calculator\n",
    "agent are defined. The role of the agent, in this case, will be to perform basic arithmetic\n",
    "operations with access to the defined tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import FunctionTool\n",
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two numbers.\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def subtract(a: int, b: int) -> int:\n",
    "    \"\"\"Subtract two numbers.\"\"\"\n",
    "    return a - b\n",
    "\n",
    "# Create agent configs\n",
    "calculator_agent = FunctionAgent(\n",
    "    name=\"calculator\",\n",
    "    description=\"Performs basic arithmetic operations\",\n",
    "    system_prompt=\"You are a calculator assistant.\",\n",
    "    tools=[\n",
    "        FunctionTool.from_defaults(fn=add),\n",
    "        FunctionTool.from_defaults(fn=subtract),\n",
    "    ],\n",
    "    llm=OpenAI(model=\"gpt-4\"),\n",
    ")"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "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** using\n",
    "LlamaIndex's `QueryEngineTool` and `FunctionAgent`. The pipeline enables\n",
    "efficient data retrieval from a structured knowledge base (Memgraph) and\n",
    "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": [
    "from llama_index.core.tools import QueryEngineTool\n",
    "\n",
    "query_engine = index.as_query_engine()\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",
    "retriever_agent = FunctionAgent(\n",
    "    name=\"retriever\",\n",
    "    description=\"Manages data retrieval\",\n",
    "    system_prompt=\"You are a retrieval assistant.\",\n",
    "    tools=[\n",
    "        budget_tool,\n",
    "    ],\n",
    "    llm=OpenAI(model=\"gpt-4\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating and running the workflow\n",
    "\n",
    "Finally, and most importantly, let's create an **AgentWorkflow** that ties together\n",
    "the previously defined agents, including the **calculator** and **retriever**\n",
    "agents. This workflow enables us to run a sequence of operations involving both\n",
    "data retrieval and arithmetic computations, allowing the agents to interact with\n",
    "one another.\n",
    "\n",
    "We define an **async function** to execute the workflow, sending a user query\n",
    "that asks for both the total amount of the 2023 Canadian federal budget and an\n",
    "additional calculation (adding 3 billion)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import (\n",
    "    AgentWorkflow,\n",
    "    FunctionAgent,\n",
    "    ReActAgent,\n",
    ")\n",
    "import asyncio\n",
    "\n",
    "# Create and run the workflow\n",
    "workflow = AgentWorkflow(\n",
    "    agents=[calculator_agent, retriever_agent], root_agent=\"calculator\"\n",
    ")\n",
    "\n",
    "# Define an async function to run the workflow\n",
    "async def run_workflow():\n",
    "    response = await workflow.run(user_msg=\"What is the total amount of the 2023 Canadian federal budget? Add 3 billion to that budget using tools\")\n",
    "    print(response)\n",
    "\n",
    "# Run the async function using asyncio\n",
    "asyncio.run(run_workflow())"
   ]
  }
 ],
 "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
}