Gemini API：使用 ChromaDB 进行文档问答

本教程面向有编程基础的开发者，深入讲解如何利用 Gemini API 的嵌入模型与 ChromaDB 向量数据库构建一个完整的文档问答系统。内容涵盖环境配置、API 密钥管理、自定义嵌入函数、向量数据库的创建与查询，以及如何将检索到的相关段落与生成模型结合，实现检索增强生成（RAG）。教程强调工程实践，包括模型选择、任务类型参数、数据流分析、常见错误排查及生产环境优化建议，帮助开发者掌握构建高效、可扩展的 AI 应用的核心技能。

技术背景

在自然语言处理（NLP）领域，传统的关键词搜索难以理解语义，而向量数据库通过将文本转换为高维向量，实现语义相似度检索。Gemini API 提供强大的嵌入模型，可将文本映射到向量空间，ChromaDB 则是一个轻量级、开源的向量数据库，支持存储、检索和过滤。本教程解决的核心问题是：如何从大量非结构化文档中快速、准确地找到与用户查询相关的信息，并利用生成模型生成自然语言答案。这种模式被称为检索增强生成（RAG），广泛应用于智能客服、知识库问答、企业文档检索等场景。其优势在于无需微调模型即可利用最新知识，且可解释性强。主要限制包括：嵌入模型的选择影响检索质量，向量数据库的规模受内存限制，以及 API 的速率限制和区域可用性。

核心原理

本教程的核心原理是检索增强生成（RAG），其数据流如下：1) 文档预处理：将原始文档切分为段落或块，作为独立的检索单元。2) 嵌入生成：使用 Gemini API 的嵌入模型（如 gemini-embedding-001）将每个文档块转换为向量表示。嵌入模型将文本映射到高维空间，语义相近的文本在向量空间中距离更近。3) 向量存储：将文档向量及其元数据（如原始文本、ID）存入 ChromaDB 集合。ChromaDB 使用 HNSW（Hierarchical Navigable Small World）算法构建索引，支持高效的最近邻搜索。4) 查询处理：用户查询同样通过嵌入模型转换为向量，然后使用 ChromaDB 的 query 方法在向量空间中查找最相似的文档向量。5) 答案生成：将检索到的相关文档作为上下文，与用户查询一起构造提示词，输入给生成模型（如 gemini-3.7-flash），生成最终答案。关键设计决策包括：使用自定义嵌入函数封装 Gemini API 调用，以便与 ChromaDB 无缝集成；设置 task_type 参数以优化嵌入效果（如 RETRIEVAL_DOCUMENT 用于文档，RETRIEVAL_QUERY 用于查询）；选择固定的嵌入模型以保证向量空间的一致性。

架构分析

系统架构分为三层：数据层、检索层和生成层。数据层负责文档的收集与预处理，本教程使用三个预定义的文档字符串，实际应用中可扩展为从文件、数据库或网络爬虫获取。检索层由 ChromaDB 和自定义嵌入函数组成。ChromaDB 客户端负责管理集合（Collection），集合是文档、向量和元数据的容器。自定义嵌入函数 GeminiEmbeddingFunction 实现了 ChromaDB 的 EmbeddingFunction 接口，其 __call__ 方法接收文档列表，调用 Gemini API 的 embed_content 方法，返回嵌入向量。该函数在创建集合时传入，使得 ChromaDB 在添加文档或查询时自动调用嵌入函数。生成层使用 Gemini API 的生成模型，将检索到的相关段落与用户问题组合成提示词，生成答案。组件关系：用户查询 -> 嵌入函数 -> ChromaDB 查询 -> 相关文档 -> 提示词构造 -> 生成模型 -> 答案。设计思想：将嵌入和生成解耦，便于替换模型或数据库；使用 ChromaDB 的持久化功能（本教程未展示，但可配置）实现数据持久化；通过自定义嵌入函数实现与 Gemini API 的紧密集成，同时保持代码的可读性和可扩展性。

完整实战教程

##### 版权所有 2026 Google LLC。

# @title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Gemini API：使用 ChromaDB 进行文档问答

注意： 此笔记本需要付费层级速率限制才能正常运行。（更多详情请参阅定价）。

概述

本教程演示了如何使用 Gemini API 创建向量数据库，并从数据库中检索问题的答案。此外，您将使用 ChromaDB{:.external}，这是一个用于创建嵌入数据库的开源 Python 工具。ChromaDB 允许您：

* 存储嵌入及其元数据

* 嵌入文档和查询

* 在嵌入数据库中进行搜索

在本教程中，您将使用嵌入从使用 ChromaDB 创建的向量数据库中检索答案。

环境设置

首先，下载并安装 ChromaDB 和 Gemini API Python 库。

%pip install -U -q "google-genai>=2.9.0" chromadb

然后导入本教程中将使用的模块。

import textwrap
import chromadb
import numpy as np
import pandas as pd

from IPython.display import Markdown
from chromadb import Documents, EmbeddingFunction, Embeddings

配置你的 API 密钥

在使用 Gemini API 之前，你必须先获取一个 API 密钥。如果你还没有，可以在 Google AI Studio 中一键创建。

<a class="button button-primary" href="https://aistudio.google.com/app/apikey" target="_blank" rel="noopener noreferrer">获取 API 密钥</a>

在 Colab 中，将密钥添加到左侧面板的“🔑”下的密钥管理器中，并将其命名为 GEMINI_API_KEY。

获得 API 密钥后，将其传递给 SDK。你可以通过两种方式实现：

* 将密钥放入 GEMINI_API_KEY 环境变量中（SDK 会自动从该变量中读取）。

* 将密钥传递给 genai.Client(api_key=...)。

from google import genai
from google.colab import userdata


GEMINI_API_KEY=userdata.get('GEMINI_API_KEY')
client = genai.Client(api_key=GEMINI_API_KEY)

关键点：接下来，你将选择一个模型。任何嵌入模型都可以用于本教程，但在实际应用中，选择一个特定的模型并坚持使用它非常重要。不同模型的输出彼此不兼容。

注意：目前，Gemini API 仅在某些地区可用。

for m in client.models.list():
  if 'embedContent' in m.supported_actions:
    print(m.name)

数据

以下是一小组文档，你将用它们来创建嵌入数据库：

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.
"""
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.
"""
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.
"""

documents = [DOCUMENT1, DOCUMENT2, DOCUMENT3]

使用 ChromaDB 创建嵌入数据库

你将创建一个自定义函数{:.external}，用于通过 Gemini API 执行嵌入。通过向此自定义函数输入一组文档，你将获得文档的向量或嵌入。

使用 gemini-embedding-001 模型对 Embeddings 的 API 变更

对于新的 embeddings 模型 embedding-001，新增了一个任务类型参数，以及一个可选的标题（仅当 task_type=RETRIEVAL_DOCUMENT 时有效）。

这些新参数仅适用于最新的 embeddings 模型。任务类型如下：

| 任务类型 | 描述 |

| --- | --- |

| RETRIEVAL_QUERY | 指定给定文本是搜索/检索设置中的查询。 |

| RETRIEVAL_DOCUMENT | 指定给定文本是搜索/检索设置中的文档。 |

| SEMANTIC_SIMILARITY | 指定给定文本将用于语义文本相似度（STS）。 |

| CLASSIFICATION | 指定嵌入将用于分类。 |

| CLUSTERING | 指定嵌入将用于聚类。 |

from google.genai import types

class GeminiEmbeddingFunction(EmbeddingFunction):
  def __call__(self, input: Documents) -> Embeddings:
    EMBEDDING_MODEL_ID = "gemini-embedding-001"  # @param ["text-embedding-004", "gemini-embedding-001"] {"allow-input": true, "isTemplate": true}
    title = "Custom query"
    response = client.models.embed_content(
        model=EMBEDDING_MODEL_ID,
        contents=input,
        config=types.EmbedContentConfig(
          task_type="retrieval_document",
          title=title
        )
    )

    return response.embeddings[0].values

现在你将创建向量数据库。在 create_chroma_db 函数中，你将实例化一个 Chroma 客户端。接着，你将创建一个集合，用于存储嵌入、文档和任何元数据。请注意，上面的嵌入函数作为参数传递给了 create_collection。

接下来，你使用 add 方法将文档添加到集合中。

def create_chroma_db(documents, name):
  chroma_client = chromadb.Client()
  db = chroma_client.create_collection(
      name=name,
      embedding_function=GeminiEmbeddingFunction()
  )

  for i, d in enumerate(documents):
    db.add(
      documents=d,
      ids=str(i)
    )
  return db

# Set up the DB
db = create_chroma_db(documents, "google-car-db")

通过查看数据库来确认数据已插入：

sample_data = db.get(include=['documents', 'embeddings'])

df = pd.DataFrame({
    "IDs": sample_data['ids'][:3],
    "Documents": sample_data['documents'][:3],
    "Embeddings": [str(emb)[:50] + "..." for emb in sample_data['embeddings'][:3]]  # Truncate embeddings
})

df

获取相关文档

db 是一个 Chroma 集合对象。你可以对其调用 query 来执行最近邻搜索，以查找相似的嵌入或文档。

def get_relevant_passage(query, db):
  passage = db.query(query_texts=[query], n_results=1)['documents'][0][0]
  return passage

# Perform embedding search
passage = get_relevant_passage("touch screen features", db)
Markdown(passage)

既然你已经在文档集中找到了相关段落，就可以利用它来构建一个提示词，并将其传递给 Gemini API。

def make_prompt(query, relevant_passage):
  escaped = relevant_passage.replace("'", "").replace('"', "").replace("\n", " ")
  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 converstional tone. If the passage is irrelevant to the answer,
    you may ignore it.
    QUESTION: '{query}'
    PASSAGE: '{relevant_passage}'

    ANSWER:
  """).format(query=query, relevant_passage=escaped)

  return prompt

将查询传递给提示词：

query = "How do you use the touchscreen in the Google car?"
prompt = make_prompt(query, passage)
Markdown(prompt)

现在使用 generate_content 方法从模型生成响应。

MODEL_ID = "gemini-3.7-flash" # @param ["gemini-3.1-pro-preview", "gemini-3.7-flash", "gemini-3.5-flash-lite", "gemini-2.5-pro"] {"allow-input":true, isTemplate: true}answer = client.models.generate_content(
    model = MODEL_ID,
    contents = prompt
)
Markdown(answer.text)

后续步骤

要了解如何使用嵌入的更多信息，请查看可用的示例。要了解如何使用 Gemini API 中的其他服务，请访问 Python 快速入门。

常见问题

常见错误及排查方法：1) API 密钥错误或未设置：确保在环境变量或代码中正确设置 GEMINI_API_KEY，并检查是否有足够的权限。2) 嵌入模型不可用：某些模型可能不支持嵌入，需通过 client.models.list() 检查 supported_actions 是否包含 'embedContent'。3) 速率限制：免费层级可能无法满足需求，需升级到付费层级或增加重试机制。4) 向量维度不匹配：不同嵌入模型输出的维度不同，若更换模型，需重新创建数据库。5) ChromaDB 内存不足：当文档量过大时，可考虑使用持久化客户端（chromadb.PersistentClient）或分布式部署。6) 查询结果不相关：检查嵌入模型是否适合领域，调整 task_type 参数，或增加 n_results 数量。7) 生成答案质量差：优化提示词，确保相关段落包含足够信息，或调整生成模型的参数（如 temperature）。生产环境优化：1) 使用批量嵌入 API 提高吞吐量。2) 对文档进行分块（chunking）时，考虑重叠和大小，以平衡检索精度和上下文长度。3) 使用缓存减少重复查询的 API 调用。4) 监控 API 使用量和延迟，设置告警。5) 考虑使用更高效的向量索引（如 IVF）或分布式向量数据库（如 Milvus）以支持大规模数据。

企业应用扩展

本教程可扩展的方向：1) 集成更多数据源：支持从 PDF、网页、数据库等导入文档，并实现自动分块。2) 实现多轮对话：将历史对话作为上下文，提升问答的连贯性。3) 添加元数据过滤：在 ChromaDB 中存储文档来源、日期等元数据，实现按条件过滤检索。4) 使用不同的嵌入模型：比较 text-embedding-004 与 gemini-embedding-001 的效果，选择最适合的模型。5) 部署为 Web 服务：使用 FastAPI 或 Flask 封装问答接口，供前端调用。6) 实现文档更新：支持增量添加和删除文档，保持数据库实时性。7) 结合 LangChain 等框架：利用现成的 RAG 链简化开发。8) 优化提示词工程：使用少样本示例或动态调整提示词，提高答案准确性。9) 评估系统性能：构建测试集，计算检索准确率和生成答案的 BLEU/ROUGE 分数。10) 探索其他向量数据库：如 Pinecone、Weaviate，对比性能和功能。

总结

本教程面向有编程基础的开发者，深入讲解如何利用 Gemini API 的嵌入模型与 ChromaDB 向量数据库构建一个完整的文档问答系统。内容涵盖环境配置、API 密钥管理、自定义嵌入函数、向量数据库的创建与查询，以及如何将检索到的相关段落与生成模型结合，实现检索增强生成（RAG）。教程强调工程实践，包括模型选择、任务类型参数、数据流分析、常见错误排查及生产环境优化建议，帮助开发者掌握构建高效、可扩展的 AI 应用的核心技能。

内容来源与验证

继续系统学习

站内还有更多完整实战教程

加入会员，解锁全部教程、完整代码与持续更新内容。

加入会员
查看全部RAG 与知识库教程