一、嵌入模型核心原理

嵌入（Embedding）是 RAG 系统的 "眼睛"，它决定了系统能否准确理解文本的语义。这是 RAG 最核心的基础技术之一，必须彻底搞懂。

1.1 为什么我们需要文本嵌入？

在学习嵌入之前，我们先思考一个问题：计算机如何理解文本的语义？

计算机只能处理数字，无法直接理解自然语言。传统的文本表示方法存在严重缺陷：

1.1.1 传统文本表示的问题：One-Hot 编码

One-Hot 编码是最简单的文本表示方法，它为每个词分配一个唯一的索引，然后用一个只有对应位置为 1、其余为 0 的向量表示该词。

示例：假设我们的词汇表是 ["我", "爱", "中国", "北京"]

• "我" → [1, 0, 0, 0]
• "爱" → [0, 1, 0, 0]
• "中国" → [0, 0, 1, 0]
• "北京" → [0, 0, 0, 1]

致命缺陷：

1. 维度灾难：如果词汇表有 10 万个词，每个词的向量就是 10 万维，计算量巨大
2. 语义缺失：任意两个词的向量都是正交的（余弦相似度为 0），无法表示语义相似性
   • "中国" 和 "北京" 的相似度 = 0
   • "猫" 和 "狗" 的相似度 = 0
   • 完全无法体现 "中国的首都是北京"、"猫和狗都是动物" 这样的语义关系

1.1.2 词嵌入（Word Embedding）的突破

2013 年 Word2Vec 的出现彻底改变了文本表示的方式。词嵌入的核心思想是：语义相近的词，在向量空间中距离更近。

示例：

• "中国" → [0.1, 0.2, 0.3, ..., 0.8]（100 维）
• "北京" → [0.12, 0.21, 0.29, ..., 0.78]（100 维）
• "美国" → [0.5, 0.6, 0.7, ..., 0.2]（100 维）
• "华盛顿" → [0.51, 0.59, 0.72, ..., 0.19]（100 维）

神奇的特性：

• 语义相似的词向量距离近：距离("中国", "北京") < 距离("中国", "美国")
• 可以进行向量运算：向量("国王") - 向量("男人") + 向量("女人") ≈ 向量("女王")

1.1.3 文本嵌入（Text Embedding）：从词到句子 / 文档

词嵌入只能表示单个词，而我们需要表示整个句子或文档。文本嵌入（也叫句子嵌入）将任意长度的文本映射为一个固定维度的稠密向量，保留整个文本的语义信息。

RAG 为什么必须用专用文本嵌入模型？

很多初学者会问：能不能用大模型（如 GPT-3.5）的隐状态作为嵌入？答案是：不推荐。

• 大模型的隐状态是为生成任务优化的，不是为语义相似度匹配优化的
• 专用嵌入模型（如 BGE）专门在海量文本对数据上训练，语义相似度匹配效果远好于大模型隐状态
• 专用嵌入模型体积更小、速度更快、成本更低

1.2 文本嵌入的核心特性

一个好的文本嵌入应该具备以下特性：

1. 语义一致性：语义相近的文本向量距离近，语义相反的文本向量距离远
2. 维度固定：无论输入文本多长，输出的向量维度都是固定的（如 768 维）
3. 归一化：向量的 L2 范数为 1，这样余弦相似度计算等价于内积，计算更高效
4. 鲁棒性：对同义词、转述句、拼写错误有一定的容忍度

1.3 相似度计算方法详解

嵌入的最终目的是计算文本之间的语义相似度，我们需要掌握三种最常用的相似度计算方法。

1.3.1 余弦相似度（Cosine Similarity）：RAG 最常用

余弦相似度衡量两个向量在方向上的相似程度，不受向量长度影响，是 RAG 检索的标准方法。

数学公式：

取值范围：[-1, 1]

• 1：两个向量方向完全相同，语义完全一致
• 0：两个向量正交，语义无关
• -1：两个向量方向完全相反，语义相反

代码实现：

import numpy as np

def cosine_similarity(a, b):
    """计算两个向量的余弦相似度"""
    dot_product = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    return dot_product / (norm_a * norm_b)

# 测试
vector1 = np.array([1, 2, 3])
vector2 = np.array([2, 4, 6])  # 与vector1方向相同
vector3 = np.array([-1, -2, -3])  # 与vector1方向相反
vector4 = np.array([1, 0, 0])  # 与vector1正交

print(f"vector1与vector2的相似度：{cosine_similarity(vector1, vector2):.2f}")  # 1.00
print(f"vector1与vector3的相似度：{cosine_similarity(vector1, vector3):.2f}")  # -1.00
print(f"vector1与vector4的相似度：{cosine_similarity(vector1, vector4):.2f}")  # 0.27

1.3.2 内积（Dot Product）

内积是向量对应位置元素相乘后求和，未归一化时受向量长度影响。

数学公式：

适用场景：当所有向量都已经归一化（L2 范数为 1）时，内积等价于余弦相似度，计算更高效。

1.3.3 欧氏距离（Euclidean Distance）

欧氏距离衡量两个向量在空间中的直线距离，距离越小越相似。

数学公式：

缺点：对向量长度非常敏感，长文本的向量模长通常更大，会导致相似度计算偏差，因此 RAG 中很少使用。

三种方法对比总结：

方法	优点	缺点	RAG 推荐度
余弦相似度	不受向量长度影响，语义匹配准确	计算稍复杂	⭐⭐⭐⭐⭐
内积	计算速度快	受向量长度影响	⭐⭐⭐（仅归一化后使用）
欧氏距离	直观易懂	受向量长度影响大	⭐

1.4 主流中文嵌入模型选型（2026 年最新）

选择合适的嵌入模型是 RAG 效果的关键。中文场景下，国产模型的效果远好于国外模型。

1.4.1 模型对比表

模型名称	维度	中文效果	速度	开源	推荐场景
BGE-large-zh-v1.5	1024	SOTA	中	✅	所有中文通用场景（首选）
BGE-base-zh-v1.5	768	优秀	快	✅	低资源场景
M3E-large	1024	优秀	中	✅	多任务场景（短句 / 长文本 / 跨语言）
gte-large-zh	1024	优秀	中	✅	专业领域（金融 / 法律）
all-MiniLM-L6-v2	384	一般	极快	✅	原型验证、超大规模数据
text-embedding-ada-002	1536	中文差	中	❌	英文场景
通义千问嵌入 v3	1536	优秀	快	❌	阿里云生态

1.4.2 选型决策树

是否需要私有化部署？
├─ 是 → 中文场景？
│  ├─ 是 → 效果优先？
│  │  ├─ 是 → BGE-large-zh-v1.5
│  │  └─ 速度优先 → BGE-base-zh-v1.5
│  └─ 否 → all-MiniLM-L6-v2
└─ 否 → 预算充足？
   ├─ 是 → 通义千问嵌入v3
   └─ 否 → 还是用开源BGE吧，免费又好用

为什么强烈推荐 BGE-large-zh-v1.5？

• 由清华大学 KEG 实验室和智谱 AI 联合开发，是目前中文语义匹配的 SOTA
• 在 MTEB 中文排行榜上连续霸榜
• 支持长文本（最大 512token）
• 开源免费，可商业使用
• 生态完善，LangChain、LlamaIndex 等框架都原生支持
• 可以手动下载完整的模型文件：
• 访问 ModelScope 的 BGE-large-zh-v1.5 页面：https://www.modelscope.cn/models/AI-ModelScope/bge-large-zh-v1.5/files
• 点击 "下载全部" 按钮，下载所有文件到本地文件夹（例如：D:\models\bge-large-zh-v1.5）

1.5 实战：文本向量化工具实现

现在我们来实现一个通用的文本向量化工具，支持批量处理第一周生成的文档分块。

1.5.1 加载 BGE 嵌入模型

from sentence_transformers import SentenceTransformer
import numpy as np
import json
from tqdm import tqdm

# 加载BGE-large-zh-v1.5模型（首次运行会自动下载，约1.3GB）
# 国内用户如果下载慢，先执行：export HF_ENDPOINT=https://hf-mirror.com
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')

def embed_text(texts, batch_size=32, normalize_embeddings=True):
    """
    批量文本向量化
    :param texts: 文本列表
    :param batch_size: 批次大小
    :param normalize_embeddings: 是否归一化向量
    :return: 向量数组，shape=(n_texts, 1024)
    """
    embeddings = model.encode(
        texts,
        batch_size=batch_size,
        show_progress_bar=True,
        normalize_embeddings=normalize_embeddings
    )
    return embeddings

# 测试单个文本向量化
test_text = "检索增强生成（RAG）是一种大模型应用技术"
embedding = embed_text([test_text])[0]
print(f"文本：{test_text}")
print(f"向量维度：{len(embedding)}")
print(f"向量前5个值：{embedding[:5]}")
print(f"向量L2范数：{np.linalg.norm(embedding):.2f}")  # 归一化后应为1.0

运行结果：

1.5.2 批量向量化上一节的文档分块

def process_chunks(input_path="processed_chunks.jsonl", output_path="embeddings.jsonl"):
    """
    批量处理文档分块，生成向量并保存
    :param input_path: 上一节生成的分块结果文件
    :param output_path: 向量输出文件
    """
    # 读取分块数据
    chunks = []
    with open(input_path, 'r', encoding='utf-8') as f:
        for line in f:
            chunk = json.loads(line)
            chunks.append(chunk)
    
    print(f"共读取到{len(chunks)}个文档分块")
    
    # 提取文本内容
    texts = [chunk["text"] for chunk in chunks]
    
    # 批量向量化
    print("开始向量化...")
    embeddings = embed_text(texts)
    
    # 将向量添加到分块数据中
    for i, chunk in enumerate(chunks):
        # 将numpy数组转换为列表，方便JSON序列化
        chunk["embedding"] = embeddings[i].tolist()
    
    # 保存结果
    with open(output_path, 'w', encoding='utf-8') as f:
        for chunk in chunks:
            f.write(json.dumps(chunk, ensure_ascii=False) + '\n')
    
    print(f"向量化完成，结果已保存到{output_path}")
    return chunks

# 执行批量向量化
if __name__ == "__main__":
    chunks_with_embeddings = process_chunks()

运行后会生成embeddings.jsonl文件，每个分块都包含了对应的向量。

2.5.3 验证语义相似度

我们来验证一下 BGE 模型的语义理解能力：

# 测试语义相似度
test_sentences = [
    "什么是RAG？",
    "检索增强生成的定义是什么？",
    "大模型的幻觉问题如何解决？",
    "如何提高RAG系统的准确率？"
]

# 向量化
test_embeddings = embed_text(test_sentences)

# 计算相似度矩阵
similarity_matrix = np.zeros((len(test_sentences), len(test_sentences)))
for i in range(len(test_sentences)):
    for j in range(len(test_sentences)):
        similarity_matrix[i][j] = cosine_similarity(test_embeddings[i], test_embeddings[j])

# 打印相似度矩阵
print("语义相似度矩阵：")
print(" " * 20, end="")
for sentence in test_sentences:
    print(f"{sentence[:10]:<10}", end="")
print()

for i, sentence in enumerate(test_sentences):
    print(f"{sentence[:20]:<20}", end="")
    for j in range(len(test_sentences)):
        print(f"{similarity_matrix[i][j]:<10.2f}", end="")
    print()

运行结果：

二、向量数据库基础与 Chroma 实战

现在我们有了文本的向量表示，接下来需要一个专门的数据库来存储和检索这些向量 —— 这就是向量数据库。

2.1 为什么我们需要向量数据库？

很多初学者会问：我把向量存在 Python 列表里，遍历计算相似度不行吗？

答案是：当向量数量超过 1 万条时，暴力检索的速度会慢到无法接受。

2.1.1 暴力检索的性能问题

假设我们有 100 万个向量，每个向量 1024 维：

• 一次查询需要计算 100 万次余弦相似度
• 每次余弦相似度计算需要 1024 次乘法和 1023 次加法
• 总共约 20 亿次运算
• 即使 CPU 每秒能执行 10 亿次运算，一次查询也需要 2 秒以上
• 当向量数量达到 1 亿条时，一次查询需要 200 秒，完全无法使用

2.1.2 向量数据库的核心价值

向量数据库通过近似最近邻（ANN）索引技术，在牺牲极小精度的情况下，将检索速度提升几个数量级。

性能对比：

向量数量	暴力检索	ANN 检索（HNSW）
1 万	10ms	1ms
100 万	2000ms	2ms
1 亿	200000ms	5ms

2.1.3 向量数据库的核心能力

1. 向量存储：高效存储高维向量和对应的元数据
2. ANN 检索：快速找到与查询向量最相似的 Top-K 个向量
3. 元数据过滤：先过滤元数据，再进行向量检索（如 "只检索 PDF 文档中的内容"）
4. 增量更新：支持动态插入、删除、更新向量
5. 持久化：将数据保存到磁盘，防止重启丢失
6. 分布式扩展：支持集群部署，处理百亿级向量

2.2 主流向量数据库选型（2026 年最新）

2.2.1 详细对比表

数据库	开源	部署方式	最大规模	混合检索	生态	学习成本	推荐场景
Chroma	✅	本地 / 嵌入式	1000 万	✅	好	极低	学习、原型验证、小型项目
FAISS	✅	本地 / 库	1 亿	❌	一般	低	离线批量检索、小规模项目
Qdrant	✅	本地 / 分布式	10 亿	✅	好	中	生产环境、高并发场景
Milvus	✅	本地 / 分布式	100 亿	✅	好	中	企业级大规模生产环境
Pinecone	❌	云托管	100 亿	✅	好	极低	云原生场景、不想运维
PGVector	✅	PostgreSQL 扩展	1000 万	✅	好	低	已有 PostgreSQL 的场景

2.2.2 学习阶段为什么选 Chroma？

Chroma 是目前最适合 RAG 入门的向量数据库，没有之一：

1. 零部署成本：嵌入式数据库，无需启动服务，直接 pip install 就能用
2. API 极其简单：Python 友好，几行代码就能完成所有操作
3. 原生支持 BGE 等嵌入模型：可以自动调用嵌入模型，无需手动向量化
4. 支持持久化：数据自动保存到磁盘，重启不丢失
5. 支持元数据过滤和混合检索：满足 RAG 的所有基础需求

2.3 向量索引原理详解（HNSW）

向量数据库的核心是索引技术，HNSW（Hierarchical Navigable Small World，分层导航小世界）是目前最主流的索引算法，也是 Chroma、Milvus、Qdrant 的默认索引。

2.3.1 HNSW 的核心思想：多层高速公路

HNSW 的灵感来自于高速公路系统：

• 第 0 层：所有节点都在这一层，相当于普通公路
• 第 1 层：只有一部分节点在这一层，相当于国道
• 第 2 层：更少的节点在这一层，相当于高速公路
• 最高层：只有极少数节点，相当于枢纽

检索过程：

1. 从最高层开始，找到离查询向量最近的节点
2. 下降到下一层，从该节点出发，找到这一层离查询向量最近的节点
3. 重复这个过程，直到到达第 0 层
4. 在第 0 层找到 Top-K 个最近的节点

这种分层结构大大减少了需要比较的节点数量，将检索复杂度从 O (n) 降低到 O (log n)。

2.3.2 HNSW 的关键参数

• M：每个节点的最大邻居数，一般设置为 16-64。M 越大，精度越高，速度越慢。
• ef_construction：构建索引时的邻居探索数量，一般设置为 100-500。越大，索引质量越高，构建时间越长。
• ef_search：检索时的邻居探索数量，一般设置为 100-500。越大，精度越高，速度越慢。

RAG 最佳实践：

• M=32
• ef_construction=200
• ef_search=200
• 这个配置在大多数场景下能平衡精度和速度

2.4 Chroma 向量数据库实战

现在我们来学习 Chroma 的所有核心操作，并将上一节生成的文档分块入库。

2.4.1 Chroma 基本概念

• Client：Chroma 客户端，连接数据库的入口
• Collection：集合，相当于关系型数据库中的表，用于存储一组相关的向量
• Document：文档，包含文本内容、向量和元数据
• ID：每个文档的唯一标识符

2.4.2 初始化 Chroma 客户端

import chromadb
from chromadb.utils import embedding_functions

# 初始化Chroma客户端，数据保存在./chroma_db目录下
client = chromadb.PersistentClient(path="./chroma_db")

# 使用BGE嵌入模型（Chroma原生支持）
bge_embedding = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="BAAI/bge-large-zh-v1.5"
)

2.4.3 创建 / 获取集合

# 创建集合，如果已存在则获取
collection = client.get_or_create_collection(
    name="rag_knowledge_base",  # 集合名称
    embedding_function=bge_embedding,  # 指定嵌入模型
    metadata={"hnsw:M": 32, "hnsw:ef_construction": 200, "hnsw:ef_search": 200},  # HNSW参数
    distance_function="cosine"  # 距离度量，使用余弦相似度
)

print(f"集合创建成功，当前文档数量：{collection.count()}")

2.4.4 批量插入文档

def add_documents_to_chroma(collection, chunks):
    """
    批量将文档分块插入Chroma
    :param collection: Chroma集合对象
    :param chunks: 文档分块列表，每个分块包含id、text、metadata
    """
    # 提取id、文本、元数据
    ids = [chunk["id"] for chunk in chunks]
    documents = [chunk["text"] for chunk in chunks]
    metadatas = [chunk["metadata"] for chunk in chunks]
    
    # 批量插入（Chroma会自动调用嵌入模型生成向量）
    collection.add(
        ids=ids,
        documents=documents,
        metadatas=metadatas
    )
    
    print(f"成功插入{len(chunks)}个文档，当前集合总文档数：{collection.count()}")

# 读取第一周生成的分块数据
chunks = []
with open("processed_chunks.jsonl", 'r', encoding='utf-8') as f:
    for line in f:
        chunk = json.loads(line)
        chunks.append(chunk)

# 插入到Chroma
add_documents_to_chroma(collection, chunks)

注意：Chroma 会自动调用我们指定的嵌入模型生成向量，无需手动向量化，非常方便。

2.4.5 基础语义检索

def semantic_search(collection, query, top_k=5):
    """
    语义检索
    :param collection: Chroma集合对象
    :param query: 查询文本
    :param top_k: 返回最相似的Top-K个结果
    :return: 检索结果
    """
    results = collection.query(
        query_texts=[query],
        n_results=top_k
    )
    
    # 格式化输出结果
    print(f"查询：{query}")
    print(f"找到{len(results['ids'][0])}个相关文档：\n")
    
    for i in range(len(results['ids'][0])):
        print(f"【结果{i+1}】")
        print(f"ID：{results['ids'][0][i]}")
        print(f"相似度得分：{1 - results['distances'][0][i]:.4f}")  # Chroma返回的是距离，1-距离=相似度
        print(f"文档内容：{results['documents'][0][i][:200]}...")  # 只显示前200个字符
        print(f"元数据：{results['metadatas'][0][i]}")
        print("-" * 100)
    
    return results

# 测试语义检索
results = semantic_search(collection, "什么是RAG？它的核心原理是什么？")

2.4.6 元数据过滤检索

元数据过滤是非常实用的功能，它允许我们只检索特定类型的文档。

示例：只检索 PDF 文档中的内容

def filtered_search(collection, query, filter_condition, top_k=5):
    """
    带元数据过滤的检索
    :param filter_condition: 过滤条件，如{"file_type": "pdf"}
    """
    results = collection.query(
        query_texts=[query],
        n_results=top_k,
        where=filter_condition
    )
    
    print(f"查询：{query}")
    print(f"过滤条件：{filter_condition}")
    print(f"找到{len(results['ids'][0])}个相关文档：\n")
    
    for i in range(len(results['ids'][0])):
        print(f"【结果{i+1}】")
        print(f"文件名：{results['metadatas'][0][i]['file_name']}")
        print(f"文档内容：{results['documents'][0][i][:200]}...")
        print("-" * 100)
    
    return results

# 测试：只检索PDF文档
filtered_search(collection, "什么是RAG？", {"file_type": "pdf"})

# 测试：只检索特定文件
filtered_search(collection, "什么是RAG？", {"file_name": "test.pdf"})

2.4.7 更新和删除文档

# 更新文档
collection.update(
    ids=["test.pdf_chunk_0"],
    documents=["更新后的文档内容"],
    metadatas=[{"file_name": "test.pdf", "updated": "2026-05-01"}]
)

# 删除文档
collection.delete(ids=["test.pdf_chunk_0"])

# 删除整个集合（谨慎使用）
# client.delete_collection(name="rag_knowledge_base")

核心知识点回顾

1. 文本嵌入原理：将不定长文本映射为固定维度稠密向量，语义相近的文本向量距离更近
2. 相似度计算：余弦相似度是 RAG 的标准方法，不受向量长度影响
3. 嵌入模型选型：中文场景首选 BGE-large-zh-v1.5，效果最好且开源免费
4. 向量数据库：解决大规模向量快速检索问题，核心是 ANN 索引技术
5. HNSW 索引：目前最主流的 ANN 索引，采用分层结构，平衡精度和速度
6. Chroma 实战：掌握了创建集合、插入文档、语义检索、元数据过滤等核心操作

import os
import json
import logging
from pathlib import Path
from typing import List, Dict, Any
from dataclasses import dataclass

# ------------- 环境配置（必须在导入依赖前设置）-------------
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
# 设置模型缓存目录，避免重复下载
os.environ["TRANSFORMERS_CACHE"] = str(Path(__file__).parent / "model_cache")

import chromadb
from chromadb.utils import embedding_functions

# ------------- 日志配置（替代print，支持分级输出）-------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger(__name__)

# ------------- 全局配置（集中管理，一键修改）-------------
@dataclass
class Config:
    """配置类：所有硬编码参数统一管理"""
    CHROMA_DB_PATH: str = "./chroma_db"          # 向量库存储路径
    COLLECTION_NAME: str = "rag_knowledge_base"  # 集合名称
    EMBEDDING_MODEL: str = "BAAI/bge-large-zh-v1.5"  # 嵌入模型
    VECTOR_DISTANCE: str = "cosine"              # 向量距离计算方式
    CHUNK_FILE_PATH: str = "processed_chunks.jsonl"  # 分块文件路径
    DEFAULT_TOP_K: int = 5                       # 默认检索数量

# ------------- 核心功能封装 -------------
def init_chroma_client(db_path: str) -> chromadb.PersistentClient:
    """
    初始化Chroma持久化客户端
    :param db_path: 向量库存储路径
    :return: Chroma客户端实例
    """
    try:
        # 自动创建目录
        Path(db_path).mkdir(parents=True, exist_ok=True)
        client = chromadb.PersistentClient(path=db_path)
        logger.info(f"Chroma客户端初始化成功，数据存储路径：{db_path}")
        return client
    except Exception as e:
        logger.error(f"Chroma客户端初始化失败：{str(e)}", exc_info=True)
        raise

def init_embedding_function(model_name: str) -> Any:
    """
    初始化嵌入模型函数
    :param model_name: 模型名称
    :return: 嵌入函数实例
    """
    try:
        embed_func = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=model_name
        )
        logger.info(f"嵌入模型 {model_name} 加载成功")
        return embed_func
    except Exception as e:
        logger.error(f"嵌入模型加载失败：{str(e)}", exc_info=True)
        raise

def get_or_create_collection(
    client: chromadb.PersistentClient,
    collection_name: str,
    embed_func: Any,
    distance_type: str
) -> chromadb.Collection:
    """
    创建/获取向量库集合
    :param client: Chroma客户端
    :param collection_name: 集合名称
    :param embed_func: 嵌入函数
    :param distance_type: 距离计算方式
    :return: 集合实例
    """
    try:
        collection = client.get_or_create_collection(
            name=collection_name,
            embedding_function=embed_func,
            metadata={"hnsw:space": distance_type}
        )
        logger.info(f"集合 {collection_name} 初始化成功，当前文档数：{collection.count()}")
        return collection
    except Exception as e:
        logger.error(f"集合初始化失败：{str(e)}", exc_info=True)
        raise

def load_document_chunks(file_path: str) -> List[Dict[str, Any]]:
    """
    从jsonl文件加载文档分块数据
    :param file_path: jsonl文件路径
    :return: 分块数据列表
    """
    chunks = []
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            for line_num, line in enumerate(f, 1):
                try:
                    chunk = json.loads(line.strip())
                    chunks.append(chunk)
                except json.JSONDecodeError:
                    logger.warning(f"第{line_num}行JSON解析失败，跳过该行")
        logger.info(f"成功加载 {len(chunks)} 个文档分块")
        return chunks
    except FileNotFoundError:
        logger.error(f"文件不存在：{file_path}")
        raise
    except Exception as e:
        logger.error(f"文件读取失败：{str(e)}", exc_info=True)
        raise

def batch_add_documents(collection: chromadb.Collection, chunks: List[Dict[str, Any]]):
    """
    批量插入文档到向量库（带异常处理）
    :param collection: Chroma集合
    :param chunks: 文档分块列表
    """
    if not chunks:
        logger.warning("无文档分块可插入")
        return

    try:
        ids = [chunk["id"] for chunk in chunks]
        documents = [chunk["text"] for chunk in chunks]
        metadatas = [chunk["metadata"] for chunk in chunks]

        collection.add(ids=ids, documents=documents, metadatas=metadatas)
        logger.info(f"成功插入 {len(chunks)} 个文档，集合总文档数：{collection.count()}")
    except KeyError as e:
        logger.error(f"分块数据缺少关键字段：{e}")
        raise
    except Exception as e:
        logger.error(f"文档插入失败：{str(e)}", exc_info=True)
        raise

def semantic_search(
    collection: chromadb.Collection,
    query: str,
    top_k: int = Config.DEFAULT_TOP_K
) -> Dict[str, Any]:
    """
    语义检索（格式化输出+优雅展示）
    :param collection: Chroma集合
    :param query: 查询语句
    :param top_k: 返回结果数量
    :return: 原始检索结果
    """
    try:
        logger.info(f"执行语义检索：{query}")
        results = collection.query(query_texts=[query], n_results=top_k)

        # 格式化输出结果
        print("\n" + "="*80)
        print(f"🔍 查询语句：{query}")
        print(f"📊 匹配到 {len(results['ids'][0])} 条相关文档")
        print("="*80 + "\n")

        for idx in range(len(results['ids'][0])):
            similarity = 1 - results['distances'][0][idx]
            doc_id = results['ids'][0][idx]
            content = results['documents'][0][idx]
            metadata = results['metadatas'][0][idx]

            print(f"📌 结果 {idx + 1}")
            print(f"ID：{doc_id}")
            print(f"相似度：{similarity:.4f}")
            print(f"元数据：{metadata}")
            print(f"内容预览：{content[:300]}...")
            print("-"*80 + "\n")

        return results
    except Exception as e:
        logger.error(f"语义检索失败：{str(e)}", exc_info=True)
        raise

def clear_collection(collection: chromadb.Collection):
    """清空集合（测试用，生产环境慎用）"""
    collection.delete()
    logger.warning(f"集合 {collection.name} 已清空")

# ------------- 主程序入口 -------------
def main():
    try:
        # 1. 初始化配置
        config = Config()

        # 2. 初始化Chroma客户端 + 嵌入模型 + 集合
        client = init_chroma_client(config.CHROMA_DB_PATH)
        embed_func = init_embedding_function(config.EMBEDDING_MODEL)
        collection = get_or_create_collection(
            client=client,
            collection_name=config.COLLECTION_NAME,
            embed_func=embed_func,
            distance_type=config.VECTOR_DISTANCE
        )

        # 3. 加载并插入文档
        chunks = load_document_chunks(config.CHUNK_FILE_PATH)
        batch_add_documents(collection, chunks)

        # 4. 测试语义检索
        semantic_search(collection, query="什么是RAG？它的核心原理是什么？")

    except Exception as e:
        logger.critical(f"程序执行失败：{str(e)}", exc_info=True)

if __name__ == '__main__':
    main()