python脚本通过docling解析文档qrdant-edge索引使用RAG-MCP填充资料
这是一个用来快速提取文档关键值的脚本，可以用来批量提取一批相同类型，但是不同型号的公共参数值，可以避免人工大量重复性的搜索和来回翻阅工作：
例如快速提取esp32_datasheet_cn.pdf的以下字段
| **型号** | **封装类型** | **工作温度(°C)** | **CPU 架构** | **核心数** | **最高频率(MHz)** | **SRAM(KB)** | **PSRAM(支持)** | **内置Flash(MB)** | **Wi-Fi 协议** | **蓝牙规格** | **典型GPIO数** | **USB接口** | **以太网MAC** | **CAN总线** | **工作电压(V)** | **深度睡眠功耗(μA)** | **安全启动/Flash加密** |可以得到：
| **型号** | **封装类型** | **工作温度(°C)** | **CPU 架构** | **核心数** | **最高频率(MHz)** | **SRAM(KB)** | **PSRAM(支持)** | **内置Flash(MB)** | **Wi-Fi 协议** | **蓝牙规格** | **典型GPIO数** | **USB接口** | **以太网MAC** | **CAN总线** | **工作电压(V)** | **深度睡眠功耗(μA)** | **安全启动/Flash加密** |
|----------|--------------|------------------|--------------|------------|-------------------|--------------|-----------------|-------------------|----------------|--------------|----------------|-------------|---------------|-------------|-----------------|----------------------|--------------------------|
| ESP32 | QFN 5*5 | -40 ~ 125 | Xtensa LX6 32-bit | 双核 | 240 | 520 | 支持（外部最大 8 MB） | | IEEE 802.11b/g/n | 蓝牙 v4.2 BR/EDR 和 Bluetooth LE | 34 | | 支持（IEEE 802.3，需外部 PHY） | 支持（TWAI，兼容 ISO 11898-1） | 2.3 ~ 3.6 | 10 | 支持（安全启动、flash 加密） |
流程图：
依赖：
[project]
name = "rag-filler"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"docling>=2.121.0",
"langchain>=1.3.16",
"langchain-core>=1.6.0",
"langchain-mcp-adapters>=0.3.2",
"langchain-openai>=1.6.0",
"langchain-text-splitters>=1.1.2",
"mcp>=1.29.0",
"qdrant-client>=1.19.0",
]
在命令行中接受三个参数，其他参数通过环境变量来输入：
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="main.py",
add_help=False,
description="Fill a reference Markdown document from a PDF manual via RAG + an LLM agent.",
)
parser.add_argument("--in", dest="pdf", required=True, metavar="PDF",
help="用来索引的PDF")
parser.add_argument("--reference", required=True, metavar="MD",
help="需要填充的markdown文件")
parser.add_argument("--out", required=True, metavar="MD",
help="输出的markdown文件")
return parser.parse_args(argv)
def load_config() -> dict:
required = ("LLM_MODEL", "EMBEDDING_MODEL")
missing = [name for name in required if not os.environ.get(name)]
if missing:
raise ConfigError(
"缺少必需的环境变量：" + ", ".join(missing)
)
return {
"llm_api_base": os.environ.get("LLM_API_BASE", DEFAULT_LLM_BASE),
"llm_api_key": os.environ.get("LLM_API_KEY", "not-needed"),
"llm_model": os.environ["LLM_MODEL"],
"embedding_api_base": os.environ.get("EMBEDDING_API_BASE", DEFAULT_EMBED_BASE),
"embedding_api_key": os.environ.get("EMBEDDING_API_KEY", "not-needed"),
"embedding_model": os.environ["EMBEDDING_MODEL"],
"embedding_dimension": int(os.environ.get("EMBEDDING_DIMENSION", DEFAULT_EMBED_DIMENSION)),
}
脚本通过mcp来暴露查询方法，而不是将文档塞入上下文，用本地模型也可以跑
这里使用FastMCP来创建mcp服务
def make_search_server(qdrant: QdrantClient,
embed: OpenAIEmbeddings) -> FastMCP:
server = FastMCP("manual-search")
@server.tool()
def search_manual(query: str, top_k: int = 5) -> list[dict]:
"""在已索引的手册中检索与 query 最相关的段落并返回（含相关度分数）。填写表格单元格之前，请调用此工具从手册中查找事实依据。"""
k = max(1, min(int(top_k), 20))
# PyCharm 解析到旧版 qdrant-client 签名会误报 query 参数类型;1.19 运行时接受 list[float]
vector = embed.embed_query(query)
hits = qdrant.query_points(collection_name=COLLECTION, query=vector, limit=k).points # type: ignore[arg-type]
return [
{
"text": (hit.payload or {}).get("text", ""),
"section": (hit.payload or {}).get("section", ""),
"score": round(float(hit.score), 4),
}
for hit in hits
]
return server
然后通过 docling库来解析PDF，注意docling这个库会下载一个小模型来解析PDF，解析到的数据通过langchain提供的方法来分块，分块后通过Qwen3-Embedding-0.6B来embedding，存入qrdant-edge
建议embedding模型在本地跑，chat模型使用api，这样可以节省大量token，性能要求也不高
async def run(args: argparse.Namespace, cfg: dict) -> int:
embed = OpenAIEmbeddings(
model=cfg["embedding_model"],
openai_api_base=cfg["embedding_api_base"],
openai_api_key=cfg["embedding_api_key"],
check_embedding_ctx_length=False,
show_progress_bar=False,
)
reference = Path(args.reference)
ref_text = reference.read_text(encoding="utf-8")
print(f"正在索引 {args.pdf} ...")
qdrant = QdrantClient(":memory:")
count = index_pdf(args.pdf, embed, qdrant, cfg["embedding_dimension"])
print(f"已将 {count} 个文本块索引到内存向量数据库")
llm = ChatOpenAI(
model=cfg["llm_model"],
base_url=cfg["llm_api_base"],
api_key=cfg["llm_api_key"],
temperature=0,
)
server = make_search_server(qdrant, embed)
try:
async with create_connected_server_and_client_session(server) as session:
tools = await load_mcp_tools(session)
if not any(getattr(tool, "name", None) == "search_manual" for tool in tools):
raise FatalError("向量检索 MCP 工具 'search_manual' 未绑定到 agent")
agent = create_agent(llm, tools, system_prompt=SYSTEM_PROMPT)
print("正在填充参考文档 ...")
response = await agent.ainvoke(
{"messages": [HumanMessage(content=(
"下面是需要填充的参考文档。请依据手册填充其中留空或占位的内容；"
"手册内容一律通过 search_manual 检索获取，不要把 PDF 全文当作已知信息。"
"输出文档的最终形式由你判断，但必须内容完整、可直接使用。\n\n"
"参考文档：\n```markdown\n" + ref_text + "\n```"
))]},
config=RunnableConfig(recursion_limit=120),
)
filled = extract_content(response["messages"][-1].content)
except FatalError:
raise
except Exception as exc:
raise FatalError(f"API 调用失败：{format_api_error(exc)}") from exc
write_output(Path(args.out), filled)
print(f"已写入 {args.out}")
return 0
实测例子：
这里RX100M3说明书.pdf是一个135页的pdf
⚡wo ❯❯ $env:LLM_MODEL="deepseek-v4-flash"; $env:EMBEDDING_MODEL="Qwen3-Embedding-0.6B"; $env:LLM_API_BASE="https://api.deepseek.com/v1"; $env:LLM_API_KEY="sk-44f7ab82118c4767b44bf89f807c418b"; uv run python main.py --in res/RX100M3说明书.pdf --reference res/ref_cam.md --out result_cam.md
正在索引 res/RX100M3说明书.pdf ...
[INFO] 2026-08-23 15:26:07,927 [RapidOCR] base.py:23: Using engine_name: torch
[INFO] 2026-08-23 15:26:07,946 [RapidOCR] device_config.py:57: Using CPU device
[INFO] 2026-08-23 15:26:07,958 [RapidOCR] download_file.py:60: File exists and is valid: E:\work\working\py\RagFiller\.venv\Lib\site-packages\rapidocr\models\PP-OCRv6_det_small.pth
[INFO] 2026-08-23 15:26:07,958 [RapidOCR] main.py:50: Using E:\work\working\py\RagFiller\.venv\Lib\site-packages\rapidocr\models\PP-OCRv6_det_small.pth
[INFO] 2026-08-23 15:26:08,114 [RapidOCR] base.py:23: Using engine_name: torch
[INFO] 2026-08-23 15:26:08,114 [RapidOCR] device_config.py:57: Using CPU device
[INFO] 2026-08-23 15:26:08,118 [RapidOCR] download_file.py:60: File exists and is valid: E:\work\working\py\RagFiller\.venv\Lib\site-packages\rapidocr\models\ch_ptocr_mobile_v2.0_cls_mobile.pth
[INFO] 2026-08-23 15:26:08,118 [RapidOCR] main.py:50: Using E:\work\working\py\RagFiller\.venv\Lib\site-packages\rapidocr\models\ch_ptocr_mobile_v2.0_cls_mobile.pth
[INFO] 2026-08-23 15:26:08,191 [RapidOCR] base.py:23: Using engine_name: torch
[INFO] 2026-08-23 15:26:08,193 [RapidOCR] device_config.py:57: Using CPU device
[INFO] 2026-08-23 15:26:08,208 [RapidOCR] download_file.py:60: File exists and is valid: E:\work\working\py\RagFiller\.venv\Lib\site-packages\rapidocr\models\PP-OCRv6_rec_small.pth
[INFO] 2026-08-23 15:26:08,208 [RapidOCR] main.py:50: Using E:\work\working\py\RagFiller\.venv\Lib\site-packages\rapidocr\models\PP-OCRv6_rec_small.pth
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 770/770 [00:00<00:00, 6633.26it/s]
E:\work\working\py\RagFiller\.venv\Lib\site-packages\torch\nn\modules\conv.py:560: UserWarning: Using padding='same' with even kernel lengths and odd dilation may require a zero-padded copy of the input be created (Triggered internally at C:\actions-runner\_work\pytorch\pytorch\aten\src\ATen\native\Convolution.cpp:1102.)
return F.conv2d(
RapidOCR returned empty result!
RapidOCR returned empty result!
RapidOCR returned empty result!
[WARNING] 2026-08-23 15:28:36,526 [RapidOCR] main.py:132: The text detection result is empty
RapidOCR returned empty result!
RapidOCR returned empty result!
RapidOCR returned empty result!
RapidOCR returned empty result!
[WARNING] 2026-08-23 15:29:07,453 [RapidOCR] main.py:132: The text detection result is empty
RapidOCR returned empty result!
已将 375 个文本块索引到内存向量数据库
E:\work\working\py\RagFiller\.venv\Lib\site-packages\pydantic_settings\sources\utils.py:47: IncompleteFieldDefinitionWarning: Field 'lifespan' has an incomplete definition: its annotation contains an unresolved forward reference, so settings sources may fail to correctly resolve its value. Call `model_rebuild()` on the model where the field is defined, once all the referenced types are defined.
warnings.warn(
[08/23/26 15:31:19] INFO Processing request of type ListToolsRequest server.py:733
正在填充参考文档 ...
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:21] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
[08/23/26 15:31:22] INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:24] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:26] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:28] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:30] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:32] INFO Processing request of type CallToolRequest server.py:733
[08/23/26 15:31:33] INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:35] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:38] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:41] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
[08/23/26 15:31:42] INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:44] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:47] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:49] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:51] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
[08/23/26 15:31:52] INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:54] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:56] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:31:59] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:32:05] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:32:07] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
[08/23/26 15:32:08] INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
[08/23/26 15:32:14] INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO Processing request of type CallToolRequest server.py:733
INFO HTTP Request: POST http://127.0.0.1:7838/v1/embeddings "HTTP/1.1 200 OK" _client.py:1085
INFO HTTP Request: POST https://api.deepseek.com/v1/chat/completions "HTTP/1.1 200 OK" _client.py:1923
已写入 result_cam.md可以得到：
| **型号** | **发布时间** | **传感器类型** | **有效像素(MP)** | **传感器尺寸** | **镜头焦距(等效)** | **最大光圈** | **光学变焦** | **对焦系统** | **连拍速度(fps)** | **视频规格** | **取景器** | **屏幕** | **机身尺寸(mm)** | **重量(g)** | **电池型号** |
|----------|--------------|----------------|------------------|----------------|---------------------|--------------|--------------|--------------|-------------------|--------------|------------|----------|------------------|--------------|--------------|
| Sony Cyber-shot DSC-RX100M3 | | Exmor R™ CMOS | 约2010万（约20.1 MP） | 13.2 mm × 8.8 mm（1.0型） | f=8.8 mm 至 25.7 mm（24 mm 至 70 mm，35 mm 胶片等效） | F1.8（W）至 F2.8（T） | 2.9倍 | 对焦模式：单次AF / 连续AF / DMF / 手动对焦；对焦区域：广域 / 中间 / 自由点 / 锁定AF；支持眼控AF、中央锁定AF、AF辅助照明 | | 文件格式：XAVC S（50p 50M、25p 50M、100p 50M，100p 为 1280×720 高速记录，可记录 100 fps 动态影像）/ AVCHD（50i 24M(FX)、50i 17M(FH)、50p 28M(PS)、25p 24M(FX)、25p 17M(FH)）/ MP4（1440×1080 12M、VGA 3M）；兼容 1080 50i / 1080 50p | 电子取景器（有机EL），弹出式；总点数 1 440 000 点，取景率 100%，倍率约 0.59 倍，眼点约 20 mm | 7.5 cm（3.0 型）TFT 液晶；总点数 1 228 800 点；可向上旋转约 180 度 | 约 101.6 × 58.1 × 41.0（符合 CIPA 标准） | 约 290 g（包括电池 NP-BX1、Memory Stick PRO Duo） | NP-BX1 |
完整代码：
查看代码
"""Rag-Filler: fill a reference Markdown document from a PDF manual via RAG + an LLM agent.
Single-file CLI (managed by uv). Usage:
python main.py --in <manual.pdf> --reference <ref.md> --out <result.md>
Environment:
LLM_API_BASE default https://api.deepseek.com/v1
LLM_API_KEY optional
LLM_MODEL required
EMBEDDING_API_BASE default http://127.0.0.1:7838/v1
EMBEDDING_API_KEY optional
EMBEDDING_MODEL required
EMBEDDING_DIMENSION optional, default 1024
"""
from __future__ import annotations
import argparse
import asyncio
import os
import sys
import tempfile
from pathlib import Path
from docling.document_converter import DocumentConverter
from langchain.agents import create_agent
from langchain_core.documents import Document
from langchain_core.messages import HumanMessage
from langchain_core.runnables.config import RunnableConfig
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import (
MarkdownHeaderTextSplitter,
RecursiveCharacterTextSplitter,
)
from mcp.server.fastmcp import FastMCP
from mcp.shared.memory import create_connected_server_and_client_session
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams
COLLECTION = "manual"
CHUNK_SIZE = 800
CHUNK_OVERLAP = 120 # ~15% of CHUNK_SIZE
EMBED_BATCH = 32
DEFAULT_LLM_BASE = "https://api.deepseek.com/v1"
DEFAULT_EMBED_BASE = "http://127.0.0.1:7838/v1"
DEFAULT_EMBED_DIMENSION = 1024
SYSTEM_PROMPT = (
"你是一个文档填充助手：根据技术手册（PDF）中的事实，填充参考文档中留空或占位的内容。\n"
"规则：\n"
"1. 填写任何内容之前，必须先调用 search_manual 工具从手册中检索证据；必要时综合多条检索结果，"
"不要对同一事实重复检索。\n"
"2. 严禁编造或猜测。手册中没有证据的内容，保持与参考文档中的原样（留空或占位符不变）。\n"
"3. 数字和单位必须与手册原文完全一致。\n"
"4. 输出文档的最终形式由你自己根据参考文档与填充内容判断（表格、段落或其它结构均可），"
"但输出必须内容完整、可直接使用。\n"
"5. 只输出最终文档本身：不要将检索过程、解释或对话内容混入输出文档。"
)
class ConfigError(Exception):
"""Raised when required environment configuration is missing."""
class FatalError(Exception):
"""Raised for expected user-facing pipeline failures."""
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="main.py",
add_help=False,
description="Fill a reference Markdown document from a PDF manual via RAG + an LLM agent.",
)
parser.add_argument("--in", dest="pdf", required=True, metavar="PDF",
help="用来索引的PDF")
parser.add_argument("--reference", required=True, metavar="MD",
help="需要填充的markdown文件")
parser.add_argument("--out", required=True, metavar="MD",
help="输出的markdown文件")
return parser.parse_args(argv)
def load_config() -> dict:
required = ("LLM_MODEL", "EMBEDDING_MODEL")
missing = [name for name in required if not os.environ.get(name)]
if missing:
raise ConfigError(
"缺少必需的环境变量：" + ", ".join(missing)
)
return {
"llm_api_base": os.environ.get("LLM_API_BASE", DEFAULT_LLM_BASE),
"llm_api_key": os.environ.get("LLM_API_KEY", "not-needed"),
"llm_model": os.environ["LLM_MODEL"],
"embedding_api_base": os.environ.get("EMBEDDING_API_BASE", DEFAULT_EMBED_BASE),
"embedding_api_key": os.environ.get("EMBEDDING_API_KEY", "not-needed"),
"embedding_model": os.environ["EMBEDDING_MODEL"],
"embedding_dimension": int(os.environ.get("EMBEDDING_DIMENSION", DEFAULT_EMBED_DIMENSION)),
}
def chunk_markdown(md: str) -> list[Document]:
"""
将 Markdown 文本按标题层级（#、##、###）分块，再按长度切分为最终片段。
返回不包含空内容的 Document 列表。
"""
# 1. 按标题分割，保留层级元数据
header_splitter = MarkdownHeaderTextSplitter(
headers_to_split_on=[
("#", "section"),
("##", "subsection"),
("###", "subsubsection"),
]
)
header_docs = header_splitter.split_text(md)
# 2. 对已分割的文档进行长度切分
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
separators=["\n\n", "\n", " ", ""], # 去掉无关的 "## " 等
)
if header_docs:
chunks = splitter.split_documents(header_docs)
else:
# 无标题时直接创建文档（保持与 split_documents 返回类型一致）
chunks = splitter.create_documents([md])
# 3. 过滤空内容块
return [c for c in chunks if c.page_content.strip()]
def embed_chunks(chunks: list[Document], embed: OpenAIEmbeddings) -> list[list[float]]:
vectors: list[list[float]] = []
for start in range(0, len(chunks), EMBED_BATCH):
batch = [c.page_content for c in chunks[start : start + EMBED_BATCH]]
vectors.extend(embed.embed_documents(batch))
if len(vectors) != len(chunks):
raise FatalError("嵌入 API 返回的向量数量与文本块数量不一致")
return vectors
def index_pdf(pdf_path: str, embed: OpenAIEmbeddings, qdrant: QdrantClient,
configured_dim: int) -> int:
try:
result = DocumentConverter().convert(str(pdf_path))
md = result.document.export_to_markdown()
except Exception as exc:
raise FatalError(f"无法解析 PDF {pdf_path}：{exc}") from exc
if not md.strip():
raise FatalError(f"无法解析 PDF {pdf_path}：未提取到任何内容")
chunks = chunk_markdown(md)
if not chunks:
raise FatalError(f"无法解析 PDF {pdf_path}：分块后无内容")
vectors = embed_chunks(chunks, embed)
actual_dim = len(vectors[0])
if actual_dim != configured_dim:
print(f"注意：嵌入模型输出维度为 {actual_dim}，EMBEDDING_DIMENSION={configured_dim} 已忽略",
file=sys.stderr)
qdrant.create_collection(
collection_name=COLLECTION,
vectors_config=VectorParams(size=actual_dim, distance=Distance.COSINE),
)
points = [
PointStruct(
id=idx,
vector=vectors[idx],
payload={
"text": chunks[idx].page_content,
"section": chunks[idx].metadata.get("section", ""),
"source": str(pdf_path),
},
)
for idx in range(len(chunks))
]
qdrant.upsert(collection_name=COLLECTION, points=points)
return len(chunks)
def make_search_server(qdrant: QdrantClient,
embed: OpenAIEmbeddings) -> FastMCP:
server = FastMCP("manual-search")
@server.tool()
def search_manual(query: str, top_k: int = 5) -> list[dict]:
"""在已索引的手册中检索与 query 最相关的段落并返回（含相关度分数）。填写表格单元格之前，请调用此工具从手册中查找事实依据。"""
k = max(1, min(int(top_k), 20))
# PyCharm 解析到旧版 qdrant-client 签名会误报 query 参数类型;1.19 运行时接受 list[float]
vector = embed.embed_query(query)
hits = qdrant.query_points(collection_name=COLLECTION, query=vector, limit=k).points # type: ignore[arg-type]
return [
{
"text": (hit.payload or {}).get("text", ""),
"section": (hit.payload or {}).get("section", ""),
"score": round(float(hit.score), 4),
}
for hit in hits
]
return server
def unwrap_fence(text: str) -> str:
stripped = text.strip()
if stripped.startswith("```"):
body = stripped.split("\n", 1)[1] if "\n" in stripped else ""
if body.endswith("```"):
body = body[:-3]
return body.strip() + "\n"
return text
def extract_content(message_content: str | list | dict) -> str:
if isinstance(message_content, list):
message_content = "".join(
part.get("text", "") for part in message_content if isinstance(part, dict)
)
elif isinstance(message_content, dict):
message_content = message_content.get("text", "")
return unwrap_fence(message_content)
def write_output(path: Path, content: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
fd, tmp = tempfile.mkstemp(prefix=".rag-filler-", suffix=".tmp", dir=str(path.parent))
try:
with os.fdopen(fd, "w", encoding="utf-8", newline="") as fh:
fh.write(content)
os.replace(tmp, path)
except BaseException:
try:
os.unlink(tmp)
except OSError:
pass
raise
async def run(args: argparse.Namespace, cfg: dict) -> int:
embed = OpenAIEmbeddings(
model=cfg["embedding_model"],
openai_api_base=cfg["embedding_api_base"],
openai_api_key=cfg["embedding_api_key"],
check_embedding_ctx_length=False,
show_progress_bar=False,
)
reference = Path(args.reference)
ref_text = reference.read_text(encoding="utf-8")
print(f"正在索引 {args.pdf} ...")
qdrant = QdrantClient(":memory:")
count = index_pdf(args.pdf, embed, qdrant, cfg["embedding_dimension"])
print(f"已将 {count} 个文本块索引到内存向量数据库")
llm = ChatOpenAI(
model=cfg["llm_model"],
base_url=cfg["llm_api_base"],
api_key=cfg["llm_api_key"],
temperature=0,
)
server = make_search_server(qdrant, embed)
try:
async with create_connected_server_and_client_session(server) as session:
tools = await load_mcp_tools(session)
if not any(getattr(tool, "name", None) == "search_manual" for tool in tools):
raise FatalError("向量检索 MCP 工具 'search_manual' 未绑定到 agent")
agent = create_agent(llm, tools, system_prompt=SYSTEM_PROMPT)
print("正在填充参考文档 ...")
response = await agent.ainvoke(
{"messages": [HumanMessage(content=(
"下面是需要填充的参考文档。请依据手册填充其中留空或占位的内容；"
"手册内容一律通过 search_manual 检索获取，不要把 PDF 全文当作已知信息。"
"输出文档的最终形式由你判断，但必须内容完整、可直接使用。\n\n"
"参考文档：\n```markdown\n" + ref_text + "\n```"
))]},
config=RunnableConfig(recursion_limit=120),
)
filled = extract_content(response["messages"][-1].content)
except FatalError:
raise
except Exception as exc:
raise FatalError(f"API 调用失败：{format_api_error(exc)}") from exc
write_output(Path(args.out), filled)
print(f"已写入 {args.out}")
return 0
def format_api_error(exc: BaseException) -> str:
"""Flatten asyncio exception groups so the real cause (e.g. 401) surfaces."""
if isinstance(exc, BaseExceptionGroup):
parts = [format_api_error(e) for e in exc.exceptions]
return "; ".join(dict.fromkeys(parts)) or str(exc)
return str(exc) or type(exc).__name__
def main(argv: list[str] | None = None) -> int:
args = parse_args(argv)
try:
cfg = load_config()
except ConfigError as exc:
print(f"错误：{exc}", file=sys.stderr)
return 2
try:
return asyncio.run(run(args, cfg))
except FatalError as exc:
print(f"错误：{exc}", file=sys.stderr)
return 1
except Exception as exc: # unexpected failure: keep output file untouched
print(f"错误：{exc}", file=sys.stderr)
return 1
if __name__ == "__main__":
sys.exit(main())
浙公网安备 33010602011771号