536 lines
18 KiB
Python
536 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
预处理模块
|
||
## 获取待处理的本地路径列表
|
||
- 从云端库指定文件夹下获取所有元数据
|
||
- 对比本地库和云端库的文档修改时间戳，就差值大于阈值的文档定义为待处理文档，返回其本地路径列表
|
||
"""
|
||
|
||
|
||
from pathlib import Path
|
||
from typing import Dict, List, Tuple, cast
|
||
|
||
from langchain_text_splitters import (
|
||
MarkdownHeaderTextSplitter,
|
||
RecursiveCharacterTextSplitter,
|
||
)
|
||
from lark_oapi import Client
|
||
from lark_oapi.api.docs.v1 import (
|
||
GetContentRequest,
|
||
GetContentResponse,
|
||
GetContentResponseBody,
|
||
)
|
||
from lark_oapi.api.wiki.v2 import (
|
||
ListSpaceNodeRequest,
|
||
ListSpaceNodeResponse,
|
||
ListSpaceNodeResponseBody,
|
||
Node,
|
||
)
|
||
from psycopg import connect
|
||
|
||
##rom pgvector.psycopg import register_vector
|
||
"""
|
||
from psycopg import connect
|
||
from pydantic import BaseModel, Field
|
||
from pydantic_ai import Embedder
|
||
from pydantic_ai.embeddings.sentence_transformers import (
|
||
SentenceTransformerEmbeddingModel,
|
||
SentenceTransformersEmbeddingSettings,
|
||
)
|
||
"""
|
||
from pathlib import Path
|
||
from paramiko import SSHClient, AutoAddPolicy
|
||
from frontmatter import load
|
||
from typing import List, Dict
|
||
import asyncio
|
||
from os import utime
|
||
|
||
|
||
LOCAL_VAULT_PATH = Path(__file__).parent / "vault"
|
||
REMOTE_DIRECTORY = "/root/obsidian/vaults/MyNewVault/RawsSources"
|
||
|
||
DATABASE_CONNECT_CONFIG = {
|
||
"host": "www.liubiren.cloud",
|
||
"port": 5201,
|
||
"user": "root",
|
||
"password": "PG198752",
|
||
"database": "database",
|
||
}
|
||
|
||
|
||
def sync_vault() -> None:
|
||
"""
|
||
同步本地库和云端库
|
||
"""
|
||
# 创建 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)
|
||
for metadata in remote_vault_metadata:
|
||
# 本地路径
|
||
local_path: Path = LOCAL_VAULT_PATH / 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:
|
||
# 下载文档
|
||
sftp.get(metadata["remote_path"], local_path.as_posix())
|
||
# 同步修改时间戳
|
||
utime(local_path, times=(modified_at, modified_at))
|
||
|
||
sftp.close()
|
||
ssh.close()
|
||
|
||
|
||
async def process_single_document(doc_uid: str, raw_content: str, file_mtime: float):
|
||
"""单文档处理：分块、硅基embedding、PG入库"""
|
||
print(f"📄 处理文档 uid={doc_uid}, mtime={file_mtime:.2f}")
|
||
pass
|
||
|
||
|
||
def process() -> None:
|
||
"""
|
||
处理
|
||
"""
|
||
with connect(**DATABASE_CONNECT_CONFIG) as conn:
|
||
cur = conn.cursor()
|
||
cur.execute("SELECT document_id FROM document")
|
||
existing_uids = set(cur[0] for cur in cur.fetchall())
|
||
for local_path in LOCAL_VAULT_PATH.rglob("*.md"):
|
||
# 读取文档
|
||
document = load(local_path)
|
||
# 唯一标识和内容（需Obsidian已安装NoteUIDGenerator插件，用于生成唯一标识）
|
||
uid, content = cast(str, document.metadata["uid"]), document.content
|
||
|
||
|
||
def main():
|
||
# 同步本地库和云端库
|
||
# sync_vault()
|
||
process()
|
||
|
||
|
||
main()
|
||
|
||
|
||
exit()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|
||
|
||
|
||
# 建立数据库连接
|
||
db_connection = connect(
|
||
host="liubiren.cloud",
|
||
port=5201,
|
||
user="root",
|
||
password="PG198752",
|
||
dbname="database",
|
||
)
|
||
register_vector(db_connection)
|
||
|
||
# 实例化飞书服务端
|
||
feishu_server_client = (
|
||
Client.builder()
|
||
.app_id("cli_a1587980be78500c")
|
||
.app_secret("vZXGZomwfmyaHXoG8s810d1YYGLsIqCA")
|
||
.build()
|
||
)
|
||
|
||
# 实例化词嵌器
|
||
embedder = Embedder(
|
||
SentenceTransformerEmbeddingModel(
|
||
(
|
||
Path(__file__).parent / "models" / "paraphrase-multilingual-MiniLM-L12-v2"
|
||
).as_posix(),
|
||
settings=SentenceTransformersEmbeddingSettings(
|
||
sentence_transformers_device="cpu",
|
||
sentence_transformers_normalize_embeddings=True,
|
||
),
|
||
),
|
||
)
|
||
|
||
# 最大分块大小
|
||
max_chunk_size: int = cast(int, embedder.max_input_tokens_sync())
|
||
# 分块重叠
|
||
chunk_overlap = max(20, min(int(round(max_chunk_size * 0.15)), 60))
|
||
|
||
# 实例化 Markdown 标题分割器
|
||
header_splitter = MarkdownHeaderTextSplitter(
|
||
headers_to_split_on=[
|
||
("#", "h1"),
|
||
("##", "h2"),
|
||
("###", "h3"),
|
||
("####", "h4"),
|
||
("#####", "h5"),
|
||
],
|
||
strip_headers=True, # 不保留标题
|
||
)
|
||
|
||
|
||
class Chunk(BaseModel):
|
||
"""
|
||
分块数据模型
|
||
"""
|
||
|
||
document_id: str = Field(..., description="文档唯一标识")
|
||
idx: int = Field(..., description="分块索引")
|
||
heading_chunk_h1: str = Field(..., description="标题块标题1")
|
||
heading_chunk_h2: str = Field(default="", description="标题块标题2")
|
||
heading_chunk_h3: str = Field(default="", description="标题块标题3")
|
||
heading_chunk_h4: str = Field(default="", description="标题块标题4")
|
||
heading_chunk_h5: str = Field(default="", description="标题块标题5")
|
||
heading_chunk_heading: str = Field(..., description="标题块标题")
|
||
heading_chunk_content: str = Field(..., description="标题块内容")
|
||
content_chunk_content: str = Field(..., description="内容块内容")
|
||
content_chunk_size: int = Field(..., description="内容块大小")
|
||
embedding_text: str = Field(..., description="词嵌文本")
|
||
|
||
|
||
def get_online_documents() -> Dict[str, int]:
|
||
"""
|
||
获取在线文档字典
|
||
:return: 在线文档字典
|
||
"""
|
||
# 在线文档字典
|
||
online_documents: Dict[str, int] = {}
|
||
|
||
# 实例化知识库
|
||
if not (wiki := feishu_server_client.wiki):
|
||
raise Exception("client.wiki is None")
|
||
|
||
# 分页标记
|
||
page_token = ""
|
||
while True:
|
||
# 构造获取知识库空间子节点列表请求实例
|
||
request: ListSpaceNodeRequest = (
|
||
ListSpaceNodeRequest.builder()
|
||
.space_id("7615153684095257820") # 默认为产品设计知识库
|
||
.page_size(50) # 分页大小
|
||
.page_token(page_token)
|
||
.build()
|
||
)
|
||
# 请求
|
||
response: ListSpaceNodeResponse = wiki.v2.space_node.list(request)
|
||
if not response.success():
|
||
raise Exception("response not success")
|
||
# 响应数据
|
||
data = cast(ListSpaceNodeResponseBody, response.data)
|
||
for item in cast(List[Node], data.items):
|
||
if cast(str, item.obj_type) == "docx":
|
||
# 文档唯一标识
|
||
id: str = cast(str, item.obj_token)
|
||
# 文档更新时间戳
|
||
updated_at: int = int(
|
||
cast(int, item.obj_edit_time)
|
||
) # Lark说明类型为 int，但实际为 str
|
||
online_documents[id] = updated_at
|
||
|
||
# 更新分页标记
|
||
page_token = cast(str, data.page_token)
|
||
# 若是否还有更多项为否则跳出循环
|
||
if not cast(bool, data.has_more):
|
||
break
|
||
return online_documents
|
||
|
||
|
||
def get_offline_documents() -> Dict[str, int]:
|
||
"""
|
||
获取离线文档字典
|
||
:return: 离线文档字典
|
||
"""
|
||
# 离线文档字典
|
||
offline_documents: Dict[str, int] = {}
|
||
with db_connection.cursor() as cursor:
|
||
cursor.execute(
|
||
"""
|
||
SELECT id, updated_at
|
||
FROM documents
|
||
"""
|
||
)
|
||
for id, updated_at in cursor.fetchall():
|
||
offline_documents[id] = updated_at
|
||
return offline_documents
|
||
|
||
|
||
def download_document(id: str) -> str:
|
||
"""
|
||
下载文档
|
||
:param id: 文档唯一标识
|
||
:return: 文档内容（Markdown格式字符串）
|
||
"""
|
||
# 实例化云文档
|
||
if not (docs := feishu_server_client.docs):
|
||
raise Exception("client.docs is None")
|
||
|
||
# 构造获取云文档内容请求实例
|
||
request: GetContentRequest = (
|
||
GetContentRequest.builder()
|
||
.doc_token(id)
|
||
.doc_type("docx")
|
||
.content_type("markdown")
|
||
.build()
|
||
)
|
||
# 请求
|
||
response: GetContentResponse = docs.v1.content.get(request)
|
||
if not response.success():
|
||
print(response.code, response.msg)
|
||
raise Exception("response not success")
|
||
# 响应数据
|
||
data = cast(GetContentResponseBody, response.data)
|
||
return cast(str, data.content)
|
||
|
||
|
||
def sync_documents1() -> None:
|
||
"""
|
||
同步文档
|
||
:return: None
|
||
"""
|
||
# 获取在线文档字典
|
||
online_documents = get_online_documents()
|
||
# 获取离线文档字典
|
||
offline_documents = get_offline_documents()
|
||
|
||
# 待删除文档唯一标识集合
|
||
pending_delete_documents = set(offline_documents.keys() - online_documents.keys())
|
||
# 待下载文档唯一标识字典
|
||
pending_download_documents: Dict[str, int] = {}
|
||
for document_id, document_updated_at in online_documents.items():
|
||
# 若在线文档唯一标识于离线文档字典不存在则将该文档唯一标识添加至待下载文档唯一标识字典
|
||
if document_id not in offline_documents:
|
||
pending_download_documents[document_id] = document_updated_at
|
||
continue
|
||
|
||
# 若在线文档唯一标识于离线文档字典存在且更新时间戳不一致则将该文档唯一标识添加至待删除/下载文档唯一标识集合
|
||
if (
|
||
document_id in offline_documents
|
||
and document_updated_at != offline_documents[document_id]
|
||
):
|
||
pending_delete_documents.add(document_id)
|
||
pending_download_documents[document_id] = document_updated_at
|
||
continue
|
||
|
||
# 批量删除待删除文档
|
||
if pending_delete_documents:
|
||
with db_connection.cursor() as cursor:
|
||
cursor.execute(
|
||
"""
|
||
DELETE FROM documents
|
||
WHERE id = ANY(%s);
|
||
""",
|
||
(list(pending_delete_documents),),
|
||
)
|
||
db_connection.commit()
|
||
|
||
for document_id, document_updated_at in pending_download_documents.items():
|
||
# 下载文档
|
||
document_content = download_document(document_id)
|
||
print(document_content)
|
||
exit()
|
||
|
||
chunks = []
|
||
# 分块索引
|
||
idx = 0
|
||
embedding_texts = []
|
||
for heading_chunk in header_splitter.split_text(document_content):
|
||
# 若标题块内容为空则跳过
|
||
if not (heading_chunk_content := heading_chunk.page_content.strip()):
|
||
continue
|
||
|
||
# 构建标题块标题
|
||
heading_chunk_heading = f"# {(heading_chunk_h1:=heading_chunk.metadata["h1"])}\n" # 必定包含 h1 标题
|
||
if heading_chunk_h2 := heading_chunk.metadata.get("h2", ""):
|
||
heading_chunk_heading += f"## {heading_chunk_h2}\n"
|
||
if heading_chunk_h3 := heading_chunk.metadata.get("h3", ""):
|
||
heading_chunk_heading += f"### {heading_chunk_h3}\n"
|
||
if heading_chunk_h4 := heading_chunk.metadata.get("h4", ""):
|
||
heading_chunk_heading += f"#### {heading_chunk_h4}\n"
|
||
if heading_chunk_h5 := heading_chunk.metadata.get("h5", ""):
|
||
heading_chunk_heading += f"##### {heading_chunk_h5}\n"
|
||
|
||
# 最大标题块内容大小
|
||
max_heading_chunk_content_size = (
|
||
max_chunk_size - embedder.count_tokens_sync(heading_chunk_heading)
|
||
)
|
||
|
||
# 若标题块内容大小超过最大标题块内容大小则再次分割
|
||
if (
|
||
embedder.count_tokens_sync(heading_chunk_content)
|
||
> max_heading_chunk_content_size
|
||
):
|
||
# 实例化递归分割器
|
||
recursive_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=max_heading_chunk_content_size,
|
||
chunk_overlap=chunk_overlap, # 分块重叠：最大分块大小的 15%，且最小为 20、最大为 60
|
||
length_function=embedder.count_tokens_sync,
|
||
separators=["\n\n", "\n", "。", "？", "！", ".", "?", "!", ""],
|
||
)
|
||
content_chunk_contents = recursive_splitter.split_text(
|
||
heading_chunk_content
|
||
)
|
||
else:
|
||
content_chunk_contents = [heading_chunk_content]
|
||
|
||
for content_chunk_content in content_chunk_contents:
|
||
# 若内容块为空则跳过
|
||
if not content_chunk_content.strip():
|
||
continue
|
||
|
||
chunks.append(
|
||
Chunk(
|
||
idx=idx,
|
||
document_id=document_id,
|
||
heading_chunk_h1=heading_chunk_h1,
|
||
heading_chunk_h2=heading_chunk_h2,
|
||
heading_chunk_h3=heading_chunk_h3,
|
||
heading_chunk_h4=heading_chunk_h4,
|
||
heading_chunk_h5=heading_chunk_h5,
|
||
heading_chunk_heading=heading_chunk_heading,
|
||
heading_chunk_content=heading_chunk_content,
|
||
content_chunk_content=content_chunk_content,
|
||
content_chunk_size=embedder.count_tokens_sync(
|
||
content_chunk_content
|
||
),
|
||
embedding_text=(
|
||
embedding_text := heading_chunk_heading
|
||
+ content_chunk_content
|
||
),
|
||
)
|
||
)
|
||
idx += 1
|
||
embedding_texts.append(embedding_text)
|
||
|
||
with db_connection.cursor() as cursor:
|
||
# 写入文档表
|
||
cursor.execute(
|
||
"""
|
||
INSERT INTO documents (id, updated_at)
|
||
VALUES (%s, %s)
|
||
""",
|
||
(document_id, document_updated_at),
|
||
)
|
||
# 写入分块表
|
||
cursor.executemany(
|
||
"""
|
||
INSERT INTO chunks (
|
||
document_id, idx,
|
||
heading_chunk_h1, heading_chunk_h2, heading_chunk_h3, heading_chunk_h4, heading_chunk_h5, heading_chunk_heading, heading_chunk_content,
|
||
content_chunk_content, content_chunk_size,
|
||
embedding_text, embedding_vector
|
||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s);
|
||
""",
|
||
[
|
||
(
|
||
chunk.document_id,
|
||
chunk.idx,
|
||
chunk.heading_chunk_h1,
|
||
chunk.heading_chunk_h2,
|
||
chunk.heading_chunk_h3,
|
||
chunk.heading_chunk_h4,
|
||
chunk.heading_chunk_h5,
|
||
chunk.heading_chunk_heading,
|
||
chunk.heading_chunk_content,
|
||
chunk.content_chunk_content,
|
||
chunk.content_chunk_size,
|
||
chunk.embedding_text,
|
||
embedding_vector,
|
||
)
|
||
for chunk, embedding_vector in zip(
|
||
chunks,
|
||
embedder.embed_documents_sync(
|
||
embedding_texts
|
||
).embeddings, # 生成词嵌向量
|
||
)
|
||
], # 构建记录
|
||
)
|
||
db_connection.commit()
|
||
|
||
|
||
def query(
|
||
question: str,
|
||
) -> List[Tuple[float, str, str]]:
|
||
"""
|
||
查询
|
||
:param question: 问题
|
||
:return: 相似度最高的前 k 个分块（每个分块包含相似度、文档 ID、标题、内容）
|
||
"""
|
||
# 生成词嵌向量
|
||
embedding_vector = embedder.embed_query_sync(query=question).embeddings[0]
|
||
|
||
with db_connection.cursor() as cursor:
|
||
cursor.execute(
|
||
"""
|
||
SELECT
|
||
1 - (embedding_vector <=> %s::vector) AS similarity_score,
|
||
heading_chunk_heading,
|
||
heading_chunk_content
|
||
FROM chunks
|
||
WHERE 1 - (embedding_vector <=> %s::vector) >= %s
|
||
ORDER BY embedding_vector <=> %s::vector
|
||
LIMIT %s;
|
||
""",
|
||
(
|
||
embedding_vector,
|
||
embedding_vector,
|
||
0.45,
|
||
embedding_vector,
|
||
5,
|
||
),
|
||
)
|
||
rows = cursor.fetchall()
|
||
return [
|
||
(similarity_score, heading_chunk_heading, heading_chunk_content)
|
||
for similarity_score, heading_chunk_heading, heading_chunk_content in rows
|
||
]
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 同步文档
|
||
# sync_documents()
|
||
print(query("需求评审有哪些流程"))
|