Introduction to Graph RAG and project overview
Hi everyone. In this video I want to show you how to perform Graph RAG with LLaMA 3.1, a local running model. So first, what is Graph? Graph is an approach to perform retrieval-augmented generation by taking the relationships of entities and documents into consideration. Key concepts are nodes and relationships. Nodes represent entities or concepts extracted from data chunks such as people, organizations, events, or locations. In a knowledge graph, each node contains attributes and properties that provide more context about the entity. We then define the connections between nodes. These can include various types of associations such as hierarchical, like parent-child, temporal, before-and-after, or causal cause-and-effect relationships. They also have properties that describe the nature and the strength of the connection. When you've got a lot of documents, you end up with a nice graph describing the relationships of all documents. Let's have a look at a very simple example. In our dataset, nodes could represent entities like Apple and Tim Cook, while a relationship might describe Tim Cook as the CEO of Apple. That approach is really powerful, but one massive drawback is that it's very expensive to compute, since you have to extract entities from each document with an LLM and then compute a relational graph. This is why it's great to use that approach with a model that runs locally like LLaMA 3.1. We're going to use LLaMA in combination with Ollama in this video and Neo4j as graph database. We will create a graph about information of a large Italian family who owns multiple restaurants in different places, so many, many relationships to model. To use LLaMA in combination with Ollama, please go to ollama.com, create an account. It's totally free, and then click on the top right on Download. Then you can download Ollama for your operating system.
Installing and running Ollama and LLaMA models
After downloading the installer — or if you use Linux then it will automatically install — then open a terminal. I use PowerShell for that, and then you can run the Ollama CLI. This will be automatically installed, and if you run ollama --help you can see that I get an output, so Ollama was correctly installed. The next step is now to run a model. So you can get models here in this model section, and in this example we use LLaMA 3.1. So we click on LLaMA 3.1 and then here is the command we can just copy. We can choose between three different models: one with 8 billion, one with 70 billion, and one with 45 billion parameters in size. Of course the size differs. The smallest one is 4.7 GB, the 70 billion is 40 GB, and the largest one is even 231 GB. So I will stick with the smallest, but to be honest, if you really want a good result then try to run it with the 70 billion parameter model, but my computer, to be honest, is a little bit too old to run that model locally. To run the model, just click on the model you want, copy the command, and go to the CLI. Copy that and this will now download the model if you don't have it yet, and then just wait a few seconds.
Basic Ollama model test and switching to LangChain
So I could now use Ollama like this: 'Hello, how are you?' and communicate with the model like this. So this is not what you want, so we're just going to close it and use it in combination with LangChain. I'm now in VS Code and here on the left you can see multiple files that we will use. First, before we go into the code, we will set up Neo4j. I created a Docker Compose YAML for you. What we're going to do is we will use the neo4j folder. Inside here is a JAR file. This is used to make use of a plugin that we need to create our graph. So to create our database, just run docker compose up and this will set everything up for you and will just work out of the box. This may take a few seconds, and after a few seconds you will see that the database is running. So we can already go to our notebook and first install the required packages.
Environment, packages, and imports
First install the required packages. So we need LangChain, of course. We need langchain-ollama because we're using Ollama. We also use an OpenAI fallback and langchain experimental because the graph solution is currently in the langchain experimental package. We also install neo4j and jupyter-graphx, which is very well suited to display graphs in a Jupyter notebook. Execute that line of code. After installing the required packages, we can now import the required classes. We use multiple classes here from LangChain like the RunnablePassthrough, ChatPromptTemplate, an OutputParser, and so on. We also import Neo4jGraph; this is in the LangChain community package here in this graphs module. We also import ChatOpenAI as a fallback model for Chat LLaMA. In the LangChain experimental package we've got a GraphTransformers module, and from there we import the LLMGraphTransformer that makes use of quite some complicated prompt to convert the data in a way that we can store it in a graph database. So we will also import the graph database from Neo4j, and not only use Neo4j as graph database but we can also use it as a normal vector database. So we make a hybrid approach: we use graph knowledge and also the more standard way to search documents, which is to use an embedding model and then search for the most similar documents to a special query. This is why we make this hybrid. We're also going to use the dotenv package and load our environment variables here in our Jupyter notebook. In the .env file there is an OpenAI API key, there is a Neo4j URI, a Neo4j username, and a Neo4j password. You can just use that as it is, but in the repository it will just be called .env.example.
Connecting to Neo4j and preparing documents
Next step is to create a connection to our database. So we instantiate the Neo4j graph class and this will set up a connection to Neo4j and now we can use this dataset dumit_text.txt. Here you can see that this describes a lot of information about this Italian family: different names, different connections here like Antonio's sister, Amato grandmother, and so on. This should all be represented in our graph later. We just are going to use a text loader to load that into memory and after loading that we are going to use a text splitter to create multiple chunks from that. So that's a very standard way to split the information into smaller chunks that an LLM can easier work with. So we're going to load that and now we're going to set up our LLMGraphTransformer. This GraphTransformer is responsible for turning the documents in this way into a way that Neo4j can work with.
Based on the environment variable here llm_type — currently I didn't set one — the default is Ollama. We are going to instantiate ChatOllama or ChatOpenAI and then we're going to pass that to the constructor of LLMGraphTransformer. The convert_to_graph_documents method will then create all of the relationships between our chunks. So we pass in the documents which we created here and this may take some time to compute. Even for this quite small example it took me like three minutes, so lean back and wait a little bit. It took me more than three minutes to compute the graph, and the first time I even ran into a KeyError. But now it worked and I'm going to show you how the graph nodes look like. So this is a graph document here. You can see that we've got a nodes attribute which is a list of different nodes with an ID. We can see that the ID is like this: 'micoS_family', type is 'family', and then we've got more and more nodes like 'love', 'concept', 'note', 'tradition'. And then we also got the relationships, and this will all be stored in Neo4j.
Storing graph in Neo4j and visualization
So currently we did not start it yet, so we have to run the add_graph_documents method first. We have to provide the graph documents, and now we store everything in Neo4j. This may also take a few seconds. After storing the documents in the database, we can now visualize them. First, we have to connect to the database. We will use the driver method where we pass in our URI; this is stored in the Neo4j_URI environment variable. We also have to provide the username and password for authentication and create our driver instance. Then we create a new session and then we can use the run method of our session to run a query against Neo4j. We will use this quite cryptic query. If you're not used to Neo4j, this actually just means that Neo4j should return all pairs of nodes connected by a relationship of type 'mentions' and we want to return s, r, t. So s is the starting node, r is the relationship, and t is the end node. Then we can run that method and actually visualize our graph. Now we can scroll a little bit down and here we can see this is the complete knowledge graph of our documents. As you can see, this is quite a lot and we can drill down just by scrolling a little bit. Here we can see some entities like Pietro — this is a person — and we can see that Pietro loves the kitchen, loves the sea, and is the parent of another person, Sophia. We can see different entities are modeled with different relationships, and at the end you get this very large knowledge graph. Even for our small dataset this is actually quite a lot of stuff. I really like this kind of graphs, to be honest. Let's now see if that's not just beautiful but actually also helpful.
Creating a vector store from Neo4j and the hybrid approach
The next step is to create a vector store from Neo4j. We will use the Neo4jVector class and use the from_existing_graph method where we only pass the embedding model and from our existing graph we will now compute the embeddings. So we can also perform vector search and then at the end we will turn that vector index into a retriever to have a standardized interface for that. Now we've got a graph database where we stored our documents and also our normal vector store. Now we can perform the retrieval-augmented generation. Since we use a graph database, we need to extract entities from a query to actually perform the retrieval step from our graph database. The graph database needs this kind of entity, so we will create a custom model called Entities which inherits from BaseModel and we want to extract the entities. This can be done by just providing this attribute which is a list of strings, and here's a description for the LLM: we want to extract all persons, organizations, and business entities that are in the text. Then we create a chat prompt template and the system message is: 'You extract organizations, persons, entities from the text.' Then we provide the user input and pipe our prompt template to the LLM with the structural output which makes use of this Entities class. I'm going to show you how that looks like.
Entity extraction, graph retriever, hybrid retriever, and final QA chain
We've got our entity chain and we can invoke it like this: we pass in the question 'Who are Nonia and Jo? One: Coro.' So we've got two names here and we can access the names attribute after performing that .run() method. We can see that the output of this chain is a list of strings with just the names, and these are used to then query the graph database. This is done here. First step in this function graph_retriever is to first extract the entities and then run a query against Neo4j. I'm going to show you how that looks like at the end. We create that function and now we've got our graph_retriever function where we only pass the question. Entities will be extracted and then DB will be queried. We ask: 'Who is Nonia?' If you run that we can see all of the nodes and connections that Nonia has: so 'influenced', 'Konata', 'taught grandchildren', 'influenced fresh pasta', 'influenced Amiko', and is the matriarch of the family. Interesting. Next step is that we create a hybrid retriever. We use the graph retriever and our vector store retriever. We define a function called full_retriever where we set in our graph_retriever function and also use the vector retriever and use the invoke method of that and we get out the most relevant documents. So we've got our relationship graph and the most relevant documents by cosine similarity. At the end we will combine all the docs and return the final dataset. This is what full_retriever will achieve. Then we create a final chain. This is a normal RAG chain; you'll find that kind of chain in almost every beginner tutorial. We've got two variables called context and question. Context is the output of a vector store or any other database, and then we've got the question. All of that will be sent to the LLM. We will create a template and then use the LangChain expression language here to create our final chain. This will create a runnable parallel and I'm going to show you the invoke method here. We just use a string input and we pass that to the full_retriever function and keep the question unchanged, and then pass context and question to our prompt to fill out these variables. After filling out these variables we will pass everything to the LLM and pass the output of the LLM to a string output parser. Let's create that shell and now we can ask: 'Who is Nonia Lucia? Did you teach anyone about restaurants or cooking?' So something about relations. We can see the answer: 'Nonia Lucia is the matriarch of the Mico family and a culinary mentor. She taught her grandchildren the art of Sicilian cooking.' That's actually correct. Okay, so that's it. This is how you can perform Graph RAG with Neo4j. I hope you liked it, and if yes please give the video a like and subscribe to my channel. See you in the next video. Bye-bye.