import os
import chromadb
from llama_index.core import (
    SimpleDirectoryReader,
    VectorStoreIndex,
    Settings,
    StorageContext,
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# ============ 1. 配置 ============
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")
if not DASHSCOPE_API_KEY:
    raise ValueError("请设置环境变量 DASHSCOPE_API_KEY")

BASE_URL = "https://llm-pz0prcj5mohy7yjg.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"

# 设置向量模型
Settings.embed_model = OpenAIEmbedding(
    mode="similarity",
    model="text-embedding-ada-002",
    model_name="qwen3.7-text-embedding",
    api_key=DASHSCOPE_API_KEY,
    api_base=BASE_URL,
    max_retries=3,
    timeout=60,
)

# 设置LLM模型
Settings.llm = OpenAI(
    model="qwen-plus",
    api_key=DASHSCOPE_API_KEY,
    api_base=BASE_URL,
    max_retries=3,
    timeout=60,
)

# 设置文本分割器
"""
配置文本切块器：把长文档拆成小片段再做向量化。

chunk_size=500：每块大约 500 个 token（不是字符数），块太大检索不准，太小上下文不够。
chunk_overlap=50：相邻块重叠约 50 个 token，避免关键句子被从中间切开后语义丢失。
例如一段 1000 token 的文档，大致会切成：[0–500]、[450–950]、[900–1000…] 这种带重叠的块，再分别入库做 embedding。
"""
Settings.text_splitter = SentenceSplitter(
    chunk_size=500,
    chunk_overlap=150,
)


# ============ 2. 加载文档 ============
def load_documents(data_dir="./data"):
    reader = SimpleDirectoryReader(
        input_dir=data_dir,
        recursive=True,
        file_metadata=lambda filename: {"source": os.path.basename(filename)} #溯源:
    )
    documents = reader.load_data()
    print(f"✅ 加载了 {len(documents)} 个文档片段")
    return documents


# ============ 3. 构建索引 ============
def build_and_save_index(documents, persist_dir="./storage"):
    # 连接 ChromaDB
    chroma_client = chromadb.PersistentClient(path=persist_dir)
    chroma_collection = chroma_client.get_or_create_collection(
        name="bank_policy_collection"
    )
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

    # 创建 StorageContext（不指定 persist_dir，避免启动时加载旧数据）
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    # 构建索引
    index = VectorStoreIndex.from_documents(
        documents,
        storage_context=storage_context,
        show_progress=True,
    )

    # 手动持久化（索引 + docstore + 其他元数据）
    index.storage_context.persist(persist_dir=persist_dir)

    # 验证向量数量
    vector_count = chroma_collection.count()
    print(f"✅ 索引已保存到 {persist_dir}，向量数: {vector_count}")
    return index


if __name__ == "__main__":
    docs = load_documents("./data")
    print(docs)
    #index = build_and_save_index(docs, "./storage")
   # print("\n🎉 索引构建完成！运行 query.py 进行问答。")

代码逐段完整讲解

这份代码是基于 LlamaIndex + 阿里云 DashScope 通义千问大模型 + Chroma 向量数据库 搭建的本地知识库 RAG 检索增强生成系统初始化脚本，整体作用：读取本地文件夹文档、文本切分、调用阿里云 Embedding 向量模型把文本转为向量存入 Chroma 向量库、持久化保存知识库索引，后续可以基于这个知识库问答。

整体流程：环境密钥校验 → 配置大语言模型 + 向量嵌入模型 → 配置文本拆分规则 → 加载本地文档 → 构建向量知识库并持久化存储。

一、前置依赖（先清楚用到的库）

运行这份代码需要提前安装依赖包：

bash

pip install llama-index llama-index-vector-stores-chroma chromadb python-dotenv

核心组件：

1. llama-index：RAG框架，封装文档加载、切分、向量化、索引、检索全套流程
2. chromadb：轻量化本地21向量数据库，用来存储向量文本
3. DashScope：阿里云通义大模型服务，兼容OpenAI接口格式

分段详细解析

1. 读取阿里云 API 密钥，环境变量校验

python

运行

DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")
if not DASHSCOPE_API_KEY:
    raise ValueError("请设置环境变量 DASHSCOPE_API_KEY")

BASE_URL = "https://llm-pz0prcj5mohy7yjg.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"

配置密钥方式： Windows：set DASHSCOPE_API_KEY=你的密钥 Mac/Linux：export DASHSCOPE_API_KEY=你的密钥

2. 配置向量嵌入模型（Embedding）

python

运行

Settings.embed_model = OpenAIEmbedding(
    mode="similarity",
    model="text-embedding-ada-002",
    model_name="qwen3.7-text-embedding",
    api_key=DASHSCOPE_API_KEY,
    api_base=BASE_URL,
    max_retries=3,
    timeout=60,
)

Embedding 向量模型作用：把文字转换成一串数字向量，语义相近的文本向量距离更近，后续用户提问时，用向量相似度召回知识库相关片段。

参数拆解：

1. os.getenv()：读取系统环境变量里存放的阿里云 DashScope 密钥，不会把密钥硬编码写在代码里，更安全；
2. 判断密钥为空就直接抛出报错，防止后续调用模型鉴权失败；
3. BASE_URL：阿里云专属兼容 OpenAI 的接口地址，因为 DashScope 支持 OpenAI 调用格式，所以 LlamaIndex 可以直接用OpenAI类对接通义系列模型。
4. mode="similarity"：相似度检索模式，RAG 标准用法
5. model：OpenAI 原生模型标识（阿里云兼容层占位参数，可不关注）
6. model_name="qwen3.7-text-embedding"：实际调用阿里云通义 3.7 向量嵌入模型
7. api_base：阿里云兼容接口地址
8. max_retries=3：接口调用失败自动重试 3 次
9. timeout=60：请求超时 60 秒，避免网络卡顿卡死

3. 配置大语言模型 LLM

python

运行

Settings.llm = OpenAI(
    model="qwen-plus",
    api_key=DASHSCOPE_API_KEY,
    api_base=BASE_URL,
    max_retries=3,
    timeout=60,
)

LLM就是最终负责理解问题、结合检索到的知识库内容、生成回答的大模型

• model="qwen-plus"：指定调用阿里云通义千问增强版大模型
• 鉴权、超时、重试配置和向量模型保持一致

4. 文本分割器配置（核心 RAG 细节）

python

运行

Settings.text_splitter = SentenceSplitter(
    chunk_size=500,
    chunk_overlap=150,
)

长文档不能直接整体向量化，必须切分成小块文本（Chunk），注释已经做了解释，补充通俗理解：

1. chunk_size=500：单个文本块最大 token 数（token≈汉字 1.5 倍），每段文字控制在 500token 以内；
   • 块太大：一段内容混杂多个主题，检索精准度变差
   • 块太小：单段信息过少，LLM 没有足够上下文作答
2. chunk_overlap=150：文本块之间重叠 150token：
   • 举例：
   • 片段 1：0 ~ 500 token
   • 片段 2：350 ~ 850 token
   • 片段 3：700 ~ 1200 token
   • 重叠目的：避免一句话、一段语义刚好被切在两块中间，导致检索时丢失完整语义。

5. 文档加载函数 load_documents ()

python

运行

def load_documents(data_dir="./data"):
    reader = SimpleDirectoryReader(
        input_dir=data_dir,
        recursive=True,
        file_metadata=lambda filename: {"source": os.path.basename(filename)}
    )
    documents = reader.load_data()
    print(f"✅ 加载了 {len(documents)} 个文档片段")
    return documents

功能：读取本地./data文件夹里的所有文档文件 参数说明：

• input_dir="./data"：文档存放目录，你把 PDF、TXT、Word 等文件放进这个文件夹即可
• recursive=True：递归读取子文件夹里的文件
• file_metadata：给每段文本打上元数据标签source:文件名 后续问答时，可以溯源回答内容来自哪个文档，方便核验信息来源
• reader.load_data()：读取全部文件内容，LlamaIndex 自动解析主流文档格式

6. 构建 + 持久化向量索引 build_and_save_index ()

这是核心函数，完成：连接向量库 → 绑定存储容器 → 文档切分 + 向量化 → 存入向量库 → 本地持久化保存

python

运行

def build_and_save_index(documents, persist_dir="./storage"):
    # 1. 初始化本地Chroma向量数据库
    chroma_client = chromadb.PersistentClient(path=persist_dir)
    # 创建/获取向量集合（一张数据表）
    chroma_collection = chroma_client.get_or_create_collection(
        name="bank_policy_collection"
    )
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

    # 2. 存储上下文：管理向量、文档、元数据整套存储资源
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    # 3. 基于文档构建向量索引
    index = VectorStoreIndex.from_documents(
        documents,
        storage_context=storage_context,
        show_progress=True, # 显示向量化进度条
    )

    # 4. 把索引、向量、文档元数据全部持久化到本地文件夹
    index.storage_context.persist(persist_dir=persist_dir)

    # 统计存入向量总数
    vector_count = chroma_collection.count()
    print(f"✅ 索引已保存到 {persist_dir}，向量数: {vector_count}")
    return index

拆解流程：

1. chromadb.PersistentClient：创建持久化本地向量库，关闭程序后向量数据不会丢失，存在./storage文件夹；
2. bank_policy_collection：向量集合名称，这里偏向银行政策知识库场景；
3. VectorStoreIndex.from_documents：LlamaIndex 一站式操作： 读取文档 → 按照上面配置的SentenceSplitter切块 → 调用阿里云 Embedding 生成向量 → 文本 + 向量 + 元数据存入 Chroma；
4. persist()：强制把所有数据落地本地磁盘；
5. chroma_collection.count()：打印一共生成了多少个文本向量块。

7. 程序入口主函数

python

运行

if __name__ == "__main__":
    docs = load_documents("./data")
    print(docs)
    #index = build_and_save_index(docs, "./storage")
    #print("\n🎉 索引构建完成！运行 query.py 进行问答。")

• 当前代码把构建索引的两行注释掉了，运行只会加载 data 目录文档并打印文档内容，用来调试查看文档是否正常读取；
• 解开注释后，就会执行完整知识库构建；
• 构建完成后，需要新建query.py脚本加载本地索引，实现提问、检索知识库、AI 回答。

---

整套代码完整运行流程

1. 在项目根目录新建 data 文件夹，放入你的知识库文档（txt、pdf 等）；
2. 配置系统环境变量 DASHSCOPE_API_KEY；
3. 运行脚本：
   • 初次调试：只会加载文档，预览内容；
   • 解开注释：生成向量知识库，存入./storage文件夹；
4. 新建问答脚本，加载 storage 里的索引，就可以基于本地文档提问。