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, an open-source Python tool that creates embedding databases. ChromaDB allows you to:
In this tutorial, you’ll use embeddings to retrieve an answer from a database of vectors created with ChromaDB.
Install the Google GenAI SDK from npm.
You can create your API key using Google AI Studio with a single click.
Remember to treat your API key like a password. Don’t accidentally save it in a notebook or source file you later commit to GitHub. In this notebook we will be storing the API key in a .env file. You can also set it as an environment variable or use a secret manager.
Here’s how to set it up in a .env file:
To load the API key from the .env file, we will use the dotenv package. This package loads environment variables from a .env file into process.env.
Then, we can load the API key in our code:
const dotenv = require("dotenv") as typeof import("dotenv");
dotenv.config({
path: "../../.env",
});
const GEMINI_API_KEY = process.env.GEMINI_API_KEY ?? "";
if (!GEMINI_API_KEY) {
throw new Error("GEMINI_API_KEY is not set in the environment variables");
}
console.log("GEMINI_API_KEY is set in the environment variables");GEMINI_API_KEY is set in the environment variablesIn our particular case the .env is is two directories up from the notebook, hence we need to use ../../ to go up two directories. If the .env file is in the same directory as the notebook, you can omit it altogether.
│
├── .env
└── examples
└── chromadb
└── Vectordb_with_chroma.ipynbWith the new SDK, now you only need to initialize a client with you API key (or OAuth if using Vertex AI). The model is now set in each call.
Now select the model you want to use in this guide, either by selecting one in the list or writing it down. Keep in mind that some models, like the 2.5 ones are thinking models and thus take slightly more time to respond (cf. thinking notebook for more details and in particular learn how to switch the thiking off).
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.
const models_1 = await ai.models.list();
let page_1 = models_1.page;
while (page_1.length > 0) {
for (const model of page_1) {
if (model.supportedActions?.includes("embedContent")) {
console.log(`- ${model.name} (${model.displayName}) | [Actions: ${model.supportedActions.join(", ")}]`);
}
}
page_1 = models_1.hasNextPage() ? await models_1.nextPage() : [];
}- models/embedding-001 (Embedding 001) | [Actions: embedContent]
- models/text-embedding-004 (Text Embedding 004) | [Actions: embedContent]
- models/gemini-embedding-exp-03-07 (Gemini Embedding Experimental 03-07) | [Actions: embedContent, countTextTokens, countTokens]
- models/gemini-embedding-exp (Gemini Embedding Experimental) | [Actions: embedContent, countTextTokens, countTokens]
- models/gemini-embedding-001 (Gemini Embedding 001) | [Actions: embedContent, countTextTokens, countTokens]Here is a small set of documents you will use to create an embedding database:
const DOCUMENT1 = `
Operating the Climate Control System Your Googlecar has a climate control
system that allows you to adjust the temperature and airflow in the car.
To operate the climate control system, use the buttons and knobs located on
the center console. Temperature: The temperature knob controls the
temperature inside the car. Turn the knob clockwise to increase the
temperature or counterclockwise to decrease the temperature.
Airflow: The airflow knob controls the amount of airflow inside the car.
Turn the knob clockwise to increase the airflow or counterclockwise to
decrease the airflow. Fan speed: The fan speed knob controls the speed
of the fan. Turn the knob clockwise to increase the fan speed or
counterclockwise to decrease the fan speed.
Mode: The mode button allows you to select the desired mode. The available
modes are: Auto: The car will automatically adjust the temperature and
airflow to maintain a comfortable level.
Cool: The car will blow cool air into the car.
Heat: The car will blow warm air into the car.
Defrost: The car will blow warm air onto the windshield to defrost it.
`;
const DOCUMENT2 = `
Your Googlecar has a large touchscreen display that provides access to a
variety of features, including navigation, entertainment, and climate
control. To use the touchscreen display, simply touch the desired icon.
For example, you can touch the "Navigation" icon to get directions to
your destination or touch the "Music" icon to play your favorite songs.
`;
const DOCUMENT3 = `
Shifting Gears Your Googlecar has an automatic transmission. To
shift gears, simply move the shift lever to the desired position.
Park: This position is used when you are parked. The wheels are locked
and the car cannot move.
Reverse: This position is used to back up.
Neutral: This position is used when you are stopped at a light or in traffic.
The car is not in gear and will not move unless you press the gas pedal.
Drive: This position is used to drive forward.
Low: This position is used for driving in snow or other slippery conditions.
`;
const documents = [DOCUMENT1, DOCUMENT2, DOCUMENT3];You will create a custom function 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.
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 Type | Description |
|---|---|
RETRIEVAL_QUERY |
Specifies the given text is a query in a search/retrieval setting. |
RETRIEVAL_DOCUMENT |
Specifies the given text is a document in a search/retrieval setting. |
SEMANTIC_SIMILARITY |
Specifies the given text will be used for Semantic Textual Similarity (STS). |
CLASSIFICATION |
Specifies that the embeddings will be used for classification. |
CLUSTERING |
Specifies that the embeddings will be used for clustering. |
import { EmbeddingFunction } from "chromadb";
class GeminiEmbeddingFunction implements EmbeddingFunction {
public async generate(texts: string[]): Promise<number[][]> {
const EMBEDDING_MODEL_ID = "models/gemini-embedding-001";
const response = await ai.models.embedContent({
model: EMBEDDING_MODEL_ID,
contents: texts,
config: {
taskType: "retrieval_document",
title: "Custom Query",
},
});
const embeddings = response.embeddings?.map((e) => e.values!) ?? [];
return embeddings;
}
}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 createCollection.
Next, you use the add method to add the documents to the collection.
To set up Chroma, you can either use the Chroma Cloud or run a local instance of Chroma. If you want to run a local instance, you can use the Chroma Docker image.
For this particular example, you can run a local instance of Chroma using the following command:
import { ChromaClient, Collection } from "chromadb";
const CHROMA_HOST = process.env.CHROMA_HOST ?? "localhost";
const CHROMA_PORT = parseInt(process.env.CHROMA_PORT ?? "8000");
async function create_chroma_db(docs: string[], name: string): Promise<Collection> {
const client = new ChromaClient({
host: CHROMA_HOST,
port: CHROMA_PORT,
});
const collection = await client.getOrCreateCollection({
name: name,
embeddingFunction: new GeminiEmbeddingFunction(),
});
const ids = docs.map((_, i) => `doc${i}`);
await collection.add({
ids: ids,
documents: docs,
metadatas: docs.map((doc, i) => ({ source: `document_${i}` })),
});
return collection;
}Confirm that the data was inserted by looking at the database:
import * as danfo from "danfojs-node";
const results = await db.get({
include: ["documents", "embeddings", "metadatas"],
});
const df = new danfo.DataFrame({
documents: results.documents,
embeddings: results.embeddings.map((e) => e.join(", ")),
metadatas: results.metadatas.map((m) => JSON.stringify(m)),
});
df.print();╔════════════╤═══════════════════╤═══════════════════╤═══════════════════╗
║ │ documents │ embeddings │ metadatas ║
╟────────────┼───────────────────┼───────────────────┼───────────────────╢
║ 0 │ │ 0.0088090375, -… │ {"source":"docu… ║
║ │ Operating th… │ │ ║
╟────────────┼───────────────────┼───────────────────┼───────────────────╢
║ 1 │ │ 0.0024602748, -… │ {"source":"docu… ║
║ │ Your Googlec… │ │ ║
╟────────────┼───────────────────┼───────────────────┼───────────────────╢
║ 2 │ │ -0.0043510622, … │ {"source":"docu… ║
║ │ Shifting Gea… │ │ ║
╚════════════╧═══════════════════╧═══════════════════╧═══════════════════╝
db is a Chroma collection object. You can call query on it to perform a nearest neighbors search to find similar embeddings or documents.
import { Collection } from "chromadb";
async function get_relevant_passage(query: string, db: Collection): Promise<string> {
const results = await db.query({
queryTexts: [query],
nResults: 1,
include: ["documents"],
});
if (results.documents.length > 0) {
return results.documents[0][0];
} else {
return "No relevant passage found.";
}
}const passage = await get_relevant_passage("touch screen features", db);
tslab.display.markdown(passage);Your Googlecar has a large touchscreen display that provides access to a variety of features, including navigation, entertainment, and climate control. To use the touchscreen display, simply touch the desired icon. For example, you can touch the “Navigation” icon to get directions to your destination or touch the “Music” icon to play your favorite songs.
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.
function make_prompt(query: string, relevant_passage: string): string {
const escaped = relevant_passage.replace(/'/g, "").replace(/"/g, "").replace(/\n/g, " ");
const prompt = `
You are a helpful and informative bot that answers questions using
text from the reference passage included below.
Be sure to respond in a complete sentence, being comprehensive,
including all relevant background information.
However, you are talking to a non-technical audience, so be sure to
break down complicated concepts and strike a friendly
and conversational tone. If the passage is irrelevant to the answer,
you may ignore it.
QUESTION: '${query}'
PASSAGE: '${escaped}'
ANSWER:
`;
return prompt;
}Pass a query to the prompt:
const prompt = make_prompt("How do you use the touchscreen in the Google car?", passage);
tslab.display.markdown(prompt);You are a helpful and informative bot that answers questions using
text from the reference passage included below.
Be sure to respond in a complete sentence, being comprehensive,
including all relevant background information.
However, you are talking to a non-technical audience, so be sure to
break down complicated concepts and strike a friendly
and conversational tone. If the passage is irrelevant to the answer,
you may ignore it.
QUESTION: 'How do you use the touchscreen in the Google car?'
PASSAGE: ' Your Googlecar has a large touchscreen display that provides access to a variety of features, including navigation, entertainment, and climate control. To use the touchscreen display, simply touch the desired icon. For example, you can touch the Navigation icon to get directions to your destination or touch the Music icon to play your favorite songs. '
ANSWER:Now use the generateContent method to to generate a response from the model.
const answer = await ai.models.generateContent({
model: MODEL_ID,
contents: prompt,
});
tslab.display.markdown(answer.text ?? "");Using the touchscreen in your Google car is quite easy! You simply touch the picture, or “icon,” for the feature you’d like to use on its big display. For instance, if you want to find your way somewhere, you can tap the Navigation icon to get directions to your destination, or if you’re in the mood for some tunes, you can touch the Music icon to play your favorite songs. This handy screen also helps you manage other features like climate control!
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 Gemini Typescript quickstart.