""" 智能切片器:按药品条目为单位切片 确保每个 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 []