Google AI：用 gemini-embedding-2 给 FAQ 做语义检索

用 gemini-embedding-2 给 FAQ 做语义检索。按照官方文档，有两个和旧模型不一样的地方：

• 不能使用 task_type 参数，任务要写进文本：查询写成 task: search result | query: ...，文档写成 title: ... | text: ...；
• 直接传入一个字符串列表，只会得到一个聚合向量。要每条文本各得到一个向量，需要把每条文本分别包成一个 Content。

先看清成品

• 本地替身上：5 条 FAQ 发出 5 个独立的请求，每个都设置了 768 维，文本格式符合文档要求；“发票 邮箱”检索到的第一条是“发票”。pytest 4 项通过。
• 复现错误写法：直接传入 ["a","b","c"]，SDK 只发出 1 个请求（包含 3 个 part），返回 1 个向量，被 check_count 拦下。
• 用无效 Key 调用真实接口：google.genai.errors.ClientError: 400 INVALID_ARGUMENT ... API key not valid。

合规提醒：Gemini API 和 Google AI Studio 只在官方支持的地区提供，官方列表中没有中国大陆和香港。不在支持地区就不要使用，不要用代理或虚假信息绕过地区限制，也不要共享账号或 API Key。按照官方条款，使用免费额度（Unpaid Services）时，提交的内容和生成的结果可能被 Google 用来改进产品，还可能有人工审阅，所以不要提交个人敏感信息或公司机密。限额按项目计算，不是按 Key 计算，多建 Key 也不会增加额度。API Key 放在环境变量 GEMINI_API_KEY 中。本文不列价格和具体限额。

运行说明：box 上没有 Gemini API Key，所以没有拿到真实的模型输出。真实运行过的部分是：用无效 Key 调用真实接口，记录真实的错误格式。成功路径用的是我自己写的本地替身（fake_gemini.py，不是 Google 官方程序），客户端是真实的 google-genai 2.32.0；替身的请求和响应格式来自本地抓取的 SDK 实际请求，以及 SDK 自带测试中的数据结构。

做完之后你会得到什么？

需要准备哪些环境？

关键事实（来自官方 embeddings 文档）

文件结构

62-google-ai-B-src/
├── faq_search.py
├── faq.json
├── fake_gemini.py
└── test_faq.py

第 1 步：检索脚本

"""FAQ 语义检索：gemini-embedding-2 + 余弦相似度。
按官方 embeddings 文档：
- gemini-embedding-2 **不能**用 task_type，要把任务写进文本：
  查询用 "task: search result | query: {q}"，文档用 "title: {title} | text: {content}"
- 直接把多个字符串放进 contents，会被**聚合成一个向量**；要逐条得到向量，每条文本要单独包成一个 Content
- output_dimensionality=768 时，gemini-embedding-2 会自动归一化"""
import json, math, os, sys
from google.genai import types

MODEL = os.environ.get("EMBED_MODEL", "gemini-embedding-2")
DIM = 768

def doc_text(d):
    return f"title: {d.get('title') or 'none'} | text: {d['text']}"

def query_text(q):
    return f"task: search result | query: {q}"

def embed(client, texts):
    contents = [types.Content(parts=[types.Part(text=t)]) for t in texts]       # 每条一个 Content → 每条一个向量
    r = client.models.embed_content(model=MODEL, contents=contents,
                                    config=types.EmbedContentConfig(output_dimensionality=DIM))
    return check_count([e.values for e in r.embeddings], len(texts))

def check_count(vecs, n):
    if len(vecs) != n:
        raise RuntimeError(f"期望 {n} 个向量，实际得到 {len(vecs)} 个（是否被聚合了？）")
    return vecs

def cos(a, b):
    na, nb = math.sqrt(sum(x * x for x in a)), math.sqrt(sum(x * x for x in b))
    return sum(x * y for x, y in zip(a, b)) / (na * nb or 1.0)

class Index:
    def __init__(self, client, docs):
        self.client, self.docs = client, docs
        self.vecs = embed(client, [doc_text(d) for d in docs])
    def search(self, q, k=2):
        qv = embed(self.client, [query_text(q)])[0]
        scored = sorted(((cos(qv, v), d) for v, d in zip(self.vecs, self.docs)), key=lambda x: -x[0])
        return scored[:k]

if __name__ == "__main__":
    from google import genai
    client = genai.Client()
    idx = Index(client, json.load(open("faq.json", encoding="utf-8")))
    for q in sys.argv[1:] or ["买的东西不想要了能退吗？"]:
        for s, d in idx.search(q):
            print(f"{s:.3f}  {d['title']}：{d['text']}")

第 2 步：测试

import json
import pytest
from google import genai
from google.genai import types
import faq_search as S, fake_gemini as F

@pytest.fixture
def client():
    F.reset(); s, u = F.start()
    yield genai.Client(api_key="good", http_options=types.HttpOptions(base_url=u)); s.shutdown()

DOCS = json.load(open("faq.json", encoding="utf-8"))

def test_one_request_per_text_with_task_format(client):
    S.Index(client, DOCS)
    path, body = F.STATE["requests"][-1]
    assert path == "/v1beta/models/gemini-embedding-2:batchEmbedContents"
    assert len(body["requests"]) == 5 and all(r["outputDimensionality"] == 768 for r in body["requests"])
    assert body["requests"][0]["content"]["parts"][0]["text"] == "title: 退货期限 | text: 商品签收后 7 天内可以申请无理由退货，定制商品除外。"
    assert "taskType" not in json.dumps(body)

def test_search_ranks_relevant_first(client):
    idx = S.Index(client, DOCS)
    assert idx.search("发票 邮箱")[0][1]["title"] == "发票"
    assert F.STATE["requests"][-1][1]["requests"][0]["content"]["parts"][0]["text"].startswith("task: search result | query: ")

def test_aggregation_pitfall_detected(client):
    # 复现错误写法：把字符串列表直接交给 SDK → 只发出 1 个请求（3 个 part），返回 1 个聚合向量
    r = client.models.embed_content(model="gemini-embedding-2", contents=["a", "b", "c"])
    assert len(F.STATE["requests"][-1][1]["requests"]) == 1 and len(r.embeddings) == 1
    with pytest.raises(RuntimeError, match="期望 3 个向量，实际得到 1 个"):
        S.check_count([e.values for e in r.embeddings], 3)

def test_bad_key(client):
    bad = genai.Client(api_key="bad", http_options=types.HttpOptions(base_url=client._api_client._http_options.base_url))
    with pytest.raises(genai.errors.APIError) as e:                           # embed 走的是老接口，能被 APIError 捕获
        S.embed(bad, ["x"])
    assert e.value.code == 400

fake_gemini.py 和 62A 相同。

第 3 步：运行

$ GEMINI_API_KEY=INVALID python faq_search.py
google.genai.errors.ClientError: 400 INVALID_ARGUMENT. {'error': {'code': 400, 'message': 'API key not valid. Please pass a valid API key.', ...
$ python -m pytest -q
4 passed

调试示例

1. 字符串列表被聚合成一个向量。 我在本地抓取 SDK 的请求时发现：contents=['a','b'] 配合 gemini-embedding-2 时，发出的是 1 个请求，里面有 2 个 part；而同样的写法配合 gemini-embedding-001 时，会发出 2 个请求。改成传入 [types.Content(parts=[types.Part(text=t)]) for t in texts] 之后，变成每条文本一个请求。随后我在官方文档中找到了对应的说明（“produces a single aggregated embedding for multiple inputs”）。为了防止以后再出这种问题，check_count 会检查向量数量是否和文本数量一致。

2. 我的 monkeypatch 写错了。 第一版测试试图用 monkeypatch 把 types.Content 替换掉，结果 SDK 报错 ValueError: Unsupported ...，测试失败的原因是我的写法，而不是被测代码。后来把数量检查拆成独立的 check_count，测试直接调用它，就不需要 monkeypatch 了。

常见问题

模板规则

• 每条文本单独包成一个 Content。
• 检查向量数量。
• 查询和文档都加上任务格式。
• 固定模型名和维度，并保存在索引的元数据里。

项目资源

向量检索的第一步是确认每条文档都有自己的向量，然后再谈检索效果。