Perform Semantic Search with Elasticsearch and OpenAI
A notebook indexing OpenAI embeddings into Elasticsearch and running kNN semantic search over them.
Maintainer of this project? Claim this page to edit the listing.
1.0.0Add to Favorites
Why it matters
Leverage Elasticsearch as a vector database to perform semantic searches on your data, powered by OpenAI embeddings.
Outcomes
What it gets done
Index vector embeddings into Elasticsearch.
Encode user queries using OpenAI's embedding models.
Execute kNN semantic search queries against Elasticsearch.
Retrieve relevant documents based on query meaning.
Install
Add it to your toolbox
Run in your project directory:
curl -fsSL https://spark.entire.vc/get/oai-elasticsearch-semantic-search | bash Steps
Steps in the chain
Overview
Semantic search using Elasticsearch and OpenAI
A notebook that indexes documents into Elasticsearch with dense_vector fields and OpenAI embeddings, then runs kNN semantic search against them. Use it as a starting point for OpenAI-embedding-based semantic search in Elasticsearch - a companion notebook builds on this same setup to add RAG if you need generated answers, not just ranked results.
What it does
This notebook indexes the OpenAI Wikipedia vector dataset into Elasticsearch, embeds a question with the OpenAI embeddings endpoint, and runs semantic search on the Elasticsearch index using that encoded question. It connects to an Elastic Cloud deployment using a Cloud ID and password (found on the deployment's page at cloud.elastic.co/deployments), downloads and extracts the Wikipedia embeddings dataset, and loads the resulting CSV into a pandas DataFrame for bulk indexing.
It creates an Elasticsearch index whose mapping uses the dense_vector field type for both a title_vector and a content_vector field, the special field type that lets Elasticsearch store and later kNN-search dense vectors. Indexing itself generates bulk actions from each DataFrame row and pushes them through the Python client's bulk-API helpers in batches of 100, since the dataset is large. After a quick sanity check with a plain match query, it encodes a search question with the same embedding model used at index time - text-embedding-3-small - requiring an OpenAI API key, then runs a k-nearest-neighbors query against the indexed vectors via Elasticsearch's kNN search option, using a small helper function to pretty-print the results for readability.
When to use - and when NOT to
Use it as a starting point for standing up semantic search over your own documents in Elasticsearch, using OpenAI embeddings for both indexing and querying, and adapt it by swapping in your own dataset and experimenting with different embedding models once the basic pipeline works. It's explicitly a stepping stone: the notebook itself points to a companion notebook that builds on this same indexing and search setup to add retrieval augmented generation (RAG) with the Chat Completions API - if you need generated answers rather than just ranked search results, that's the next step, not this one.
Inputs and outputs
Input is a document dataset (demonstrated with OpenAI's Wikipedia embeddings) and an OpenAI API key for query-time embedding. Output is a populated Elasticsearch index with dense-vector fields and a ranked list of kNN search results for a given query.
Integrations
It integrates Elasticsearch (via Elastic Cloud, using its Python client's bulk helpers and kNN query support) with the OpenAI embeddings endpoint, specifically the text-embedding-3-small model.
Who it's for
Developers setting up vector-based semantic search over their own data in Elasticsearch using OpenAI embeddings, as a foundation to build further retrieval or RAG applications on top of, without having to design the indexing and kNN query mechanics from scratch.
Source README
Semantic search using Elasticsearch and OpenAI
This notebook demonstrates how to:
- Index the OpenAI Wikipedia vector dataset into Elasticsearch
- Embed a question with the OpenAI
embeddingsendpoint - Perform semantic search on the Elasticsearch index using the encoded question
Install packages and import modules
Connect to Elasticsearch
ℹ️ We're using an Elastic Cloud deployment of Elasticsearch for this notebook.
If you don't already have an Elastic deployment, you can sign up for a free Elastic Cloud trial.
To connect to Elasticsearch, you need to create a client instance with the Cloud ID and password for your deployment.
Find the Cloud ID for your deployment by going to https://cloud.elastic.co/deployments and selecting your deployment.
Download the dataset
In this step we download the OpenAI Wikipedia embeddings dataset, and extract the zip file.
Read CSV file into a Pandas DataFrame
Next we use the Pandas library to read the unzipped CSV file into a DataFrame. This step makes it easier to index the data into Elasticsearch in bulk.
Create index with mapping
Now we need to create an Elasticsearch index with the necessary mappings. This will enable us to index the data into Elasticsearch.
We use the dense_vector field type for the title_vector and content_vector fields. This is a special field type that allows us to store dense vectors in Elasticsearch.
Later, we'll need to target the dense_vector field for kNN search.
Index data into Elasticsearch
The following function generates the required bulk actions that can be passed to Elasticsearch's Bulk API, so we can index multiple documents efficiently in a single request.
For each row in the DataFrame, the function yields a dictionary representing a single document to be indexed.
As the dataframe is large, we will index data in batches of 100. We index the data into Elasticsearch using the Python client's helpers for the bulk API.
Let's test the index with a simple match query.
Encode a question with OpenAI embedding model
To perform semantic search, we need to encode queries with the same embedding model used to encode the documents at index time.
In this example, we need to use the text-embedding-3-small model.
You'll need your OpenAI API key to generate the embeddings.
Run semantic search queries
Now we're ready to run queries against our Elasticsearch index using our encoded question. We'll be doing a k-nearest neighbors search, using the Elasticsearch kNN query option.
First, we define a small function to pretty print the results.
Now let's run our kNN query.
Next steps
Success! Now you know how to use Elasticsearch as a vector database to store embeddings, encode queries by calling the OpenAI embeddings endpoint, and run semantic search.
Play around with different queries, and if you want to try with your own data, you can experiment with different embedding models.
ℹ️ Check out our other notebook Retrieval augmented generation using Elasticsearch and OpenAI. That notebook builds on this example to demonstrate how to use Elasticsearch together with the OpenAI chat completions API for retrieval augmented generation (RAG).
FAQ
Common questions
Discussion
Questions & comments · 0
Sign In Sign in to leave a comment.