Creating a Persitant-Dynamic Backend

Now that we have our graph triplets, we would like to do inference with them! But hold on. In real world applications our knowledge graph may change over time. In this case we will want to be able to handle triplets being both added and deleted. Additionally, a persitent database for this information will be crucial in the case of crashes or other unforseen issues!

This notebook will show you how to connect a simple knowledge-graph RAG agent to a database that is being actively updated, and we will do this without sacrificing performance! Let's get started.

The first thing we will need to do are a great many imports.

Setting-up our Backend.

Now that we have all of those imports done, we can start with the good stuff. We will be using ArangoDB as a backend. ArangoDB is a graph database that works very well with NetworkX, a popular Python library for graph analysis. Better yet, ArangoDB has a cuGraph persistance layer available, which will allow us to accelerate NetworkX on GPU! If you want to learn more about this check out the blog HERE.

We will get started by launching the database! The following command will launch an arangodb instance on port 8530 with the username "root" and password "ilovekgrag".

Now we need to populate our database. I hope you will forgive us, but we took the liberty of creating some convenient CSV's for reading in the triples from the previous notebook. The details of this process can be found in the file "data/getcsv.py" if you are interested! We will first need to set up the structure of our graph database, and then populate with our CSV files. Let's do that.

Now the backend is set up to accept our CSVs! Let's place our edges and vertices into the database. These commands get big, and we have variables in Python already defined that will be useful, so let's just define a way to run these commands in Python. This will also be useful later.

Awesome, now we have a database we can really work with! But we still have a problem. This database is not on GPU. If we want to serve end users concurrently this is going to be a problem, and really limit our number of feasible inferences. So we need to read the data from this backend into NetworkX, and then use cuGraph as our backend to perform RAG tasks on GPU! Let's do that.

You may be curious about the naming convention here. Why are we calling them cpu_graph and background_graph? Since our ArangoDB backend is not stored on the GPU, it's NetworkX graph goes through an intermediate CPU phase, and is then placed on the GPU. This explains the cpu_graph name, but why background_graph? Well, it takes time to move memory around. And since we are trying to create a dynamic knowledge graph, it is beneficial to keep two copies of our graph at a time. One copy can be continuously updated as updates are streamed in, while the other can be used for inference in that time. This will require some multi-threading, but we will keep it as simple as possible, and let you imagine more complicated work loads with continuous-asynchronous user queries and updates!

To make this all work we will need a couple functions. Mainly, a function to handle database alterations, and a function to "swap" our working_graph and out background_graph. Let's build those.

Now let's test it!

Great it works! We will add those edges back later when we try to simulate a real workload. Now the bulk of our backend functionality is all prepared for us. We should probably get around to adding our working_graph as well as a function for exchaning it with background_graph when the time comes!

The last step we need for our dynamic backend to work is the ability to update the backend after making Arango updates! Let's make a function for it, and we can call that from an independent thread later.

And that's that! Notice that our working_graph now has less nodes than the background_graph since we removed a bunch! We will get a chance to use this more shortly, but first we need to make an agent to actualy perform our inference with.

Building our RAG Agent

Now we want to put a RAG agent on top of our backend to allow for user inferences. We will use some NIM endpoints for this and connect to them through LangChain.

Since this is just a simple example, our agent doesn't need to do too much. Given a knowledge graph and a user query, we want it to retrieve relevent relationships from the graph to return to the user. Fortunately there is a way to do this using the GraphQAChain in LangChain. Since we are using the GPU, and not the CPU, some small changes to the QA chain need to be altered. This has already been done, and the changes can be found within the qa_chain_overrides.py file. We have actally already used some of these earlier in the notebook!

Testing the system.

Now that we have an agent for chat completion, and an agent for handling dynamic information, let's see how this comes together in a test enfironment. We will simulate this by looping, and performing queries as often as well can. First let's get some baseline times for inference and background updates!

Wow! Our backend changes are pretty fast! Maybe we don't need to multithread anything. To make sure, let's see what happens when we increase the number of edges we are adding and removing.

Unfortunately, having many seconds of down-time where we aren't serving user-queries is not very good. So, let's parallelize this solution. The following code uses a mixture of threading and multiprocessing to achieve parallelism!

Awesome! That is way better! We are getting slightly longer update times, but achieving the same inference times we were getting before! Users will be far less upset than when we had long dead periods.