Using Python OpenAI Embeddings for Semantic Search

Python OpenAI embeddings turn text into dense numeric vectors that capture semantic meaning. The OpenAI Python client exposes the embeddings endpoint through a small, consistent API, and once you have the vectors you can compare documents, rank search results, or cluster related content.

Setting Up the OpenAI Client for Embeddings

To call the embeddings API, install the openai package and create a client. The client reads the API key from the OPENAI_API_KEY environment variable, which keeps the key out of source code.

pip install openai

from openai import OpenAI

client = OpenAI()

If the key is not in the environment, pass it explicitly:

client = OpenAI(api_key="sk-...")

The client is reusable across requests. Creating a new client per request adds no value and makes authentication setup harder to maintain.

Creating Your First Embedding

The embeddings endpoint accepts a model name and an input string. The minimal call looks like this:

response = client.embeddings.create(
    model="text-embedding-3-small",
    input="The quick brown fox jumps over the lazy dog",
)

The input parameter can be a single string or a list of strings. When you pass a list, the API returns one embedding per item, in the same order.

The returned object contains the model name, usage statistics, and the actual vectors:

embedding = response.data[0].embedding
print(len(embedding))  # 1536 for text-embedding-3-small by default

response.data is a list of Embedding objects. Each has an embedding field holding a list of floats, plus an index that matches the position of the input item.

Understanding the Embedding Response

The response structure matters because you will index vectors and look them up later.

print(response.model)                    # model used
print(response.usage.prompt_tokens)       # tokens consumed by the input
print(response.usage.total_tokens)

Each item in response.data has three fields:

• embedding: a list of floats, the vector representation
• index: the position of the input that produced this vector
• object: always "embedding"

When you send multiple inputs, match results by index rather than assuming order. In practice the order matches the input order, but relying on index is safer when you later expand to batch pipelines.

Choosing a Model and Dimension

OpenAI offers several embedding models. The text-embedding-3 family is the current default choice for new projects:

The v3 models accept a dimensions parameter that reduces the vector size:

response = client.embeddings.create(
    model="text-embedding-3-small",
    input="A short document",
    dimensions=512,
)

Reducing dimensions lowers storage and compute cost. For simple similarity search, 512 or 768 dimensions often retain most of the signal. There is no universal correct value; test on your own data if quality matters.

Computing Similarity Between Embeddings

Embeddings are compared with cosine similarity. The cosine of the angle between two vectors ranges from -1 to 1, where 1 means the vectors point in the same direction.

import math

def cosine_similarity(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    return dot / (norm_a * norm_b)

The text-embedding-3 models return unit-length vectors, so the dot product equals the cosine similarity. You can skip the norm calculation when you know both vectors are normalized:

def dot_product(a, b):
    return sum(x * y for x, y in zip(a, b))

If you store normalized vectors, the dot product is faster and numerically equivalent to cosine similarity. If you ever mix sources, keep the normalization explicit to avoid subtle mismatches.

Running Semantic Search Over a Document Set

A common use for embeddings is retrieval: embed a set of documents, then rank them against a query by similarity.

documents = [
    "PostgreSQL supports JSONB for flexible document storage.",
    "Redis is an in-memory key-value store with optional persistence.",
    "RabbitMQ provides reliable message delivery between services.",
]

doc_embeddings = []
for doc in documents:
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=doc,
    )
    doc_embeddings.append(response.data[0].embedding)

query = "Which database stores JSON documents?"
query_response = client.embeddings.create(
    model="text-embedding-3-small",
    input=query,
)
query_embedding = query_response.data[0].embedding

scores = [
    (doc, cosine_similarity(query_embedding, doc_vec))
    for doc, doc_vec in zip(documents, doc_embeddings)
]
scores.sort(key=lambda item: item[1], reverse=True)

for doc, score in scores:
    print(f"{score:.4f}  {doc}")

This pattern is the core of a retrieval pipeline. For larger collections, store the vectors in a dedicated vector index rather than scanning a Python list on every query.

Batching and Token Usage

The embeddings API accepts multiple inputs in a single call, which reduces request overhead:

response = client.embeddings.create(
    model="text-embedding-3-small",
    input=documents,
)

Batching is useful when you index a corpus. The API counts tokens per request, and each model has a maximum input length per request. For the text-embedding-3 models, the limit is 8,191 tokens per request. If a document exceeds that, split it into chunks before embedding.

Token usage is reported per request in response.usage. The cost scales with the number of tokens, so batching does not reduce the token count, but it does reduce the number of HTTP round trips.

Handling Rate Limits and Errors

The embeddings endpoint can return HTTP 429 when you exceed the rate limit for your account. The openai Python client raises openai.RateLimitError in that case. A simple retry with exponential backoff handles transient limits:

import time

for attempt in range(5):
    try:
        response = client.embeddings.create(
            model="text-embedding-3-small",
            input=text,
        )
        break
    except openai.RateLimitError:
        time.sleep(2 ** attempt)

Other common failures:

• openai.AuthenticationError: the API key is invalid or missing
• openai.APIConnectionError: network problem reaching the API
• openai.BadRequestError: the model name is wrong or the input is malformed

Catch the specific exception type instead of a bare except, so auth failures and rate limits are handled differently. Retrying an auth error is pointless; retrying a rate limit is correct.