SingleStore Notebooks
New
Semantic Search with OpenAI QA
Notebook
Note
This notebook can be run on a Free Starter Workspace. To create a Free Starter Workspace navigate to Start using the left nav. You can also use your existing Standard or Premium workspace with this Notebook.
Question answering using embeddings
GPT models are trained on a fixed snapshot of text and then frozen. Ask one about something that happened after its training data ends, or about anything in your own private documents, and it has two ways to fail: admit it doesn't know, or invent a confident-sounding answer.
This notebook demonstrates the standard fix, usually called retrieval-augmented generation (RAG). Instead of hoping the model already knows the answer, you:
Search a body of text you control for the passages most relevant to the question.
Ask the model the question with those passages pasted into the prompt.
The model never learns anything new. You are simply putting the right reference material in front of it before asking. This is dramatically cheaper and faster than fine-tuning, and it has a property fine-tuning lacks: you can see exactly which source text produced the answer.
The example corpus is roughly 6,000 sections of Wikipedia articles about the 2022 Winter Olympics — events that postdate the model's training data, so we know any correct answer had to come from retrieval rather than memory.
In this Notebook you will use a combination of Semantic Search and a Large Langauge Model (LLM) to build a basic Retrieval Augmented Generation (RAG) application. For a great introduction into what RAG is, please read A Beginner's Guide to Retrieval Augmented Generation (RAG)
Prerequisites for interacting with ChatGPT
Install OpenAI package
Let's start by installing the openai Python package via running the command below.
In [1]:
1import sys2!{sys.executable} -m pip install --quiet --upgrade "openai<4"
In [2]:
1!pip check
Running !pip check should result in something like: "jupyter-resource-usage 1.1.0 has requirement psutil~=5.6, but you have psutil 7.2.2." If you have to perform "restart Kernel" then you'll need to start from the top of the notebook and run the python code in order, otherwise it will likely throw an error.
In [3]:
1import openai, sys2 3print(openai.__version__, sys.version.split()[0])4 5# Report whether this image's brotli is one of the affected ones. Informational only --6# the OpenAI client cell below sidesteps the decoder either way.7try:8 import brotli9 10 brotli.Decompressor().process(brotli.compress(b'probe'), output_buffer_limit=1 << 20)11 print(f'brotli {brotli.__version__}: accepts output_buffer_limit')12except TypeError:13 print(f'brotli {brotli.__version__}: too old for httpx2 (expected on this image) -- '14 f'the OpenAI client below will not request brotli-encoded responses')15except ImportError:16 print('brotli not installed -- nothing to work around')
Connect to ChatGPT and display the response
In [4]:
1import openai2 3EMBEDDING_MODEL = "text-embedding-ada-002"4GPT_MODEL = "gpt-3.5-turbo"
Before you run this notebook
You will need an OpenAI API key to use the openai Python library, and your account
will need credits on it. Three steps:
1. Create an account. Go to platform.openai.com and sign up. Note that this is the API platform and is billed separately from a ChatGPT Plus subscription — paying for ChatGPT does not give you API credits.
2. Add credits. Go to
Settings → Billing and add
a payment method, then purchase credits. API access is prepaid: with a zero balance,
every call fails with a 429 error reading insufficient_quota, even though your key
is perfectly valid.
A full run of this notebook costs a fraction of a cent — the corpus ships with its embeddings already computed, so only a handful of small API calls are actually billed. The practical minimum is OpenAI's smallest credit purchase, which is far more than this notebook will consume.
3. Create an API key. Go to API keys and click Create new secret key. Copy it immediately — the full key is displayed only once, and you cannot retrieve it later. If you lose it, delete that key and create another.
The cell below will prompt you for the key when you run it. Paste it into the prompt (Cmd+V or Ctrl+V) and press Enter. The input is masked, so you will not see the characters appear.
Do not paste your key directly into a code cell. Anything typed into a cell is saved inside the
.ipynbfile and travels with it to anyone you share, commit, or export the notebook to. The prompt below keeps the key in memory for this session only. If a key does get committed or shared, delete it on the API keys page and create a new one.
In [5]:
1import getpass2import os3 4import httpx25 6os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')7 8# This image ships a brotli older than 1.2.0, but httpx2 2.12 calls the brotli decoder9# with `output_buffer_limit=`, a keyword only accepted from 1.2.0 onwards. A10# brotli-encoded response then dies inside the decoder with `TypeError: process() takes11# no keyword arguments`, logged directly beneath a `200 OK` -- so it reads like a network12# or API-key fault when it is neither. Upgrading brotli from a cell does not help,13# because this kernel has already imported it. Simply not advertising brotli means the14# broken decoder is never reached, and gzip is perfectly good for JSON.15http_client = httpx2.Client(headers={'Accept-Encoding': 'gzip, deflate'})16 17client = openai.OpenAI(http_client=http_client)
In [6]:
1e = client.embeddings.create(model=EMBEDDING_MODEL, input="test")2print("embedding dims:", len(e.data[0].embedding))3r = client.chat.completions.create(model=GPT_MODEL, messages=[{"role": "user", "content": "say ok"}])4print("chat:", r.choices[0].message.content)
Test the connection.
In [7]:
1response = client.chat.completions.create(2 model=GPT_MODEL,3 messages=[4 {"role": "system", "content": "You are a helpful assistant."},5 {"role": "user", "content": "Who won the gold medal for curling in Olymics 2022?"},6 ]7)8 9print(response.choices[0].message.content)
Get the data about Winter Olympics and provide the information to ChatGPT as context
1. Install and import libraries
In [8]:
1!pip install tabulate tiktoken wget --quiet
In [9]:
1import json2import numpy as np3import os4import pandas as pd5import wget
2. Fetch the CSV data and read it into a DataFrame
Download pre-chunked text and pre-computed embeddings. This file is ~200 MB, so may take a minute depending on your connection speed.
In [10]:
1embeddings_url = "https://cdn.openai.com/API/examples/data/winter_olympics_2022.csv"2embeddings_path = "winter_olympics_2022.csv"3 4if not os.path.exists(embeddings_path):5 wget.download(embeddings_url, embeddings_path)6 print("File downloaded successfully.")7else:8 print("File already exists in the local file system.")
Here we are using the converters= parameter of the pd.read_csv function to convert the JSON
array in the CSV file to numpy arrays.
In [11]:
1def json_to_numpy_array(x: str | None) -> np.ndarray | None:2 """Convert JSON array string into numpy array."""3 return np.array(json.loads(x)) if x else None4 5df = pd.read_csv(embeddings_path, converters=dict(embedding=json_to_numpy_array))6df
In [12]:
1df.info(show_counts=True)
3. Set up the database
Action Required
If you have a Free Starter Workspace deployed already, select the database from drop-down menu at the top of this notebook. It updates the connection_url to connect to that database.
Create the database.
In [13]:
1shared_tier_check = %sql show variables like 'is_shared_tier'2if not shared_tier_check or shared_tier_check[0][1] == 'OFF':3 %sql DROP DATABASE IF EXISTS winter_wikipedia;4 %sql CREATE DATABASE winter_wikipedia;
Action Required
Make sure to select the winter_wikipedia database from the drop-down menu at the top of this notebook. It updates the connection_url which is used by the %%sql magic command and SQLAlchemy to make connections to the selected database.
In [14]:
1%%sql2CREATE TABLE IF NOT EXISTS winter_olympics_2022 /* Creating table for sample data. */(3 id INT PRIMARY KEY,4 text TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci,5 embedding BLOB6);
4. Populate the table with our DataFrame
Create a SQLAlchemy connection.
In [15]:
1import singlestoredb as s22 3conn = s2.create_engine().connect()
Use the to_sql method of the DataFrame to upload the data to the requested table.
In [16]:
1df.to_sql('winter_olympics_2022', con=conn, index=True, index_label='id', if_exists='append', chunksize=1000)
5. Do a semantic search with the same question from above and use the response to send to OpenAI again
Step 1: Search
strings_ranked_by_relatedness() embeds your query as a vector, compares it against
the precomputed embedding of every section in the corpus, and returns the closest
matches ranked by cosine similarity.
This is not keyword search. Nothing here is matching the literal strings "curling" or "gold". The query and the documents are both mapped into a space where proximity means similar meaning, so a passage can rank highly without sharing a single word with the query. That property is the entire reason embeddings are useful for this.
In [17]:
1import sqlalchemy as sa2 3 4def get_embedding(text: str, model: str = 'text-embedding-ada-002') -> str:5 """Return the embeddings."""6 return [x.embedding for x in client.embeddings.create(input=[text], model=model).data][0]7 8 9def strings_ranked_by_relatedness(10 query: str,11 df: pd.DataFrame,12 table_name: str,13 relatedness_fn=lambda x, y: 1 - spatial.distance.cosine(x, y),14 top_n: int=100,15) -> tuple:16 """Returns a list of strings and relatednesses, sorted from most related to least."""17 18 # Get the embedding of the query.19 query_embedding_response = get_embedding(query, EMBEDDING_MODEL)20 21 # Create the SQL statement.22 stmt = sa.text(f"""23 SELECT24 text,25 DOT_PRODUCT_F64(JSON_ARRAY_PACK_F64(:embedding), embedding) AS score26 FROM {table_name}27 ORDER BY score DESC28 LIMIT :limit29 """)30 31 # Execute the SQL statement.32 results = conn.execute(stmt, dict(embedding=json.dumps(query_embedding_response), limit=top_n))33 34 strings = []35 relatednesses = []36 37 for row in results:38 strings.append(row[0])39 relatednesses.append(row[1])40 41 # Return the results.42 return strings[:top_n], relatednesses[:top_n]
In [18]:
1from tabulate import tabulate2 3strings, relatednesses = strings_ranked_by_relatedness(4 "curling gold medal",5 df,6 "winter_olympics_2022",7 top_n=58)9 10for string, relatedness in zip(strings, relatednesses):11 print(f"{relatedness=:.3f}")12 print(tabulate([[string]], headers=['Result'], tablefmt='fancy_grid'))13 print('\n\n')
Reading these results
Three things in the output above are worth noticing.
The retrieved text is raw and messy. These sections are unprocessed Wikipedia
markup — {{Medals table, [[Niklas Edin]], flagIOC template calls, HTML <br>
tags. Nobody cleaned it. This matters more than it looks: retrieval does not require a
tidy, curated knowledge base, and the language model is perfectly capable of reading
structured junk and pulling the answer out of it. A common mistake is to over-invest in
preprocessing before checking whether it changes the answer.
The relatedness scores are clustered very tightly — 0.872 down to 0.867 across all five results. Cosine similarity between text embeddings is compressed near the top of the range, so the absolute value tells you little. A score of 0.87 does not mean "87% relevant." What's useful is the ranking, and even that should be held loosely: a spread of 0.005 means the order among these five is close to arbitrary. Retrieve several passages rather than trusting the single top hit.
The results answer a question you did not quite ask. The query "curling gold medal" surfaced the medal table, the medalist list, and all three gold medal game recaps — because curling awarded three golds in 2022: Sweden (men's), Great Britain (women's), and Italy (mixed doubles). The retrieval step doesn't resolve that ambiguity, it just hands over everything nearby. Watch how the model deals with it in the next step.
In [19]:
1import tiktoken2 3 4def num_tokens(text: str, model: str=GPT_MODEL) -> int:5 """Return the number of tokens in a string."""6 encoding = tiktoken.encoding_for_model(model)7 return len(encoding.encode(text))8 9 10def query_message(11 query: str,12 df: pd.DataFrame,13 model: str,14 token_budget: int15) -> str:16 """Return a message for GPT, with relevant source texts pulled from SingleStoreDB."""17 strings, relatednesses = strings_ranked_by_relatedness(query, df, "winter_olympics_2022")18 introduction = 'Use the below articles on the 2022 Winter Olympics to answer the subsequent question. If the answer cannot be found in the articles, write "I could not find an answer."'19 question = f"\n\nQuestion: {query}"20 message = introduction21 for string in strings:22 next_article = f'\n\nWikipedia article section:\n"""\n{string}\n"""'23 if (24 num_tokens(message + next_article + question, model=model)25 > token_budget26 ):27 break28 else:29 message += next_article30 return message + question31 32 33def ask(34 query: str,35 df: pd.DataFrame=df,36 model: str=GPT_MODEL,37 token_budget: int=4096 - 500,38 print_message: bool=False,39) -> str:40 """Answers a query using GPT and a table of relevant texts and embeddings in SingleStoreDB."""41 message = query_message(query, df, model=model, token_budget=token_budget)42 if print_message:43 print(message)44 messages = [45 {"role": "system", "content": "You answer questions about the 2022 Winter Olympics."},46 {"role": "user", "content": message},47 ]48 response = client.chat.completions.create(49 model=model,50 messages=messages,51 temperature=052 )53 response_message = response.choices[0].message.content54 return response_message
Step 2: Ask
ask() runs the search above, then packs as many top-ranked sections as fit into the
prompt — the token budget is 3,596, leaving room for the model's reply — along with an
instruction to answer using only that text and to say "I could not find an answer" if
it isn't there.
That last instruction is the part that makes this trustworthy. Without it, a model handed irrelevant context will often answer anyway from vague recollection. Compare the answer below against the passages you just retrieved: every fact in it should be traceable to text on your screen.
In [20]:
1print(ask('Who won the gold medal for curling in Olymics 2022?'))
What you just built, and why it matters
A handful of cells did the whole thing: turn a question into a vector, rank stored text by similarity to it, pack the best passages into a prompt, and answer from them. That pattern is retrieval-augmented generation, and it is how you get a language model to answer from data it was never trained on.
The real product is grounding, not search. This model knows nothing about the 2022
Winter Olympics beyond what you handed it. That constraint is the feature: the answer is
traceable to rows you can read yourself, it improves when you fix the underlying text
rather than when you retrain anything, and — because of the instruction built in
query_message — it can decline instead of inventing. Every one of those properties
comes from the retrieval step, not from the model.
Why it matters that the search ran in the database
Look again at strings_ranked_by_relatedness. The only work Python did was embed the
query string. Every comparison against every stored embedding happened inside
SingleStore, in one statement, and only the top rows travelled back.
The vectors sit beside the text they describe. They are a column in an ordinary table, not a parallel copy living in a second system. One schema, one set of permissions, one backup, one version of the truth.
Similarity, filtering and joining are a single operation. Want only sections about a
particular event, or the medal table joined in, or the most relevant rows added this
week? Those are a WHERE clause and a JOIN bolted onto the query you already have,
planned and executed together. Stand a dedicated vector service next to your database and
each of them becomes an application problem instead: query the index, ship a list of IDs
back, join them yourself, and keep two stores in step on every write. Every filter turns
into a round trip, and consistency becomes something you maintain by hand.
Brute force is faster than it sounds. SingleStore compiles the query to machine code and runs the distance calculation as a vectorized operation over compressed columnstore data, using SIMD instructions across every core and every partition at once. Comparing your question against the entire corpus really is brute force, and it still returns immediately.
You are not stuck with brute force, either. An exact scan is linear — ten times the
rows is ten times the work. When a corpus outgrows that, this table moves to the native
VECTOR column type with a vector index and the same query becomes approximate
nearest-neighbour search. That is a change to a table definition, not to your
architecture and not to your application code.
Clean up
Action Required
If you created a new database in your Standard or Premium Workspace, you can drop the database by running the cell below. Note: this will not drop your database for Free Starter Workspaces. To drop a Free Starter Workspace, terminate the Workspace using the UI.
In [21]:
1shared_tier_check = %sql show variables like 'is_shared_tier'2if not shared_tier_check or shared_tier_check[0][1] == 'OFF':3 %sql DROP DATABASE IF EXISTS winter_wikipedia;
Details
About this Template
Provide context to chatGPT using data stored in SingleStoreDB.
This Notebook can be run in Shared Tier, Standard and Enterprise deployments.
Tags
starteropenaivectordbgenai
See Notebook in action
Launch this notebook in SingleStore and start executing queries instantly.
License
This Notebook has been released under the Apache 2.0 open source license.