RAG技术实战指南：让大模型拥有”记忆”的秘密#

1
回答 = LLM(问题 + 检索到的相关文档)

1
from langchain.text_splitter import RecursiveCharacterTextSplitter
2
from langchain.text_splitter import MarkdownHeaderTextSplitter
3

4
class SmartTextSplitter:
5
    """智能文档切分器"""
6

7
    def __init__(
8
        self,
9
        chunk_size: int = 500,
10
        chunk_overlap: int = 50,
11
        separators: list = None
12
    ):
13
        self.chunk_size = chunk_size
14
        self.chunk_overlap = chunk_overlap
15

16
        # 默认分隔符优先级：从强到弱
17
        self.separators = separators or [
18
            "\n\n\n",  # 章节分隔
19
            "\n\n",    # 段落分隔
20
            "\n",      # 行分隔
21
            "。",      # 中文句号
22
            ".",       # 英文句号
23
            " ",       # 空格
24
            ""         # 最后强制切分
25
        ]
26

27
        self.splitter = RecursiveCharacterTextSplitter(
28
            chunk_size=chunk_size,
29
            chunk_overlap=chunk_overlap,
30
            separators=self.separators,
31
            length_function=self._count_tokens
32
        )
33

34
    def _count_tokens(self, text: str) -> int:
35
        """更准确的token计数（中文友好）"""
36
        # 简化版：中文约1.5字符/token，英文约4字符/token
37
        chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
38
        other_chars = len(text) - chinese_chars
39
        return int(chinese_chars * 0.7 + other_chars * 0.25)
40

41
    def split_text(self, text: str) -> list[str]:
42
        """切分文本"""
43
        return self.splitter.split_text(text)
44

45
    def split_markdown(self, text: str) -> list[dict]:
46
        """切分Markdown文档，保留标题元数据"""
47
        # 先按标题切分
48
        headers_to_split_on = [
49
            ("#", "header1"),
50
            ("##", "header2"),
51
            ("###", "header3"),
52
        ]
53

54
        markdown_splitter = MarkdownHeaderTextSplitter(
55
            headers_to_split_on=headers_to_split_on
56
        )
57
        md_splits = markdown_splitter.split_text(text)
58

59
        # 再对每个片段进行细切分
60
        final_chunks = []
61
        for split in md_splits:
62
            content = split.page_content
63
            metadata = split.metadata
64

65
            # 如果片段太长，继续切分
66
            if len(content) > self.chunk_size:
67
                sub_chunks = self.splitter.split_text(content)
68
                for i, sub_chunk in enumerate(sub_chunks):
69
                    final_chunks.append({
70
                        "content": sub_chunk,
71
                        "metadata": {
72
                            **metadata,
73
                            "chunk_index": i
74
                        }
75
                    })
76
            else:
77
                final_chunks.append({
78
                    "content": content,
79
                    "metadata": metadata
80
                })
81

82
        return final_chunks
83

84
# 使用示例
85
splitter = SmartTextSplitter(chunk_size=500, chunk_overlap=50)
86
chunks = splitter.split_markdown(long_markdown_text)

1
from sentence_transformers import SentenceTransformer
2
import numpy as np
3
from typing import List, Optional
4

5
class EmbeddingManager:
6
    """Embedding管理器，支持多种模型切换"""
7

8
    def __init__(self, model_name: str = "BAAI/bge-large-zh-v1.5"):
9
        self.model_name = model_name
10
        self.model = None
11
        self._load_model()
12

13
    def _load_model(self):
14
        """加载模型"""
15
        # 支持本地模型
16
        local_models = {
17
            "bge-large-zh": "BAAI/bge-large-zh-v1.5",
18
            "bge-m3": "BAAI/bge-m3",
19
            "m3e-large": "moka-ai/m3e-large",
20
        }
21

22
        if self.model_name in local_models:
23
            self.model = SentenceTransformer(local_models[self.model_name])
24
        else:
25
            # 使用OpenAI API
26
            self.model = None  # 将使用API
27

28
    def embed_texts(self, texts: List[str]) -> np.ndarray:
29
        """批量生成向量"""
30
        if self.model:
31
            # 本地模型
32
            return self.model.encode(texts, normalize_embeddings=True)
33
        else:
34
            # OpenAI API
35
            return self._embed_with_openai(texts)
36

37
    def embed_query(self, query: str) -> np.ndarray:
38
        """生成查询向量（可添加特殊处理）"""
39
        # 对于某些模型，查询需要添加前缀
40
        if "bge" in self.model_name.lower():
41
            query = "为这个句子生成表示以用于检索相关文章：" + query
42

43
        return self.embed_texts([query])[0]
44

45
    def _embed_with_openai(self, texts: List[str]) -> np.ndarray:
46
        """使用OpenAI API"""
47
        import openai
48

49
        response = openai.embeddings.create(
50
            model="text-embedding-3-small",
51
            input=texts
52
        )
53

54
        embeddings = [item.embedding for item in response.data]
55
        return np.array(embeddings)
56

57
# 使用示例
58
embedding_manager = EmbeddingManager("bge-large-zh")
59
vectors = embedding_manager.embed_texts(["这是第一段文本", "这是第二段文本"])
60
query_vector = embedding_manager.embed_query("搜索这段相关内容")

1
import chromadb
2
from chromadb.config import Settings
3

4
class ChromaVectorStore:
5
    """Chroma向量存储封装"""
6

7
    def __init__(self, collection_name: str = "documents", persist_dir: str = "./chroma_db"):
8
        self.client = chromadb.PersistentClient(path=persist_dir)
9
        self.collection = self.client.get_or_create_collection(
10
            name=collection_name,
11
            metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
12
        )
13

14
    def add_documents(
15
        self,
16
        documents: List[str],
17
        embeddings: List[List[float]],
18
        metadatas: List[dict] = None,
19
        ids: List[str] = None
20
    ):
21
        """添加文档"""
22
        if not ids:
23
            ids = [f"doc_{i}" for i in range(len(documents))]
24

25
        if not metadatas:
26
            metadatas = [{"source": "unknown"} for _ in documents]
27

28
        self.collection.add(
29
            documents=documents,
30
            embeddings=embeddings,
31
            metadatas=metadatas,
32
            ids=ids
33
        )
34

35
    def search(
36
        self,
37
        query_embedding: List[float],
38
        top_k: int = 5,
39
        where_filter: dict = None
40
    ) -> dict:
41
        """相似度检索"""
42
        results = self.collection.query(
43
            query_embeddings=[query_embedding],
44
            n_results=top_k,
45
            where=where_filter  # 元数据过滤
46
        )
47

48
        return {
49
            "documents": results["documents"][0],
50
            "metadatas": results["metadatas"][0],
51
            "distances": results["distances"][0],
52
            "ids": results["ids"][0]
53
        }
54

55
    def delete_by_metadata(self, metadata_filter: dict):
56
        """按元数据删除"""
57
        self.collection.delete(where=metadata_filter)
58

59
# 使用示例
60
store = ChromaVectorStore("my_knowledge_base")
61
store.add_documents(chunks, embeddings, metadatas)
62
results = store.search(query_vector, top_k=5)

1
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
2

3
class MilvusVectorStore:
4
    """Milvus向量存储（生产级）"""
5

6
    def __init__(
7
        self,
8
        collection_name: str = "documents",
9
        host: str = "localhost",
10
        port: str = "19530",
11
        dimension: int = 1024
12
    ):
13
        self.collection_name = collection_name
14
        self.dimension = dimension
15

16
        # 连接Milvus
17
        connections.connect(host=host, port=port)
18

19
        # 创建或获取collection
20
        if utility.has_collection(collection_name):
21
            self.collection = Collection(collection_name)
22
        else:
23
            self._create_collection()
24

25
        # 创建索引（首次需要）
26
        self._create_index()
27

28
    def _create_collection(self):
29
        """创建collection schema"""
30
        fields = [
31
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
32
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.dimension),
33
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=2000),
34
            FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=200),
35
            FieldSchema(name="chunk_index", dtype=DataType.INT64),
36
        ]
37

38
        schema = CollectionSchema(fields=fields, description="知识库向量")
39
        self.collection = Collection(name=self.collection_name, schema=schema)
40

41
    def _create_index(self):
42
        """创建向量索引"""
43
        index_params = {
44
            "metric_type": "COSINE",
45
            "index_type": "HNSW",
46
            "params": {"M": 8, "efConstruction": 64}
47
        }
48

49
        self.collection.create_index(
50
            field_name="embedding",
51
            index_params=index_params
52
        )
53

54
    def insert(self, data: List[dict]):
55
        """批量插入"""
56
        embeddings = [item["embedding"] for item in data]
57
        contents = [item["content"][:2000] for item in data]
58
        sources = [item.get("source", "unknown")[:200] for item in data]
59
        chunk_indices = [item.get("chunk_index", 0) for item in data]
60

61
        self.collection.insert([
62
            embeddings,
63
            contents,
64
            sources,
65
            chunk_indices
66
        ])
67

68
        # 刷新以确保写入
69
        self.collection.flush()
70

71
    def search(
72
        self,
73
        query_vector: List[float],
74
        top_k: int = 10,
75
        expr: str = None
76
    ) -> List[dict]:
77
        """相似度检索"""
78
        # 加载collection到内存
79
        self.collection.load()
80

81
        search_params = {"metric_type": "COSINE", "params": {"ef": 64}}
82

83
        results = self.collection.search(
84
            data=[query_vector],
85
            anns_field="embedding",
86
            param=search_params,
87
            limit=top_k,
88
            expr=expr  # 过滤表达式
89
        )
90

91
        # 格式化结果
92
        formatted_results = []
93
        for hit in results[0]:
94
            formatted_results.append({
95
                "id": hit.id,
96
                "distance": hit.distance,
97
                "content": hit.entity.get("content"),
98
                "source": hit.entity.get("source"),
99
            })
100

101
        return formatted_results

1
"""
2
完整的RAG系统实现
3
"""
4

5
from typing import List, Dict, Optional
6
from dataclasses import dataclass
7
import chromadb
8

9
@dataclass
10
class RAGConfig:
11
    """RAG配置"""
12
    embedding_model: str = "BAAI/bge-large-zh-v1.5"
13
    vector_db: str = "chroma"
14
    chunk_size: int = 500
15
    chunk_overlap: int = 50
16
    top_k: int = 5
17
    rerank: bool = True
18
    llm_model: str = "qwen-plus"
19

20

21
class RAGSystem:
22
    """完整RAG系统"""
23

24
    def __init__(self, config: RAGConfig = None):
25
        self.config = config or RAGConfig()
26

27
        # 初始化各组件
28
        self.text_splitter = SmartTextSplitter(
29
            chunk_size=self.config.chunk_size,
30
            chunk_overlap=self.config.chunk_overlap
31
        )
32

33
        self.embedding_manager = EmbeddingManager(self.config.embedding_model)
34

35
        self.vector_store = ChromaVectorStore("knowledge_base")
36

37
        if self.config.rerank:
38
            self.reranker = ReRanker()
39

40
    def ingest_documents(self, documents: List[Dict]):
41
        """文档入库"""
42
        all_chunks = []
43
        all_embeddings = []
44
        all_metadatas = []
45
        all_ids = []
46

47
        for doc_idx, doc in enumerate(documents):
48
            # 切分
49
            chunks = self.text_splitter.split_text(doc["content"])
50

51
            # 生成向量
52
            embeddings = self.embedding_manager.embed_texts(chunks)
53

54
            # 构建元数据
55
            for chunk_idx, chunk in enumerate(chunks):
56
                all_chunks.append(chunk)
57
                all_embeddings.append(embeddings[chunk_idx].tolist())
58
                all_metadatas.append({
59
                    "source": doc.get("source", "unknown"),
60
                    "doc_idx": doc_idx,
61
                    "chunk_idx": chunk_idx,
62
                    "title": doc.get("title", "")
63
                })
64
                all_ids.append(f"doc_{doc_idx}_chunk_{chunk_idx}")
65

66
        # 存入向量库
67
        self.vector_store.add_documents(
68
            documents=all_chunks,
69
            embeddings=all_embeddings,
70
            metadatas=all_metadatas,
71
            ids=all_ids
72
        )
73

74
        return len(all_chunks)
75

76
    def query(
77
        self,
78
        question: str,
79
        top_k: int = None,
80
        rerank: bool = None,
81
        return_sources: bool = True
82
    ) -> Dict:
83
        """查询"""
84
        top_k = top_k or self.config.top_k
85
        rerank = rerank if rerank is not None else self.config.rerank
86

87
        # 生成查询向量
88
        query_vector = self.embedding_manager.embed_query(question)
89

90
        # 检索（多取一些用于重排序）
91
        retrieve_k = top_k * 3 if rerank else top_k
92
        results = self.vector_store.search(query_vector, top_k=retrieve_k)
93

94
        # 重排序
95
        if rerank:
96
            ranked = self.reranker.rerank(question, results["documents"], top_k=top_k)
97
            final_docs = [item[0] for item in ranked]
98
            # 找到对应的元数据
99
            final_metadatas = []
100
            for doc in final_docs:
101
                idx = results["documents"].index(doc)
102
                final_metadatas.append(results["metadatas"][idx])
103
        else:
104
            final_docs = results["documents"][:top_k]
105
            final_metadatas = results["metadatas"][:top_k]
106

107
        # 构建Prompt
108
        context = "\n\n".join([f"[文档{i+1}] {doc}" for i, doc in enumerate(final_docs)])
109

110
        prompt = f"""基于以下文档内容回答用户问题。如果文档中没有相关信息，请明确说明。
111

112
参考文档：
113
{context}
114

115
用户问题：{question}
116

117
请给出准确、简洁的回答，并注明信息来源。"""
118

119
        # 调用LLM
120
        answer = self._call_llm(prompt)
121

122
        # 构建返回结果
123
        result = {
124
            "question": question,
125
            "answer": answer,
126
        }
127

128
        if return_sources:
129
            result["sources"] = [
130
                {
131
                    "content": doc,
132
                    "metadata": meta
133
                }
134
                for doc, meta in zip(final_docs, final_metadatas)
135
            ]
136

137
        return result
138

139
    def _call_llm(self, prompt: str) -> str:
140
        """调用LLM"""
141
        # 这里可以根据config使用不同的LLM
142
        # 示例使用OpenAI格式API
143
        import openai
144

145
        client = openai.OpenAI(
146
            api_key="your-api-key",
147
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
148
        )
149

150
        response = client.chat.completions.create(
151
            model=self.config.llm_model,
152
            messages=[{"role": "user", "content": prompt}],
153
            temperature=0.1  # 低温度减少幻觉
154
        )
155

156
        return response.choices[0].message.content
157

158

159
# 使用示例
160
def main():
161
    # 配置
162
    config = RAGConfig(
163
        embedding_model="bge-large-zh",
164
        chunk_size=500,
165
        top_k=5,
166
        rerank=True
167
    )
168

169
    # 初始化系统
170
    rag = RAGSystem(config)
171

172
    # 入库文档
173
    documents = [
174
        {
175
            "content": "长文档内容...",
176
            "source": "产品手册",
177
            "title": "产品使用指南"
178
        },
179
        {
180
            "content": "另一篇文档...",
181
            "source": "FAQ",
182
            "title": "常见问题解答"
183
        }
184
    ]
185

186
    rag.ingest_documents(documents)
187

188
    # 查询
189
    result = rag.query("如何使用这个产品？")
190

191
    print(f"问题: {result['question']}")
192
    print(f"回答: {result['answer']}")
193
    print(f"来源: {len(result['sources'])} 个文档片段")
194

195

196
if __name__ == "__main__":
197
    main()

1
import asyncio
2
from concurrent.futures import ThreadPoolExecutor
3

4
class AsyncRAGSystem(RAGSystem):
5
    """异步RAG系统，提升吞吐量"""
6

7
    def __init__(self, config: RAGConfig = None, max_workers: int = 4):
8
        super().__init__(config)
9
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
10

11
    async def async_ingest(self, documents: List[Dict]) -> int:
12
        """异步入库"""
13
        loop = asyncio.get_event_loop()
14

15
        # 并行处理文档
16
        tasks = [
17
            loop.run_in_executor(
18
                self.executor,
19
                self._process_single_document,
20
                doc
21
            )
22
            for doc in documents
23
        ]
24

25
        results = await asyncio.gather(*tasks)
26

27
        # 批量写入
28
        total_chunks = sum(r["chunk_count"] for r in results)
29

30
        # 合并所有结果写入向量库
31
        all_data = []
32
        for r in results:
33
            all_data.extend(r["data"])
34

35
        self.vector_store.add_documents(
36
            documents=[d["content"] for d in all_data],
37
            embeddings=[d["embedding"] for d in all_data],
38
            metadatas=[d["metadata"] for d in all_data],
39
            ids=[d["id"] for d in all_data]
40
        )
41

42
        return total_chunks
43

44
    def _process_single_document(self, doc: Dict) -> Dict:
45
        """处理单个文档"""
46
        chunks = self.text_splitter.split_text(doc["content"])
47
        embeddings = self.embedding_manager.embed_texts(chunks)
48

49
        data = []
50
        for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
51
            data.append({
52
                "content": chunk,
53
                "embedding": emb.tolist(),
54
                "metadata": {
55
                    "source": doc.get("source", "unknown"),
56
                    "chunk_idx": i
57
                },
58
                "id": f"{doc.get('source', 'doc')}_{i}"
59
            })
60

61
        return {"chunk_count": len(chunks), "data": data}
62

63
    async def async_query(self, question: str) -> Dict:
64
        """异步查询"""
65
        loop = asyncio.get_event_loop()
66

67
        # 并行执行：向量生成 + 初步检索
68
        query_vector_task = loop.run_in_executor(
69
            self.executor,
70
            self.embedding_manager.embed_query,
71
            question
72
        )
73

74
        query_vector = await query_vector_task
75

76
        # 检索
77
        results = await loop.run_in_executor(
78
            self.executor,
79
            self.vector_store.search,
80
            query_vector.tolist(),
81
            self.config.top_k * 3
82
        )
83

84
        # 重排序 + LLM生成并行
85
        if self.config.rerank:
86
            rerank_task = loop.run_in_executor(
87
                self.executor,
88
                self.reranker.rerank,
89
                question,
90
                results["documents"],
91
                self.config.top_k
92
            )
93
            ranked = await rerank_task
94
            final_docs = [item[0] for item in ranked]
95
        else:
96
            final_docs = results["documents"][:self.config.top_k]
97

98
        # LLM生成
99
        context = "\n\n".join(final_docs)
100
        prompt = f"基于以下内容回答：\n{context}\n问题：{question}"
101

102
        answer = await loop.run_in_executor(
103
            self.executor,
104
            self._call_llm,
105
            prompt
106
        )
107

108
        return {"question": question, "answer": answer, "sources": final_docs}

1
class RAGMonitor:
2
    """RAG系统监控"""
3

4
    def __init__(self):
5
        self.metrics = {
6
            "query_count": 0,
7
            "avg_latency": 0,
8
            "avg_retrieval_latency": 0,
9
            "avg_llm_latency": 0,
10
            "retrieval_accuracy": []  # 人工标注
11
        }
12

13
    def log_query(
14
        self,
15
        question: str,
16
        retrieval_latency: float,
17
        llm_latency: float,
18
        answer: str,
19
        sources: List[str]
20
    ):
21
        """记录查询日志"""
22
        self.metrics["query_count"] += 1
23

24
        # 更新平均延迟
25
        n = self.metrics["query_count"]
26
        old_avg = self.metrics["avg_latency"]
27
        new_latency = retrieval_latency + llm_latency
28
        self.metrics["avg_latency"] = old_avg + (new_latency - old_avg) / n
29

30
        # 单独统计
31
        self.metrics["avg_retrieval_latency"] = (
32
            self.metrics["avg_retrieval_latency"] +
33
            (retrieval_latency - self.metrics["avg_retrieval_latency"]) / n
34
        )
35
        self.metrics["avg_llm_latency"] = (
36
            self.metrics["avg_llm_latency"] +
37
            (llm_latency - self.metrics["avg_llm_latency"]) / n
38
        )
39

40
    def get_report(self) -> Dict:
41
        """获取监控报告"""
42
        return {
43
            "total_queries": self.metrics["query_count"],
44
            "avg_total_latency_ms": round(self.metrics["avg_latency"] * 1000, 2),
45
            "avg_retrieval_ms": round(self.metrics["avg_retrieval_latency"] * 1000, 2),
46
            "avg_llm_ms": round(self.metrics["avg_llm_latency"] * 1000, 2),
47
            "retrieval_accuracy": self._calculate_accuracy()
48
        }
49

50
    def _calculate_accuracy(self) -> float:
51
        """计算检索准确率（需要人工标注）"""
52
        if not self.metrics["retrieval_accuracy"]:
53
            return None
54

55
        return sum(self.metrics["retrieval_accuracy"]) / len(self.metrics["retrieval_accuracy"])