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