动手实践：构建小型语义搜索应用

让我们把之前学习的知识点结合起来，构建一个小型、实用的语义搜索应用。本练习整合一个嵌入 (embedding)模型、一个向量 (vector)数据库客户端和一个简单的网页框架，体现一个完整但简易的搜索流程，从用户查询到相关结果。我们将使用在本地容易配置的组件，让您可以专注于它们之间的配合。

在本次实践中，我们将使用：

1. 嵌入模型：使用 sentence-transformers 库，搭配 all-MiniLM-L6-v2 等预训练 (pre-training)模型。这个模型高效，能为句子和短段落提供高质量的嵌入。
2. 向量数据库：ChromaDB。我们将使用其 Python 客户端进行本地持久化存储，这简化了本示例的配置。
3. 网页框架：FastAPI。一个现代的 Python 网页框架，易于使用并能自动生成交互式 API 文档。

1. 配置与依赖

首先，请确保您已安装所需的库。您可以使用 pip 进行安装：

pip install sentence-transformers chromadb fastapi uvicorn[standard] python-multipart Jinja2

• sentence-transformers：用于加载嵌入 (embedding)模型和生成向量 (vector)。
• chromadb：用于与 Chroma 向量数据库交互的客户端库。
• fastapi：用于创建我们 API 端点的网页框架。
• uvicorn：一个用于运行我们的 FastAPI 应用的 ASGI 服务器。
• python-multipart：FastAPI 处理表单数据所需（尽管我们可能使用 JSON）。
• Jinja2：FastAPI 用于可选 HTML 模板，如果需要，通常包含在 FastAPI 的依赖项中。

2. 数据准备与索引脚本

让我们创建一个脚本（index_data.py）来准备一些示例数据，生成嵌入 (embedding)，并将它们索引到 ChromaDB 中。

# index_data.py
import chromadb
from sentence_transformers import SentenceTransformer

# --- Configuration ---
MODEL_NAME = 'all-MiniLM-L6-v2'
COLLECTION_NAME = "docs_collection"
PERSIST_DIRECTORY = "./chroma_db_persist" # 存储数据库数据的目录

# --- Sample Data ---
# 文档的简单列表（此处为句子）
documents = [
    "The quick brown fox jumps over the lazy dog.", # 敏捷的棕色狐狸跳过懒惰的狗。
    "Artificial intelligence is transforming many industries.", # 人工智能正在改变许多行业。
    "Vector databases are optimized for similarity search.", # 向量数据库为相似性搜索而优化。
    "Natural language processing enables computers to understand text.", # 自然语言处理使计算机能够理解文本。
    "The capital of France is Paris.", # 法国的首都是巴黎。
    "Apples are a type of fruit, often red or green.", # 苹果是一种水果，通常是红色或绿色。
    "Machine learning algorithms learn from data.", # 机器学习算法从数据中学习。
    "Semantic search provides results based on meaning, not just keywords.", # 语义搜索提供基于含义而非仅基于关键词的结果。
]

# --- Initialization ---
print("Initializing embedding model...") # 正在初始化嵌入模型...
# 加载预训练的句子 Transformer 模型
# 这个模型将句子和段落映射到 384 维稠密向量空间
# 如果模型不存在，它将自动下载
model = SentenceTransformer(MODEL_NAME)

print("Initializing ChromaDB client...") # 正在初始化 ChromaDB 客户端...
# 初始化具有持久化功能的 ChromaDB 客户端
# 这会将数据库状态保存到指定目录
client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)

print(f"Getting or creating collection: {COLLECTION_NAME}") # 正在获取或创建集合：{COLLECTION_NAME}
# 获取或创建集合。如果存在，则加载它。
# 根据我们的 SentenceTransformer 模型指定嵌入函数
collection = client.get_or_create_collection(
    name=COLLECTION_NAME,
    embedding_function=chromadb.utils.embedding_functions.SentenceTransformerEmbeddingFunction(model_name=MODEL_NAME)
    # 如果需要，您也可以显式传递 metadata={'hnsw:space': 'cosine'}，
    # 但 SentenceTransformerEmbeddingFunction 通常会默认选择合适的。
)

# --- Indexing ---
print("Generating IDs and preparing data for indexing...") # 正在生成 ID 并准备数据进行索引...
# 为本示例生成简单的顺序 ID
doc_ids = [f"doc_{i}" for i in range(len(documents))]

# 检查数据是否需要索引（基于预期数量的简单检查）
# 在实际应用中，您可能需要更完善的方式来跟踪已索引数据
if collection.count() < len(documents):
    print(f"Indexing {len(documents)} documents...") # 正在索引 {len(documents)} 个文档...
    try:
        # 将文档添加到集合
        # ChromaDB 的 SentenceTransformerEmbeddingFunction 在此处自动处理嵌入生成
        collection.add(
            documents=documents,
            ids=doc_ids
        )
        print("Documents indexed successfully.") # 文档索引成功。
    except Exception as e:
        print(f"Error indexing documents: {e}") # 索引文档时出错：{e}
else:
    print("Documents seem to be already indexed.") # 文档似乎已经索引。

print(f"Collection '{COLLECTION_NAME}' now contains {collection.count()} documents.") # 集合 '{COLLECTION_NAME}' 现在包含 {collection.count()} 个文档。
print("Indexing script finished.") # 索引脚本完成。

说明：

• 我们定义了示例 documents 和配置参数 (parameter)。
• 我们加载 SentenceTransformer 模型。首次运行时，它将下载模型权重 (weight)。
• 我们为 ChromaDB 初始化一个 PersistentClient，并指定一个目录（./chroma_db_persist）来存储数据库文件。这使我们的索引在多次运行后仍能持久保存。
• 我们使用 client.get_or_create_collection。这很方便，因为它会在集合不存在时创建它，或在集合存在时加载现有集合。我们通过 embedding_function 将 SentenceTransformer 模型与该集合关联。当我们添加文档或执行查询时，ChromaDB 会自动使用此函数。
• 我们为每个文档生成简单的唯一 ID（doc_ids）。
• 我们使用 collection.add 将文档及其对应的 ID 添加到集合中。由于我们配置了 embedding_function，ChromaDB 会在内部调用模型，在存储每个文档之前获取其向量 (vector)。我们包含了一个基本检查，以避免每次都重复索引。

运行此脚本一次来填充您的本地 ChromaDB：

python index_data.py

您应该会看到指示初始化和成功索引的输出，并且会创建一个 chroma_db_persist 目录。

3. 使用 FastAPI 构建搜索 API

现在，让我们创建将处理搜索请求的网页应用（main.py）。

# main.py
import chromadb
from fastapi import FastAPI, Query, HTTPException
from sentence_transformers import SentenceTransformer
import uvicorn # 用于运行应用

# --- Configuration ---
MODEL_NAME = 'all-MiniLM-L6-v2'
COLLECTION_NAME = "docs_collection"
PERSIST_DIRECTORY = "./chroma_db_persist"
N_RESULTS = 3 # 返回的搜索结果数量

# --- Application Initialization ---
app = FastAPI(
    title="Simple Semantic Search API", # 简单语义搜索 API
    description="An API that uses a vector database for semantic search.", # 一个使用向量数据库进行语义搜索的 API。
    version="0.1.0"
)

# --- Global Variables / Resources ---
# 应用启动时一次性初始化资源
try:
    print("Loading embedding model...") # 正在加载嵌入模型...
    embedding_model = SentenceTransformer(MODEL_NAME)
    print("Model loaded successfully.") # 模型加载成功。

    print("Connecting to ChromaDB...") # 正在连接到 ChromaDB...
    db_client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
    collection = db_client.get_collection(name=COLLECTION_NAME)
    # 验证集合中是否有项目（可选但良好实践）
    if collection.count() == 0:
         print(f"Warning: Collection '{COLLECTION_NAME}' is empty. Did you run index_data.py?") # 警告：集合 '{COLLECTION_NAME}' 为空。您运行过 index_data.py 吗？
    print("ChromaDB connection successful.") # ChromaDB 连接成功。

except Exception as e:
    print(f"Error during initialization: {e}") # 初始化时出错：{e}
    # 适当处理初始化失败，可能是退出或抛出特定错误
    embedding_model = None
    collection = None

# --- API Endpoints ---
@app.get("/search/")
async def perform_search(
    query: str = Query(..., min_length=3, description="The search query text.") # 搜索查询文本。
):
    """
    对已索引文档执行语义搜索。

    接收一个查询字符串，生成其嵌入，并搜索向量
    数据库以查找最相似的文档。
    """
    if not embedding_model or not collection:
        raise HTTPException(status_code=503, detail="Search service is not available due to initialization error.") # 由于初始化错误，搜索服务不可用。

    print(f"Received query: '{query}'") # 收到查询：'{query}'

    try:
        # 1. 为查询生成嵌入
        print("Generating query embedding...") # 正在生成查询嵌入...
        query_embedding = embedding_model.encode(query).tolist()
        print("Query embedding generated.") # 查询嵌入已生成。

        # 2. 查询向量数据库
        print(f"Querying collection '{COLLECTION_NAME}'...") # 正在查询集合 '{COLLECTION_NAME}'...
        results = collection.query(
            query_embeddings=[query_embedding], # 注意：query_embeddings 期望一个嵌入列表
            n_results=N_RESULTS,
            include=['documents', 'distances'] # 请求 ChromaDB 返回文档和距离
        )
        print("Query executed successfully.") # 查询执行成功。

        # 3. 格式化并返回结果
        # 结果结构可能有点嵌套，我们来简化它
        if results and results.get('ids') and results['ids'][0]:
             formatted_results = []
             ids = results['ids'][0]
             distances = results['distances'][0]
             documents = results['documents'][0]

             for i in range(len(ids)):
                 formatted_results.append({
                     "id": ids[i],
                     "document": documents[i],
                     "distance": distances[i] # 距离越小意味着对于余弦/欧几里得距离越相似
                 })
             return {"results": formatted_results}
        else:
             return {"results": []} # 如果未找到结果，返回空列表

    except Exception as e:
        print(f"Error during search for query '{query}': {e}") # 查询 '{query}' 时搜索出错：{e}
        raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}") # 搜索失败：{str(e)}

@app.get("/")
async def read_root():
    """ 一个简单的根端点，用于检查 API 是否正在运行。 """
    return {"message": "Semantic Search API is running. Use the /search/ endpoint."} # 语义搜索 API 正在运行。请使用 /search/ 端点。

# --- Main Execution ---
# 此代码块允许直接使用 `python main.py` 运行应用
if __name__ == "__main__":
    print("Starting FastAPI server...") # 正在启动 FastAPI 服务器...
    uvicorn.run(app, host="0.0.0.0", port=8000)

说明：

• 我们初始化 FastAPI。
• 全局资源：我们在应用启动时只进行一次加载 SentenceTransformer 模型并连接到持久化的 ChromaDB 集合。这避免了在每次请求时重新加载模型或重新连接数据库，否则会非常低效。为提高稳定性，还添加了错误处理。
• /search/ 端点：
  • 它接受一个 query 参数 (parameter)（一个字符串）。
  • 它使用与索引时相同的模型为输入的 query 生成嵌入 (embedding)。这对于有意义地比较向量 (vector)很重要。
  • 它使用 collection.query 查找与 query_embedding 最相似的 N_RESULTS 个文档嵌入。我们请求 ChromaDB 在响应中包含原始 documents 和 distances。
  • 它将 ChromaDB 的结果格式化为更整洁的字典列表，并以 JSON 形式返回。
  • 包含错误处理以应对嵌入生成或数据库查询期间可能出现的问题。
• 根端点：一个简单的 / 端点确认 API 正在运行。
• 运行应用：if __name__ == "__main__": 代码块允许您直接使用 python main.py 运行服务器。另外，您也可以使用 uvicorn main:app --reload --host 0.0.0.0 --port 8000。--reload 标志在开发期间很有用，因为它会在您保存更改时自动重启服务器。

4. 运行与测试

1. 索引数据：如果您尚未操作，请运行 python index_data.py。

2. 启动 API 服务器：运行 uvicorn main:app --reload --port 8000。

3. 测试 API：打开您的网页浏览器或使用 curl 等工具向搜索端点发送请求：

   • 浏览器：访问 http://localhost:8000/search/?query=what+is+AI
   • curl：

     curl "http://localhost:8000/search/?query=information%20about%20databases"
     

     curl "http://localhost:8000/search/?query=tell%20me%20about%20animals"
     

您应该会收到 JSON 响应，其中包含来自小型数据集的最相关文档，基于语义相似性，并附带它们的距离。例如，查询“数据库”应返回与向量 (vector)数据库和可能的机器学习 (machine learning)相关的结果。查询“动物”应检索到关于狐狸的句子。

搜索应用流程图

该图说明了语义搜索应用的请求流程。用户向 API 端点发送查询，该端点使用嵌入 (embedding)模型将查询转换为向量 (vector)。然后，此向量用于在 ChromaDB 集合中搜索相似的文档向量。结果被格式化并返回给用户。

进一步思考

本示例提供了一个主要结构。您可以通过多种方式进行扩展：

• 更大规模的数据集：索引一个更大规模的数据集（例如，文章、产品描述）。请记住之前讨论过的索引策略以提高效率。
• 元数据过滤：在索引期间添加元数据（例如，类别、时间戳），并使用 ChromaDB 的过滤功能（query 方法中的 where 子句）来精炼搜索结果。
• 不同数据库/模型：通过根据 Pinecone、Weaviate 或 Milvus 各自的 Python 客户端修改客户端初始化和查询逻辑，将 ChromaDB 替换为它们。尝试不同的嵌入 (embedding)模型。
• 用户界面：使用 Jinja2 模板与 FastAPI 或独立的网页框架（如 React、Vue）构建一个简单的 HTML 前端，使其与此 API 交互。
• 混合搜索：将关键词搜索（例如，使用 Whoosh 或 Elasticsearch）与向量 (vector)搜索结合，以可能提高相关性。
• 评估：使用带有标签的数据集实现评估指标（如 Recall@K），以衡量搜索结果的质量。

这个动手实践说明了本课程中讨论的组件（嵌入模型、向量数据库和搜索逻辑）如何结合起来，创建能够理解用户查询背后含义的应用。