Semantic search in Python

What you’re building

A search box over eight handwritten notes. You type “where did I put my keys?” and it returns the note about the spare set in the kitchen drawer — a note that contains no form of the word “key” anywhere in it.

The whole thing is one file and no database. That’s deliberate: everything a vector database does, you’re about to do by hand in three lines, and it’s much easier to appreciate the database once you’ve felt what it’s replacing.

Setup

Free, runs locally

pip install sentence-transformers numpy

No API key needed. The first run downloads roughly 90 MB.

OpenAI API

pip install openai numpy
export OPENAI_API_KEY="sk-..."

Grab a key from platform.openai.com/api-keys if you don’t have one.

Keep keys out of your source

OpenAI() picks the key up from the environment on its own. Hard-coding it means it goes into git, and once a key is in git history it is compromised — rotating it is the only fix. If you like .env files, python-dotenv works nicely here.

The code

Save this as notes_search.py.

Free, runs locally

notes_search.py

import numpy as np
from sentence_transformers import SentenceTransformer

NOTES = [
    "The spare set for the front door is in the kitchen drawer.",
    "The dog needs his flea treatment on the 3rd.",
    "Mum's birthday is in April, she likes tulips.",
    "Renew the car insurance before it lapses.",
    "The standing desk needs a 4mm hex key to assemble.",
    "Pizza dough: 500g flour, 300ml water, rest overnight.",
    "The wifi password is taped under the router.",
    "Physio said ice the knee for twenty minutes.",
]

model = SentenceTransformer("all-MiniLM-L6-v2")


def embed(texts: list[str]) -> np.ndarray:
    """Turn a list of strings into a matrix of vectors, one row per string."""
    return model.encode(texts, normalize_embeddings=True)


def search(query: str, notes: list[str], matrix: np.ndarray, top_k: int = 3):
    query_vector = embed([query])[0]
    scores = matrix @ query_vector           # cosine similarity, see note below
    ranked = np.argsort(scores)[::-1][:top_k]
    return [(notes[i], float(scores[i])) for i in ranked]


if __name__ == "__main__":
    note_matrix = embed(NOTES)               # do this once, up front

    for question in ["where did I put my keys?", "what do I cook on Friday?"]:
        print(f"\nQ: {question}")
        for note, score in search(question, NOTES, note_matrix):
            print(f"  {score:.3f}  {note}")

OpenAI API

notes_search.py

import numpy as np
from openai import OpenAI

NOTES = [
    "The spare set for the front door is in the kitchen drawer.",
    "The dog needs his flea treatment on the 3rd.",
    "Mum's birthday is in April, she likes tulips.",
    "Renew the car insurance before it lapses.",
    "The standing desk needs a 4mm hex key to assemble.",
    "Pizza dough: 500g flour, 300ml water, rest overnight.",
    "The wifi password is taped under the router.",
    "Physio said ice the knee for twenty minutes.",
]

client = OpenAI()
MODEL = "text-embedding-3-small"


def embed(texts: list[str]) -> np.ndarray:
    """One API call for the whole list. Rows come back in the order you sent them."""
    response = client.embeddings.create(model=MODEL, input=texts)
    vectors = np.array([item.embedding for item in response.data])
    return vectors / np.linalg.norm(vectors, axis=1, keepdims=True)


def search(query: str, notes: list[str], matrix: np.ndarray, top_k: int = 3):
    query_vector = embed([query])[0]
    scores = matrix @ query_vector           # cosine similarity, see note below
    ranked = np.argsort(scores)[::-1][:top_k]
    return [(notes[i], float(scores[i])) for i in ranked]


if __name__ == "__main__":
    note_matrix = embed(NOTES)               # one call, eight notes

    for question in ["where did I put my keys?", "what do I cook on Friday?"]:
        print(f"\nQ: {question}")
        for note, score in search(question, NOTES, note_matrix):
            print(f"  {score:.3f}  {note}")

Run it:

python notes_search.py

Your scores will differ depending on which model you chose. You should see something like this:

Q: where did I put my keys?
  0.492  The spare set for the front door is in the kitchen drawer.
  0.271  The standing desk needs a 4mm hex key to assemble.
  0.144  The wifi password is taped under the router.

Q: what do I cook on Friday?
  0.518  Pizza dough: 500g flour, 300ml water, rest overnight.
  0.201  The dog needs his flea treatment on the 3rd.
  0.169  The spare set for the front door is in the kitchen drawer.

Those numbers are illustrative, not captured from a run — exact scores move between model versions, and any tutorial quoting them to three decimal places is showing you one moment in time. The ranking is the part that matters.

That first result is the whole point of this guide. The winning note shares not one word with the question. Meanwhile the hex key note — the only note that literally contains “key” — came second, and a keyword search would have put it first. The model understood that a 4mm hex key is a tool, and that a spare set for the front door is what you actually meant.

Line by line

1. NOTES — your data. In a real app this is rows from a database. Eight strings in a list keeps the moving parts visible.

2. embed(texts) — the only part that talks to a model. Give it a list of strings, get back a matrix: one row per string, every row the same width. Note that we send all eight notes in a single call. Looping and calling the API once per note would work, and would also be eight times slower and eight times more rate-limit-prone. Batch by default.

3. Normalising — normalize_embeddings=True locally, or the / np.linalg.norm(...) division for OpenAI. This scales every vector to length 1. Do this once here and the similarity calculation later becomes a plain dot product, which is why the next line is so short. The cosine similarity page explains why this works.

4. matrix @ query_vector — this is the search. The @ is NumPy’s matrix multiply. It takes the dot product of the query against all eight notes at once and hands back eight scores. Because the vectors are normalised, each score is the cosine similarity, between roughly -1 and 1, where higher is more similar. If this line feels too small to be doing the work: it is genuinely the whole search.

5. np.argsort(scores)[::-1][:top_k] — argsort gives indices from lowest to highest, [::-1] flips it to highest-first, and [:top_k] takes the top few. Slightly cryptic, completely standard NumPy.

6. note_matrix = embed(NOTES) runs once — embedding your documents is the expensive bit, so you do it up front and reuse it. Only the query gets embedded per search. In production you’d store that matrix in a database rather than rebuilding it every time the process starts.

Three things that will go wrong

openai.AuthenticationError: Incorrect API key provided The key isn’t reaching the process. export only applies to the terminal you ran it in — new tab, new export. On Windows, setx needs a fresh terminal to take effect. Check with python -c "import os; print(os.environ.get('OPENAI_API_KEY', 'NOT SET'))".

RateLimitError on your very first call Almost always means no billing set up rather than genuine rate limiting. New accounts need a payment method on file before the API responds, even for the tiny amounts embeddings cost. Check platform.openai.com/settings/organization/billing.

Every result scores about the same, and the ranking looks random Two usual suspects. Either your query is far shorter or far more generic than your documents (“stuff” will match everything weakly), or you’ve embedded your notes with one model and your query with another. That second one is silent and nasty — see the warning about mixing models.

Batching has a ceiling

Batch, but don’t get carried away. OpenAI’s embeddings endpoint accepts at most 2,048 inputs per request, and caps each request at 300,000 tokens across all of them. Exceed either and you get an HTTP 400, not a partial result.

So for a large corpus you want batches, plural — split your list into groups of a few hundred and loop over the groups. Still vastly better than one call per document, and it won’t walk into the limit.

Scores are relative, not absolute

Don’t hard-code a threshold like if score > 0.8. What counts as a good score varies wildly between models and even between kinds of query. 0.49 was an excellent match above. On a different model it might be mediocre. Rank your results and take the top few; if you need a cut-off, measure one against your own data rather than borrowing a number from a blog post.

Making it real

The version above holds everything in memory and re-embeds on every startup. That’s correct for eight notes and hopeless for eighty thousand. What changes as you scale:

• Persist the vectors. pgvector adds a vector column to Postgres you already run, which is the least disruptive option for most teams. Chroma is the easiest to start with. Qdrant and FAISS are strong when you outgrow the easy options.
• Split long documents. A whole PDF in one embedding turns into meaningless average. See chunking.
• Cache aggressively. Same text, same model, same vector, every time. Hash the text and skip the API call.
• Handle the API failing. The OpenAI SDK retries automatically by default; make sure you’re not swallowing the exception it eventually raises.

Next

• How cosine similarity works — what that @ was really doing
• Chunking and vector databases — going from eight notes to eighty thousand
• The same app in C#, Java or TypeScript

References

• OpenAI — Embeddings guide
• OpenAI — Embeddings API reference
• openai-python on GitHub — including the retry and timeout options
• sentence-transformers documentation — encode() parameters in full
• NumPy — matmul and argsort
• pgvector · Chroma · Qdrant · FAISS