RAG增强检索：知识库驱动的智能体

1

pip install pypdf chromadb openai

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

# document_loader.py —— 多格式文档加载器
from pathlib import Path
from pypdf import PdfReader
import re
import requests

def load_pdf(file_path: str) -> str:
    """加载PDF文件，提取全部文本"""
    reader = PdfReader(file_path)
    text = ""
    for page in reader.pages:
        # 逐页提取文本，用换行分隔
        text += page.extract_text() + "\n"
    return text

def load_markdown(file_path: str) -> str:
    """加载Markdown文件，保留原始文本"""
    return Path(file_path).read_text(encoding="utf-8")

def load_webpage(url: str) -> str:
    """抓取网页，去除HTML标签提取纯文本"""
    resp = requests.get(url, timeout=10)
    resp.raise_for_status()
    html = resp.text
    # 先去掉script/style整段，避免JS代码混入正文
    html = re.sub(r"<(script|style)[^>]*>.*?</\1>", "", html, flags=re.DOTALL)
    # 再去掉其他HTML标签
    text = re.sub(r"<[^>]+>", " ", html)
    # 压缩多余空白
    return re.sub(r"\s+", " ", text).strip()

def load_document(source: str) -> str:
    """根据来源自动选择加载器"""
    if source.startswith("http"):
        return load_webpage(source)
    elif source.endswith(".pdf"):
        return load_pdf(source)
    elif source.endswith((".md", ".markdown", ".txt")):
        return load_markdown(source)
    else:
        raise ValueError(f"不支持的文件格式: {source}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

# text_splitter.py —— 递归字符分块器
from typing import List

class RecursiveTextSplitter:
    """递归字符文本分块器：按分隔符层级递归切分"""

    def __init__(
        self,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        separators: List[str] = None,
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        # 分隔符按优先级从高到低：段落 > 换行 > 句号 > 空格
        self.separators = separators or ["\n\n", "\n", "。", ".", " ", ""]

    def split_text(self, text: str) -> List[str]:
        """将长文本切分为多个分块"""
        # 找出文本中存在的最高优先级分隔符
        separator = self.separators[-1]
        for sep in self.separators:
            if sep == "":
                continue
            if sep in text:
                separator = sep
                break

        # 按选定的分隔符切分
        if separator:
            splits = text.split(separator)
        else:
            splits = list(text)  # 实在没分隔符就按字符切

        # 拼接成不超过chunk_size的分块，带overlap
        chunks = []
        current = ""
        for piece in splits:
            candidate = current + separator + piece if current else piece
            if len(candidate) > self.chunk_size and current:
                # 当前块已满，保存并开启新块
                chunks.append(current.strip())
                # 重叠部分：保留尾部chunk_overlap字符，避免关键信息被切断
                overlap_text = current[-self.chunk_overlap:] if self.chunk_overlap > 0 else ""
                current = overlap_text + separator + piece
            else:
                current = candidate

        if current.strip():
            chunks.append(current.strip())

        return chunks

1
2
3
4
5

splitter = RecursiveTextSplitter(chunk_size=500, chunk_overlap=50)
text = load_document("company_handbook.pdf")
chunks = splitter.split_text(text)
print(f"文档切分为 {len(chunks)} 个分块")
print(f"第一个分块示例: {chunks[0][:100]}...")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

# embedding.py —— Embedding生成封装
from openai import OpenAI

# OpenAI SDK会自动读取环境变量里的密钥，无需手动传入
client = OpenAI()
EMBEDDING_MODEL = "text-embedding-3-small"

def get_embedding(text: str) -> list:
    """生成单条文本的Embedding向量"""
    text = text.replace("\n", " ")  # API建议去掉换行，避免影响向量质量
    resp = client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=text,
    )
    return resp.data[0].embedding

def get_embeddings_batch(texts: list) -> list:
    """批量生成Embedding，效率更高"""
    cleaned = [t.replace("\n", " ") for t in texts]
    resp = client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=cleaned,
    )
    # 按index排序确保返回顺序与输入一致
    sorted_data = sorted(resp.data, key=lambda x: x.index)
    return [d.embedding for d in sorted_data]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

# vector_store.py —— 基于ChromaDB的向量存储与检索
import chromadb

class VectorStore:
    """向量数据库封装：存储与检索"""

    def __init__(self, persist_path: str = "./chroma_db", collection_name: str = "knowledge"):
        # 创建持久化客户端，数据写入本地目录
        self.client = chromadb.PersistentClient(path=persist_path)
        # 获取或创建集合（类似数据库中的表）
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},  # 使用余弦相似度
        )

    def add_documents(self, chunks: list, metadatas: list = None, batch_size: int = 100):
        """将文本分块批量写入向量库"""
        for i in range(0, len(chunks), batch_size):
            batch = chunks[i:i + batch_size]
            # 为每个分块生成唯一ID
            ids = [f"doc_{i + j}" for j in range(len(batch))]
            # 生成向量（这里手动调用，便于演示流程全貌）
            embeddings = get_embeddings_batch(batch)
            meta = metadatas[i:i + batch_size] if metadatas else [{}] * len(batch)
            self.collection.add(
                embeddings=embeddings,
                documents=batch,
                metadatas=meta,
                ids=ids,
            )
        print(f"已写入 {len(chunks)} 个分块到向量库")

    def search(self, query: str, top_k: int = 5) -> list:
        """语义检索：返回与query最相似的top_k个分块"""
        query_embedding = get_embedding(query)
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            include=["documents", "metadatas", "distances"],
        )
        # 整理结果为字典列表，方便上层使用
        docs = results["documents"][0]
        metas = results["metadatas"][0]
        dists = results["distances"][0]
        return [
            {"content": d, "metadata": m, "distance": dist}
            for d, m, dist in zip(docs, metas, dists)
        ]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

# 构建知识库
store = VectorStore(persist_path="./my_kb")

# 加载并分块文档
raw = load_document("产品手册.pdf")
chunks = RecursiveTextSplitter(chunk_size=500, chunk_overlap=50).split_text(raw)

# 为每个分块打上来源元数据，方便后续溯源
metas = [{"source": "产品手册.pdf", "chunk_index": i} for i in range(len(chunks))]

# 写入向量库
store.add_documents(chunks, metadatas=metas)

# 检索测试
results = store.search("如何申请退款？", top_k=3)
for r in results:
    print(f"[相似度: {1 - r['distance']:.4f}] {r['content'][:80]}...")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

# multi_query.py —— 多查询检索
from openai import OpenAI

client = OpenAI()  # 自动读取环境变量OPENAI_API_KEY
LLM_MODEL = "gpt-5.4"

def generate_multi_queries(question: str, n: int = 3) -> list:
    """让LLM把原问题改写成n个不同表述的子问题"""
    prompt = f"""你是一个检索查询改写器。请把下面的问题改写成{n}个语义相同但表述不同的版本，
用于从向量库中检索相关文档。每行一个，不要编号。

原问题：{question}

改写结果："""
    resp = client.chat.completions.create(
        model=LLM_MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7,
    )
    queries = resp.choices[0].message.content.strip().split("\n")
    # 原问题也加入检索队列，保底
    return [question] + [q.strip() for q in queries if q.strip()][:n]

def multi_query_search(store, question: str, top_k: int = 3) -> list:
    """多查询检索：合并各子问题的结果并去重"""
    queries = generate_multi_queries(question, n=3)
    print(f"生成 {len(queries)} 个查询: {queries}")

    all_results = []
    seen_contents = set()  # 用于去重
    for q in queries:
        for r in store.search(q, top_k=top_k):
            # 用内容前50字符作为去重键
            key = r["content"][:50]
            if key not in seen_contents:
                seen_contents.add(key)
                all_results.append(r)

    # 按距离升序（距离越小越相似）
    all_results.sort(key=lambda x: x["distance"])
    return all_results[:top_k * 2]  # 返回合并后的top结果

1

pip install sentence-transformers

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

# reranker.py —— 交叉编码器重排序
from sentence_transformers import CrossEncoder

# 加载预训练交叉编码器，支持中英文
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def rerank(query: str, candidates: list, top_k: int = 3) -> list:
    """对候选分块用交叉编码器重新打分排序"""
    # 构造(query, doc)对
    pairs = [(query, c["content"]) for c in candidates]
    scores = reranker.predict(pairs)
    # 把分数附加到候选上
    for c, s in zip(candidates, scores):
        c["rerank_score"] = float(s)
    # 按重排序分数降序
    candidates.sort(key=lambda x: x["rerank_score"], reverse=True)
    return candidates[:top_k]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

# hybrid_search.py —— 混合检索
import math

class BM25Index:
    """简易BM25关键词索引（演示用，生产环境可用rank_bm25库）"""

    def __init__(self, documents: list, k1: float = 1.5, b: float = 0.75):
        self.docs = documents
        self.k1 = k1
        self.b = b
        # 分词（中文简单按字切，生产建议用jieba）
        self.tokenized = [list(d) for d in documents]
        self.doc_len = [len(t) for t in self.tokenized]
        self.avg_len = sum(self.doc_len) / len(self.doc_len) if self.doc_len else 0
        # 统计词频
        self.df = {}  # 文档频率：某词出现在多少篇文档
        self.tf = []  # 各文档的词频
        for tokens in self.tokenized:
            freq = {}
            for t in tokens:
                freq[t] = freq.get(t, 0) + 1
            self.tf.append(freq)
            for t in freq:
                self.df[t] = self.df.get(t, 0) + 1
        self.N = len(documents)

    def search(self, query: str, top_k: int = 5) -> list:
        """返回[(index, score)]列表"""
        q_tokens = list(query)
        scores = []
        for i in range(self.N):
            s = 0.0
            for t in q_tokens:
                if t not in self.df:
                    continue
                # BM25公式：IDF * TF归一化
                idf = math.log((self.N - self.df[t] + 0.5) / (self.df[t] + 0.5) + 1)
                tf = self.tf[i].get(t, 0)
                norm = tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * self.doc_len[i] / self.avg_len))
                s += idf * norm
            scores.append((i, s))
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]


def hybrid_search(store, bm25_index, chunks, query: str, top_k: int = 5, alpha: float = 0.5) -> list:
    """混合检索：语义(alpha) + 关键词(1-alpha)"""
    # 语义检索
    semantic_results = store.search(query, top_k=top_k * 2)
    # 关键词检索
    bm25_results = bm25_index.search(query, top_k=top_k * 2)

    # 归一化分数并融合
    score_map = {}  # key: chunk内容前50字符
    # 语义分数归一化（距离转相似度）
    max_sim = max([1 - r["distance"] for r in semantic_results]) or 1
    for r in semantic_results:
        key = r["content"][:50]
        score_map[key] = score_map.get(key, 0) + alpha * (1 - r["distance"]) / max_sim

    # BM25分数归一化
    max_bm = max([s for _, s in bm25_results]) or 1
    for idx, s in bm25_results:
        key = chunks[idx][:50]
        score_map[key] = score_map.get(key, 0) + (1 - alpha) * s / max_bm

    # 取分数最高的top_k个，回填原文
    ranked = sorted(score_map.items(), key=lambda x: x[1], reverse=True)[:top_k]
    content_map = {c[:50]: c for c in chunks}
    return [{"content": content_map.get(k, k), "score": s} for k, s in ranked]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

# rag_agent.py —— 完整的RAG增强Agent
import json
from openai import OpenAI

client = OpenAI()  # 自动读取环境变量OPENAI_API_KEY
LLM_MODEL = "gpt-5.4"

# 工具定义：让Agent可以调用知识库检索
tools = [
    {
        "type": "function",
        "function": {
            "name": "search_knowledge_base",
            "description": "从知识库中检索与用户问题相关的文档片段。当用户询问公司政策、产品文档、内部流程等专属知识时调用此工具。",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "检索查询词，应聚焦于用户问题的核心关键词",
                    },
                    "top_k": {
                        "type": "integer",
                        "description": "返回的文档片段数量，默认5",
                        "default": 5,
                    },
                },
                "required": ["query"],
            },
        },
    }
]

def execute_tool(tool_name: str, args: dict, store, reranker) -> str:
    """执行工具调用，返回结果字符串"""
    if tool_name == "search_knowledge_base":
        query = args["query"]
        top_k = args.get("top_k", 5)
        # 先粗检索，召回较多候选
        candidates = store.search(query, top_k=top_k * 3)
        # 再重排序精炼
        top_results = rerank(query, candidates, top_k=top_k)
        # 拼成文本供Agent阅读
        lines = []
        for i, r in enumerate(top_results, 1):
            source = r.get("metadata", {}).get("source", "未知")
            lines.append(f"[{i}] (来源: {source})\n{r['content']}")
        return "\n\n".join(lines) if lines else "未检索到相关文档。"
    return "未知工具"

def run_rag_agent(question: str, store, reranker, history: list = None) -> str:
    """运行RAG增强Agent：自主决定是否检索"""
    system_prompt = """你是一个知识库问答助手。你可以调用search_knowledge_base工具检索知识库。
规则：
1. 当问题涉及公司专属知识（政策、产品、流程、文档内容）时，必须先检索再回答。
2. 回答必须基于检索到的文档内容，不要编造。
3. 如果检索结果与问题无关，诚实告知用户知识库中没有相关信息。
4. 回答时引用来源，格式如：(来源: xxx.pdf)"""

    messages = [{"role": "system", "content": system_prompt}]
    if history:
        messages.extend(history)
    messages.append({"role": "user", "content": question})

    # Agent循环：最多3轮工具调用
    for _ in range(3):
        resp = client.chat.completions.create(
            model=LLM_MODEL,
            messages=messages,
            tools=tools,
            temperature=0.3,
        )
        msg = resp.choices[0].message

        # 如果没有工具调用，说明Agent已准备好最终回答
        if not msg.tool_calls:
            return msg.content

        # 处理工具调用
        messages.append(msg)  # 把assistant消息（含tool_calls）加入历史
        for tc in msg.tool_calls:
            args = json.loads(tc.function.arguments)
            print(f"🔧 Agent调用工具: {tc.function.name}({args})")
            result = execute_tool(tc.function.name, args, store, reranker)
            # 把工具结果喂回给Agent
            messages.append({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": result,
            })

    return "Agent处理轮次超限，请简化问题后重试。"

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

if __name__ == "__main__":
    # 1. 初始化向量库并灌入文档
    store = VectorStore(persist_path="./company_kb")
    if len(store.collection.get()["ids"]) == 0:
        # 首次运行：加载文档
        raw = load_document("员工手册.pdf")
        chunks = RecursiveTextSplitter(500, 50).split_text(raw)
        metas = [{"source": "员工手册.pdf", "chunk_index": i} for i in range(len(chunks))]
        store.add_documents(chunks, metadatas=metas)

    # 2. 初始化重排序器
    reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

    # 3. 提问
    answer = run_rag_agent(
        question="公司年假最多能休几天？需要提前多久申请？",
        store=store,
        reranker=reranker,
    )
    print("🤖 回答:", answer)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

# evaluation.py —— 检索质量评估
def evaluate_retrieval(store, test_cases: list, top_k: int = 5):
    """
    test_cases: [{"question": "...", "relevant_ids": ["产品手册.pdf#3", "产品手册.pdf#7"]}, ...]
    """
    total_recall = 0
    total_precision = 0
    total_mrr = 0

    for case in test_cases:
        results = store.search(case["question"], top_k=top_k)
        # 拼出检索到的文档标识：来源#chunk序号
        retrieved_ids = [r["metadata"].get("source", "") + f"#{r['metadata'].get('chunk_index', '')}" for r in results]
        relevant = set(case["relevant_ids"])
        retrieved = set(retrieved_ids)

        # 召回率
        hits = relevant & retrieved
        recall = len(hits) / len(relevant) if relevant else 0
        total_recall += recall

        # 准确率
        precision = len(hits) / len(retrieved) if retrieved else 0
        total_precision += precision

        # MRR：第一个相关结果的排名倒数
        mrr = 0
        for i, rid in enumerate(retrieved_ids):
            if rid in relevant:
                mrr = 1 / (i + 1)
                break
        total_mrr += mrr

    n = len(test_cases)
    return {
        "recall": total_recall / n,
        "precision": total_precision / n,
        "mrr": total_mrr / n,
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

def evaluate_generation(question: str, answer: str, context: str) -> dict:
    """用LLM对生成回答打分（1-5分）"""
    judge_prompt = f"""请对以下RAG系统的回答打分（1-5分），评估三个维度：
1. 相关性：回答是否切题
2. 准确性：回答是否忠于参考文档，无编造
3. 完整性：是否完整回答了问题

问题：{question}
参考文档：{context}
回答：{answer}

请输出JSON格式：{{"relevance": x, "faithfulness": x, "completeness": x, "reason": "..."}}"""

    resp = client.chat.completions.create(
        model=LLM_MODEL,
        messages=[{"role": "user", "content": judge_prompt}],
        response_format={"type": "json_object"},
        temperature=0,
    )
    return json.loads(resp.choices[0].message.content)