1

agent/application/preprocessing.py

@@ -6,34 +6,34 @@
 2）该文档存在但更新时间戳不一致则先再数据库中删除该文档所有块向量记录，再重新分块、向量化、存储到数据库中
 """
 
-import asyncio
-from pathlib import Path
-from typing import Dict, List, cast, Set
 
+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 lark_oapi.api.docs.v1 import (
-    GetContentRequest,
-    GetContentResponse,
-    GetContentResponseBody,
-)
 from 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,
 )
-import tiktoken
-from langchain_text_splitters import (
-    MarkdownHeaderTextSplitter,
-    RecursiveCharacterTextSplitter,
-)
 
 # 建立数据库连接
 db_connection = connect(
@@ -65,6 +65,8 @@ embedder = Embedder(
 )
 # 最大分块大小
 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(
@@ -72,19 +74,30 @@ header_splitter = MarkdownHeaderTextSplitter(
         ("#", "h1"),
         ("##", "h2"),
         ("###", "h3"),
+        ("####", "h4"),
+        ("#####", "h5"),
     ],
-    strip_headers=False,  # 保留标题
+    strip_headers=True,  # 不保留标题
 )
 
-# 实例化 Markdown 递归分割器
+
-recursive_splitter = RecursiveCharacterTextSplitter(
+class Chunk(BaseModel):
-    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,
+    document_id: str = Field(..., description="文档唯一标识")
-    separators=["\n\n", "\n", "。", "？", "！", ".", "?", "!"],
+    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]:
@@ -153,63 +166,9 @@ def get_offline_documents() -> Dict[str, int]:
     return offline_documents
 
 
-def get_pending_insert_documents() -> Set[str]:
+def download_document(id: str) -> 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格式字符串）
     """
@@ -235,52 +194,210 @@ def get_content(id: str) -> str:
     return cast(str, data.content)
 
 
-# 获取待插入文档列表
+def sync_documents() -> None:
-# pending_insert_documents = get_pending_insert_documents()
+    """
-pending_insert_documents = set(["1"])
+    同步文档
+    :return: None
+    """
+    # 获取在线文档字典
+    online_documents = get_online_documents()
+    # 获取离线文档字典
+    offline_documents = get_offline_documents()
 
-for id in pending_insert_documents:
+    # 待删除文档唯一标识集合
-    # 获取云文档内容
+    pending_delete_documents = set(offline_documents.keys() - online_documents.keys())
-    # content = get_content(id)
+    # 待下载文档唯一标识字典
-    content = """
+    pending_download_documents: Dict[str, int] = {}
-        # 这个例子展示了如何处理一个包含多级标题、代码块和不同内容类型的长Markdown文档。它首先按标题分割，然后对每个部分进行字符级分割，确保每个chunk的大小适合后续处理（如嵌入或向量化）。
+    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
 
-        text
+        # 若在线文档唯一标识于离线文档字典存在且更新时间戳不一致则将该文档唯一标识添加至待删除/下载文档唯一标识集合
+        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
 
-        ## 5. 常见问题和解决方案
+    # 批量删除待删除文档
+    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()
 
-        1. **问题**：分割后的chunk太大或太小。
+    for document_id, document_updated_at in pending_download_documents.items():
-        **解决方案**：调整RecursiveCharacterTextSplitter的chunk_size参数。
+        # 下载文档
+        document_content = download_document(document_id)
+        print(document_content)
+        exit()
 
-        2. **问题**：重要的上下文信息丢失。
+        chunks = []
-        **解决方案**：增加chunk_overlap参数，允许相邻chunk之间有一定的重叠。
+        # 分块索引
+        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
 
-        3. **问题**：代码块被分割。
+            # 构建标题块标题
-        **解决方案**：考虑使用自定义分割策略，或在预处理阶段标记代码块。
+            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"
 
-        4. **问题**：某些地区可能无法直接访问某些API。
+            # 最大标题块内容大小
-        **解决方案**：使用API代理服务，如示例中的`http://api.wlai.vip`。
+            max_heading_chunk_content_size = (
+                max_chunk_size - embedder.count_tokens_sync(heading_chunk_heading)
+            )
 
-        ## 6. 总结和进一步学习资源
+            # 若标题块内容大小超过最大标题块内容大小则再次分割
+            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]
 
-        本文介绍了如何使用Python和LangChain库实现基于Markdown标题的智能分块策略。这种方法不仅保持了文档的结构完整性，还提供了灵活的配置选项，适用于各种文档处理场景。
+            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)
 
-        - LangChain官方文档：[Text Splitters](https://python.langchain.com/docs/modules/data_connection/document_transformers/)
+        with db_connection.cursor() as cursor:
-        - Python Markdown库：[Python-Markdown](https://python-markdown.github.io/)
+            # 写入文档表
-        - 向量数据库应用：[Pinecone](https://www.pinecone.io/)
+            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()
 
-        ## 参考资料
 
-        1. LangChain Documentation. (2023). Text Splitters. https://python.langchain.com/docs/modules/data_connection/document_transformers/
+def query(
-        2. Gruber, J. (2004). Markdown. https://daringfireball.net/projects/markdown/
+    question: str,
-        3. Pinecone. (2023). Chunking Strategies for LLM Applications. https://www.pinecone.io/learn/chunking-strategies/
+) -> 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
+    ]
 
-        ---END---
+
-        """
+if __name__ == "__main__":
-    for section in header_splitter.split_text(content):
+    # 同步文档
-        print("-" * 50)
+    sync_documents()
-        print(section)
+    # print(query("需求评审有哪些流程"))