This commit is contained in:
liubiren 2026-09-28 20:15:54 +08:00
parent e6c7adf1ac
commit cb63e7bc6d
5 changed files with 1000608 additions and 541 deletions

File diff suppressed because one or more lines are too long

View File

@ -0,0 +1,17 @@
{
"add_prefix_space": true,
"backend": "tokenizers",
"bos_token": "<s>",
"clean_up_tokenization_spaces": true,
"cls_token": "<s>",
"eos_token": "</s>",
"is_local": false,
"local_files_only": false,
"mask_token": "<mask>",
"model_max_length": 8192,
"pad_token": "<pad>",
"sep_token": "</s>",
"sp_model_kwargs": {},
"tokenizer_class": "XLMRobertaTokenizer",
"unk_token": "<unk>"
}

View File

@ -0,0 +1,419 @@
# -*- coding: utf-8 -*-
"""
预处理模块
"""
from os import utime
from pathlib import Path
from re import sub, MULTILINE
from hashlib import blake2s
from frontmatter import load
from langchain_text_splitters import (
MarkdownHeaderTextSplitter,
RecursiveCharacterTextSplitter,
)
from paramiko import AutoAddPolicy, SSHClient
from pgvector.psycopg import register_vector
from psycopg import connect
from pydantic import BaseModel, Field
from transformers import AutoTokenizer
class Chunk(BaseModel):
"""
分块数据模型
"""
index: int = Field(..., description="分块索引")
document_id: str = Field(..., description="文档唯一标识")
embedding_text: str = Field(..., description="分块嵌入文本")
return_text: str = Field(..., description="分块返回文本")
return_text_fingerprint: str = Field(..., description="分块返回文本指纹")
LOCAL_DIRECTORY = Path(__file__).parent / "vault" # 本地库文件夹
REMOTE_DIRECTORY = "/root/obsidian/vaults/MyNewVault/RawsSources" # 云端库文件夹
CHUNK_SIZE = 8192 # 分块大小,BAAI/BGE-M3 标值
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_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:
# 下载文档
sftp.get(metadata["remote_path"], local_path.as_posix())
# 同步修改时间戳
utime(local_path, times=(modified_at, modified_at))
sftp.close()
ssh.close()
def process_documents() -> None:
"""
处理文档
"""
# 实例化标题分割器
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)
) # 首尾不添加模型专用特殊标记
with connect(
host="www.liubiren.cloud",
port=5201,
user="root",
password="PG198752",
dbname="database",
) as connection:
document_sync_records: dict[str, float] = {} # 同步记录字典
with connection.cursor() as cursor:
cursor.execute("SELECT uid, modified_at FROM document_sync_records;")
for uid, modified_at in cursor.fetchall():
document_sync_records[uid] = modified_at
for local_path in LOCAL_DIRECTORY.rglob("*.md"): # 递归遍历本地库所有文档
# 读取文档
document = load(local_path)
# 解析文档唯一标识(需Obsidian已安装NoteUIDGenerator插件,用于生成唯一标识)
document_id = str(document.metadata["uid"])
if modified_at := document_sync_records.get(document_id):
# 若本地文档修改时间戳与数据库文档修改时间戳差异小于阈值则跳过
if abs(local_path.stat().st_mtime - modified_at) < 60:
continue
# 若本地文档修改时间戳与数据库文档修改时间戳差异大于阈值则删除删除该文档所有文档块
cursor.execute(
"DELETE FROM document_chunks WHERE document_id = %s",
(document_id,),
)
connection.commit()
# 文档标题
document_title = local_path.stem
index = 0
chunks: 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,}", r"\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)
# 初始化分块内容列表
contents = [content]
# 若文本词元数超过限制则使用递归字符分割器进一步分割分块内容
if count_tokens(return_text) > CHUNK_SIZE:
# 实例化递归字符分割器
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
chunks.append(
Chunk(
index=index,
document_id=document_id,
embedding_text=header + content,
return_text=return_text,
return_text_fingerprint=blake2s(
return_text.encode(), digest_size=16
).hexdigest(),
)
)
index += 1
for chunk in chunks:
print(chunk.return_text)
print("\n")
exit()
def main():
# 同步本地库和云端库
sync_vault()
# 处理文档
process_documents()
main()
exit()
register_vector(db_connection)
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()
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("需求评审有哪些流程"))

View File

@ -1,535 +0,0 @@
# -*- 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("需求评审有哪些流程"))

View File

@ -857,9 +857,3 @@ AI产品可观测性分三层:
| Q110 | 可观测性 | AI产品运维核心,原文未覆盖监控体系 | | Q110 | 可观测性 | AI产品运维核心,原文未覆盖监控体系 |
| Q111 | AI出海 | 国际化是大厂高频话题,原文完全缺失 | | Q111 | AI出海 | 国际化是大厂高频话题,原文完全缺失 |
| Q112 | 2026 Agent趋势 | 时效性补充,展示行业洞察力 | | Q112 | 2026 Agent趋势 | 时效性补充,展示行业洞察力 |
---
## 相关标签
#AI产品经理 #面试准备 #大模型 #Agent #RAG #AI商业化 #Prompt工程 #AI安全