liubiren
/
Python
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Activity

1

This commit is contained in:
liubiren 2026-09-30 16:12:04 +08:00
parent 6f9f68af44
commit ffdcee2953
5 changed files with 140 additions and 110 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
agent/document/vault/HelloWorld.md
View File

@ -1,4 +0,0 @@
---
uid: 58ea90c4-59e1-4796-be8e-838928eecc21
---
20260901开始使用Obsidian

0
agent/document/baai_bge_m3_tokenizer/tokenizer.json → agent/document_preprocess/baai_bge_m3_tokenizer/tokenizer.json
View File

0
agent/document/baai_bge_m3_tokenizer/tokenizer_config.json → agent/document_preprocess/baai_bge_m3_tokenizer/tokenizer_config.json
View File

246
agent/document/main.py → agent/document_preprocess/main.py
View File

@ -23,111 +23,132 @@ from transformers import AutoTokenizer
class Chunk(BaseModel): class Chunk(BaseModel):
""" """
分块数据模型 分片数据模型
""" """
document_id: str = Field(..., description="文档唯一标识") document_id: str = Field(..., description="文档唯一标识")
index: int = Field(..., description="分块索引") index: int = Field(..., description="分片索引")
text: str = Field(..., description="分块文本") text: str = Field(..., description="分片文本")
text_vector: list[float] | None = Field(default=None, description="分块文本向量") text_vector: list[float] | None = Field(default=None, description="分片文本向量")
text_fingerprint: str = Field(..., description="分块文本文本指纹") text_fingerprint: str = Field(..., description="分片文本文本指纹")
text_tokens: int = Field(..., description="分块文本词元数") text_tokens: int = Field(..., description="分片文本词元数")
return_text: str = Field(..., description="返回文本") return_text: str = Field(..., description="分片返回文本")
LOCAL_DIRECTORY = Path(__file__).parent / "vault" # 本地库文件夹 REMOTE_DIRECTORY_PATH = (
REMOTE_DIRECTORY = "/root/obsidian/vaults/MyNewVault/RawsSources" # 云端库文件夹 "/root/obsidian/vaults/MyNewVault/RawsSources" # Obsidian云端库文件夹路径
)
LOCAL_DIRECTORY_PATH = Path(__file__).parent / "vault" # 本地库文件夹路径
CHUNK_SIZE = 8192 # 分块大小，BAAI/BGE-M3标值
CHUNK_SIZE = 8192 # 分片大小，BAAI/BGE-M3标值
TPM = 500_000 * 0.8 # 每分钟词元数限额，硅基流动标值，安全冗余20% TPM = 500_000 * 0.8 # 每分钟词元数限额，硅基流动标值，安全冗余20%
WINDOW_PERIOD = 60 # 窗口期 WINDOW_PERIOD = 60 # 窗口期
def sync_vault() -> None: def sync_to_local_vault() -> None:
""" """
同步本地库和云端库 同步Obsidian云端库至本地库
""" """
print("===== 开始同步本地库和云端库 =====") print("===== 开始同步Obsidian云端库至本地库 =====")
# 创建 SSHClient # 创建SSHClient
ssh = SSHClient() ssh_client = SSHClient()
ssh.set_missing_host_key_policy(AutoAddPolicy()) ssh_client.set_missing_host_key_policy(AutoAddPolicy())
ssh.connect( ssh_client.connect(
hostname="www.liubiren.cloud", hostname="www.liubiren.cloud",
port=22, port=22,
username="root", username="root",
password="Te198752.", password="Te198752.",
) )
# 创建 SFTP # 创建SFTP会话
sftp = ssh.open_sftp() sftp = ssh_client.open_sftp()
def get_remote_vault_metadata(folder: str) -> list[dict]: def recursive_scan_remote_vault(directory_path: str) -> list[dict]:
""" """
获取云端库指定文件夹下所有元数据（文档即MD文件，每个文档包含云端路径和修改时间戳） 递归扫描云端库指定文件夹路径，返回该文件夹内所有MD文档元数据
:param folder: 文件夹 :param directory_path: 云端库指定文件夹路径
:return: [ {"remote_path":str, "modified_at":float}, ... ] :return: [{"remote_path":str, "modified_at":float}, ... ]
""" """
metadata = [] metadatas = [] # 所有MD文档元数据列表
# 获取指定文件夹下所有实体列表 # 列出云端库指定文件夹下所有文件和文件夹
entries = sftp.listdir_attr(folder) entries = sftp.listdir_attr(directory_path)
for entry in entries: for entry in entries:
# 若实体模式掩码为空则跳过 # 若实体模式掩码为空则跳过
if (entry_mode_mask := entry.st_mode) is None: if (entry_mode_mask := entry.st_mode) is None:
continue continue
# 构建云端路径 # 构建实体云端库路径
remote_path = f"{folder}/{entry.filename}" remote_path = f"{directory_path}/{entry.filename}"
# 若实体为文件夹则递归获取该文件夹下所有元数据 # 若实体为文件夹则递归扫描该文件夹路径，并将该文件夹内所有MD文档元数据添加至所有MD文档元数据列表
if entry_mode_mask & 0o040000: if entry_mode_mask & 0o040000:
metadata.extend(get_remote_vault_metadata(remote_path)) metadatas.extend(recursive_scan_remote_vault(remote_path))
# 若实体文件为MD文件则构建文档源数据并添加到列表 # 若实体文件为MD文件则构建文档元数据并添加到所有MD文档元数据列表
elif entry.filename.lower().endswith(".md"): elif entry.filename.lower().endswith(".md"):
metadata.append( metadatas.append(
{"remote_path": remote_path, "modified_at": entry.st_mtime or 0} {"remote_path": remote_path, "modified_at": entry.st_mtime or 0}
) )
return metadata return metadatas
# 获取云端库指定文件夹下所有元数据 print(f"🔍 正在扫描云端库...")
remote_vault_metadata = get_remote_vault_metadata(REMOTE_DIRECTORY) # 递归扫描云端库指定文件夹路径，返回该文件夹内所有MD文档元数据
print(f"📁 云端库扫描完成，文档数：{len(remote_vault_metadata)}") metadatas = recursive_scan_remote_vault(REMOTE_DIRECTORY_PATH)
# 统计文档数
remote_document_count = len(metadatas)
print(f"✅ 扫描完成，云端库共有MD文档 {remote_document_count} 份")
download_count = 0 print(f"🧹 正在清理MD文档...")
for metadata in remote_vault_metadata: local_paths = list(LOCAL_DIRECTORY_PATH.rglob("*.md"))
# 本地路径 deleted_document_count = 0 # 删除文档数
local_path: Path = LOCAL_DIRECTORY / Path(metadata["remote_path"]).relative_to( for local_path in local_paths:
REMOTE_DIRECTORY # 构建文档云端库路径
) 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) local_path.parent.mkdir(parents=True, exist_ok=True)
is_download_needed = True # 需要下载 is_download_needed = True # 需要下载
modified_at = metadata["modified_at"] # 若文档本地库路径存在且修改时间戳差异小于阈值则无需下载
# 若本地文档存在且修改时间戳与云端文档修改时间戳差异小于阈值则无需下载
if local_path.exists(): if local_path.exists():
if abs(local_path.stat().st_mtime - modified_at) < 60: if abs(local_path.stat().st_mtime - metadata["modified_at"]) < 60:
is_download_needed = False is_download_needed = False
if is_download_needed: if is_download_needed:
print(f"⬇️ 下载文档：{metadata['remote_path']}")
# 下载文档 # 下载文档
sftp.get(metadata["remote_path"], local_path.as_posix()) sftp.get(metadata["remote_path"], local_path.as_posix())
# 同步修改时间戳 print(f" > 📥 {local_path.stem} 下载完成")
utime(local_path, times=(modified_at, modified_at)) # 修改文档本地库路径时间戳
download_count += 1 utime(local_path, times=(metadata["modified_at"], metadata["modified_at"]))
download_document_count += 1
print(f"✅ 同步完成，下载文档数：{download_count}") print(f"✅ 下载完成，本次共下载MD文档 {download_document_count} 份")
sftp.close() sftp.close()
ssh.close() ssh_client.close()
print("===== 本地库和云端库同步完成 =====\n") print("===== 同步Obsidian云端库至本地库完成 =====\n")
def process_documents() -> None: def process_documents() -> None:
""" """
处理文档 处理文档
""" """
print("===== 开始文档处理 =====") print("===== 开始处理文档 =====")
# 实例化标题分割器 # 实例化标题切分器
header_splitter = MarkdownHeaderTextSplitter( header_splitter = MarkdownHeaderTextSplitter(
headers_to_split_on=[ headers_to_split_on=[
("#", "h1"), ("#", "h1"),
@ -177,54 +198,53 @@ def process_documents() -> None:
password="PG198752", password="PG198752",
dbname="database", dbname="database",
) as connection: ) as connection:
print("📖 加载文档同步记录...") print("📥 正在加载文档同步记录...")
document_sync_records: dict[str, float] = {} # 文档同步记录字典 document_sync_records: dict[str, float] = {} # 文档同步记录字典
with connection.cursor() as cursor: with connection.cursor() as cursor:
cursor.execute("SELECT id, modified_at FROM document_sync_records;") cursor.execute("SELECT id, modified_at FROM document_sync_records;")
for id, modified_at in cursor.fetchall(): for id, modified_at in cursor.fetchall():
document_sync_records[id] = modified_at document_sync_records[id] = modified_at
print(f"✅ 加载完成，数据库中文档数：{len(document_sync_records)}") # 统计记录数
record_count = len(document_sync_records)
print(f"✅ 加载完成，数据库共有MD文档 {record_count} 份")
for idx, local_path in enumerate( # 列出本地库所有MD文档路径
local_paths := list(LOCAL_DIRECTORY.rglob("*.md")), 1 local_paths = list(LOCAL_DIRECTORY_PATH.rglob("*.md"))
): # 统计文档数
# 文档标题 document_count = len(local_paths)
document_title = local_path.stem document_ids = set()
print(f"📄 [{idx}/{len(local_paths)}] 正在处理文档：{document_title}\n") for idx, local_path in enumerate(local_paths, 1):
print(f"🧩 [{idx}/{document_count}] | 正在切片 {local_path.stem} ...")
# 读取文档 # 读取文档
document = load(local_path) document = load(local_path)
# 解析文档唯一标识（需Obsidian已安装NoteUIDGenerator插件，用于生成唯一标识） # 解析文档唯一标识（需Obsidian已安装NoteUIDGenerator插件，用于生成唯一标识）
document_id = str(document.metadata["uid"]) document_id = str(document.metadata["uid"])
# 本地文档修改时间戳 document_ids.add(document_id)
# 文档本地库路径的修改时间戳
modified_at = local_path.stat().st_mtime modified_at = local_path.stat().st_mtime
operation_type = "create" # 操作类型 operation_type = "create" # 操作类型
if document_id in document_sync_records: if document_id in document_sync_records:
# 若本地文档修改时间戳与数据库文档修改时间戳差异小于阈值则跳过 # 若本地文档修改时间戳与数据库文档修改时间戳差异小于阈值则跳过
if abs(modified_at - document_sync_records[document_id]) < 60: if abs(modified_at - document_sync_records[document_id]) < 60:
print(f"⏭️ 该文档未变动，无需处理") print(f"⏩ 该文档未修改，无需切片和向量化")
continue continue
operation_type = "update" operation_type = "update"
print(f"🔄 该文档需要更新")
else:
print(f"🔄 该文档需要创建")
index = 0 index = 0
current_batch: list[Chunk] = [] # 当前批次列表 current_batch: list[Chunk] = [] # 当前批次列表
current_batch_text_tokens: int = 0 # 当前批次列表的文本词元数 current_batch_text_tokens: int = 0 # 当前批次列表的文本词元数
batches: list[list[Chunk]] = [] # 批次列表 batches: list[list[Chunk]] = [] # 批次列表
# 先根据标题分割文档，若分块词元大小超过阈值则使用递归字符分割器进一步分割 # 先根据标题切分文档，若分片词元大小超过阈值则使用递归字符切分器进一步切分
for chunk in header_splitter.split_text(document.content): for chunk in header_splitter.split_text(document.content):
# --- 构建分块标题 --- # --- 构建分片标题 ---
headers = [] headers = []
# 若h1标题为空则使用文档标题 if h1 := chunk.metadata.get("h1"):
if not (h1 := chunk.metadata.get("h1")):
headers.append(document_title.strip())
else:
headers.append(h1.strip()) headers.append(h1.strip())
# 若h1标题为空则使用文档名称
else:
headers.append(local_path.stem.strip())
if h2 := chunk.metadata.get("h2"): if h2 := chunk.metadata.get("h2"):
headers.append(h2.strip()) headers.append(h2.strip())
if h3 := chunk.metadata.get("h3"): if h3 := chunk.metadata.get("h3"):
@ -237,8 +257,8 @@ def process_documents() -> None:
headers.append(h6.strip()) headers.append(h6.strip())
header = f"标题层级：{" > ".join(headers)}\n\n" header = f"标题层级：{" > ".join(headers)}\n\n"
# --- 构建分块内容 --- # --- 构建分片文本 ---
# 若分块内容为空则跳过 # 若分片内容为空则跳过
if not (content := chunk.page_content.strip()): if not (content := chunk.page_content.strip()):
continue continue
@ -250,10 +270,10 @@ def process_documents() -> None:
content = sub(r" +", " ", content) content = sub(r" +", " ", content)
# 压缩换行符 # 压缩换行符
content = sub(r"\n{3,}", "\n\n", content) content = sub(r"\n{3,}", "\n\n", content)
# 删除分割线标记 # 删除切分线标记
content = content.replace("---", "") content = content.replace("---", "")
# 返回文本 # 构建返回文本
return_text = header + content return_text = header + content
# 删除斜体标记 # 删除斜体标记
@ -269,16 +289,15 @@ def process_documents() -> None:
# 删除普通连接标记（删除URL，保留文本） # 删除普通连接标记（删除URL，保留文本）
content = sub(r"\[(.*?)\]\(.*?\)", r"\1", content) content = sub(r"\[(.*?)\]\(.*?\)", r"\1", content)
# 若文本词元数未超过限制则以分块内容作为分块内容列表 # 若文本（标题+内容）词元数未超过限制则以分片内容作为分片内容列表，否则使用递归字符切分器进一步切分分片内容
if count_tokens(header + content) < CHUNK_SIZE: if count_tokens(header + content) < CHUNK_SIZE:
contents = [content] contents = [content]
# 若文本词元数超过限制则使用递归字符分割器进一步分割分块内容
else: else:
# 实例化递归字符分割器 # 实例化递归字符切分器
recursive_splitter = RecursiveCharacterTextSplitter( recursive_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE - count_tokens(header), chunk_size=CHUNK_SIZE - count_tokens(header),
chunk_overlap=60, # 分块重叠 chunk_overlap=60, # 分片重叠
length_function=count_tokens, length_function=count_tokens,
separators=["\n\n", "\n", "。", "？", "！", ".", "?", "!", ""], separators=["\n\n", "\n", "。", "？", "！", ".", "?", "!", ""],
) )
@ -299,7 +318,7 @@ def process_documents() -> None:
text_tokens=count_tokens(text), text_tokens=count_tokens(text),
return_text=return_text, return_text=return_text,
) )
# 若当前批次列表的文本词元数加上当前分块的文本词元数超过阈值则将当前批次列表添加至批次列表 # 若当前批次列表的文本词元数加上当前分片的文本词元数超过阈值则将当前批次列表添加至批次列表
if current_batch_text_tokens + chunk.text_tokens > TPM: if current_batch_text_tokens + chunk.text_tokens > TPM:
batches.append(current_batch) batches.append(current_batch)
current_batch = [chunk] current_batch = [chunk]
@ -314,23 +333,24 @@ def process_documents() -> None:
# 若当前批次列表不为空则将当前批次列表添加至批次列表 # 若当前批次列表不为空则将当前批次列表添加至批次列表
if current_batch: if current_batch:
batches.append(current_batch) batches.append(current_batch)
print(f"📝 该文档切块完成，分块数：{index}") # 统计批次数
batch_count = len(batches)
print(f"✅ 切片完成，共切片 {batch_count} 批次、{index} 片")
# 按批次请求计算文本向量，若某批次请求发生异常则跳过该文档 # 按批次请求计算文本向量，若某批次请求发生异常则跳过该文档
try: try:
print(f"🔢 [{idx}/{document_count}] | 正在向量化 {local_path.stem} ...")
with connection.cursor() as cursor: with connection.cursor() as cursor:
# 更新场景：先删除当前文档全部旧chunk # 若操作类型为更新则删除该文档所有分片
if operation_type == "update": if operation_type == "update":
cursor.execute( cursor.execute(
"DELETE FROM document_chunks WHERE document_id = %s;", "DELETE FROM document_chunks WHERE document_id = %s;",
(document_id,), (document_id,),
) )
print(f"🗑️ 删除原有分块")
for i, batch in enumerate(batches, 1): for batch in batches:
# 统计批次的文本词元数 # 统计批次的文本词元数
batch_text_tokens = sum(chunk.text_tokens for chunk in batch) batch_text_tokens = sum(chunk.text_tokens for chunk in batch)
print(f" > 📦 Batch[{i}/{len(batches)}] 正在处理批次分块")
# 若窗口队列的文本词元数加上批次的文本词元数超过阈值则等待窗口队列的最早请求时间加上窗口期 # 若窗口队列的文本词元数加上批次的文本词元数超过阈值则等待窗口队列的最早请求时间加上窗口期
while True: while True:
# 获取窗口队列的文本词元数 # 获取窗口队列的文本词元数
@ -342,10 +362,10 @@ def process_documents() -> None:
max((window_deque[0][0] + WINDOW_PERIOD) - time(), 0.2) max((window_deque[0][0] + WINDOW_PERIOD) - time(), 0.2)
) )
request_counts = 0 # 请求次数 request_count = 0 # 请求次数
response_payload: dict | None = None response_payload: dict | None = None
while request_counts < 3: while request_count < 3:
request_counts += 1 request_count += 1
response: Response | None = None response: Response | None = None
try: try:
response = post( response = post(
@ -369,7 +389,7 @@ def process_documents() -> None:
except exceptions.HTTPError as exception: except exceptions.HTTPError as exception:
# 触发限速 # 触发限速
if response is not None and response.status_code == 429: if response is not None and response.status_code == 429:
sleep(2**request_counts) sleep(2**request_count)
continue continue
raise RuntimeError( raise RuntimeError(
@ -378,7 +398,7 @@ def process_documents() -> None:
except Exception as exception: except Exception as exception:
print(f"Exception: {exception}") print(f"Exception: {exception}")
sleep(2**request_counts) sleep(2**request_count)
if response_payload is None: if response_payload is None:
raise RuntimeError("response_payload is None") raise RuntimeError("response_payload is None")
@ -390,7 +410,7 @@ def process_documents() -> None:
): ):
chunk.text_vector = text_vector chunk.text_vector = text_vector
# 批量创建文档分块记录 # 批量创建文档分片记录
cursor.executemany( cursor.executemany(
""" """
INSERT INTO document_chunks INSERT INTO document_chunks
@ -410,7 +430,7 @@ def process_documents() -> None:
for chunk in batch for chunk in batch
], ],
) )
print(f" ✅ 该批次分块处理完成")
# 创建文档同步记录 # 创建文档同步记录
cursor.execute( cursor.execute(
""" """
@ -422,22 +442,36 @@ def process_documents() -> None:
(document_id, modified_at), (document_id, modified_at),
) )
connection.commit() connection.commit()
# 更新内存缓存 print(f"✅ 向量化完成并已同步至数据库")
document_sync_records[document_id] = modified_at
print(f"✅ 该文档处理完成\n")
except Exception: except Exception:
connection.rollback() connection.rollback()
continue continue
print("===== 全部文档处理完成 =====") # 若待删除文档唯一标识列表非空则删除相应文档同步记录和文档分片记录
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")
def main(): if __name__ == "__main__":
# 同步本地库和云端库 # 同步Obsidian云端库至本地库
sync_vault() sync_to_local_vault()
# 处理文档 # 处理文档
process_documents() process_documents()
main()

0
agent/document/vault/AI产品经理大厂面试100+题.md → agent/document_preprocess/vault/AI产品经理大厂面试100+题.md
View File