使用嵌入向量进行搜索重排序：Google Colab 教程

本教程面向有编程基础的开发者，深入讲解如何利用 Gemini 嵌入向量对搜索结果进行重排序。通过结合 Gemini 的函数调用与 Wikipedia API，实现从查询生成、文档检索到基于语义相似度排序的完整流程。教程涵盖环境配置、API 密钥管理、模型选择、工具定义、自动与手动函数调用，以及基于余弦相似度的重排序技术，并引入 HyDE 方法提升排序效果。旨在帮助开发者掌握构建高效搜索重排序系统的工程实践。

技术背景

传统搜索引擎依赖关键词匹配，难以理解查询语义，导致结果相关性不足。嵌入向量将文本映射到高维向量空间，使语义相近的文本在向量空间中距离更近，从而支持基于向量相似度的语义搜索。本教程利用 Gemini 的嵌入模型和函数调用能力，结合 Wikipedia API，构建一个轻量级的搜索重排序系统。该方法适用于需要快速实现语义搜索、文档检索、问答系统等场景，但受限于 Wikipedia 数据源和 API 速率限制，不适合大规模生产环境。

核心原理

核心原理是将文本转换为嵌入向量，通过计算向量间的余弦相似度来评估语义相关性。流程为：用户查询经 Gemini 生成多个辅助查询，通过函数调用触发 Wikipedia 搜索，获取候选文档摘要；然后使用嵌入模型将查询和文档摘要编码为向量，计算相似度分数，选择最高分文档作为最终答案。HyDE 方法则先生成假设性答案，再与文档比较，以提升排序鲁棒性。

架构分析

系统由三部分组成：1) 查询生成模块，利用 Gemini 生成多样化搜索查询；2) 检索模块，通过函数调用访问 Wikipedia API，获取候选文档并提取相关信息；3) 重排序模块，使用嵌入模型计算相似度，对候选文档排序。关键设计决策包括：使用函数调用实现外部 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 嵌入进行搜索重排序

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

本笔记本演示了如何使用嵌入来重新排序搜索结果。本演练将重点介绍以下目标：

  1. 设置开发环境和 API 访问权限以使用 Gemini。
  2. 使用 Gemini 的函数调用支持来访问 Wikipedia API。
  3. 通过 Gemini API 嵌入内容。
  4. 重新排序搜索结果。

这就是你实现搜索重排序的方法：

  1. 用户将进行搜索查询。
  2. 你将使用 Wikipedia API 返回相关的搜索结果。
  3. 搜索结果将被嵌入，并通过计算余弦相似度等距离度量来评估其相关性。
  4. 最相关的搜索结果将作为最终答案返回。
本笔记本中的非源代码材料采用知识共享-署名-相同方式共享 CC-BY-SA 4.0 许可，https://creativecommons.org/licenses/by-sa/4.0/legalcode。

环境设置

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

%pip install -U -q google-genai

同时安装本教程中将用到的 wikipedia 包。

%pip install -U -q wikipedia

注意：wikipedia 包 指出它“为易用性和简单性而设计，并非用于高级用途”，并且生产环境或大量使用应改为“使用 Pywikipediabot 或其他更高级的 Python MediaWiki API 包装器”。

import json
import textwrap

from google import genai
from google.genai import types

import wikipedia
from wikipedia.exceptions import DisambiguationError, PageError

import numpy as np

from IPython.display import Markdown

def to_markdown(text):
  text = text.replace('•', '  *')
  return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))

获取 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.colab import userdata

# 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)

选择要使用的模型

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}

定义工具

如前所述，本教程使用 Gemini 的函数调用功能来访问 Wikipedia API。请参阅文档以了解更多关于函数调用的信息。

定义搜索函数

为了满足搜索引擎的需求，您将按以下方式设计此函数：

* 对于每个搜索查询，搜索引擎将使用 wikipedia.search 方法获取相关主题。

* 从相关主题中，引擎将选择 n_topics(int) 个最佳候选，并使用 gemini-3.7-flash 从页面中提取相关信息。

* 引擎将通过维护搜索历史来避免重复条目。

def wikipedia_search(search_queries: list[str]) -> list[str]:
  """Search wikipedia for each query and summarize relevant docs."""
  n_topics=3
  search_history = set() # tracking search history
  search_urls = []
  summary_results = []

  for query in search_queries:
    print(f'Searching for "{query}"')
    search_terms = wikipedia.search(query)

    print(f"Related search terms: {search_terms[:n_topics]}")
    for search_term in search_terms[:n_topics]: # select first `n_topics` candidates
      if search_term in search_history: # check if the topic is already covered
        continue

      print(f'Fetching page: "{search_term}"')
      search_history.add(search_term) # add to search history

      try:
        # extract the relevant data by using `gemini-3.7-flash` model
        page = wikipedia.page(search_term, auto_suggest=False)
        url = page.url
        print(f"Information Source: {url}")
        search_urls.append(url)
        page = page.content
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=textwrap.dedent(f"""\
                Extract relevant information
                about user's query: {query}
                From this source:

                {page}

                Note: Do not summarize. Only Extract and return the relevant information
        """))

        urls = [url]
        if response.candidates[0].citation_metadata:
          extra_citations = response.candidates[0].citation_metadata.citation_sources
          extra_urls = [source.url for source in extra_citations]
          urls.extend(extra_urls)
          search_urls.extend(extra_urls)
          print("Additional citations:", response.candidates[0].citation_metadata.citation_sources)
        try:
          text = response.text
        except ValueError:
          pass
        else:
          summary_results.append(text + "\n\nBased on:\n  " + ',\n  '.join(urls))

      except DisambiguationError:
        print(f"""Results when searching for "{search_term}" (originally for "{query}")
        were ambiguous, hence skipping""")
        continue

      except PageError:
        print(f'{search_term} did not match with any page id, hence skipping.')
        continue
        
      except:
        print(f'{search_term} did not match with any page id, hence skipping.')
        continue

  print(f"Information Sources:")
  for url in search_urls:
    print('    ', url)

  return summary_results

example = wikipedia_search(["What are LLMs?"])

以下是搜索结果的样子：

from IPython.display import display

for e in example:
  display(to_markdown(e))

将工具传递给模型

如果你将一个函数列表传递给 GenerativeModel 的 tools 参数，它会从函数的签名和类型提示中提取一个 schema，然后将该 schema 传递给 API 调用。作为响应，模型可能会返回一个 FunctionCall 对象，要求调用该函数。

注意：这种方法仅处理 AllowedTypes = int | float | str | dict | list['AllowedTypes'] 的注解。

对 Gemini 模型的请求会在自身保留对该函数的引用，以便稍后可以在本地执行该函数。

生成辅助搜索查询

为了让用户原始查询有多个辅助搜索查询，你需要让模型生成更多这样的查询。这将有助于搜索引擎在全面层面上覆盖所提出的问题。

instructions = """You have access to the Wikipedia API which you will be using
to answer a user's query. Your job is to generate a list of search queries which
might answer a user's question. Be creative by using various key-phrases from
the user's query. To generate variety of queries, ask questions which are
related to  the user's query that might help to find the answer. The more
queries you generate the better are the odds of you finding the correct answer.
Here is an example:

user: Tell me about Cricket World cup 2023 winners.

function_call: wikipedia_search(['What is the name of the team that
won the Cricket World Cup 2023?', 'Who was the captain of the Cricket World Cup
2023 winning team?', 'Which country hosted the Cricket World Cup 2023?', 'What
was the venue of the Cricket World Cup 2023 final match?', 'Cricket World cup 2023',
'Who lifted the Cricket World Cup 2023 trophy?'])

The search function will return a list of article summaries, use these to
answer the  user's question.

Here is the user's query: {query}
"""

为了生成更具创意和更多样化的问题，你可以将模型的温度参数设置为较高的值。该值的范围在[0.0,1.0]之间（包含端点）。值越接近1.0，生成的回答就越多样、越有创意；而值越接近0.0，模型的回答通常会更直接、更保守。

启用自动函数调用并调用 API

现在，使用 enable_automatic_function_calling=True 开启一个新的对话。启用后，genai.ChatSession 将自动处理调用函数所需的来回交互，并返回最终响应：

tools = [wikipedia_search]

config = types.GenerateContentConfig(
    temperature=0.6,
    tools=tools
)

chat = client.chats.create(
    model="gemini-3.7-flash",
    config=config
)

query = "Explain how deep-sea life survives."

res = chat.send_message(instructions.format(query=query))

to_markdown(res.text)

看起来成功了。你可以浏览聊天记录，查看函数调用中发送和接收的详细信息。

for content in chat._comprehensive_history:
  part = content.parts[0]

  print(f'{content.role} -> ', end='')
  print(json.dumps(part.to_json_dict(), indent=2))
  print('---' * 20)

在聊天历史中，你可以看到全部4个步骤：

  1. 用户发送了查询。
  2. 模型回复了一个 genai.protos.FunctionCall，调用 wikipedia_search 并附带一些相关的搜索。
  3. 因为你在创建 genai.ChatSession 时设置了 enable_automatic_function_calling=True，它执行了搜索函数并将文章摘要列表返回给模型。
  4. 按照提示中的指令，模型基于这些摘要生成了最终答案。

[可选] 手动执行函数调用

如果你想了解幕后发生了什么，本节将手动执行 FunctionCall 来进行演示。

chat = client.chats.create(
    model=MODEL_ID,
    config=config
)

result = chat.send_message(instructions.format(query=query))

最初，模型返回一个FunctionCall：

fc = result.automatic_function_calling_history[1].parts[0].function_call
fc = fc.to_json_dict()
print(json.dumps(fc, indent=2))

fc['name']

调用函数并传入生成的参数以获取结果。

summaries = wikipedia_search(**fc['args'])

现在将 FunctionResult 发送给模型。

response = chat.send_message(
    [genai.types.Part.from_function_response(
        name='wikipedia_search',
        response={'result': summaries}
    )]
)

to_markdown(response.text)

对搜索结果进行重新排序

用于嵌入内容的辅助函数：

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

def get_embeddings(content: list[str]) -> np.ndarray:
  embeddings = client.models.embed_content(
    model=EMBEDDINGS_MODEL_ID,
    contents=content,
    config=types.EmbedContentConfig(
      task_type='SEMANTIC_SIMILARITY'
    )
  )
  return embeddings.embeddings[0].values

有关嵌入的更多信息，请参阅嵌入指南。

下一步，你需要定义一些函数，用于计算两个嵌入向量之间的相似度分数。这些分数将帮助你确定哪个嵌入向量与用户的查询最相关。

现在你将实现余弦相似度作为你的度量标准。这里返回的嵌入向量是单位长度的，因此它们的 L1 范数（np.linalg.norm()）将约为 1。因此，计算余弦相似度本质上等同于计算它们的点积分数。

def dot_product(a: np.ndarray, b: np.ndarray):
  return (np.array(a) @ np.array(b).T)

与用户查询的相似度

现在，是时候找出维基百科 API 返回的最相关的搜索结果了。

使用 Gemini API 获取用户查询和搜索结果的嵌入向量。

search_res = get_embeddings(summaries)
embedded_query = get_embeddings([query])

计算相似度分数：

sim_value = dot_product(search_res, embedded_query)

使用 np.argmax 选择最佳候选。

用户输入： 解释深海生物如何生存。

回答：

Markdown(summaries[np.argmax(sim_value)])

假设文档嵌入相似度（HyDE）

受 Gao 等人 的启发，这里的目标是利用 gemini-3.7-flash 的内部知识生成用户查询的模板答案。这个假设答案将作为基线，用于计算所有搜索结果的相关性。

res = client.models.generate_content(
    model=MODEL_ID,
    contents=f"""Generate a hypothetical answer
to the user's query by using your own knowledge. Assume that you know everything
about the said topic. Do not use factual information, instead use placeholders
to complete your answer. Your answer should feel like it has been written by a human.

query: {query}""")

to_markdown(res.text)

使用 Gemini API 获取基线答案的嵌入向量，并将其与搜索结果进行比较。

hypothetical_ans = get_embeddings([res.text])

计算相似度分数，以对搜索结果进行排名

sim_value = dot_product(search_res, hypothetical_ans)

sim_value

使用 np.argmax 选择最佳候选。

用户输入： 解释深海生物如何生存。

回答：

to_markdown(summaries[np.argmax(sim_value)])

现在你已经使用嵌入向量创建了一个搜索重排序引擎！

后续步骤

希望这个示例对你有帮助！查看 Gemini API cookbook 中的更多示例，以了解更多信息。

常见问题

常见错误包括：API 密钥未正确设置导致认证失败；模型选择不当导致速率限制或功能不支持；Wikipedia 页面歧义或不存在导致异常；嵌入模型返回结果格式错误。排查方法：检查 API 密钥是否有效，确认模型 ID 正确，处理 DisambiguationError 和 PageError 异常，验证嵌入向量维度。生产环境优化：使用缓存减少 API 调用，采用批量嵌入，考虑使用更高效的索引结构（如 ANN），并监控 API 使用量。

企业应用扩展

可扩展方向：1) 集成更多数据源（如新闻、学术论文）；2) 使用更先进的嵌入模型（如 gemini-embedding-2-preview）；3) 实现混合检索（关键词+语义）；4) 引入学习排序（LTR）模型；5) 优化查询生成策略，如使用温度调节控制多样性；6) 部署为 REST API 服务。

总结

本教程面向有编程基础的开发者，深入讲解如何利用 Gemini 嵌入向量对搜索结果进行重排序。通过结合 Gemini 的函数调用与 Wikipedia API，实现从查询生成、文档检索到基于语义相似度排序的完整流程。教程涵盖环境配置、API 密钥管理、模型选择、工具定义、自动与手动函数调用，以及基于余弦相似度的重排序技术，并引入 HyDE 方法提升排序效果。旨在帮助开发者掌握构建高效搜索重排序系统的工程实践。

内容来源与验证

继续系统学习

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

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

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