返回 RAG 工程实战目录

核心模块开发实现

把架构落成可测试、可配置的核心模块

生命周期 · 阶段四 开发实现(编码落地) | 上游:设计(03)+ 流水线(05)+ 查询链路(06)| 下游:测试验收(08~09)

本文站在开发工程师角度,按 03 篇的设计给出工程结构与各核心模块的参考实现(Python + FastAPI + Celery + Qdrant)。

一、工程结构

rag-kb/
├── app/
│   ├── api/                  # 路由层(薄,只做参数校验与调用 service)
│   │   ├── documents.py
│   │   ├── chat.py
│   │   └── feedback.py
│   ├── core/
│   │   ├── config.py         # 配置(pydantic-settings)
│   │   ├── logging.py        # 结构化日志
│   │   └── security.py       # JWT、权限
│   ├── models/               # SQLAlchemy ORM(对应 03 篇 DDL)
│   ├── services/             # 业务服务
│   │   ├── document_service.py
│   │   └── chat_service.py
│   ├── pipeline/             # 入库流水线(Celery 任务链)
│   │   ├── tasks.py          # 任务编排:chain(parse → ocr → chunk → embed → index)
│   │   ├── parser/           # factory.py + pdf_parser.py + docx_parser.py ...
│   │   ├── ocr.py
│   │   ├── chunker.py
│   │   ├── embedder.py
│   │   └── indexer.py
│   ├── retrieval/            # 查询链路
│   │   ├── rewriter.py       # 查询清洗/改写
│   │   ├── retriever.py      # 混合检索 + RRF
│   │   ├── reranker.py
│   │   └── generator.py      # Prompt 组装 + 流式生成
│   └── adapters/             # 模型适配器(OpenAI 兼容协议)
│       ├── llm_adapter.py
│       └── embed_adapter.py
├── tests/                    # pytest:unit/ + e2e/
│   ├── test_chunker.py
│   ├── test_rrf.py
│   └── eval/                 # 评测集回归(见 09 篇)
├── docker/
│   ├── Dockerfile
│   └── docker-compose.yml
├── .env.example
└── pyproject.toml

二、配置管理

# app/core/config.py
from pydantic_settings import BaseSettings

class Settings(BaseSettings):
    # 服务
    app_name: str = "rag-kb"
    debug: bool = False
    # 存储
    mysql_dsn: str  # 从环境变量或 Secret 注入,不在代码中提供默认密码
    redis_url: str = "redis://localhost:6379/0"
    qdrant_url: str = "http://localhost:6333"
    minio_endpoint: str = "localhost:9000"
    # 模型(OpenAI 兼容协议,可切 vLLM/Ollama/商用 API)
    llm_base_url: str = "http://localhost:11434/v1"
    llm_model: str = "qwen2.5:7b"
    embed_base_url: str = "http://localhost:8080"
    embed_model: str = "bge-large-zh-v1.5"
    embed_dim: int = 1024
    # 入库限制
    max_file_size_mb: int = 200
    allowed_types: list[str] = ["pdf", "docx", "xlsx", "pptx", "html", "md", "txt", "png", "jpg"]
    # 检索参数
    retrieval_top_k: int = 20
    rerank_top_k: int = 5

    class Config:
        env_file = ".env"

settings = Settings()

要点:模型名、地址、维度全部配置化embed_dim 与向量库强绑定,改配置必须走重建流程。

三、文档接收:上传接口(含幂等)

# app/api/documents.py
import hashlib
from fastapi import APIRouter, UploadFile, File, HTTPException, Header
from app.core.config import settings
from app.services import document_service

router = APIRouter(prefix="/api/v1/documents", tags=["documents"])

@router.post("")
async def upload_document(
    file: UploadFile = File(...),
    kb_id: int = 1,
    idempotency_key: str | None = Header(default=None),
):
    # 1. 类型与大小校验(类型看扩展名 + 魔数双重校验)
    ext = file.filename.rsplit(".", 1)[-1].lower()
    if ext not in settings.allowed_types:
        raise HTTPException(400, detail="1002 不支持的文件格式")
    size = 0
    chunks = []
    while block := await file.read(1024 * 1024):
        size += len(chunks.append(block) or block)
    if size > settings.max_file_size_mb * 1024 * 1024:
        raise HTTPException(400, detail="1003 文件超过大小限制")

    # 2. 内容哈希(幂等依据)
    file_hash = hashlib.sha256(b"".join(chunks)).hexdigest()
    result = await document_service.create_ingest_task(
        kb_id=kb_id, filename=file.filename, ext=ext,
        content=b"".join(chunks), file_hash=file_hash,
        idempotency_key=idempotency_key,
    )
    return {"code": 0, "data": result}
    # result: {doc_id, task_id, status, deduplicated}
# app/services/document_service.py(核心逻辑摘要)
async def create_ingest_task(kb_id, filename, ext, content, file_hash, idempotency_key):
    # 幂等:同 kb + 同 hash + 未删除 → 直接返回已有任务
    existed = await doc_repo.find_by_hash(kb_id, file_hash)
    if existed:
        return {"doc_id": existed.id, "task_id": existed.last_task_id,
                "status": existed.status, "deduplicated": True}
    # 存 MinIO → 建 kb_document + kb_ingest_task(PENDING) → 发 Celery 队列
    storage_path = f"kb/{kb_id}/{file_hash}.{ext}"
    minio_client.put_object(storage_path, content)
    doc = await doc_repo.create(kb_id=kb_id, doc_name=filename, file_hash=file_hash,
                                storage_path=storage_path, file_type=ext)
    task = await task_repo.create(doc_id=doc.id)
    pipeline_tasks.run_pipeline.delay(task.id)
    return {"doc_id": doc.id, "task_id": task.id, "status": "PENDING", "deduplicated": False}

四、解析器:工厂 + 策略模式

# app/pipeline/parser/factory.py
from .pdf_parser import PdfParser
from .docx_parser import DocxParser
from .xlsx_parser import XlsxParser

PARSERS = {"pdf": PdfParser, "docx": DocxParser, "xlsx": XlsxParser}

def get_parser(ext: str):
    parser_cls = PARSERS.get(ext)
    if not parser_cls:
        raise ValueError(f"暂不支持格式: {ext}")
    return parser_cls()
# app/pipeline/parser/pdf_parser.py(关键:文本层检测 + 表格保护)
import fitz  # PyMuPDF
from dataclasses import dataclass

@dataclass
class Section:
    page: int
    section_path: str      # 如 "3.2 处置流程"
    content: str
    is_table: bool = False

class PdfParser:
    def parse(self, file_bytes: bytes) -> list[Section]:
        sections, needs_ocr_pages = [], []
        with fitz.open(stream=file_bytes, filetype="pdf") as doc:
            for page_no, page in enumerate(doc, start=1):
                text = page.get_text("text").strip()
                if len(text) < 20:                       # 文本层过少 → 转 OCR 分支
                    needs_ocr_pages.append(page_no)
                    continue
                tables = page.find_tables()              # 表格单独抽取,防切断
                table_spans = [t.bbox for t in tables.tables]
                sections.append(Section(page_no, "", self._remove_noise(text, page)))
        return sections, needs_ocr_pages

    def _remove_noise(self, text, page) -> str:
        # 剔除页眉页脚:位置在页面顶部/底部 5% 区域内的文本行
        ...

开发约定:所有 parser 统一返回 list[Section];表格走独立通道(整表一个切片);噪声(页眉页脚/水印)在 parser 层剔除,不流入切片。

五、切片实现

# app/pipeline/chunker.py
from dataclasses import dataclass

@dataclass
class Chunk:
    doc_id: int
    chunk_index: int
    content: str
    page: int
    section: str
    meta: dict          # department/security 等随切片下发的文档级元数据

class SemanticChunker:
    """结构感知切片:优先按章节边界,超长再递归细分;表格整表不切"""
    def __init__(self, max_size=500, overlap=50):
        self.max_size, self.overlap = max_size, overlap

    def split(self, sections, doc_meta: dict) -> list[Chunk]:
        chunks, idx = [], 0
        for sec in sections:
            if sec.is_table:
                pieces = [sec.content]                       # 整表一片
            elif len(sec.content) <= self.max_size:
                pieces = [sec.content]
            else:                                            # 递归细分:段落→句子
                pieces = self._recursive_split(sec.content)
            for p in pieces:
                chunks.append(Chunk(doc_meta["doc_id"], idx, p, sec.page,
                                    sec.section_path, doc_meta.copy()))
                idx += 1
        return chunks

    def _recursive_split(self, text: str) -> list[str]:
        seps = ["\n\n", "\n", "。", ";", ",", " "]
        return self._recursive(text, seps)

    def _recursive(self, text, seps):
        if len(text) <= self.max_size:
            return [text]
        sep, rest = seps[0], seps[1:]
        parts = [p for p in text.split(sep) if p.strip()]
        out, buf = [], ""
        for p in parts:                                      # 贪心合并 + 重叠窗口
            candidate = (buf + sep + p) if buf else p
            if len(candidate) > self.max_size and buf:
                out.append(buf)
                buf = buf[-self.overlap:] + sep + p          # 尾部重叠防断句
            else:
                buf = candidate
        if buf:
            out.append(buf)
        return out

六、向量化:批量编码 + 重试

# app/pipeline/embedder.py
from tenacity import retry, stop_after_attempt, wait_exponential
from app.adapters.embed_adapter import EmbedClient

class Embedder:
    def __init__(self, client: EmbedClient, batch_size=64):
        self.client, self.batch_size = client, batch_size

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10))
    def _embed_batch(self, texts):
        return self.client.embed(texts, normalize=True)

    def embed_all(self, chunks, on_progress=None):
        vectors = []
        for i in range(0, len(chunks), self.batch_size):
            batch = chunks[i:i + self.batch_size]
            vecs = self._embed_batch([c.content for c in batch])
            vectors.extend(vecs)
            if on_progress:                                  # 进度回写任务表
                on_progress(int((i + len(batch)) / len(chunks) * 100))
        assert len(vectors) == len(chunks), "向量化丢片,禁止静默入库"
        return vectors

七、索引写入(Qdrant)

# app/pipeline/indexer.py
from qdrant_client import QdrantClient, models

class Indexer:
    def __init__(self, url, collection="kb_chunks"):
        self.client = QdrantClient(url=url)
        self.collection = collection

    def upsert(self, chunks, vectors):
        points = [
            models.PointStruct(
                id=f"{c.doc_id}:{c.chunk_index}",           # 天然幂等:重跑即覆盖
                vector=v.tolist(),
                payload={
                    "text": c.content, "doc_id": c.doc_id,
                    "doc_name": c.meta["doc_name"], "page": c.page,
                    "section": c.section, "kb_id": c.meta["kb_id"],
                    "department": c.meta["department"], "security": c.meta["security"],
                    "version": c.meta["version"],
                },
            ) for c, v in zip(chunks, vectors)
        ]
        for i in range(0, len(points), 256):                 # 分批 upsert
            self.client.upsert(self.collection, points[i:i + 256])

    def delete_by_doc(self, doc_id: int):
        self.client.delete(self.collection,
            points_selector=models.FilterSelector(
                filter=models.Filter(must=[models.FieldCondition(
                    key="doc_id", match=models.MatchValue(value=doc_id))])))

八、查询链路:混合检索 + RRF + 重排

# app/retrieval/retriever.py
def rrf_fuse(result_lists: list[list[str]], k: int = 60) -> dict[str, float]:
    """多路结果融合:score = Σ 1/(k + rank)"""
    scores = {}
    for results in result_lists:
        for rank, point_id in enumerate(results, start=1):
            scores[point_id] = scores.get(point_id, 0) + 1 / (k + rank)
    return dict(sorted(scores.items(), key=lambda x: -x[1]))

class HybridRetriever:
    def retrieve(self, query_vec, query_text, perm_filter, top_k=20):
        vec_hits = self.qdrant.search(                        # 向量路(带权限过滤)
            self.collection, query_vec, limit=top_k,
            query_filter=perm_filter)                          # department/security 强制过滤
        bm25_hits = self.es.search(query_text, top_k=top_k)    # 关键词路
        fused = rrf_fuse([[h.id for h in vec_hits], [h.id for h in bm25_hits]])
        return self._load_points(list(fused)[:30])
# app/retrieval/generator.py(Prompt 组装 + 拒答)
PROMPT = """你是不良资产业务知识库助手。请严格遵守:
1. 仅根据【参考资料】回答,禁止使用资料之外的知识
2. 资料中没有答案时,直接回答"根据现有资料无法回答该问题"
3. 回答末尾列出引用编号与页码
4. 金额、日期、比例等数字必须与资料完全一致

【参考资料】
{context}

【问题】
{question}"""

def build_prompt(question, chunks, max_chars=6000):
    ctx, total = [], 0
    for i, c in enumerate(chunks, 1):                          # 按相关度装填,超预算截断
        piece = f"[{i}] 《{c.payload['doc_name']}》P{c.payload['page']}:{c.payload['text']}"
        if total + len(piece) > max_chars:
            break
        ctx.append(piece); total += len(piece)
    return PROMPT.format(context="\n\n".join(ctx), question=question)

九、流式问答(SSE)

# app/api/chat.py
import json
from fastapi.responses import StreamingResponse

@router.post("/api/v1/chat")
async def chat(req: ChatRequest, user=Depends(current_user)):
    return StreamingResponse(event_stream(req, user), media_type="text/event-stream")

async def event_stream(req, user):
    perm = await get_user_filter(user)                         # 权限过滤条件
    q = rewriter.rewrite(req.question)                         # 清洗/改写
    chunks = retriever.retrieve(embed(q), q, perm)             # 混合检索
    chunks = reranker.rerank(q, chunks, top_k=5)               # 精排
    if not chunks or chunks[0].score < REFUSE_THRESHOLD:       # 拒答兜底
        yield f"data: {json.dumps({'type': 'refuse', 'reason': '知识库中未找到相关内容'})}\n\n"
        return
    yield f"data: {json.dumps({'type': 'meta', 'retrieved': summarize(chunks)})}\n\n"
    prompt = build_prompt(req.question, chunks)
    async for delta in llm_adapter.stream(prompt, temperature=0.2):
        yield f"data: {json.dumps({'type': 'delta', 'content': delta})}\n\n"
    yield f"data: {json.dumps({'type': 'done'})}\n\n"
    await save_message(req, q, chunks, user)                   # 检索明细落库(排障用)

十、日志与异常规范

  • 结构化日志:JSON 格式,字段 ts/level/trace_id/module/msg/extra;trace_id 由网关生成、中间件注入,贯穿到 LLM 调用
  • 分级:入库失败 ERROR(含 doc_id + stage)、检索异常 WARNING、拒答 INFO
  • 统一异常处理:业务异常带错误码(见 03 篇);未捕获异常兜底返回 5000,日志留堆栈,不外泄堆栈给前端
  • 重试规范:模型调用 tenacity 指数退避 ≤3 次;Celery 任务 acks_late=True + 任务表状态机支持续跑

十一、单元测试示例

# tests/test_chunker.py
import pytest

def test_short_section_not_split(chunker, doc_meta):
    sections = [Section(1, "1 概述", "短内容" * 10, False)]
    chunks = chunker.split(sections, doc_meta)
    assert len(chunks) == 1

def test_long_section_recursive_and_overlap(chunker, doc_meta):
    text = "句子一。句子二。句子三。" * 100
    chunks = chunker.split([Section(2, "2 流程", text, False)], doc_meta)
    assert all(len(c.content) <= chunker.max_size + 100 for c in chunks)
    assert any("句子" in c.content[-chunker.overlap - 10:] for c in chunks[:-1])  # 尾部重叠

def test_table_never_split(chunker, doc_meta):
    table = "|a|b|\n" + "|1|2|\n" * 300
    chunks = chunker.split([Section(3, "3 表", table, is_table=True)], doc_meta)
    assert len(chunks) >= 1 and table in [c.content for c in chunks]

# tests/test_rrf.py
def test_rrf_prefers_multi_list_hit():
    fused = rrf_fuse([["a", "b", "c"], ["b", "d"]])
    assert fused["b"] > fused["a"]      # 两路都命中的排最前

测试要求:优先覆盖 pipeline/retrieval/ 的高风险分支、失败恢复和数据对账;行覆盖率可作为辅助指标,不以统一的 80% 代替风险判断。模型/存储依赖在单元测试中隔离,并保留上传→问答主链路的集成或 E2E 冒烟。

十二、开发自查清单(提测前)

  • [ ] 幂等:同一文件重复上传不产生重复切片(有单测)
  • [ ] 一致性:删除文档后向量库切片同步物理删除
  • [ ] 断点:流水线失败后重试从失败 stage 续跑,不重复已完成的阶段
  • [ ] 校验:文件类型魔数校验、大小限制、空文件拦截
  • [ ] 权限:检索 filter 强制生效,无 filter 的代码路径无法通过 CR
  • [ ] 拒答:低相关度/空结果走拒答分支,有测试覆盖
  • [ ] 流式:SSE 无缓冲、中断有 done/error 事件
  • [ ] 配置:无硬编码模型名/地址/密钥;embed_dim 变更有启动校验(防错配)
  • [ ] 日志:trace_id 贯穿;检索明细落库
  • [ ] 单测:新增模块有测试,覆盖率达标,CI 绿

本篇交付与下一步

本篇代码用于表达模块边界和关键实现模式。复制前应按锁定依赖版本运行,并补齐认证、事务、重试、超时、资源释放和异常映射。下一步进入 08-测试策略与用例设计