DEV Community

Cover image for Stop Treating Your Vector Database Like a Black Box: Visualize ChromaDB
Rijul Rajesh
Rijul Rajesh

Posted on

Stop Treating Your Vector Database Like a Black Box: Visualize ChromaDB

Hello, I'm Rijul. I'm building git-lrc, a micro AI code reviewer that runs on every commit. It's free and source-available on GitHub. Star git-lrc to help more developers discover the project. Do give it a try and share your feedback

Have you built your first RAG system and set up ChromaDB?

If so, you might have started thinking of ChromaDB as a bit of a black box.

You put documents into it, embeddings come out, searches happen, and relevant chunks are returned.

But what is actually happening inside?

We can make that easier to understand by visualizing the embeddings stored in ChromaDB.

You can get something like this:

Let's see how we can build it.


Before We Start

If you haven't built a RAG system yet, I recommend checking out my RAG article and ChromaDB Article first.

It will give you an idea of how ChromaDB and vector databases work.

Once you're familiar with that, you can come back here and visualize what's happening inside the vector database.


How Does the Visualization Work?

The basic idea is quite simple.

Every chunk in our vector store is represented as a 3072-dimensional vector.

In simple terms, that's a list of 3072 numbers that captures the semantic meaning of that chunk.

The problem is that we can't visualize 3072 dimensions directly.

So we need to reduce those 3072 dimensions down to just 3.

That's where PCA, or Principal Component Analysis, comes in.

What Is PCA?

PCA is a technique used to reduce a high-dimensional vector to a smaller number of dimensions while preserving as much information as possible.

Instead of simply removing random numbers, PCA finds the most important patterns across the dimensions and represents those patterns using fewer dimensions.

In our case:

3072 dimensions
       ↓
      PCA
       ↓
  3 dimensions
       ↓
  3D scatter plot
Enter fullscreen mode Exit fullscreen mode

We can then represent every chunk as a single point in a 3D space.

Chunks with similar embeddings may end up closer together, allowing clusters to appear naturally.


The Five Steps

The visualization script essentially performs five steps.

1. Connect to the Vector Store

The open_store() function opens the local ChromaDB database and gets the rag_demo collection.

I created this RAG demo in my RAG article, so you can check that out if you want to see how the data was created.

2. Pull Out the Records

The collect_entries() function fetches everything we need from the collection:

  • Embedding vectors
  • Chunk IDs
  • Source documents
  • Chunk text

We will also use the source document information to group the points in the visualization.

3. Compress 3072 Dimensions Into 3

The compress_to_3d() function uses PCA to reduce each 3072-dimensional embedding down to 3 dimensions.

These three values become the X, Y, and Z coordinates of our visualization.

4. Build the Tooltips

The build_hover_labels() function creates the information that appears when you hover over a point.

This lets us see things such as:

  • Chunk ID
  • Source document
  • Part of the chunk text

5. Draw the Scene

Finally, render_scene() creates the interactive 3D visualization using Plotly.

Each point represents a chunk from our vector database.

We also create a separate trace for each source document, which allows us to distinguish where the chunks came from.


The Code

Here is the complete visualization script:

from pathlib import Path

import chromadb
import plotly.graph_objects as go
from sklearn.decomposition import PCA

COLLECTION_NAME = "rag_demo"
DB_PATH = Path(__file__).parent / "chroma_db"
TARGET_AXES = 3


def open_store() -> chromadb.Collection:
    """Return a handle to the local vector collection."""
    client = chromadb.PersistentClient(path=str(DB_PATH))
    return client.get_collection(COLLECTION_NAME)


def collect_entries(store: chromadb.Collection) -> tuple[list, list, list, list]:
    """Pull every stored record out of the collection.

    Returns (vectors, ids, source_tags, chunk_texts) in lock-step order.
    """
    records = store.get(include=["embeddings", "documents", "metadatas"])
    vectors = records["embeddings"]
    ids = records["ids"]
    source_tags = [
        (meta or {}).get("source", "unknown") for meta in records["metadatas"]
    ]
    chunk_texts = records["documents"] or []
    return vectors, ids, source_tags, chunk_texts


def compress_to_3d(vectors: list) -> tuple[object, object]:
    """Fit a PCA model on the vectors and return (model, projected_points)."""
    model = PCA(n_components=TARGET_AXES)
    projected_points = model.fit_transform(vectors)
    return model, projected_points


def build_hover_labels(ids, source_tags, chunk_texts) -> list[str]:
    """Compose the rich tooltip shown when hovering over a point."""
    return [
        f"id: {rid}<br>source: {src}<br>chunk: {text[:120]!r}"
        for rid, src, text in zip(ids, source_tags, chunk_texts)
    ]


def render_scene(
    projected_points,
    source_tags,
    hover_labels,
    variance_explained: float,
) -> None:
    """Build the interactive 3D figure and open it in the browser."""
    figure = go.Figure()
    figure.update_layout(
        title=(
            f"{COLLECTION_NAME} embeddings projected into 3D "
            f"(explained variance {variance_explained:.0%})"
        ),
        scene={
            "xaxis_title": "PCA Axis 1",
            "yaxis_title": "PCA Axis 2",
            "zaxis_title": "PCA Axis 3",
        },
    )

    # One trace per source document so each cluster can be toggled in the legend.
    for source in sorted(set(source_tags)):
        mask = [tag == source for tag in source_tags]
        figure.add_trace(
            go.Scatter3d(
                x=projected_points[mask, 0],
                y=projected_points[mask, 1],
                z=projected_points[mask, 2],
                mode="markers",
                name=source,
                customdata=[
                    hover_labels[i] for i in range(len(mask)) if mask[i]
                ],
                hovertemplate="%{customdata}<extra></extra>",
            )
        )

    figure.show()


def main() -> None:
    store = open_store()

    vectors, ids, source_tags, chunk_texts = collect_entries(store)

    if len(vectors) == 0:
        print("No embeddings found. Run `python ingest.py` first.")
        return

    model, projected_points = compress_to_3d(vectors)
    hover_labels = build_hover_labels(ids, source_tags, chunk_texts)
    variance_explained = model.explained_variance_ratio_.sum()

    render_scene(
        projected_points,
        source_tags,
        hover_labels,
        variance_explained,
    )

    print(
        f"Plotted {len(vectors)} chunks in 3D "
        f"({variance_explained:.0%} of variance preserved)."
    )


if __name__ == "__main__":
    main()
Enter fullscreen mode Exit fullscreen mode

Try It Yourself

Now that you have an idea of how the visualization works, you can try it yourself using my demo repository.

Clone the repository and run:

python ingest.py
python visualize.py
Enter fullscreen mode Exit fullscreen mode

ingest.py loads the data into ChromaDB if it hasn't already been added.

Then visualize.py reads the stored embeddings and launches the 3D visualization.

You should get something like this:

Each dot represents a different chunk in the vector database.

The color represents the source document.

For example, all the red points might represent chunks that came from policies.md.

This gives us a much more intuitive way to look at what's happening inside our vector database.

Instead of just seeing a collection of embeddings, we can actually see how the chunks are distributed in the embedding space.


Wrapping Up

When learning something new, it is easy to stay at the theoretical level.

We can read about embeddings, vector databases, similarity search, and RAG without ever really developing an intuition for what is happening underneath.

Visualization can help bridge that gap.

By reducing our high-dimensional embeddings to three dimensions, we can actually see the data and start developing an intuition for how similar chunks are positioned relative to one another.

It won't show us everything happening inside ChromaDB, but it gives us a useful way to look inside what can otherwise feel like a black box.

See you in another article.

AI agents write code fast. They also silently remove logic, change behavior, and introduce bugs -- without telling you. You often find out in production.

git-lrc fixes this. It hooks into git commit and reviews every diff before it lands. 60-second setup. Completely free.

Any feedback or contributors are welcome! It's online, source-available, and ready for anyone to use.

Give it a ⭐ star on Github

Top comments (0)