Embeddings

Generate vector embeddings for semantic search and similarity

The embed and embed_many functions create vector representations of text that can be used for semantic search, similarity comparisons, clustering, and other machine learning tasks.

Basic Usage

import asyncio
from ai_query import embed
from ai_query.providers import openai

async def main():
    result = await embed(
        model=openai.embedding("text-embedding-3-small"),
        value="Hello world"
    )
    print(len(result.embedding))  # 1536
    print(result.usage.tokens)  # Token count

asyncio.run(main())

Embedding Multiple Values

For efficiency, use embed_many to embed multiple texts in a single API call:

from ai_query import embed_many
from ai_query.providers import openai

documents = [
    "Python is a programming language",
    "JavaScript runs in browsers",
    "Rust is known for memory safety"
]

result = await embed_many(
    model=openai.embedding("text-embedding-3-small"),
    values=documents
)

print(len(result.embeddings))  # 3
print(len(result.embeddings[0]))  # 1536 (vector dimension)

Supported Providers

OpenAI

from ai_query import embed
from ai_query.providers import openai

# text-embedding-3-small (1536 dimensions, cheaper)
result = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="Hello world"
)

# text-embedding-3-large (3072 dimensions, higher quality)
result = await embed(
    model=openai.embedding("text-embedding-3-large"),
    value="Hello world"
)

# With custom dimensions (OpenAI v3 models only)
result = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="Hello world",
    dimensions=256  # Reduce dimensions for storage efficiency
)

Google

from ai_query import embed
from ai_query.providers import google

# gemini-embedding-001 (latest model, flexible dimensions)
result = await embed(
    model=google.embedding("gemini-embedding-001"),
    value="Hello world"
)

# With custom dimensions (supports 128-3072, default is 3072)
result = await embed(
    model=google.embedding("gemini-embedding-001"),
    value="Hello world",
    provider_options={
        "google": {
            "outputDimensionality": 768  # Smaller for storage efficiency
        }
    }
)

Parameters

Parameter	Type	Description
model	EmbeddingModel	The embedding model (from openai.embedding() or google.embedding())
value	str	The text to embed (for embed)
values	list[str]	List of texts to embed (for embed_many)
provider_options	dict	Provider-specific options

Response Object

The embed function returns an EmbedResult:

result = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="Hello world"
)

# The original input
print(result.value)  # "Hello world"

# The embedding vector
print(result.embedding)  # [0.123, -0.456, ...]
print(len(result.embedding))  # 1536

# Token usage
print(result.usage.tokens)  # Number of tokens used

The embed_many function returns an EmbedManyResult:

result = await embed_many(
    model=openai.embedding("text-embedding-3-small"),
    values=["Hello", "World"]
)

# All embeddings (in same order as input)
print(len(result.embeddings))  # 2

# Each embedding vector
print(len(result.embeddings[0]))  # 1536

Common Use Cases

Semantic Search

import numpy as np
from ai_query import embed, embed_many
from ai_query.providers import openai

# Create embeddings for documents
documents = [
    "How to train a neural network",
    "Best practices for Python coding",
    "Introduction to machine learning"
]

doc_result = await embed_many(
    model=openai.embedding("text-embedding-3-small"),
    values=documents
)

# Create embedding for query
query_result = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="machine learning tutorial"
)

# Calculate cosine similarity
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# Find most similar document
similarities = [
    cosine_similarity(query_result.embedding, doc_emb)
    for doc_emb in doc_result.embeddings
]

most_similar_idx = np.argmax(similarities)
print(f"Most similar: {documents[most_similar_idx]}")

Text Clustering

from sklearn.cluster import KMeans
from ai_query import embed_many
from ai_query.providers import openai

texts = ["...", "...", "..."]  # Your texts

result = await embed_many(
    model=openai.embedding("text-embedding-3-small"),
    values=texts
)

# Cluster embeddings
kmeans = KMeans(n_clusters=3)
clusters = kmeans.fit_predict(result.embeddings)

Duplicate Detection

from ai_query import embed_many
from ai_query.providers import openai

texts = ["Hello world", "Hello World!", "Goodbye world"]

result = await embed_many(
    model=openai.embedding("text-embedding-3-small"),
    values=texts
)

# Find near-duplicates using cosine similarity threshold
threshold = 0.95
for i in range(len(texts)):
    for j in range(i + 1, len(texts)):
        sim = cosine_similarity(result.embeddings[i], result.embeddings[j])
        if sim > threshold:
            print(f"Near-duplicate: '{texts[i]}' and '{texts[j]}'")

Provider Options

OpenAI Options

result = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="Hello world",
    dimensions=512,  # Reduce dimensions (v3 models only)
    provider_options={
        "openai": {
            "encoding_format": "float"  # or "base64"
        }
    }
)

Google Options

result = await embed(
    model=google.embedding("gemini-embedding-001"),
    value="Hello world",
    provider_options={
        "google": {
            "taskType": "RETRIEVAL_DOCUMENT",  # Optimize for retrieval
            "title": "Document Title",  # Optional document title
            "outputDimensionality": 768  # Custom dimensions (128-3072)
        }
    }
)

Available Models

OpenAI

Model	Dimensions	Max Tokens	Description
text-embedding-3-large	3072	8192	Highest quality, best for accuracy-critical tasks
text-embedding-3-small	1536	8192	Good balance of quality and cost
text-embedding-ada-002	1536	8192	Legacy model

Google

Model	Dimensions	Max Tokens	Description
gemini-embedding-001	128-3072	2048	Latest model with flexible dimensions via Matryoshka learning

Next Steps

embed() Reference

Full API reference for embed functions

RAG Agent

Build a retrieval-augmented generation agent