Notebooks
G
Google Gemini Vectordb With Chroma

Vectordb With Chroma

ChromaDBgemini-cookbookgemini-apiexamplesgemini
alph-notebooks/gemini-cookbook / Vectordb_with_chroma.ipynb
Copyright 2025 Google LLC.
[1]

Gemini API: Document Q&A with ChromaDB

⚠️

This notebook requires paid tier rate limits to run properly.
(cf. pricing for more details).

Overview

This tutorial demonstrates how to use the Gemini API to create a vector database and retrieve answers to questions from the database. Moreover, you will use ChromaDB{:.external}, an open-source Python tool that creates embedding databases. ChromaDB allows you to:

  • Store embeddings as well as their metadata
  • Embed documents and queries
  • Search through the database of embeddings

In this tutorial, you'll use embeddings to retrieve an answer from a database of vectors created with ChromaDB.

Setup

First, download and install ChromaDB and the Gemini API Python library.

[ ]

Then import the modules you'll use in this tutorial.

[2]

Configure your API key

Before you can use the Gemini API, you must first obtain an API key. If you don't already have one, create a key with one click in Google AI Studio.

Get an API key

In Colab, add the key to the secrets manager under the "🔑" in the left panel. Give it the name GEMINI_API_KEY.

Once you have the API key, pass it to the SDK. You can do this in two ways:

  • Put the key in the GEMINI_API_KEY environment variable (the SDK will automatically pick it up from there).
  • Pass the key to genai.Client(api_key=...)
[ ]

Key Point: Next, you will choose a model. Any embedding model will work for this tutorial, but for real applications it's important to choose a specific model and stick with it. The outputs of different models are not compatible with each other.

Note: At this time, the Gemini API is only available in certain regions.

[4]
models/embedding-001
models/text-embedding-004
models/gemini-embedding-exp-03-07
models/gemini-embedding-exp
models/gemini-embedding-001

Data

Here is a small set of documents you will use to create an embedding database:

[5]

Creating the embedding database with ChromaDB

You will create a custom function{:.external} for performing embedding using the Gemini API. By inputting a set of documents into this custom function, you will receive vectors, or embeddings of the documents.

API changes to Embeddings with model gemini-embedding-001

For the new embeddings model, embedding-001, there is a new task type parameter and the optional title (only valid with task_type=RETRIEVAL_DOCUMENT).

These new parameters apply only to the newest embeddings models.The task types are:

Task TypeDescription
RETRIEVAL_QUERYSpecifies the given text is a query in a search/retrieval setting.
RETRIEVAL_DOCUMENTSpecifies the given text is a document in a search/retrieval setting.
SEMANTIC_SIMILARITYSpecifies the given text will be used for Semantic Textual Similarity (STS).
CLASSIFICATIONSpecifies that the embeddings will be used for classification.
CLUSTERINGSpecifies that the embeddings will be used for clustering.
[6]
EMBEDDING_MODEL_ID

Now you will create the vector database. In the create_chroma_db function, you will instantiate a Chroma client. From there, you will create a collection, which is where you store your embeddings, documents, and any metadata. Note that the embedding function from above is passed as an argument to the create_collection.

Next, you use the add method to add the documents to the collection.

[7]
[11]
/tmp/ipykernel_83887/781126645.py:5: DeprecationWarning: The class GeminiEmbeddingFunction does not implement __init__. This will be required in a future version.
  embedding_function=GeminiEmbeddingFunction()

Confirm that the data was inserted by looking at the database:

[13]

Getting the relevant document

db is a Chroma collection object. You can call query on it to perform a nearest neighbors search to find similar embeddings or documents.

[14]
[15]

Now that you have found the relevant passage in your set of documents, you can use it make a prompt to pass into the Gemini API.

[16]

Pass a query to the prompt:

[17]

Now use the generate_content method to to generate a response from the model.

[18]
MODEL_ID

Next steps

To learn more about how you can use the embeddings, check out the examples available. To learn how to use other services in the Gemini API, visit the Python quickstart.