478 lines
20 KiB
Python
478 lines
20 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="分片返回文本")
|
||
|
||
|
||
REMOTE_DIRECTORY_PATH = (
|
||
"/root/obsidian/vaults/MyNewVault/RawsSources" # Obsidian云端库文件夹路径
|
||
)
|
||
LOCAL_DIRECTORY_PATH = Path(__file__).parent / "vault" # 本地库文件夹路径
|
||
|
||
|
||
CHUNK_SIZE = 8192 # 分片大小，BAAI/BGE-M3标值
|
||
|
||
TPM = 500_000 * 0.8 # 每分钟词元数限额，硅基流动标值，安全冗余20%
|
||
WINDOW_PERIOD = 60 # 窗口期
|
||
|
||
|
||
def sync_to_local_vault() -> None:
|
||
"""
|
||
同步Obsidian云端库至本地库
|
||
"""
|
||
print("===== 开始同步Obsidian云端库至本地库 =====")
|
||
# 创建SSHClient
|
||
ssh_client = SSHClient()
|
||
ssh_client.set_missing_host_key_policy(AutoAddPolicy())
|
||
ssh_client.connect(
|
||
hostname="www.liubiren.cloud",
|
||
port=22,
|
||
username="root",
|
||
password="Te198752.",
|
||
)
|
||
# 创建SFTP会话
|
||
sftp = ssh_client.open_sftp()
|
||
|
||
def recursive_scan_remote_vault(directory_path: str) -> list[dict]:
|
||
"""
|
||
递归扫描云端库指定文件夹路径，返回该文件夹内所有MD文档元数据
|
||
:param directory_path: 云端库指定文件夹路径
|
||
:return: [{"remote_path":str, "modified_at":float}, ... ]
|
||
"""
|
||
metadatas = [] # 所有MD文档元数据列表
|
||
# 列出云端库指定文件夹下所有文件和文件夹
|
||
entries = sftp.listdir_attr(directory_path)
|
||
for entry in entries:
|
||
# 若实体模式掩码为空则跳过
|
||
if (entry_mode_mask := entry.st_mode) is None:
|
||
continue
|
||
|
||
# 构建实体云端库路径
|
||
remote_path = f"{directory_path}/{entry.filename}"
|
||
# 若实体为文件夹则递归扫描该文件夹路径，并将该文件夹内所有MD文档元数据添加至所有MD文档元数据列表
|
||
if entry_mode_mask & 0o040000:
|
||
metadatas.extend(recursive_scan_remote_vault(remote_path))
|
||
|
||
# 若实体文件为MD文件则构建文档元数据并添加到所有MD文档元数据列表
|
||
elif entry.filename.lower().endswith(".md"):
|
||
metadatas.append(
|
||
{"remote_path": remote_path, "modified_at": entry.st_mtime or 0}
|
||
)
|
||
return metadatas
|
||
|
||
print(f"🔍 正在扫描云端库...")
|
||
# 递归扫描云端库指定文件夹路径，返回该文件夹内所有MD文档元数据
|
||
metadatas = recursive_scan_remote_vault(REMOTE_DIRECTORY_PATH)
|
||
# 统计文档数
|
||
remote_document_count = len(metadatas)
|
||
print(f"✅ 扫描完成，云端库共有MD文档 {remote_document_count} 份")
|
||
|
||
print(f"🧹 正在清理MD文档...")
|
||
local_paths = list(LOCAL_DIRECTORY_PATH.rglob("*.md"))
|
||
deleted_document_count = 0 # 删除文档数
|
||
for local_path in local_paths:
|
||
# 构建文档云端库路径
|
||
remote_path = (
|
||
Path(REMOTE_DIRECTORY_PATH) / local_path.relative_to(LOCAL_DIRECTORY_PATH)
|
||
).as_posix()
|
||
if remote_path not in [metadata["remote_path"] for metadata in metadatas]:
|
||
local_path.unlink()
|
||
deleted_document_count += 1
|
||
|
||
print(f"✅ 清理完成，本次共清理MD文档 {deleted_document_count} 份")
|
||
|
||
print(f"📥 正在下载MD文档...")
|
||
download_document_count = 0 # 下载文档数
|
||
for metadata in metadatas:
|
||
# 构建文档本地库路径
|
||
local_path: Path = LOCAL_DIRECTORY_PATH / Path(
|
||
metadata["remote_path"]
|
||
).relative_to(REMOTE_DIRECTORY_PATH)
|
||
# 检查文档本地库路径是否存在，若不存在则创建
|
||
local_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
is_download_needed = True # 需要下载
|
||
# 若文档本地库路径存在且修改时间戳差异小于阈值则无需下载
|
||
if local_path.exists():
|
||
if abs(local_path.stat().st_mtime - metadata["modified_at"]) < 60:
|
||
is_download_needed = False
|
||
|
||
if is_download_needed:
|
||
# 下载文档
|
||
sftp.get(metadata["remote_path"], local_path.as_posix())
|
||
print(f" > 📥 {local_path.stem} 下载完成")
|
||
# 修改文档本地库路径时间戳
|
||
utime(local_path, times=(metadata["modified_at"], metadata["modified_at"]))
|
||
download_document_count += 1
|
||
|
||
print(f"✅ 下载完成，本次共下载MD文档 {download_document_count} 份")
|
||
sftp.close()
|
||
ssh_client.close()
|
||
print("===== 同步Obsidian云端库至本地库完成 =====\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
|
||
# 统计记录数
|
||
record_count = len(document_sync_records)
|
||
print(f"✅ 加载完成，数据库共有MD文档 {record_count} 份")
|
||
|
||
# 列出本地库所有MD文档路径
|
||
local_paths = list(LOCAL_DIRECTORY_PATH.rglob("*.md"))
|
||
# 统计文档数
|
||
document_count = len(local_paths)
|
||
document_ids = set()
|
||
for idx, local_path in enumerate(local_paths, 1):
|
||
print(f"🧩 [{idx}/{document_count}] | 正在切片 {local_path.stem} ...")
|
||
# 读取文档
|
||
document = load(local_path)
|
||
# 解析文档唯一标识（需Obsidian已安装NoteUIDGenerator插件，用于生成唯一标识）
|
||
document_id = str(document.metadata["uid"])
|
||
document_ids.add(document_id)
|
||
# 文档本地库路径的修改时间戳
|
||
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"
|
||
|
||
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 = []
|
||
if h1 := chunk.metadata.get("h1"):
|
||
headers.append(h1.strip())
|
||
# 若h1标题为空则使用文档名称
|
||
else:
|
||
headers.append(local_path.stem.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)
|
||
# 统计批次数
|
||
batch_count = len(batches)
|
||
print(f"✅ 切片完成，共切片 {batch_count} 批次、{index} 片")
|
||
|
||
# 按批次请求计算文本向量，若某批次请求发生异常则跳过该文档
|
||
try:
|
||
print(f"🔢 [{idx}/{document_count}] | 正在向量化 {local_path.stem} ...")
|
||
with connection.cursor() as cursor:
|
||
# 若操作类型为更新则删除该文档所有分片
|
||
if operation_type == "update":
|
||
cursor.execute(
|
||
"DELETE FROM document_chunks WHERE document_id = %s;",
|
||
(document_id,),
|
||
)
|
||
|
||
for batch in batches:
|
||
# 统计批次的文本词元数
|
||
batch_text_tokens = sum(chunk.text_tokens for chunk in batch)
|
||
# 若窗口队列的文本词元数加上批次的文本词元数超过阈值则等待窗口队列的最早请求时间加上窗口期
|
||
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_count = 0 # 请求次数
|
||
response_payload: dict | None = None
|
||
while request_count < 3:
|
||
request_count += 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_count)
|
||
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_count)
|
||
|
||
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
|
||
],
|
||
)
|
||
|
||
# 创建文档同步记录
|
||
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()
|
||
print(f"✅ 向量化完成并已同步至数据库")
|
||
|
||
except Exception:
|
||
connection.rollback()
|
||
continue
|
||
|
||
# 若待删除文档唯一标识列表非空则删除相应文档同步记录和文档分片记录
|
||
if pending_delete_document_ids := list(
|
||
set(document_sync_records.keys()) - document_ids
|
||
):
|
||
with connection.cursor() as cursor:
|
||
cursor.execute(
|
||
"""
|
||
WITH deleted_documents AS (
|
||
DELETE FROM document_sync_records
|
||
WHERE id = ANY(%s)
|
||
RETURNING id
|
||
)
|
||
DELETE FROM document_chunks
|
||
WHERE document_id IN (SELECT id FROM deleted_documents);
|
||
""",
|
||
(pending_delete_document_ids,),
|
||
)
|
||
connection.commit()
|
||
|
||
print("===== 处理文档完成 =====\n")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 同步Obsidian云端库至本地库
|
||
sync_to_local_vault()
|
||
# 处理文档
|
||
process_documents()
|