AIyaodianzhushou/data-pipeline/processors/chunker.py

"""
智能切片器：按药品条目为单位切片
确保每个 chunk 是完整的药品信息，不截断关键内容
"""
import logging
from typing import Optional
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


@dataclass
class Chunk:
    content: str
    metadata: dict = field(default_factory=dict)


class DrugChunker:
    def __init__(
        self,
        chunk_size: int = 1500,
        chunk_overlap: int = 200,
        min_chunk_size: int = 200,
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.min_chunk_size = min_chunk_size

    def chunk_drug_entry(self, drug_entry: dict) -> list[Chunk]:
        """
        将结构化药品条目切成 chunk。
        策略：按 section 切分，保持每个 section 完整。
        若单个 section 过长，再按段落细分。
        """
        chunks = []
        drug_name = drug_entry.get("name", "")
        source = drug_entry.get("source", {})

        for section_key, section_text in drug_entry.get("sections", {}).items():
            if not section_text or len(section_text.strip()) < 10:
                continue

            if len(section_text) <= self.chunk_size:
                chunks.append(Chunk(
                    content=f"【{drug_name} - {section_key}】\n{section_text}",
                    metadata={
                        "drug_name": drug_name,
                        "section": section_key,
                        "source": source.get("version", "") + " "
                                + source.get("volume", "") + " "
                                + source.get("page", ""),
                        "category": drug_entry.get("category", ""),
                    }
                ))
            else:
                sub_chunks = self._split_long_section(
                    section_text, drug_name, section_key, source
                )
                chunks.extend(sub_chunks)

        return chunks

    def _split_long_section(
        self, text: str, drug_name: str, section: str, source: dict
    ) -> list[Chunk]:
        """对过长 section 按段落切分"""
        paragraphs = text.split("\n")
        chunks = []
        buffer = ""
        for para in paragraphs:
            if len(buffer) + len(para) > self.chunk_size and len(buffer) >= self.min_chunk_size:
                chunks.append(Chunk(
                    content=f"【{drug_name} - {section}】\n{buffer}",
                    metadata={
                        "drug_name": drug_name,
                        "section": section,
                        "source": f"{source.get('version', '')} {source.get('volume', '')} {source.get('page', '')}",
                        "category": "",
                    }
                ))
                buffer = para
            else:
                buffer += ("\n" if buffer else "") + para

        if buffer:
            chunks.append(Chunk(
                content=f"【{drug_name} - {section}】\n{buffer}",
                metadata={
                    "drug_name": drug_name,
                    "section": section,
                    "source": f"{source.get('version', '')} {source.get('volume', '')} {source.get('page', '')}",
                    "category": "",
                }
            ))

        return chunks

    def chunk_regulation(self, text: str, metadata: dict) -> list[Chunk]:
        """凡例/通则按小标题切分"""
        return []

    def chunk_exam_knowledge(self, knowledge_point: dict) -> list[Chunk]:
        """考试知识点按大纲章节切分"""
        return []