ZhouYI/03_RAG优化/agent/rag/splitter.py

"""文本分块策略。"""

from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.documents import Document
from langchain_experimental.text_splitter import SemanticChunker
from langchain_text_splitters import RecursiveCharacterTextSplitter

from agent.config import load_config


def recursive_split(
    documents: list[Document],
    chunk_size: int = 500,
    chunk_overlap: int = 100,
) -> list[Document]:
    """使用适合中文文本的递归字符分块。"""
    splitter = RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", "。", "！", "？", ".", " ", ""],
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
    )
    return splitter.split_documents(documents)


def semantic_split(text: str) -> list[str]:
    """使用 Embedding 的语义断点进行分块。"""
    config = load_config()
    embedding_model = DashScopeEmbeddings(
        model="text-embedding-v4",
        dashscope_api_key=config.api_key,
    )
    splitter = SemanticChunker(
        embeddings=embedding_model,
        breakpoint_threshold_type="percentile",
        breakpoint_threshold_amount=85,
    )
    return splitter.split_text(text)