Retrieval augmented generation with OpenAI and OpenSearch^®

Retrieval augmented generation (RAG) is an AI technique for retrieving facts from an external knowledge base to provide large language models (LLMs) on the most accurate, up-to-date context to be able to craft better replies and reduce the risk of hallucination.

This tutorial guides you through an example of using Aiven for OpenSearch® as a backend vector database for OpenAI embeddings and how to perform text, semantic, or mixed search which can serve as basis for a RAG system. We'll use the set of Wikipedia articles as base knowledge to influence the replies of a chatbot.

Why use OpenSearch as a backend vector database?

OpenSearch is a widely adopted open source search/analytics engine. It allows to storing, querying and transforming of documents in a variety of shapes. It also provides fast and scalable functionality to perform both accurate and fuzzy text search. Using OpenSearch as vector database enables you to mix and match semantic and text search queries on top of a performant and scalable engine.

Prerequisites

Before you begin, have the following:

1. An Aiven Account. You can create an account and start a free trial with Aiven by navigating to the signup page and creating a user.
2. An Aiven for OpenSearch service. You can spin up an Aiven for OpenSearch service in minutes in the Aiven Console with the following steps
   • Click on Create service
   • Select OpenSearch
   • Choose the Cloud Provider and Region
   • Select the Service plan (the hobbyist plan is enough for the notebook)
   • Provide the Service name
   • Click on Create service
3. The OpenSearch Connection String. The connection string is visible as Service URI in the Aiven for OpenSearch service overview page.
4. Your OpenAI API key
5. Python and pip installed locally.

Installing dependencies

The tutorial requires the following Python packages:

• openai
• pandas
• wget
• python-dotenv
• opensearch-py

You can install the above packages with:

pip install openai pandas wget python-dotenv opensearch-py
Copy to clipboard

OpenAI key settings

We'll use OpenAI to create embeddings starting from a set of documents, so we need an API key. You can get one from the OpenAI API Key page after logging in.

To avoid leaking the OpenAI key, you can store it as an environment variable named OPENAI_API_KEY.

To store safely the information, create a .env file in the same folder where the notebook is located and add the following line, replacing the <INSERT_YOUR_API_KEY_HERE> with your OpenAI API Key.

OPENAI_API_KEY=<INSERT_YOUR_API_KEY_HERE>
Copy to clipboard

Connect to Aiven for OpenSearch

Once the Aiven for OpenSearch service is in the RUNNING state, we can retrieve the connection string from the Aiven for OpenSearch service page.

Copy the Service URI paramete and store it in the same .env file created above, after replacing the https://USER:PASSWORD@HOST:PORT string with the Service URI.

OPENSEARCH_URI=https://USER:PASSWORD@HOST:PORT
Copy to clipboard

We can now connect to Aiven for OpenSearch by adding the following code in Python:

import os
from opensearchpy import OpenSearch
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

connection_string = os.getenv("OPENSEARCH_URI")

# Create the client with SSL/TLS enabled, but hostname verification disabled.
client = OpenSearch(connection_string, use_ssl=True, timeout=100)
Copy to clipboard

The code above reads the OpenSearch connection string from the .env file (os.getenv("OPENSEARCH_URI")) and creates a client connection using SSL with a timeout of 100 seconds.

Download the dataset

In theory we could use any dataset for this purpose, and you are more than welcome to bring your own. However, for simplicity's sake and to avoid the need to calculate embeddings on a huge dataset of documents, we'll use a set with pre-calculated OpenAI embeddings which score Wikipedia articles. We can get the file and unzip it with:

import wget
import zipfile

embeddings_url = 'https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip'
wget.download(embeddings_url)

with zipfile.ZipFile("vector_database_wikipedia_articles_embedded.zip",
"r") as zip_ref:
    zip_ref.extractall("data")
Copy to clipboard

Let's load the file in a pandas dataframe and check its content with:

import pandas as pd

wikipedia_dataframe = pd.read_csv("data/vector_database_wikipedia_articles_embedded.csv")

wikipedia_dataframe.head()
Copy to clipboard

The file contains:

• id a unique Wikipedia article identifier
• url the Wikipedia article URL
• title the title of the Wikipedia page
• text the text of the article
• title_vector and content_vector the embedding calculated on the title and content of the wikipedia article respectively
• vector_id the id of the vector

Define the OpenSearch mapping to store the OpenAI embeddings

To properly store and query all the fields included in the dataset, we need to define an OpenSearch index settings and mapping optimized for the storage of the information, including the embeddings. For this purpose we can define the settings and the mappings via:

index_settings ={
    "index": {
      "knn": True,
      "knn.algo_param.ef_search": 100
    }
  }

index_mapping= {
    "properties": {
      "title_vector": {
          "type": "knn_vector",
          "dimension": 1536,
          "method": {
            "name": "hnsw",
            "space_type": "l2",
            "engine": "faiss"
        }
      },
      "content_vector": {
          "type": "knn_vector",
          "dimension": 1536,
          "method": {
            "name": "hnsw",
            "space_type": "l2",
            "engine": "faiss"
        },
      },
      "text": {"type": "text"},
      "title": {"type": "text"},
      "url": { "type": "keyword"},
      "vector_id": {"type": "long"}
      
    }
}
Copy to clipboard

The code above:

• Defines an index with knn search enabled. The k-nearest neighbors (k-NN) search searches a vector space (generated by embeddings) in order to retrieve the k closest vectors. You can read more in the OpenSearch k-NN documentation.
• Defines a mapping with:
  • title_vector and content_vector of type knn_vector and 1536 dimension (vector with 1536 entries)
  • text, containing the article text as a text field
  • title, containing the article title as a text field
  • url, containing the article text as a keyword field
  • vector_id, containing the id of the vector as long field

With the settings and mappings defined, we can now create the openai_wikipedia_index index in Aiven for OpenSearch with:

index_name = "openai_wikipedia_index"
client.indices.create(
        index=index_name, 
        body={"settings": index_settings, "mappings":index_mapping}
        )
Copy to clipboard

Load data into OpenSearch

With the index created, the next step is to parse the the pandas dataframe and load the data into OpenSearch using the Bulk APIs. The following function loads a set of rows in the dataframe:

def dataframe_to_bulk_actions(df):
    for index, row in df.iterrows():
        yield {
            "_index": index_name,
            "_id": row['id'],
            "_source": {
                'url' : row["url"],
                'title' : row["title"],
                'text' : row["text"],
                'title_vector' : json.loads(row["title_vector"]),
                'content_vector' : json.loads(row["content_vector"]),
                'vector_id' : row["vector_id"]
            }
        }
Copy to clipboard

To speed up ingestion we can load the data in batches of 200 rows.

from opensearchpy import helpers
import json

start = 0
end = len(wikipedia_dataframe)
batch_size = 200
for batch_start in range(start, end, batch_size):
    batch_end = min(batch_start + batch_size, end)
    batch_dataframe = wikipedia_dataframe.iloc[batch_start:batch_end]
    actions = dataframe_to_bulk_actions(batch_dataframe)
    helpers.bulk(client, actions)
Copy to clipboard

Once all the documents are loaded, we can try a query to retrieve the documents containing Pizza:

res = client.search(index=index_name, body={
    "_source": {
        "excludes": ["title_vector", "content_vector"]
    },
    "query": {
        "match": {
            "text": {
                "query": "Pizza"
            }
        }
    }
})

print(res["hits"]["hits"][0]["_source"]["text"])
Copy to clipboard

The result is the Wikipedia article talking about Pizza:

Pizza is an Italian food that was created in Italy (The Naples area). It is made with different toppings. Some of the most common toppings are cheese, sausages, pepperoni, vegetables, tomatoes, spices and herbs and basil. These toppings are added over a piece of bread covered with sauce. The sauce is most often tomato-based, but butter-based sauces are used, too. The piece of bread is usually called a "pizza crust". Almost any kind of topping can be put over a pizza. The toppings used are different in different parts of the world. Pizza comes from Italy from Neapolitan cuisine. However, it has become popular in many parts of the world.

History 
The origin of the word Pizza is uncertain. The food was invented in Naples about 200 years ago. It is the name for a special type of flatbread, made with special dough. The pizza enjoyed a second birth as it was taken to the United States in the late 19th century.

...
Copy to clipboard

Encode chatbot questions with OpenAI text-embedding-ada-002 model

To perform a semantic search, we need to calculate question encodings with the same embedding model used to encode the documents at index time. In this example, we need to use the text-embedding-ada-002 model.

from openai import OpenAI

# Define model
EMBEDDING_MODEL = "text-embedding-ada-002"

# Define the Client
openaiclient = OpenAI(
    # This is the default and can be omitted
    api_key=os.getenv("OPENAI_API_KEY"),
)

# Define question
question = 'is Pineapple a good ingredient for Pizza?'

# Create embedding
question_embedding = openaiclient.embeddings.create(input=question, model=EMBEDDING_MODEL)
Copy to clipboard

Run semantic search queries with OpenSearch

With the above embedding calculated, we can now run semantic searches against the OpenSearch index to retrieve the necessary context for the retrieval-augmented generation. We're using knn as query type and scan the content of the content_vector field.

response = client.search(
  index = index_name,
  body = {
      "size": 15,
      "query" : {
        "knn" : {
          "content_vector":{
          "vector":  question_embedding.data[0].embedding,
          "k": 3
        }
      }
    }
  }
)

for result in response["hits"]["hits"]:
  print("Id:" + str(result['_id']))
  print("Score: " + str(result["_score"]))
  print("Title: " + str(result["_source"]["title"]))
  print("Text: " + result["_source"]["text"][0:100])
Copy to clipboard

The result is the list of articles ranked by score:

Id:13967
Score: 13.94602
Title: Pizza
Text: Pizza is an Italian food that was created in Italy (The Naples area). It is made with different topp
Id:90918
Score: 13.754393
Title: Pizza Hut
Text: Pizza Hut is an American pizza restaurant, or pizza parlor. Pizza Hut also serves salads, pastas and
Id:66079
Score: 13.634726
Title: Pizza Pizza
Text: Pizza Pizza Limited (PPL), doing business as Pizza Pizza (), is a franchised Canadian pizza fast foo
Id:85932
Score: 11.388243
Title: Margarita
Text: Margarita may mean:
 The margarita, a cocktail made with tequila and triple sec
Margarita Island, a 
Id:13968
Score: 10.576359
Title: Pepperoni
Text: Pepperoni is a meat food that is sometimes sliced thin and put on pizza. It is a kind of salami, whi
Id:87088
Score: 9.424156
Title: Margherita of Savoy
...
Copy to clipboard

Use OpenAI Chat Completions API to generate a RAG reply

The step above retrieves the content semantically similar to the question. Now let's use OpenAI chat completions function to return a retrieval-augmented generated reply based on the information retrieved.

# Retrieve the text of the first result in the above dataset
top_hit_summary = response['hits']['hits'][0]['_source']['text']

# Craft a reply
response = openaiclient.chat.completions.create(
    model="gpt-3.5-turbo",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Answer the following question:" 
            + question 
            + "by using the following text:" 
            + top_hit_summary
        }
        ]
    )

choices = response.choices

for choice in choices:
    print(choice.message.content)
Copy to clipboard

The result is going to be similar to the below:

Pineapple is a contentious ingredient for pizza, and it is commonly used at Pizza Pizza Limited (PPL), a Canadian fast-food restaurant with locations throughout Ontario. The menu includes a variety of toppings such as pepperoni, pineapples, mushrooms, and other non-exotic produce. Pizza Pizza has been a staple in the area for over 30 years, with over 500 locations in Ontario and expanding across the nation, including recent openings in Montreal and British Columbia.
Copy to clipboard

Conclusion

OpenSearch is a powerful tool providing both text and vector search capabilities. Used alongside OpenAI APIs allows you to craft personalized AI applications able to augment the context based on semantic search, and return responses to queries that are augmented by AI. A logical next step would be to pair OpenSearch with another storage system to, for example, store the responses that your customers find useful and train the model further. Building an end to end system including a databases like PostgreSQL® and a streaming integration with Apache Kafka® could provide the resiliency of a relational database and the hybrid search capability of OpenSearch with data feeding in near real time.

You can try Aiven for OpenSearch, or any of the other Open Source tools, in the Aiven platform free trial by signing up.