Embedding、Rerank与向量数据库概念详解与应用教程

# 创建并激活虚拟环境
conda create -n rag-kb python=3.10
conda activate rag-kb

# 安装核心依赖
pip install -r requirements.txt

# 打开 config.py，将这一行中的 "your-api-key-here" 替换为你的智谱 API Key
ZHIPU_API_KEY = "your-api-key-here"      # 请替换为你的智谱 API Key

rag-knowledge-base/
├── requirements.txt        # 依赖清单
├── config.py               # 配置文件（模型路径、API 密钥等）
├── rag_core.py             # 核心模块（Embedding、Rerank、ChromaDB 封装）
├── cli.py                  # 命令行知识库创建工具
├── main.py                 # FastAPI 应用（知识检索 API + 前端页面）
└── static/
    └── index.html          # 前端界面

# ChromaDB 持久化存储路径
CHROMA_PERSIST_DIR = "./chroma_data"

# Embedding 模型路径
EMBEDDING_MODEL_PATH = r"C:\models\bge-small-zh-v1.5"

# Rerank 模型路径
RERANK_MODEL_PATH = r"C:\models\bge-reranker-base"

# 智谱 API 配置（请替换为你的 API Key）
ZHIPU_API_KEY = "your-api-key-here"
ZHIPU_API_BASE = "https://open.bigmodel.cn/api/paas/v4"
LLM_MODEL_NAME = "glm-4.7-flash"

# 向量检索参数
VECTOR_SEARCH_TOP_K = 30
RERANK_TOP_N = 5

from sentence_transformers import SentenceTransformer

class EmbeddingService:
    """Embedding 编码服务"""

    def __init__(self, model_name_or_path: str):
        self.model = SentenceTransformer(model_name_or_path)

    @property
    def dimension(self) -> int:
        return self.model.get_sentence_embedding_dimension()

    def encode(self, texts: list[str]) -> list[list[float]]:
        embeddings = self.model.encode(
            texts, normalize_embeddings=True, show_progress_bar=False
        )
        return embeddings.tolist()

    def encode_query(self, query: str) -> list[float]:
        return self.encode([query])[0]

import chromadb

class VectorStore:
    """ChromaDB 向量数据库封装"""

    def __init__(self, persist_dir: str, embedding_service: EmbeddingService):
        self.client = chromadb.PersistentClient(path=persist_dir)
        self.embedding_service = embedding_service

    def create_collection(self, name: str) -> chromadb.Collection:
        try:
            self.client.delete_collection(name)
        except ValueError:
            pass
        return self.client.create_collection(name)

    def add_documents(self, collection_name: str, documents: list[str],
                      metadatas: list[dict] | None = None) -> list[str]:
        collection = self.get_collection(collection_name)
        ids = [f"doc_{i}_{hash(doc) % (10**8)}" for i, doc in enumerate(documents)]
        embeddings = self.embedding_service.encode(documents)
        collection.add(
            embeddings=embeddings, documents=documents,
            metadatas=metadatas or [{"source": "unknown"} for _ in documents],
            ids=ids
        )
        return ids

    def search(self, collection_name: str, query_vector: list[float],
               top_k: int = 10) -> list[dict]:
        collection = self.get_collection(collection_name)
        results = collection.query(
            query_embeddings=[query_vector],
            n_results=min(top_k, collection.count()),
            include=["documents", "metadatas", "distances"]
        )
        hits = []
        if results["documents"] and results["documents"][0]:
            for i in range(len(results["documents"][0])):
                hits.append({
                    "id": results["ids"][0][i],
                    "text": results["documents"][0][i],
                    "metadata": results["metadatas"][0][i] if results["metadatas"] else {},
                    "score": 1 - results["distances"][0][i] if results["distances"] else 0
                })
        return hits

from sentence_transformers import CrossEncoder

class RerankService:
    """Rerank 重排序服务"""

    def __init__(self, model_name_or_path: str):
        self.model = CrossEncoder(model_name_or_path)

    def rerank(self, query: str, candidates: list[dict], top_n: int = 5) -> list[dict]:
        if not candidates:
            return []
        pairs = [(query, cand["text"]) for cand in candidates]
        scores = self.model.predict(pairs)
        for i, score in enumerate(scores):
            candidates[i]["rerank_score"] = float(score)
        reranked = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
        return reranked[:top_n]

from openai import OpenAI

class LLMService:
    """LLM 回答生成服务"""

    def __init__(self, api_key: str, api_base: str, model_name: str):
        self.client = OpenAI(api_key=api_key, base_url=api_base)
        self.model_name = model_name

    def generate_answer(self, query: str, context_docs: list[dict]) -> str:
        context_parts = []
        for i, doc in enumerate(context_docs):
            source = doc.get("metadata", {}).get("source", "未知来源")
            context_parts.append(f"[文档 {i+1}](来源: {source})\n{doc['text']}")
        context_text = "\n\n".join(context_parts)

        system_prompt = (
            "你是一个基于知识库的问答助手..."
        )
        user_prompt = f"## 上下文信息\n\n{context_text}\n\n## 用户问题\n\n{query}"

        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            temperature=0.3,
            max_tokens=1024
        )
        return response.choices[0].message.content

import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from config import *
from rag_core import EmbeddingService, VectorStore


def chunk_text(text: str, chunk_size: int = 256, overlap: int = 32) -> list[tuple[str, int]]:
    """将长文本分割为有重叠的文本块"""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append((text[start:end], start))
        start += chunk_size - overlap
    return chunks


def cmd_create(args):
    """处理 create 子命令：创建知识库"""
    embedding_service = EmbeddingService(EMBEDDING_MODEL_PATH)
    vector_store = VectorStore(CHROMA_PERSIST_DIR, embedding_service)
    vector_store.create_collection(args.name)
    # 读取文件 -> 分块 -> 编码 -> 存入向量数据库
    for file_path in args.files:
        with open(file_path, "r", encoding="utf-8") as f:
            text = f.read()
        chunks = chunk_text(text, CHUNK_SIZE, CHUNK_OVERLAP)
        # 编码并存入...
    # ...


def main():
    parser = argparse.ArgumentParser(description="知识库管理工具")
    subparsers = parser.add_subparsers(dest="command", required=True, help="可用命令")

    create_parser = subparsers.add_parser("create", help="创建知识库并导入文档")
    create_parser.add_argument("--name", "-n", required=True, help="知识库名称")
    create_parser.add_argument("--files", "-f", nargs="+", required=True,
                               help="要导入的文档文件路径")

    delete_parser = subparsers.add_parser("delete", help="删除知识库")
    delete_parser.add_argument("--name", "-n", required=True, help="要删除的知识库名称")

    args = parser.parse_args()
    if args.command == "create":
        cmd_create(args)
    elif args.command == "delete":
        cmd_delete(args)


if __name__ == "__main__":
    main()

# 创建名为 "company_handbook" 的知识库，导入两个文档
python cli.py create --name company_handbook --files handbook.txt policies.txt

# 删除知识库
python cli.py delete --name company_handbook

# 查看帮助
python cli.py --help

from fastapi import FastAPI, Query, HTTPException
from fastapi.responses import HTMLResponse
from config import *
from rag_core import (
    EmbeddingService, VectorStore, RerankService, LLMService,
)
import uvicorn

# FastAPI 应用初始化
app = FastAPI(
    title="RAG Knowledge Base API",
    description="基于 Embedding + Rerank + ChromaDB + LLM 的知识检索服务",
    version="1.0.0"
)

# 全局服务实例
embedding_service = EmbeddingService(EMBEDDING_MODEL_PATH)
vector_store = VectorStore(CHROMA_PERSIST_DIR, embedding_service)
rerank_service = RerankService(RERANK_MODEL_PATH)
llm_service = LLMService(ZHIPU_API_KEY, ZHIPU_API_BASE, LLM_MODEL_NAME)


@app.get("/api/collections")
def list_collections():
    """获取所有知识库名称列表"""
    return {"collections": vector_store.list_collections()}


@app.get("/api/search")
def search_knowledge_base(
    query: str = Query(..., description="用户查询"),
    collection: str = Query(..., description="知识库名称"),
    top_k: int = Query(VECTOR_SEARCH_TOP_K, description="初筛候选数"),
    top_n: int = Query(RERANK_TOP_N, description="精选数"),
):
    """在指定知识库中执行 RAG 检索流水线"""
    # 阶段1: 查询编码
    query_vector = embedding_service.encode_query(query)
    # 阶段2: 向量检索（初筛）
    candidates = vector_store.search(collection, query_vector, top_k)
    if not candidates:
        return {"query": query, "answer": "未在知识库中找到相关信息。"}
    # 阶段3: Rerank 重排序
    reranked = rerank_service.rerank(query, candidates, top_n)
    # 阶段4: LLM 回答生成
    answer = llm_service.generate_answer(query, reranked)
    return {"query": query, "candidates": reranked, "answer": answer}


@app.get("/", response_class=HTMLResponse)
def index():
    """返回前端交互页面"""
    html_path = os.path.join(os.path.dirname(__file__), "static", "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

# 设置 API 密钥
# 在 config.py 中设置 ZHIPU_API_KEY 为你的 API Key

# 确保在 rag-knowledge-base/ 目录下
cd rag-knowledge-base/

# 启动 FastAPI 服务
python main.py

# 创建知识库，导入 txt 文档
python cli.py create --name company_handbook --files ./docs/handbook.txt ./docs/policies.txt

# 创建多个知识库
python cli.py create --name product_docs --files ./docs/product_v1.txt ./docs/product_v2.txt

# 删除知识库
python cli.py delete --name company_handbook