核心模块开发实现
把架构落成可测试、可配置的核心模块
生命周期 · 阶段四 开发实现(编码落地) | 上游:设计(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-测试策略与用例设计。