首页
学习
活动
专区
圈层
工具
发布
社区首页 >专栏 >AI 应用实战:构建一个完整的 RAG 知识库问答系统 在 AI 应用开发中,RAG 是最常见也最实用的架构之一。下面从工程实现角度,用 Python 完整写

AI 应用实战:构建一个完整的 RAG 知识库问答系统 在 AI 应用开发中,RAG 是最常见也最实用的架构之一。下面从工程实现角度,用 Python 完整写

原创
作者头像
用户12583550
发布于 2026-11-13 16:49:56
发布于 2026-11-13 16:49:56
330
举报

在 AI 应用开发中,RAG 是最常见也最实用的架构之一。下面从工程实现角度,用 Python 完整写一个 RAG 系统,包括文档加载、文本分块、Embedding、向量检索、BM25 检索、混合重排序、Prompt 构建、答案生成、缓存与评测。

代码语言:javascript
复制
import os
import re
import json
import math
import time
import hashlib
import logging
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional, Any
from collections import Counter, defaultdict

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("rag_system")

一、文档加载模块

代码语言:javascript
复制
class Document: 15015.h.kuaisou.com 
    """表示一个待处理的文档对象。"""

    def __init__(
        self,
        doc_id: str,
        text: str,
        source: str = "",
        metadata: Optional[Dict[str, Any]] = None,
    ):
        self.doc_id = doc_id
        self.text = text
        self.source = source
        self.metadata = metadata or {}

    def __repr__(self) -> str: 15025.h.kuaisou.com 
        return f"Document(doc_id={self.doc_id!r}, source={self.source!r}, len={len(self.text)})"


class DocumentLoader: 15027.h.kuaisou.com 
    """
    文档加载器,支持从文本字符串、JSON、TXT 文件加载文档。
    实际项目中可以继续扩展 PDF、Markdown、Word 等格式。
    """

    def load_from_text(self, text: str, doc_id: str = "doc_0", source: str = "") -> Document: 15029.h.kuaisou.com 
        return Document(doc_id=doc_id, text=text, source=source)

    def load_from_json(self, file_path: str) -> List[Document]: 15030.h.kuaisou.com 
        if not os.path.exists(file_path): 15031.h.kuaisou.com 
            raise FileNotFoundError(f"JSON file not found: {file_path}")

        with open(file_path, "r", encoding="utf-8") as f: 15032.h.kuaisou.com 
            data = json.load(f)

        documents = []
        for i, item in enumerate(data): 15033.h.kuaisou.com 
            doc = Document(
                doc_id=item.get("doc_id", f"doc_{i}"),
                text=item.get("text", ""),
                source=item.get("source", file_path),
                metadata=item.get("metadata", {}),
            )
            documents.append(doc)
        return documents

    def load_from_txt(self, file_path: str, doc_id: str = "doc_0") -> Document: 15034.h.kuaisou.com
        if not os.path.exists(file_path): 15035.h.kuaisou.com 
            raise FileNotFoundError(f"TXT file not found: {file_path}")

        with open(file_path, "r", encoding="utf-8") as f: 15036.h.kuaisou.com
        return Document(
            doc_id=doc_id,
            text=text,
            source=file_path,
            metadata={"file_type": "txt"},
        )

二、文本分块模块

代码语言:javascript
复制
@dataclass
class TextChunk: 15037.h.kuaisou.com 
    """文本块对象,用于后续 Embedding 和检索。"""

    chunk_id: str
    text: str
    doc_id: str = ""
    source: str = ""
    start_char: int = 0
    end_char: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __len__(self) -> int: 15038.h.kuaisou.com 
        return len(self.text)


class RecursiveCharacterTextSplitter: 15039.h.kuaisou.com 
    """
    递归文本分块器。
    优先按照段落、句子、标点进行切分,尽量保持语义完整。
    """

    DEFAULT_SEPARATORS = [
        "\n\n",
        "\n",
        "。",
        ";",
        "!",
        "?",
        ". ",
        "! ",
        "? ",
        ",",
        ",",
        " ",
        "",
    ]

    def __init__(
        self,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        separators: Optional[List[str]] = None,
        length_function=len,
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.separators = separators or self.DEFAULT_SEPARATORS
        self.length_function = length_function

    def split_text(self, text: str) -> List[str]: 16001.h.kuaisou.com 
        return self._split_recursive(text, self.separators)

    def _split_recursive(self, text: str, separators: List[str]) -> List[str]: 16002.h.kuaisou.com 
        if not text: 16003.h.kuaisou.com 
            return []

        separator = separators[-1]
        new_separators = []

        for i, sep in enumerate(separators): 16004.h.kuaisou.com 
            if sep == "" or sep in text: 16005.h.kuaisou.com 
                separator = sep
                new_separators = separators[i + 1:]
                break

        if separator: 16006.h.kuaisou.com 
            splits = text.split(separator)
        else: 16007.h.kuaisou.com 
            splits = list(text)

        chunks = []
        current_chunk: List[str] = []
        current_length = 0

        for split in splits: 16008.h.kuaisou.com 
            split_len = self.length_function(split)
            sep_len = self.length_function(separator)

            if current_length + split_len + sep_len <= self.chunk_size: 16009.h.kuaisou.com
                current_chunk.append(split)
                current_length += split_len + sep_len
            else: 16010.h.kuaisou.com 
                if current_chunk: 16011.h.kuaisou.com 
                    chunk_text = separator.join(current_chunk)
                    if chunk_text: 16012.h.kuaisou.com 
                        chunks.append(chunk_text)

                    overlap_text = self._build_overlap(current_chunk, separator)
                    if overlap_text: 16013.h.kuaisou.com 
                        current_chunk = [overlap_text, split]
                        current_length = self.length_function(overlap_text) + split_len + sep_len
                    else: 16014.h.kuaisou.com 
                        current_chunk = [split]
                        current_length = split_len + sep_len
                else: 16015.h.kuaisou.com 
                    if split_len > self.chunk_size and new_separators: 16016.h.kuaisou.com 
                        sub_chunks = self._split_recursive(split, new_separators)
                        chunks.extend(sub_chunks)
                    else: 16017.h.kuaisou.com 
                        chunks.append(split)
                        current_chunk = []
                        current_length = 0

        if current_chunk: 16018.h.kuaisou.com 
            chunk_text = separator.join(current_chunk)
            if chunk_text: 16019.h.kuaisou.com 
                chunks.append(chunk_text)

        return chunks

    def _build_overlap(self, parts: List[str], separator: str) -> str: 16020.h.kuaisou.com 
        if not parts or self.chunk_overlap <= 0: 16021.h.kuaisou.com 
            return ""

        overlap_parts = []
        overlap_len = 0

        for part in reversed(parts): 16022.h.kuaisou.com 
            part_len = self.length_function(part)
            if overlap_len + part_len <= self.chunk_overlap: 16023.h.kuaisou.com 
                overlap_parts.insert(0, part)
                overlap_len += part_len
            else: 16024.h.kuaisou.com 
                remaining = self.chunk_overlap - overlap_len
                if remaining > 0: 16025.h.kuaisou.com 
                    overlap_parts.insert(0, part[-remaining:])
                break

        return separator.join(overlap_parts)

    def create_chunks(self, text: str, doc_id: str = "", source: str = "") -> List[TextChunk]:
        raw_chunks = self.split_text(text)
        chunks = [] 
        pos = 0

        for i, chunk_text in enumerate(raw_chunks): 16026.h.kuaisou.com 
            start = text.find(chunk_text, pos)
            if start == -1: 16027.h.kuaisou.com 
                start = pos
            end = start + len(chunk_text)
            pos = end

            chunks.append(
                TextChunk(
                    chunk_id=f"{doc_id}:chunk_{i}",
                    text=chunk_text,
                    doc_id=doc_id,
                    source=source,
                    start_char=start,
                    end_char=end,
                    metadata={"index": i, "doc_id": doc_id, "source": source},
                )
            )
        return chunks

三、Embedding 与向量索引

代码语言:javascript
复制
class SimpleEmbeddingModel: 16028.h.kuaisou.com 
    """
    简易 Embedding 模型。
    生产环境应替换为 OpenAI、BGE、SentenceTransformer 等真实向量模型。
    """

    def __init__(self, dim: int = 128): 16029.h.kuaisou.com 
        self.dim = dim

    def _hash_word(self, word: str) -> int: 16030.h.kuaisou.com 
        return int(hashlib.md5(word.lower(16032.h.kuaisou.com).encode(16033.h.kuaisou.com)).hexdigest(), 16)

    def embed(self, text: str) -> List[float]: 16031.h.kuaisou.com 
        words = re.findall(r'\w+|[^\w\s]', text.lower(16034.h.kuaisou.com))
        vector = [0.0] * self.dim

        for word in words: 16035.h.kuaisou.com 
            h = self._hash_word(word)
            idx = h % self.dim
            vector[idx] += 1.0

        norm = math.sqrt(sum(v * v for v in vector))
        if norm > 0: 16036.h.kuaisou.com 
            vector = [v / norm for v in vector]
        return vector

    def embed_batch(self, texts: List[str]) -> List[List[float]]: 16037.h.kuaisou.com 
        return [self.embed(t) for t in texts]


def cosine_similarity(a: List[float], b: List[float]) -> float: 16038.h.kuaisou.com 
    if len(a) != len(b): 16039.h.kuaisou.com 
        raise ValueError("Vector dimensions must match")

    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(y * y for y in b))

    if norm_a == 0 or norm_b == 0: 16040.h.kuaisou.com 
        return 0.0

    return dot / (norm_a * norm_b)


class VectorIndex: 16041.h.kuaisou.com 
    """内存向量索引,支持添加、检索、删除。"""

    def __init__(self, dim: int = 128): 16042.h.kuaisou.com 
        self.dim = dim
        self.vectors: Dict[str, List[float]] = {}
        self.metadata: Dict[str, Dict[str, Any]] = {}

    def add(self, chunk_id: str, vector: List[float], metadata: Optional[Dict[str, Any]] = None):
        if len(vector) != self.dim: 16043.h.kuaisou.com 
            raise ValueError(f"Expected dim {self.dim}, got {len(vector)}")

        self.vectors[chunk_id] = vector
        self.metadata[chunk_id] = metadata or {}

    def add_batch(
        self,
        chunk_ids: List[str],
        vectors: List[List[float]],
        metadatas: Optional[List[Dict[str, Any]]] = None,
    ):
        for i, cid in enumerate(chunk_ids): 16044.h.kuaisou.com 
            meta = metadatas[i] if metadatas else {}
            self.add(cid, vectors[i], meta)

    def search(self, query_vector: List[float], k: int = 5) -> List[Tuple[str, float, Dict[str, Any]]]:
        scores = []
        for cid, vec in self.vectors.items(16046.h.kuaisou.com): 16045.h.kuaisou.com 
            score = cosine_similarity(query_vector, vec)
            scores.append((cid, score, self.metadata.get(cid, {})))

        scores.sort(key=lambda x: -x[1])
        return scores[:k]

    def search_with_threshold(
        self,
        query_vector: List[float],
        k: int = 5,
        threshold: float = 0.3,
    ) -> List[Tuple[str, float, Dict[str, Any]]]: 16047.h.kuaisou.com 
        candidates = self.search(query_vector, k=k * 3)
        filtered = [(cid, s, m) for cid, s, m in candidates if s >= threshold]
        return filtered[:k]

    def remove(self, chunk_id: str) -> bool: 16048.h.kuaisou.com 
        self.vectors.pop(chunk_id, None)
        self.metadata.pop(chunk_id, None)
        return True

    def clear(self): 16049.h.kuaisou.com 
        self.vectors.clear(16050.h.kuaisou.com)
        self.metadata.clear(16051.h.kuaisou.com)

    def size(self) -> int: 16052.h.kuaisou.com 
        return len(self.vectors)

四、BM25 关键词检索

代码语言:javascript
复制
class BM25Retriever: 16053.h.kuaisou.com 
    """ 
    BM25 关键词检索器。
    用于补充向量检索,提升精确词匹配能力。
    """

    def __init__(self, k1: float = 1.5, b: float = 0.75): 16054.h.kuaisou.com 
        self.k1 = k1
        self.b = b
        self.doc_freq: Dict[str, int] = defaultdict(int)
        self.doc_lengths: Dict[str, int] = {}
        self.doc_terms: Dict[str, Counter] = {}
        self.num_docs = 0
        self.avg_doc_length = 0.0

    def add_document(self, doc_id: str, text: str): 16055.h.kuaisou.com 
        terms = self._tokenize(text)
        term_counts = Counter(terms)

        self.doc_terms[doc_id] = term_counts
        self.doc_lengths[doc_id] = len(terms)

        for term in set(terms): 16056.h.kuaisou.com 
            self.doc_freq[term] += 1

        self.num_docs += 1
        total_length = sum(self.doc_lengths.values(16057.h.kuaisou.com))
        self.avg_doc_length = total_length / self.num_docs if self.num_docs > 0 else 0.0

    def add_documents(self, doc_ids: List[str], texts: List[str]): 16059.h.kuaisou.com 
        for doc_id, text in zip(doc_ids, texts): 16058.h.kuaisou.com 
            self.add_document(doc_id, text)

    def search(self, query: str, k: int = 5) -> List[Tuple[str, float]]: 16060.h.kuaisou.com 
        query_terms = self._tokenize(query)
        scores: Dict[str, float] = defaultdict(float)

        for term in query_terms: 16061.h.kuaisou.com 
            df = self.doc_freq.get(term, 0)
            if df == 0: 16062.h.kuaisou.com 
                continue

            idf = math.log((self.num_docs - df + 0.5) / (df + 0.5) + 1.0)

            for doc_id, term_counts in self.doc_terms.items(16064.h.kuaisou.com): 16063.h.kuaisou.com 
                tf = term_counts.get(term, 0)
                doc_len = self.doc_lengths[doc_id]

                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (
                    1 - self.b + self.b * doc_len / self.avg_doc_length
                )
                scores[doc_id] += idf * numerator / denominator

        ranked = sorted(scores.items(17001.h.kuaisou.com), key=lambda x: -x[1])
        return ranked[:k]

    def _tokenize(self, text: str) -> List[str]: 17002.h.kuaisou.com 
        return [t.lower() for t in re.findall(r'\w+|[^\w\s]', text)]

五、混合检索与重排序

代码语言:javascript
复制
class RRFMerger: 17003.h.kuaisou.com 
    """
    RRF:Reciprocal Rank Fusion。
    用于合并向量检索和 BM25 检索结果。
    """

    def __init__(self, k: int = 60):
        self.k = k

    def merge(
        self,
        result_lists: List[List[Tuple[str, float]]],
        top_n: int = 5,
    ) -> List[Tuple[str, float]]: 17004.h.kuaisou.com 
        rrf_scores: Dict[str, float] = {}

        for result_list in result_lists: 17005.h.kuaisou.com 
            for rank, (chunk_id, _) in enumerate(result_list, start=1): 17006.h.kuaisou.com 
                rrf_scores[chunk_id] = rrf_scores.get(chunk_id, 0.0) + 1.0 / (self.k + rank)

        merged = sorted(rrf_scores.items(17007.h.kuaisou.com), key=lambda x: -x[1])
        return merged[:top_n]


class CrossEncoderReranker: 17008.h.kuaisou.com 
    """
    简易重排序器。
    生产环境可替换为 BGE-reranker、Cohere Rerank 等。
    """

    def __init__(self, embedding_model: SimpleEmbeddingModel): 17009.h.kuaisou.com 
        self.embedding_model = embedding_model

    def rerank(
        self,
        query: str,
        candidates: List[Tuple[str, str, float]],
        top_n: int = 5,
    ) -> List[Tuple[str, str, float]]: 17010.h.kuaisou.com 
        query_vec = self.embedding_model.embed(query)
        scored = []

        for chunk_id, chunk_text, init_score in candidates: 17011.h.kuaisou.com 
            chunk_vec = self.embedding_model.embed(chunk_text)
            sim = cosine_similarity(query_vec, chunk_vec)
            final_score = 0.4 * init_score + 0.6 * sim
            scored.append((chunk_id, chunk_text, final_score))

        scored.sort(key=lambda x: -x[2])
        return scored[:top_n]


class HybridRetriever: 17012.h.kuaisou.com 
    """
    混合检索器:
    1. 向量检索
    2. BM25 检索
    3. RRF 融合
    4. 重排序
    """

    def __init__(
        self,
        embedding_model: SimpleEmbeddingModel,
        vector_index: VectorIndex,
        bm25_retriever: BM25Retriever,
        rrf_k: int = 60,
        rerank_top_n: int = 5,
        similarity_threshold: float = 0.3,
    ):
        self.embedding_model = embedding_model
        self.vector_index = vector_index
        self.bm25_retriever = bm25_retriever
        self.rrf_merger = RRFMerger(k=rrf_k)
        self.reranker = CrossEncoderReranker(embedding_model)
        self.rerank_top_n = rerank_top_n
        self.similarity_threshold = similarity_threshold

    def retrieve(self, query: str, k: int = 5) -> List[Tuple[str, str, float]]:
        query_vec = self.embedding_model.embed(query)

        vec_results = self.vector_index.search_with_threshold(
            query_vec,
            k=k * 3,
            threshold=self.similarity_threshold,
        )
        vec_ranked = [(cid, score) for cid, score, _ in vec_results]

        bm25_results = self.bm25_retriever.search(query, k=k * 3)

        merged = self.rrf_merger.merge(
            [vec_ranked, bm25_results],
            top_n=k * 2,
        )

        candidates = []
        for cid, rrf_score in merged: 17013.h.kuaisou.com 
            text = self.vector_index.metadata.get(cid, {}).get("text", "")
            if text: 17014.h.kuaisou.com 
                candidates.append((cid, text, rrf_score))

        reranked = self.reranker.rerank(query, candidates, top_n=self.rerank_top_n)
        return reranked

六、RAG Prompt 构建与答案生成

代码语言:javascript
复制
class RAGPromptBuilder: 17015.h.kuaisou.com 
    """
    构建防幻觉 RAG Prompt。
    明确要求模型只基于上下文回答,避免编造。
    """

    SYSTEM_PROMPT = """你是一个专业的知识库问答助手。请严格遵循以下规则:

1. 只根据【上下文】中的内容回答问题。
2. 如果上下文不足以回答问题,请回复:“根据现有知识库无法解答该问题。”
3. 不要编造上下文之外的信息。
4. 回答末尾标注引用来源。
5. 回答尽量简洁、准确、专业。"""

    USER_TEMPLATE = """【上下文】
{context}

【用户问题】
{question}

请基于上述上下文回答问题:"""

    def __init__(self, max_context_chunks: int = 5): 17016.h.kuaisou.com 
        self.max_context_chunks = max_context_chunks

    def build_prompt(self, question: str, retrieved_chunks: List[Dict[str, str]]) -> str: 17017.h.kuaisou.com 
        context_parts = []

        for i, chunk in enumerate(retrieved_chunks[: self.max_context_chunks], 1): 17018.h.kuaisou.com 
            context_parts.append(
                f"[片段{i}] (来源: {chunk.get('source', 'unknown')}, "
                f"ID: {chunk['chunk_id']})\n{chunk['text']}\n"
            )

        context = "\n".join(context_parts)
        user_prompt = self.USER_TEMPLATE.format(context=context, question=question)
        return self.SYSTEM_PROMPT + "\n\n" + user_prompt


class RAGGenerator: 17019.h.kuaisou.com 
    """
    RAG 答案生成器。
    实际项目中应接入 OpenAI、DeepSeek、Qwen 等 LLM。
    """

    def __init__(self, prompt_builder: Optional[RAGPromptBuilder] = None): 17020.h.kuaisou.com 
        self.prompt_builder = prompt_builder or RAGPromptBuilder(17021.h.kuaisou.com)

    def generate(
        self,
        question: str,
        retrieved_chunks: List[Dict[str, str]],
    ) -> Dict[str, Any]: 17022.h.kuaisou.com 
        if not retrieved_chunks: 17023.h.kuaisou.com 
            return {
                "answer": "根据现有知识库无法解答该问题。",
                "sources": [],
                "question": question,
            }

        prompt = self.prompt_builder.build_prompt(question, retrieved_chunks)
        answer = self._call_llm(prompt)

        sources = [
            {"chunk_id": c["chunk_id"], "source": c.get("source", "")}
            for c in retrieved_chunks
        ]

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

    def _call_llm(self, prompt: str) -> str: 17024.h.kuaisou.com 
        """
        占位方法,实际项目中替换为真实 LLM 调用。
        示例:OpenAI、DeepSeek、Qwen、本地部署模型。
        """
        return f"[LLM模拟回答] 基于 {len(prompt)} 字符上下文生成的答案。"

七、缓存模块

代码语言:javascript
复制
class RAGCache: 17121.h.kuaisou.com 
    """
    简单缓存层,用于缓存相同问题的检索和生成结果。
    生产环境可替换为 Redis、Memcached 等。
    """

    def __init__(self, ttl: int = 3600): 17025.h.kuaisou.com 
        self.cache: Dict[str, Dict[str, Any]] = {}
        self.timestamps: Dict[str, float] = {}
        self.ttl = ttl

    def _make_key(self, question: str) -> str: 17026.h.kuaisou.com 
        return hashlib.sha256(question.strip().encode("utf-8")).hexdigest()

    def get(self, question: 17117.h.kuaisou.com 

原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。

如有侵权,请联系 cloudcommunity@tencent.com 删除。

目录
  • 在 AI 应用开发中,RAG 是最常见也最实用的架构之一。下面从工程实现角度,用 Python 完整写一个 RAG 系统,包括文档加载、文本分块、Embedding、向量检索、BM25 检索、混合重排序、Prompt 构建、答案生成、缓存与评测。
  • 一、文档加载模块
  • 二、文本分块模块
  • 三、Embedding 与向量索引
  • 四、BM25 关键词检索
  • 五、混合检索与重排序
  • 六、RAG Prompt 构建与答案生成
  • 七、缓存模块
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档