# -*- 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 """ 文档同步记录表建表SQL CREATE TABLE IF NOT EXISTS document_sync_records ( id TEXT PRIMARY KEY, modified_at FLOAT NOT NULL ); 文档分片表建表SQL CREATE EXTENSION IF NOT EXISTS vector; CREATE EXTENSION IF NOT EXISTS pg_search; CREATE TABLE IF NOT EXISTS document_chunks ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), document_id TEXT NOT NULL, index INT NOT NULL, text TEXT NOT NULL, text_vector vector(1024) NOT NULL, text_fingerprint TEXT NOT NULL, text_tokens INT NOT NULL, return_text TEXT ); CREATE INDEX idx_document_chunks_text ON document_chunks USING paradedb ( id, (text::pdb.jieba) ); CREATE INDEX idx_document_chunks_text_vector ON document_chunks USING hnsw (text_vector vector_cosine_ops); CREATE INDEX idx_document_chunks_text_fingerprint ON document_chunks(text_fingerprint); """ 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/MyVault/RawsSources" # Obsidian云端库文件夹路径 ) LOCAL_DIRECTORY_PATH = Path(__file__).parent / "vault" # 本地库文件夹路径 DATABASE_CONNECTION_CONFIG = { "host": "www.liubiren.cloud", "port": 5201, "user": "root", "password": "PG198752", "dbname": "database", } CHUNK_SIZE = 8192 # 分片大小，BAAI/BGE-M3标值 TPM = 500_000 * 0.8 # 每分钟词元数限额，硅基流动标值，安全冗余 WINDOW_PERIOD = 60 # 窗口期 def get_text_vectors(input: list[Chunk] | str) -> list[list[float]]: """ 获取文本嵌入向量 :param input: 输入 :return: 文本嵌入向量列表 """ 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 input] if isinstance(input, list) else [input] ), }, timeout=60, ) response.raise_for_status() response_payload = response.json() 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") return [item["embedding"] for item in response_payload["data"]] 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(**DATABASE_CONNECTION_CONFIG) 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) ) # 获取文本嵌入向量 text_vectors = get_text_vectors(batch) # 请求成功则将本次请求添加至窗口队列 window_deque.append((time(), batch_text_tokens)) # 写入文本向量 for chunk, text_vector in zip( batch, text_vectors, ): 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") def hybrid_retrieve( text: str, keywords: list[str] | None = None, ) -> list[dict]: """ 混合检索 :param text: 检索文本 :param keywords: 检索关键词列表 :return: 检索结果 """ candidates: list[dict] = [] with connect(**DATABASE_CONNECTION_CONFIG) as connection: with connection.cursor() as cursor: cursor.execute( """ WITH bm25_retrieve_results AS ( SELECT id, document_id, index, text, return_text, text_fingerprint, ROW_NUMBER() OVER (ORDER BY paradedb.score(id) DESC) AS bm25_retrieve_rank FROM document_chunks WHERE text @@@ %s ORDER BY paradedb.score(id) DESC LIMIT 50 ), vector_retrieve_results AS ( SELECT id, document_id, index, text, return_text, text_fingerprint, ROW_NUMBER() OVER (ORDER BY text_vector <-> %s::vector(1024)) AS vector_retrieve_rank FROM document_chunks ORDER BY text_vector <-> %s::vector(1024) LIMIT 50 ), merged_results AS ( SELECT COALESCE(bm25_retrieve_results.id, vector_retrieve_results.id) AS id, COALESCE(bm25_retrieve_results.document_id, vector_retrieve_results.document_id) AS document_id, COALESCE(bm25_retrieve_results.index, vector_retrieve_results.index) AS index, COALESCE(bm25_retrieve_results.text, vector_retrieve_results.text) AS text, COALESCE(bm25_retrieve_results.return_text, vector_retrieve_results.return_text) AS return_text, COALESCE(bm25_retrieve_results.text_fingerprint, vector_retrieve_results.text_fingerprint) AS text_fingerprint, COALESCE(1.0/(bm25_retrieve_rank + 60), 0) + COALESCE(1.0/(vector_retrieve_rank + 60), 0) AS rank_score FROM bm25_retrieve_results FULL OUTER JOIN vector_retrieve_results USING(id) ), deduplication AS ( SELECT *, ROW_NUMBER() OVER (PARTITION BY text_fingerprint ORDER BY rank_score DESC) AS duplicate_index FROM merged_results ) SELECT id, document_id, index, text, return_text, rank_score FROM deduplication WHERE duplicate_index = 1 ORDER BY rank_score DESC LIMIT 100; """, ( " ".join(keywords) if keywords is not None else text, text_vector := get_text_vectors(text)[0], # 获取文本嵌入向量 text_vector, ), ) rows = cursor.fetchall() if cursor.description is None: return candidates candidates = [ dict( zip( ([metadata[0] for metadata in cursor.description]), row, ) ) for row in rows ] try: 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/rerank", headers={ "Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz", "Content-Type": "application/json", }, json={ "model": "BAAI/bge-reranker-v2-m3", "query": text, "documents": [item["return_text"] for item in candidates], "return_documents": False, # 不返回文档内容 "top_n": 10, # 返回结果数 }, timeout=60, ) response.raise_for_status() response_payload = response.json() 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") # 若发生异常则返回候选前10个分片 except Exception: return candidates[:10] return [ candidates[result["index"]] for result in response_payload["results"] if result["relevance_score"] >= 0.85 ] if __name__ == "__main__": # 同步Obsidian云端库至本地库 sync_to_local_vault() # 处理文档 process_documents() """ results = hybrid_retrieve("如何实现事实校验") for result in results: print(result["return_text"]) """