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 按段落切分，相邻 chunk 之间保留 overlap 以避免上下文断裂"""
        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": "",
                    }
                ))
                # overlap: 从上一个 buffer 尾部保留 chunk_overlap 字符作为下一段的上下文
                if self.chunk_overlap > 0 and len(buffer) > self.chunk_overlap:
                    buffer = buffer[-self.chunk_overlap:] + "\n" + para
                else:
                    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]:
        """凡例/通则按编号或段落切分，生成可供检索的 chunk。"""
        if not text or len(text.strip()) < 10:
            return []

        chunks = []
        source_info = (
            f"{metadata.get('source', '')} {metadata.get('version', '')} "
            f"{metadata.get('volume', '')}".strip()
        )
        reg_name = metadata.get("name", "通则")

        # 尝试按编号拆分（一、二、三... 或 1. 2. 3.）
        import re
        cn_nums = "一二三四五六七八九十百千"
        patterns = [
            re.compile(r'(?:^|\n)([' + cn_nums + r']{1,3})[、，,.]'),
            re.compile(r'(?:^|\n)（([' + cn_nums + r']{1,3})）'),
            re.compile(r'(?:^|\n)(\d{1,2})[\.\)）]'),
        ]

        sections = None
        for pattern in patterns:
            matches = list(pattern.finditer(text))
            if len(matches) >= 2:
                sections = []
                for i, match in enumerate(matches):
                    start = match.start()
                    end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
                    section_text = text[start:end].strip()
                    if len(section_text) > 20:
                        sections.append(section_text)
                break

        if sections:
            for sec_text in sections:
                if len(sec_text) <= self.chunk_size:
                    chunks.append(Chunk(
                        content=f"【{reg_name}】\n{sec_text}",
                        metadata={
                            "regulation": reg_name,
                            "source": source_info,
                            **{k: v for k, v in metadata.items() if k not in ("name", "source")},
                        }
                    ))
                else:
                    # 长段落实用 _split_long_section 逻辑
                    sub = self._split_long_section(sec_text, reg_name, "通则", {
                        "version": metadata.get("version", ""),
                        "volume": metadata.get("volume", ""),
                        "page": metadata.get("page", ""),
                    })
                    chunks.extend(sub)
        else:
            # 无编号结构，按段落切分
            paras = [p.strip() for p in text.split("\n") if p.strip()]
            buffer = ""
            for para in paras:
                if len(buffer) + len(para) > self.chunk_size and len(buffer) >= self.min_chunk_size:
                    chunks.append(Chunk(
                        content=f"【{reg_name}】\n{buffer}",
                        metadata={
                            "regulation": reg_name,
                            "source": source_info,
                            **{k: v for k, v in metadata.items() if k not in ("name", "source")},
                        }
                    ))
                    # overlap
                    if self.chunk_overlap > 0 and len(buffer) > self.chunk_overlap:
                        buffer = buffer[-self.chunk_overlap:] + "\n" + para
                    else:
                        buffer = para
                else:
                    buffer += ("\n" if buffer else "") + para
            if buffer:
                chunks.append(Chunk(
                    content=f"【{reg_name}】\n{buffer}",
                    metadata={
                        "regulation": reg_name,
                        "source": source_info,
                        **{k: v for k, v in metadata.items() if k not in ("name", "source")},
                    }
                ))

        return chunks

    def chunk_exam_knowledge(self, knowledge_point: dict) -> list[Chunk]:
        """考试知识点按大纲章节切分，生成可向量化的 chunk。"""
        if not knowledge_point:
            return []

        title = knowledge_point.get("title", "")
        content = knowledge_point.get("content", "")
        subject = knowledge_point.get("subject", "")
        chapter_id = knowledge_point.get("chapter_id", "")
        difficulty = knowledge_point.get("difficulty", "")
        frequency = knowledge_point.get("frequency", "")
        key_points = knowledge_point.get("key_points", [])

        if not content or len(content.strip()) < 10:
            return []

        chunks = []
        header = f"【{subject or '考试'} - {chapter_id or '未知章节'} - {title or '知识点'}】"

        if len(content) <= self.chunk_size:
            chunks.append(Chunk(
                content=f"{header}\n{content}",
                metadata={
                    "subject": subject,
                    "chapter_id": chapter_id,
                    "title": title,
                    "difficulty": difficulty,
                    "frequency": frequency,
                    "key_points": ",".join(key_points) if key_points else "",
                    "content_type": "exam_knowledge",
                }
            ))
        else:
            paras = content.split("\n")
            buffer = ""
            for para in paras:
                if len(buffer) + len(para) > self.chunk_size and len(buffer) >= self.min_chunk_size:
                    chunks.append(Chunk(
                        content=f"{header}\n{buffer}",
                        metadata={
                            "subject": subject,
                            "chapter_id": chapter_id,
                            "title": title,
                            "difficulty": difficulty,
                            "frequency": frequency,
                            "content_type": "exam_knowledge",
                        }
                    ))
                    if self.chunk_overlap > 0 and len(buffer) > self.chunk_overlap:
                        buffer = buffer[-self.chunk_overlap:] + "\n" + para
                    else:
                        buffer = para
                else:
                    buffer += ("\n" if buffer else "") + para
            if buffer:
                chunks.append(Chunk(
                    content=f"{header}\n{buffer}",
                    metadata={
                        "subject": subject,
                        "chapter_id": chapter_id,
                        "title": title,
                        "difficulty": difficulty,
                        "frequency": frequency,
                        "content_type": "exam_knowledge",
                    }
                ))

        return chunks