文档搜索教程：使用 Gemini API 进行语义搜索

本教程面向有编程基础的开发者，深入讲解如何利用 Gemini API 的嵌入模型实现语义文档搜索。通过构建嵌入数据库、计算点积相似度，并结合生成模型构建问答系统，读者将掌握从文本嵌入到检索增强生成（RAG）的完整流程。教程涵盖环境配置、API 密钥管理、模型选择、任务类型区分、代码实现及优化建议，帮助开发者快速落地 AI 搜索应用。

技术背景

传统关键词搜索基于字面匹配，难以理解语义，导致检索结果不准确。语义搜索通过将文本映射到高维向量空间，使语义相近的文本在向量空间中距离更近，从而提升检索质量。Gemini API 提供预训练嵌入模型，支持多种任务类型，可灵活应用于文档检索、问答系统等场景。本教程旨在帮助开发者掌握利用 Gemini API 构建语义搜索应用的核心技术。

核心原理

核心原理是将文本转换为向量表示（嵌入），通过计算向量间的点积（余弦相似度）衡量语义相似性。Gemini 嵌入模型输出归一化向量，点积值范围在 [-1, 1] 之间，值越大表示语义越接近。教程中，文档和查询分别使用 RETRIEVAL_DOCUMENT 和 RETRIEVAL_QUERY 任务类型生成嵌入，确保模型针对检索场景优化。问答阶段，通过检索最相关文档，结合生成模型（如 gemini-3.7-flash）生成自然语言答案，实现检索增强生成（RAG）。

架构分析

系统架构分为三部分：1) 嵌入生成模块：调用 Gemini API 的 embed_content 方法，将文档和查询转换为向量；2) 检索模块：使用 NumPy 计算查询向量与文档向量的点积，选取最高分文档；3) 生成模块：利用 Gemini 生成模型，基于检索到的文档和用户问题生成回答。数据流为：文档 → 嵌入 → 存储（DataFrame）→ 查询嵌入 → 点积排序 → 检索结果 → 提示词构造 → 生成回答。关键设计决策包括：使用 DataFrame 存储嵌入便于操作，区分任务类型以优化嵌入质量，以及通过点积而非余弦相似度简化计算（因向量已归一化）。

完整实战教程

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

# @title Licensed under the Apache License, Version 2.0 (the "License");
#
# 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.

使用嵌入向量的文档搜索

<a target="_blank" href="https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/document_search.ipynb">

</a>

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

概述

本示例演示了如何使用 Gemini API 创建嵌入，以便执行文档搜索。您将使用 Python 客户端库构建一个词嵌入，使您能够将搜索字符串或问题与文档内容进行比较。

在本教程中，您将使用嵌入对一组文档执行文档搜索，以提出与 Google Car 相关的问题。

前提条件

您可以在 Google Colab 中运行此快速入门。

要在您自己的开发环境中完成此快速入门，请确保您的环境满足以下要求：

设置

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

%pip install -U -q google-genai

import textwrap
import numpy as np
import pandas as pd

from google import genai
from google.genai import types

# Used to securely store your API key
from google.colab import userdata

from IPython.display import Markdown

获取 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=...)。

# Or use `os.getenv('GEMINI_API_KEY')` to fetch an environment variable.
GEMINI_API_KEY=userdata.get('GEMINI_API_KEY')

client = genai.Client(api_key=GEMINI_API_KEY)

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

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

选择要使用的模型

MODEL_ID = "gemini-embedding-001" # @param ["gemini-embedding-2-preview", "gemini-embedding-001"] {"allow-input":true, isTemplate: true}

生成嵌入

在本节中，您将看到如何使用Gemini API中的嵌入为数据框中的不同文本生成嵌入。

Gemini嵌入模型支持多种任务类型，每种类型都针对特定目标而定制。以下是可用类型及其应用的总体概述：

| 任务类型 | 描述 |

| --- | --- |

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

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

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

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

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

sample_text = ("Title: The next generation of AI for developers and Google Workspace"
    "\n"
    "Full article:\n"
    "\n"
    "Gemini API & Google AI Studio: An approachable way to explore and prototype with generative AI applications")

embedding = client.models.embed_content(model=MODEL_ID,
                                contents=sample_text,
                                config=types.EmbedContentConfig(
                                  task_type="RETRIEVAL_DOCUMENT"))

print(embedding.embeddings)

构建嵌入数据库

以下是三个示例文本，用于构建嵌入数据库。你将使用 Gemini API 为每个文档创建嵌入。将它们转换为数据框以便更好地可视化。

DOCUMENT1 = {
    "title": "Operating the Climate Control System",
    "content": "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 = {
    "title": "Touchscreen",
    "content": "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 = {
    "title": "Shifting Gears",
    "content": "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]

将字典的内容整理成数据框，以便更好地可视化。

df = pd.DataFrame(documents)
df.columns = ['Title', 'Text']
df

获取这些文本中每一段的嵌入向量。将此信息添加到数据框中。

# Get the embeddings of each text and add to an embeddings column in the dataframe
def embed_fn(text):
  return client.models.embed_content(model=MODEL_ID,
                             contents=text,
                             config=types.EmbedContentConfig(
                               task_type="RETRIEVAL_DOCUMENT")
                             )

df['Embeddings'] = df.apply(lambda row: embed_fn(row['Text']).embeddings[0].values, axis=1)
df

带问答的文档搜索

现在嵌入已生成，让我们创建一个问答系统来搜索这些文档。您将提出一个关于超参数调优的问题，创建该问题的嵌入，并将其与数据框中的嵌入集合进行比较。

问题的嵌入将是一个向量（浮点值列表），将使用点积与文档的向量进行比较。API 返回的此向量已经归一化。点积表示两个向量在方向上的相似性。

点积的值范围在 -1 到 1 之间（含）。如果两个向量之间的点积为 1，则向量方向相同。如果点积值为 0，则这些向量正交或彼此无关。最后，如果点积为 -1，则向量指向相反方向，彼此不相似。

注意，使用新的嵌入模型（embedding-001）时，请将任务类型指定为 QUERY 用于用户查询，并在嵌入文档文本时指定为 DOCUMENT。

| 任务类型 | 描述 |

| --- | --- |

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

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

query = "How do you shift gears in the Google car?"
model = MODEL_ID

request = client.models.embed_content(model=model,
                                contents=query,
                                config=types.EmbedContentConfig(
                                  task_type="RETRIEVAL_QUERY"))

使用 find_best_passage 函数计算点积，然后对数据框按点积值从大到小排序，以从数据库中检索相关段落。

def find_best_passage(query, dataframe):
  """
  Compute the distances between the query and each document in the dataframe
  using the dot product.
  """
  query_embedding = client.models.embed_content(model=model,
                                                contents=query,
                                                config=types.EmbedContentConfig(
                                                  task_type="RETRIEVAL_QUERY")).embeddings[0].values
  dot_products = np.dot(np.stack(dataframe['Embeddings']), query_embedding)
  idx = np.argmax(dot_products)
  return dataframe.iloc[idx]['Text'] # Return text from index with max value

从数据库中查看最相关的文档：

passage = find_best_passage(query, df)
passage

问答应用

让我们尝试使用文本生成 API 来创建一个问答系统。在下面输入你自己的自定义数据，以创建一个简单的问答示例。你仍然会使用点积作为相似度的度量。

def make_prompt(query, relevant_passage):
  escaped = relevant_passage.replace("'", "").replace('"', "").replace("\n", " ")
  prompt = textwrap.dedent("""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

prompt = make_prompt(query, passage)
print(prompt)

选择一个 Gemini 内容生成模型来寻找你的查询的答案。

answer = client.models.generate_content(
    model="gemini-3.7-flash",
    contents=prompt
)

print(answer.text)

Markdown(answer.text)

后续步骤

要了解如何使用 Gemini API 中的其他服务，请参阅入门指南。

常见问题

常见错误：1) API 密钥未正确设置，导致认证失败；解决：确保密钥已存入环境变量或显式传入。2) 模型选择不一致，导致嵌入不兼容；解决：固定使用同一模型。3) 任务类型使用错误，影响检索效果；解决：查询用 RETRIEVAL_QUERY，文档用 RETRIEVAL_DOCUMENT。4) 点积计算维度不匹配；解决：检查嵌入向量维度是否一致。5) 生成模型输出不相关；解决：优化提示词，确保检索文档相关。生产环境优化：使用向量数据库（如 Pinecone）存储嵌入，支持大规模检索；对文档进行分块处理，提高检索精度；缓存嵌入结果，减少 API 调用；使用异步请求提升吞吐量。

企业应用扩展

可扩展方向：1) 支持多种文档格式（PDF、Word）的解析与嵌入；2) 集成向量数据库实现持久化存储和高效检索；3) 实现多轮对话式搜索，结合对话历史优化查询；4) 添加重排序模型（如 Cross-Encoder）提升检索精度；5) 部署为 REST API 服务，供前端应用调用；6) 结合用户反馈进行模型微调，优化嵌入效果。

总结

本教程面向有编程基础的开发者，深入讲解如何利用 Gemini API 的嵌入模型实现语义文档搜索。通过构建嵌入数据库、计算点积相似度，并结合生成模型构建问答系统，读者将掌握从文本嵌入到检索增强生成（RAG）的完整流程。教程涵盖环境配置、API 密钥管理、模型选择、任务类型区分、代码实现及优化建议，帮助开发者快速落地 AI 搜索应用。

内容来源与验证

继续系统学习

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

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

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