AI Engineer Portal

Embedding and vector storage

Why this matters

Embeddings are the translation layer between human language and mathematical retrieval. Every time you call a RAG pipeline and ask "find me the most relevant chunks for this query," that question is answered in embedding space — not in keyword space. If your embeddings are bad, or if you choose the wrong similarity metric, or if your vector database is misconfigured, retrieval will fail in ways that are extremely hard to debug because everything looks like it is working.

For a full-stack engineer, the practical question is: which embedding model do I choose, which vector database do I choose, and how do I wire them together without painting myself into a corner? This lesson gives you the mental model to answer those questions on a real project.

Core concepts

What an embedding is. An embedding model takes a string of text and returns a fixed-length float vector — typically 768, 1024, or 1536 dimensions. The geometry of the space is trained so that semantically similar texts land close together. "The contract was signed in March" and "The agreement was executed in Q1" will have vectors that are close to each other even though they share no words.

Embedding model choices. OpenAI's text-embedding-3-small and text-embedding-3-large are the easiest starting point — one API call, consistent quality, no infrastructure. Cohere's Embed v3 supports embedding types (search_document vs. search_query), which means the model can optimize separately for what goes in the index versus what a query looks like. Open-source options like BAAI/bge-large-en-v1.5 or sentence-transformers/all-MiniLM-L6-v2 run locally or on-device, which matters for privacy-sensitive workloads.

Key factors when choosing:

Dimensionality vs. quality tradeoff. Higher dimensions capture more semantic nuance but cost more storage and are slower to search. text-embedding-3-small (1536-d, can be reduced to 256) outperforms the old ada-002 at half the cost.
Matryoshka embeddings. Some modern models (OpenAI's v3 series, BGE) support dimensionality reduction while preserving most quality. You can store 256-d vectors in development and 1536-d in production, using the same model.
Domain fit. A general-purpose model will underperform on domain-specific content (medical, legal, code). Consider fine-tuning or using a domain-specialized model.

Vector databases. The market has converged on a few major options:

Chroma: Local-first, embedded (no server required), ideal for prototypes and single-user applications. Write vectors to disk in seconds.
Pinecone: Managed cloud service. No infrastructure to operate. Strong filtering and metadata support. Best for production at small-to-medium scale.
pgvector: A PostgreSQL extension. If your application already runs Postgres, adding vector search means one fewer service to operate. Query performance is slightly lower than specialized databases at large scale.
Weaviate, Qdrant, Milvus: Open-source specialized databases with self-hosted and managed options. Better for high-throughput or complex filtering scenarios.

Similarity metrics. The three common options are:

Cosine similarity: Measures angle between vectors, ignoring magnitude. Most embedding models are trained for cosine. Use this by default.
Dot product: Cosine similarity multiplied by vector magnitude. Some models like Cohere's Embed v3 are trained for dot product. Faster to compute.
Euclidean (L2) distance: Measures straight-line distance. Sensitive to vector magnitude. Less common for text embeddings.

Always normalize vectors before indexing if your vector DB uses dot product internally but your model is trained for cosine. Mismatched metrics silently degrade retrieval quality.

Working example

from __future__ import annotations
import os
import json
from typing import Optional
from dataclasses import dataclass

import chromadb
from openai import OpenAI


@dataclass
class SearchResult:
    text: str
    metadata: dict
    distance: float


class RAGIndex:
    """
    A simple RAG index backed by Chroma (local) and OpenAI embeddings.
    Drop-in replaceable: swap embed_text for a different model,
    swap the collection for a different vector DB.
    """

    def __init__(
        self,
        collection_name: str = "knowledge",
        persist_dir: str = "./chroma_db",
        embedding_model: str = "text-embedding-3-small",
        embedding_dims: int = 256,  # Matryoshka reduction
    ):
        self.openai = OpenAI()
        self.model = embedding_model
        self.dims = embedding_dims

        # Persistent local Chroma
        self.client = chromadb.PersistentClient(path=persist_dir)
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},  # Explicit metric
        )

    def embed_text(self, text: str) -> list[float]:
        """Embed a single string. In production, batch these calls."""
        response = self.openai.embeddings.create(
            model=self.model,
            input=text,
            dimensions=self.dims,
        )
        return response.data[0].embedding

    def add_documents(self, chunks: list[dict]) -> None:
        """
        Add chunks to the index.
        Each chunk: {"id": str, "text": str, "metadata": dict}
        """
        if not chunks:
            return

        # Batch embedding is cheaper and faster
        texts = [c["text"] for c in chunks]
        response = self.openai.embeddings.create(
            model=self.model,
            input=texts,
            dimensions=self.dims,
        )
        embeddings = [item.embedding for item in response.data]

        self.collection.add(
            ids=[c["id"] for c in chunks],
            embeddings=embeddings,
            documents=texts,
            metadatas=[c.get("metadata", {}) for c in chunks],
        )
        print(f"Indexed {len(chunks)} chunks.")

    def search(
        self,
        query: str,
        top_k: int = 5,
        where: Optional[dict] = None,  # Metadata filter
    ) -> list[SearchResult]:
        """Semantic search over the index."""
        query_embedding = self.embed_text(query)
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            where=where,
        )
        output = []
        for text, meta, distance in zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0],
        ):
            output.append(SearchResult(text=text, metadata=meta, distance=distance))
        return output


# --- Usage ---
index = RAGIndex(collection_name="docs_v1")

# Add chunks (normally these come from your chunking pipeline)
chunks = [
    {"id": "doc1-0", "text": "Vector databases store high-dimensional embeddings.", "metadata": {"source": "intro.md", "section": "overview"}},
    {"id": "doc1-1", "text": "Cosine similarity measures the angle between two vectors.", "metadata": {"source": "intro.md", "section": "metrics"}},
    {"id": "doc2-0", "text": "pgvector is a PostgreSQL extension for vector search.", "metadata": {"source": "databases.md", "section": "options"}},
]

index.add_documents(chunks)

# Search with optional metadata filter
results = index.search("How do I measure similarity between embeddings?", top_k=3)
for r in results:
    print(f"[dist={r.distance:.3f}] {r.text}")

# Search with metadata filter (only docs from intro.md)
filtered = index.search("similarity", where={"source": "intro.md"})

Common mistakes

Embedding queries the same way you embed documents. Some models (Cohere Embed v3) have different modes for search_document vs. search_query. Using the wrong mode for queries silently degrades recall. Read the model card.
Not batching embeddings. Calling embed_text once per chunk in a loop is 10–100x slower and more expensive than batching. OpenAI supports up to 2048 inputs per batch call.
Changing the embedding model mid-project without re-indexing. Vector spaces are model-specific. If you switch from ada-002 to text-embedding-3-small, old vectors are meaningless in the new space. You must re-embed everything.
Storing embeddings but losing the original text. Always store the source text alongside the vector. Chroma stores it in the documents field. Without the original text, you cannot serve it to the LLM.
Ignoring index size at production scale. 1 million chunks × 1536 dimensions × 4 bytes = 6GB of raw vectors. Matryoshka reduction to 256 dimensions brings this to 1GB. Plan storage before you discover it at launch.

Try it yourself

Set up a local Chroma collection. Embed 10 sentences on a topic you know well. Query it with paraphrases of those sentences. Then query it with something completely unrelated and look at the distances. What distance threshold would you use to discard irrelevant results? Now try changing embedding_dims from 1536 to 256 — does retrieval quality change for your queries?