444 lines
18 KiB
Python
444 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
主模块
|
||
"""
|
||
from collections import deque
|
||
from hashlib import blake2s
|
||
from os import utime
|
||
from pathlib import Path
|
||
from re import MULTILINE, sub
|
||
from time import sleep, time
|
||
|
||
from frontmatter import load
|
||
from langchain_text_splitters import (
|
||
MarkdownHeaderTextSplitter,
|
||
RecursiveCharacterTextSplitter,
|
||
)
|
||
from paramiko import AutoAddPolicy, SSHClient
|
||
from psycopg import connect
|
||
from pydantic import BaseModel, Field
|
||
from requests import Response, exceptions, post
|
||
from transformers import AutoTokenizer
|
||
|
||
|
||
class Chunk(BaseModel):
|
||
"""
|
||
分块数据模型
|
||
"""
|
||
|
||
document_id: str = Field(..., description="文档唯一标识")
|
||
index: int = Field(..., description="分块索引")
|
||
text: str = Field(..., description="分块文本")
|
||
text_vector: list[float] | None = Field(default=None, description="分块文本向量")
|
||
text_fingerprint: str = Field(..., description="分块文本文本指纹")
|
||
text_tokens: int = Field(..., description="分块文本词元数")
|
||
return_text: str = Field(..., description="返回文本")
|
||
|
||
|
||
LOCAL_DIRECTORY = Path(__file__).parent / "vault" # 本地库文件夹
|
||
REMOTE_DIRECTORY = "/root/obsidian/vaults/MyNewVault/RawsSources" # 云端库文件夹
|
||
|
||
CHUNK_SIZE = 8192 # 分块大小，BAAI/BGE-M3标值
|
||
|
||
TPM = 500_000 * 0.8 # 每分钟词元数限额，硅基流动标值，安全冗余20%
|
||
WINDOW_PERIOD = 60 # 窗口期
|
||
|
||
|
||
def sync_vault() -> None:
|
||
"""
|
||
同步本地库和云端库
|
||
"""
|
||
print("===== 开始同步本地库和云端库 =====")
|
||
# 创建 SSHClient
|
||
ssh = SSHClient()
|
||
ssh.set_missing_host_key_policy(AutoAddPolicy())
|
||
ssh.connect(
|
||
hostname="www.liubiren.cloud",
|
||
port=22,
|
||
username="root",
|
||
password="Te198752.",
|
||
)
|
||
# 创建 SFTP
|
||
sftp = ssh.open_sftp()
|
||
|
||
def get_remote_vault_metadata(folder: str) -> list[dict]:
|
||
"""
|
||
获取云端库指定文件夹下所有元数据（文档即MD文件，每个文档包含云端路径和修改时间戳）
|
||
:param folder: 文件夹
|
||
:return: [ {"remote_path":str, "modified_at":float}, ... ]
|
||
"""
|
||
metadata = []
|
||
# 获取指定文件夹下所有实体列表
|
||
entries = sftp.listdir_attr(folder)
|
||
for entry in entries:
|
||
# 若实体模式掩码为空则跳过
|
||
if (entry_mode_mask := entry.st_mode) is None:
|
||
continue
|
||
|
||
# 构建云端路径
|
||
remote_path = f"{folder}/{entry.filename}"
|
||
# 若实体为文件夹则递归获取该文件夹下所有元数据
|
||
if entry_mode_mask & 0o040000:
|
||
metadata.extend(get_remote_vault_metadata(remote_path))
|
||
|
||
# 若实体文件为MD文件则构建文档源数据并添加到列表
|
||
elif entry.filename.lower().endswith(".md"):
|
||
metadata.append(
|
||
{"remote_path": remote_path, "modified_at": entry.st_mtime or 0}
|
||
)
|
||
return metadata
|
||
|
||
# 获取云端库指定文件夹下所有元数据
|
||
remote_vault_metadata = get_remote_vault_metadata(REMOTE_DIRECTORY)
|
||
print(f"📁 云端库扫描完成，文档数：{len(remote_vault_metadata)}")
|
||
|
||
download_count = 0
|
||
for metadata in remote_vault_metadata:
|
||
# 本地路径
|
||
local_path: Path = LOCAL_DIRECTORY / Path(metadata["remote_path"]).relative_to(
|
||
REMOTE_DIRECTORY
|
||
)
|
||
local_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
is_download_needed = True # 需要下载
|
||
modified_at = metadata["modified_at"]
|
||
# 若本地文档存在且修改时间戳与云端文档修改时间戳差异小于阈值则无需下载
|
||
if local_path.exists():
|
||
if abs(local_path.stat().st_mtime - modified_at) < 60:
|
||
is_download_needed = False
|
||
|
||
if is_download_needed:
|
||
print(f"⬇️ 下载文档：{metadata['remote_path']}")
|
||
# 下载文档
|
||
sftp.get(metadata["remote_path"], local_path.as_posix())
|
||
# 同步修改时间戳
|
||
utime(local_path, times=(modified_at, modified_at))
|
||
download_count += 1
|
||
|
||
print(f"✅ 同步完成，下载文档数：{download_count}")
|
||
sftp.close()
|
||
ssh.close()
|
||
print("===== 本地库和云端库同步完成 =====\n")
|
||
|
||
|
||
def process_documents() -> None:
|
||
"""
|
||
处理文档
|
||
"""
|
||
print("===== 开始文档处理 =====")
|
||
|
||
# 实例化标题分割器
|
||
header_splitter = MarkdownHeaderTextSplitter(
|
||
headers_to_split_on=[
|
||
("#", "h1"),
|
||
("##", "h2"),
|
||
("###", "h3"),
|
||
("####", "h4"),
|
||
("#####", "h5"),
|
||
("######", "h6"),
|
||
],
|
||
strip_headers=True, # 不保留标题
|
||
)
|
||
|
||
# 实例化分词器（已保存预训练模型，本地加载）
|
||
tokenizer = AutoTokenizer.from_pretrained(
|
||
Path(__file__).parent / "baai_bge_m3_tokenizer"
|
||
)
|
||
|
||
def count_tokens(text: str) -> int:
|
||
"""
|
||
统计词元数
|
||
:param text: 文本
|
||
:return: token 数量
|
||
"""
|
||
return len(
|
||
tokenizer.encode(text, add_special_tokens=False)
|
||
) # 首尾不添加模型专用特殊标记
|
||
|
||
# 实例化窗口队列（用于保存窗口期内请求）
|
||
window_deque: deque[tuple[float, int]] = deque()
|
||
|
||
def get_window_deque_text_tokens() -> int:
|
||
"""
|
||
获取窗口队列的文本词元数
|
||
:return: 文本词元数
|
||
"""
|
||
# 当前时间戳
|
||
current_time = time()
|
||
# 清理过期请求
|
||
while window_deque and window_deque[0][0] < current_time - WINDOW_PERIOD:
|
||
window_deque.popleft()
|
||
return sum(text_tokens for _, text_tokens in window_deque)
|
||
|
||
with connect(
|
||
host="www.liubiren.cloud",
|
||
port=5201,
|
||
user="root",
|
||
password="PG198752",
|
||
dbname="database",
|
||
) as connection:
|
||
print("📖 加载文档同步记录...")
|
||
document_sync_records: dict[str, float] = {} # 文档同步记录字典
|
||
with connection.cursor() as cursor:
|
||
cursor.execute("SELECT id, modified_at FROM document_sync_records;")
|
||
for id, modified_at in cursor.fetchall():
|
||
document_sync_records[id] = modified_at
|
||
print(f"✅ 加载完成，数据库中文档数：{len(document_sync_records)}")
|
||
|
||
for idx, local_path in enumerate(
|
||
local_paths := list(LOCAL_DIRECTORY.rglob("*.md")), 1
|
||
):
|
||
# 文档标题
|
||
document_title = local_path.stem
|
||
print(f"📄 [{idx}/{len(local_paths)}] 正在处理文档：{document_title}\n")
|
||
# 读取文档
|
||
document = load(local_path)
|
||
# 解析文档唯一标识（需Obsidian已安装NoteUIDGenerator插件，用于生成唯一标识）
|
||
document_id = str(document.metadata["uid"])
|
||
# 本地文档修改时间戳
|
||
modified_at = local_path.stat().st_mtime
|
||
|
||
operation_type = "create" # 操作类型
|
||
if document_id in document_sync_records:
|
||
# 若本地文档修改时间戳与数据库文档修改时间戳差异小于阈值则跳过
|
||
if abs(modified_at - document_sync_records[document_id]) < 60:
|
||
print(f"⏭️ 该文档未变动，无需处理")
|
||
continue
|
||
|
||
operation_type = "update"
|
||
print(f"🔄 该文档需要更新")
|
||
|
||
else:
|
||
print(f"🔄 该文档需要创建")
|
||
|
||
index = 0
|
||
current_batch: list[Chunk] = [] # 当前批次列表
|
||
current_batch_text_tokens: int = 0 # 当前批次列表的文本词元数
|
||
batches: list[list[Chunk]] = [] # 批次列表
|
||
# 先根据标题分割文档，若分块词元大小超过阈值则使用递归字符分割器进一步分割
|
||
for chunk in header_splitter.split_text(document.content):
|
||
# --- 构建分块标题 ---
|
||
|
||
headers = []
|
||
# 若h1标题为空则使用文档标题
|
||
if not (h1 := chunk.metadata.get("h1")):
|
||
headers.append(document_title.strip())
|
||
else:
|
||
headers.append(h1.strip())
|
||
if h2 := chunk.metadata.get("h2"):
|
||
headers.append(h2.strip())
|
||
if h3 := chunk.metadata.get("h3"):
|
||
headers.append(h3.strip())
|
||
if h4 := chunk.metadata.get("h4"):
|
||
headers.append(h4.strip())
|
||
if h5 := chunk.metadata.get("h5"):
|
||
headers.append(h5.strip())
|
||
if h6 := chunk.metadata.get("h6"):
|
||
headers.append(h6.strip())
|
||
header = f"标题层级：{" > ".join(headers)}\n\n"
|
||
|
||
# --- 构建分块内容 ---
|
||
# 若分块内容为空则跳过
|
||
if not (content := chunk.page_content.strip()):
|
||
continue
|
||
|
||
# 删除行尾空格
|
||
content = sub(r"[ \t]+$", "", content, flags=MULTILINE)
|
||
# 将全角空格替换为半角空格
|
||
content = content.replace(" ", " ")
|
||
# 压缩空格
|
||
content = sub(r" +", " ", content)
|
||
# 压缩换行符
|
||
content = sub(r"\n{3,}", "\n\n", content)
|
||
# 删除分割线标记
|
||
content = content.replace("---", "")
|
||
|
||
# 返回文本
|
||
return_text = header + content
|
||
|
||
# 删除斜体标记
|
||
content = sub(r"\*(.*?)\*", r"\1", content)
|
||
# 删除加粗标记
|
||
content = sub(r"\*\*(.*?)\*\*", r"\1", content)
|
||
# 删除代码块标记
|
||
content = sub(r"`(.*?)`", r"\1", content)
|
||
# 删除文档链接标记
|
||
content = sub(r"\[\[(.*?)\]\]", r"\1", content)
|
||
# 删除图片标记（直接删除）
|
||
content = sub(r"!\[.*?\]\(.*?\)", "", content)
|
||
# 删除普通连接标记（删除URL，保留文本）
|
||
content = sub(r"\[(.*?)\]\(.*?\)", r"\1", content)
|
||
|
||
# 若文本词元数未超过限制则以分块内容作为分块内容列表
|
||
if count_tokens(header + content) < CHUNK_SIZE:
|
||
contents = [content]
|
||
|
||
# 若文本词元数超过限制则使用递归字符分割器进一步分割分块内容
|
||
else:
|
||
# 实例化递归字符分割器
|
||
recursive_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=CHUNK_SIZE - count_tokens(header),
|
||
chunk_overlap=60, # 分块重叠
|
||
length_function=count_tokens,
|
||
separators=["\n\n", "\n", "。", "？", "！", ".", "?", "!", ""],
|
||
)
|
||
contents = recursive_splitter.split_text(content)
|
||
|
||
for content in contents:
|
||
# 若内容为空则跳过
|
||
if not content.strip():
|
||
continue
|
||
|
||
chunk = Chunk(
|
||
document_id=document_id,
|
||
index=index,
|
||
text=(text := header + content),
|
||
text_fingerprint=blake2s(
|
||
text.encode(), digest_size=16
|
||
).hexdigest(),
|
||
text_tokens=count_tokens(text),
|
||
return_text=return_text,
|
||
)
|
||
# 若当前批次列表的文本词元数加上当前分块的文本词元数超过阈值则将当前批次列表添加至批次列表
|
||
if current_batch_text_tokens + chunk.text_tokens > TPM:
|
||
batches.append(current_batch)
|
||
current_batch = [chunk]
|
||
current_batch_text_tokens = chunk.text_tokens
|
||
|
||
else:
|
||
current_batch.append(chunk)
|
||
current_batch_text_tokens += chunk.text_tokens
|
||
|
||
index += 1
|
||
|
||
# 若当前批次列表不为空则将当前批次列表添加至批次列表
|
||
if current_batch:
|
||
batches.append(current_batch)
|
||
print(f"📝 该文档切块完成，分块数：{index}")
|
||
|
||
# 按批次请求计算文本向量，若某批次请求发生异常则跳过该文档
|
||
try:
|
||
with connection.cursor() as cursor:
|
||
# 更新场景：先删除当前文档全部旧chunk
|
||
if operation_type == "update":
|
||
cursor.execute(
|
||
"DELETE FROM document_chunks WHERE document_id = %s;",
|
||
(document_id,),
|
||
)
|
||
print(f"🗑️ 删除原有分块")
|
||
|
||
for i, batch in enumerate(batches, 1):
|
||
# 统计批次的文本词元数
|
||
batch_text_tokens = sum(chunk.text_tokens for chunk in batch)
|
||
print(f" > 📦 Batch[{i}/{len(batches)}] 正在处理批次分块")
|
||
# 若窗口队列的文本词元数加上批次的文本词元数超过阈值则等待窗口队列的最早请求时间加上窗口期
|
||
while True:
|
||
# 获取窗口队列的文本词元数
|
||
window_deque_text_tokens = get_window_deque_text_tokens()
|
||
if window_deque_text_tokens + batch_text_tokens < TPM:
|
||
break
|
||
|
||
sleep(
|
||
max((window_deque[0][0] + WINDOW_PERIOD) - time(), 0.2)
|
||
)
|
||
|
||
request_counts = 0 # 请求次数
|
||
response_payload: dict | None = None
|
||
while request_counts < 3:
|
||
request_counts += 1
|
||
response: Response | None = None
|
||
try:
|
||
response = post(
|
||
"https://api.siliconflow.cn/v1/embeddings",
|
||
headers={
|
||
"Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json={
|
||
"model": "BAAI/bge-m3",
|
||
"input": [chunk.text for chunk in batch],
|
||
},
|
||
timeout=60,
|
||
)
|
||
response.raise_for_status()
|
||
response_payload = response.json()
|
||
# 请求成功则将本次请求添加至窗口队列
|
||
window_deque.append((time(), batch_text_tokens))
|
||
break
|
||
|
||
except exceptions.HTTPError as exception:
|
||
# 触发限速
|
||
if response is not None and response.status_code == 429:
|
||
sleep(2**request_counts)
|
||
continue
|
||
|
||
raise RuntimeError(
|
||
f"HTTPError: {exception}, resp={response.text if response is not None else 'response is None'}"
|
||
)
|
||
|
||
except Exception as exception:
|
||
print(f"Exception: {exception}")
|
||
sleep(2**request_counts)
|
||
|
||
if response_payload is None:
|
||
raise RuntimeError("response_payload is None")
|
||
|
||
# 写入文本向量
|
||
for chunk, text_vector in zip(
|
||
batch,
|
||
[item["embedding"] for item in response_payload["data"]],
|
||
):
|
||
chunk.text_vector = text_vector
|
||
|
||
# 批量创建文档分块记录
|
||
cursor.executemany(
|
||
"""
|
||
INSERT INTO document_chunks
|
||
(document_id, index, text, text_vector, text_fingerprint, text_tokens, return_text)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
||
""",
|
||
[
|
||
(
|
||
chunk.document_id,
|
||
chunk.index,
|
||
chunk.text,
|
||
chunk.text_vector,
|
||
chunk.text_fingerprint,
|
||
chunk.text_tokens,
|
||
chunk.return_text,
|
||
)
|
||
for chunk in batch
|
||
],
|
||
)
|
||
print(f" ✅ 该批次分块处理完成")
|
||
# 创建文档同步记录
|
||
cursor.execute(
|
||
"""
|
||
INSERT INTO document_sync_records (id, modified_at)
|
||
VALUES (%s, %s)
|
||
ON CONFLICT (id) DO UPDATE
|
||
SET modified_at = EXCLUDED.modified_at;
|
||
""",
|
||
(document_id, modified_at),
|
||
)
|
||
connection.commit()
|
||
# 更新内存缓存
|
||
document_sync_records[document_id] = modified_at
|
||
print(f"✅ 该文档处理完成\n")
|
||
|
||
except Exception:
|
||
connection.rollback()
|
||
continue
|
||
|
||
print("===== 全部文档处理完成 =====")
|
||
|
||
|
||
def main():
|
||
# 同步本地库和云端库
|
||
sync_vault()
|
||
# 处理文档
|
||
process_documents()
|
||
|
||
|
||
main()
|