Query Kusto with OpenAI Embeddings for Semantic Search
A notebook using Azure Data Explorer (Kusto) as a vector database for OpenAI embeddings and cosine similarity search.
Why it matters
Leverage Azure Data Explorer (Kusto) as a vector database to perform semantic similarity searches on AI embeddings generated by OpenAI.
Outcomes
What it gets done
Store precomputed OpenAI embeddings in Kusto.
Convert text queries into embeddings using OpenAI.
Execute cosine similarity searches within Kusto.
Utilize the series-cosine-similarity-fl UDF for efficient search.
Install
Add it to your toolbox
Run in your project directory:
curl -fsSL https://spark.entire.vc/get/oai-gettingstartedwithkustoandopenaiembeddings | bash Steps
Steps in the chain
Overview
Kusto as a Vector database for AI embeddings
A notebook using Azure Data Explorer (Kusto) as a vector database - storing OpenAI embeddings, embedding search queries with the same model, and running cosine similarity search via a Kusto UDF. Use it when you already have data in Azure Data Explorer and want vector similarity search on top of it - requires creating the cosine-similarity UDF in your database first, or the search step fails.
What it does
This notebook shows Azure Data Explorer (Kusto) used as a vector database for OpenAI embeddings, end to end: loading precomputed OpenAI embeddings (of Wikipedia articles, so you don't have to recompute them yourself), storing them in a Kusto table, converting a raw text search query into an embedding via the OpenAI API, and running a cosine similarity search over the stored vectors in Kusto to find the closest matches.
Storing the vectors uses a Spark option (named CreakeIfNotExists in the source notebook) to create the target table automatically from a dataframe's contents if it doesn't already exist. Vectorizing both the stored documents and the search query must use the same embedding model - the notebook specifically calls out using text-embedding-3-small, matching the model the precomputed embeddings were generated with, and supports either an Azure OpenAI or a standard OpenAI key for this step. The similarity search itself relies on Kusto's series-cosine-similarity-fl function, a user-defined function that must be created in the target database before running the search.
When to use - and when NOT to
Use it when you already have (or plan to store) data in Azure Data Explorer and want to add vector similarity search on top of it, rather than standing up a separate dedicated vector database. It depends on a Kusto server instance and either Azure OpenAI credentials or an OpenAI API key being available before you start, and requires manually creating the series-cosine-similarity-fl UDF in your database first - the search step will fail without it.
Inputs and outputs
Input is precomputed OpenAI embeddings (or documents to embed), plus a raw text search query. Output is the set of stored records whose embeddings are closest, by cosine similarity, to the query's embedding.
Integrations
It combines Azure Data Explorer (Kusto) for vector storage and the series-cosine-similarity-fl UDF for similarity search, with OpenAI or Azure OpenAI's text-embedding-3-small model for generating embeddings.
Who it's for
Teams already using Azure Data Explorer who want to add embedding-based semantic search over their existing Kusto data without introducing a separate vector database.
Source README
Kusto as a Vector database for AI embeddings
This notebook provides step-by-step instructions on using Azure Data Explorer (Kusto) as a vector database with OpenAI embeddings.
This notebook presents an end-to-end process of:
- Using precomputed embeddings created by OpenAI API.
- Storing the embeddings in Kusto.
- Converting raw text query to an embedding with OpenAI API.
- Using Kusto to perform cosine similarity search in the stored embeddings
Prerequisites
For the purposes of this exercise we need to prepare a couple of things:
- Azure Data Explorer(Kusto) server instance. https://azure.microsoft.com/en-us/products/data-explorer
- Azure OpenAI credentials or OpenAI API key.
Download precomputed Embeddings
In this section we are going to load prepared embedding data, so you don't have to recompute the embeddings of Wikipedia articles with your own credits.
Store vectors in a Kusto table
Create a table & load the vectors in Kusto based on the contents in the dataframe. The spark option CreakeIfNotExists will automatically create a table if it doesn't exist
Prepare your OpenAI API key
The OpenAI API key is used for vectorization of the documents and queries. You can follow the instructions to create and retrieve your Azure OpenAI key and endpoint. https://learn.microsoft.com/en-us/azure/cognitive-services/openai/tutorials/embeddings
Please make sure to use the text-embedding-3-small model. Since the precomputed embeddings were created with text-embedding-3-small model we also have to use it during search.
If using Azure Open AI
If using Open AI
Only run this cell if you plan to use Open AI for embedding
Generate embedding for the search term
Semantic search in Kusto
We will search the Kusto table for the closest vectors.
We will be using the series-cosine-similarity-fl UDF for similarity search.
Please create the function in your database before proceeding -
https://learn.microsoft.com/en-us/azure/data-explorer/kusto/functions-library/series-cosine-similarity-fl?tabs=query-defined
FAQ
Common questions
Discussion
Questions & comments · 0
Sign In Sign in to leave a comment.