Using LlamaIndex for AI-Assisted Knowledge Management

Introduction

One of the hottest questions of the moment for a lot of strategy and decision makers across most industries worldwide is:

How can AI help my business?

Afterall, with great disruption also comes great opportunity. A sound business strategy should not ignore emerging changes in the market. We’re still at the early stages of understanding AI and I’m not going to provide a definitive answer to this question in this article, but the good news is that this article should provide a part of the answer.

Knowledge is power, right?

And yet we all know how it is to struggle trying to retain and efficiently re-use the knowledge we gather.

We strive to learn from our successes and mistakes and we invest a lot of time and money in building fancy knowledge bases just to discover later that unfortunately we keep repeating the same mistakes and reinventing the wheel.

In my experience as a consultant, the biggest issue (especially for medium and large companies) is not a lack of knowledge but on the contrary, it’s too much knowledge and an inability to use it in a timely and effective manner.

The solution

This article presents a very simple yet effective way of indexing large quantities of pre-existing knowledge that can later be retrieved by natural language queries or integrated with chatbot systems.

As usual, take it as a starting point. The code example is trivial and lacks any error handling, but provides the building blocks to work from.

My example builds on your existing knowledge base and leverages LlamaIndex and the power of Large Language Models (in this case GPT 3.5 Turbo from OpenAI).

Why LlamaIndex? Well, created by Jerry Liu, LlamaIndex is a robust open-source resource that empowers you to organize and search your data for a variety of applications, including answering questions, summarizing information or serving as a part of a chatbot system. It provides data connectors to ingest your existing data sources in many different formats (such as text files, PDF, docs, SQL, etc.). It then allows you to structure your data (via indices or graphs) so that this data can be easily used with LLMs. In many ways, it is similar to Langchain but more focused on data storage and retrieval instead of automated AI agents.

In short, this article will show you how, with just a few lines of code, you can index your enterprise's knowledge base and then have the ability to query and retrieve information from GPT 3.5 Turbo with your own knowledge base on top of that in the most natural way: plain English.

Logic diagram

Creating the Index

Retrieving the knowledge

Prerequisites

Make sure you check these points before you start writing the code:

1. Make sure you store your OpenAI API key in a local environment variable for secure and efficient access. The code works on the assumption that the API key is stored on your local environment (OPENAI_API_KEY).
2. I’ve used Python v3.11. If you’re running an older version, an update is recommended to make sure you don’t run into any compatibility issues.
3. Install the requirements:
   
   pip install openai
pip install llama-index
   
    
4. Create a subfolder in your .PY file’s location (in my example the name of the subfolder is ‘stories’). You will store your knowledge base in .TXT files in that location. If your knowledge articles are in different formats (e.g., PDF or DOCX) you will have to:
   1. Change the code to use a different LlamaIndex data connector (https://gpt-index.readthedocs.io/en/latest/reference/readers.html) – this is my recommended solution – or:
   2. Convert all your documents in .TXT format and use the code as it is.

For my demo, I have created (with the help of GPT-4) three fictional stories that will represent our proprietary ‘knowledge’ base:

Your ‘stories’ folder should now look like this:

The Code

For the sake of simplicity, I’ve split the functionality into two different scripts:

• Index_stories.py (responsible for reading the ‘stories’ folder, creating an index and saving it for later queries)
• Query_stories.py (demonstrating how to query GPT 3.5 and then filter the AI response through our own knowledge base)

Let’s begin with Index_stories.py:

from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader

# Loading from a directory
documents = SimpleDirectoryReader('stories').load_data()

# Construct a vector store index
index = GPTVectorStoreIndex.from_documents(documents)

# Save your index to a .json file
index.storage_context.persist()

As you can see, the code is using SimpleDirectoryReader from LlamaIndex to read all .TXT files from the ‘stories’ folder. It then creates a simple vector index that can later be used to run queries over the content of these documents.

In case you’re wondering what a vector index represents, imagine you're in a library with thousands of books, and you're looking for a specific book. Instead of having to go through each book one by one, this index acts in a similar way to a library catalog.  It helps you find the book you're looking for quickly.

In the context of this code, GPTVectorStoreIndex is like that library catalog. It's a tool that helps organize and find specific pieces of information (like documents or stories) quickly and efficiently. When you ask a question, it looks through all the information it has and finds the most relevant answer for you. It's like a super-efficient librarian that knows exactly where everything is.

The last line of the code saves the index in a sub-folder called ‘storage’ so that we do not have to recreate it every time and we are able to reuse it in the future.

Now, for the querying part. Here’s the second script: Query_stories.py:

from llama_index import GPTVectorStoreIndex, StorageContext, load_index_from_storage

import openai
import os

openai.api_key = os.getenv('OPENAI_API_KEY')

def prompt_chatGPT(task):
    response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": task}
                ]
                )
    AI_response = response['choices'][0]['message']['content'].strip()
   
    return AI_response

# rebuild storage context
storage_context = StorageContext.from_defaults(persist_dir="storage")

# load index
index = load_index_from_storage(storage_context)

# Querying GPT 3.5 Turbo
prompt = "Tell me how Tortellini Macaroni's brother managed to conquer Rome."
answer = prompt_chatGPT(prompt)
print('Original AI answer: ' + answer +'\n\n')

# Refining the answer in the context of our knowledge base
query_engine = index.as_query_engine()
response = query_engine.query(f'The answer to the following prompt: "{prompt}" is :"answer". If the answer is aligned to our knowledge, return the answer. Otherwise return a corrected answer')

print('Custom knowledge answer: ' + str(response))

How it works

After indexing the ‘stories’ folder, once you run the query_stories.py script the code will first load the index from the ‘storage’ sub-folder. It then prompts the GPT 3.5 Turbo model with a hard-coded question: “Tell me how Tortellini Macaroni’s brother managed to conquer Rome”. After the response is received, it queries our ‘stories’ to see if the answer aligns with our ‘knowledge’. Then, you’ll receive two answers.

The first one is the original answer from GPT 3.5 Turbo:

As expected, the AI model identified Mr. Spaghetti as a potentially fictional character and could not find any historical references of him conquering Rome.

The second answer though, checks with our ‘knowledge’ and, because we have different information in our ‘stories’, it modifies the answer into:

If you’ve read the three GPT-4-created stories you’ve noticed that Story1.txt mentions Biscotti as a fictional conqueror of Rome but not his brother and Story2.txt mentions Tortellini and his farm adventures but does not mention any relationship with Biscotti. Only the third story (Story3.txt) describes the nature of their relationship.

This shows not only that the vector index managed to correctly record the knowledge from the individual stories but also proves the query function managed to provide a contextual response to our question.

In addition to the Vector Store Index, there are several other types of indexes that can be used depending on the specific needs of your project.

For instance, the List Index simply stores Nodes as a sequential chain, making it a straightforward and efficient choice for certain applications, such as where the order of data matters and where you frequently need to access all the data in the order it was added. An example might be a timeline of events or a log of transactions, where you often want to retrieve all entries in chronological order.

Another option is the Tree Index which builds a hierarchical tree from a set of Nodes, which can be particularly useful when dealing with complex, nested data. For instance, if you're building a file system explorer, a tree index would be a good choice because files and directories naturally form a tree-like structure.

There’s also the Keyword Table Index, which extracts keywords from each Node and builds a mapping from each keyword to the corresponding Nodes. This can be a powerful tool for text-based queries, allowing for quick and precise retrieval of relevant information.

Each of these indexes offers unique advantages, and the choice between them would depend on the nature of your data and the specific requirements of your use case.

Conclusion

Now, think about the possibilities. Instead of fictional stories, we could have a collection of our standard operating procedures, process descriptions, knowledge articles, disaster recovery plans, change schedules and so on and so forth. Or, as another example we could build a chatbot that can solve generic users requests simply using GPT3.5 knowledge but forwards more specific issues (indexed from our knowledge base) to a support team.

This brings unlimited potential in automation of our business processes and improvement of the decision-making process. You get the best of both worlds: the power of Large Language Models combined with the value of your own knowledge base.

Security considerations

Working on this article made me realize that we cannot really trust our interactions with an AI model unless we are in full control of the entire technology stack. Just because the interface might look familiar, it doesn’t necessary mean that bad actors cannot compromise the integrity of the data by injecting false or censored responses to our queries. But that’s a story for another time!

Final note

This article barely scratches the surface of the full capabilities of LlamaIndex. It is not meant to be a comprehensive guide into this topic but rather serve as an example start point for integrating AI technologies in our day-to-day business processes. I encourage you to have an in-depth study of LlamaIndex’s capabilities (https://gpt-index.readthedocs.io/en/latest/) if you want to take advantage of its full capabilities.

About the Author

Andrei Gheorghiu is an experienced trainer with a passion for helping learners achieve their maximum potential. He always strives to bring a high level of expertise and empathy to his teaching.

With a background in IT audit, information security, and IT service management, Andrei has delivered training to over 10,000 students across different industries and countries. He is also a Certified Information Systems Security Professional and Certified Information Systems Auditor, with a keen interest in digital domains like Security Management and Artificial Intelligence.

In his free time, Andrei enjoys trail running, photography, video editing and exploring the latest developments in technology.

You can connect with Andrei on:

• LinkedIn: https://www.linkedin.com/in/gheorghiu/
• Twitter: https://twitter.com/aqg8017