Skip to main content

Introduction

Retrival Augmented Generation (RAG for short) has reached a high demand in different industries and use cases. In a RAG pipeline your large language model (LLM) of choice consults with the most relevant data you bring (augment) to guide the LLM generation. The augmented data can be of different forms or shape, from your firm’s knowledge base, internal documents, customer calls, or even news documents! Using this arrangement, the final generated response is more than just a generic answer and is informed by the data you add. In this tutorial, we walk through building your RAG using news retrieved from Bigdata, and you can chat with the news!

Setup

The two major components of a RAG pipeline are the Retrieval which finds the most relevant piece of data (documents) to the query (prompt) and the Generation which is all about the language model to generate the final response. Bigdata SDK addresses two challenges, first, providing the data (documents), second, finding the most relevant data document (news chunks). For our tutorial we use OpenAI’s GPT-4o-mini as the generative model. Let’s install the required libraries and dependencies:
Then import the required libraries, initialize a few constants, and instantiate the classes:
Make sure you have entered your Bigdata.com’s credential in order to authenticate with the Bigdata SDK:

Retrieve Documents

We will use Bigdata to augment our RAG pipeline with the retrieved documents. The news data is stored in chunks, and we can search and get the chunks using the Bigdata SDK. We first run a search for the given question (query). Based on similarity of the query and the news documents, the SDK return the news chunks matching our query. With all being said, let’s define a function to get the documents:
In the get_news_data function, we pass the query text and the limit (maximum number) of news chunks we want to retrieve. The function returns a list of news documents. Let’s test the function with a sample query:
From the retrieved headlines, we can say there are significant mentions of the “climate change” in the news related to the US election. As we now have our data ready, we can move on to preparing the LLM to generate responses. But before that, we can tag our data for the LLM to have a better understanding of the nature of each data part. Here is a function to generate XML from the retrieved news documents:
Let’s test the function with the search results we got earlier:

Generate Response

Now that we have our data ready, we can generate the response using the LLM. The function generate_response takes the LLM, query, and the chunks of news as input and returns the response generated by the LLM along with the usage statistics:
We are now able to generate a response for our query and the search results we have:
To calculate the cost of using the OpenAI API, we easily can do the math using the pricing details provided by OpenAI as of early October 2024. The cost is calculated based on the number of tokens used in the generation process. The cost per token is \$0.15 for a million input tokens and \$0.60 for a million output tokens. Here is a function to calculate the cost:
Extremely affordable, right! With having the response generated, we can display it in Markdown format right here in the notebook:
The generation looks good except the citations are not yet linked. Let’s process the citations and add the links to the references at the end of the response:
Let’s process the citations and display the response with the references:
The response is now well-structured and includes the references at the end. The references provide a way to access the original news documents on Bigdata website for further reading or verification.

Conclusion

In this tutorial, we have successfully built an AI RAG pipeline using documents retrieved via the Bigdata SDK. The pipeline is capable of generating responses to questions using the retrieved documents and providing citations to the sources. You can further enhance this pipeline by customizing the retrieval process or using different language models. Happy Searching! 🚀