"""Provide service-layer logic for text chunker.""" import re from typing import List, Dict, Optional, Tuple from dataclasses import dataclass, field from loguru import logger # Keep service responsibilities explicit so downstream behavior stays predictable. @dataclass class ChunkMetadata: """Represent the Chunk Metadata type.""" doc_id: str = "" doc_name: str = "" chunk_id: str = "" section_number: str = "" # Keep service responsibilities explicit so downstream behavior stays predictable. section_title: str = "" # Keep service responsibilities explicit so downstream behavior stays predictable. clause_number: str = "" # Keep service responsibilities explicit so downstream behavior stays predictable. page_number: int = 0 start_position: int = 0 # Keep service responsibilities explicit so downstream behavior stays predictable. end_position: int = 0 # Keep service responsibilities explicit so downstream behavior stays predictable. regulation_type: str = "" # Keep service responsibilities explicit so downstream behavior stays predictable. version: str = "" @dataclass class TextChunk: """Represent the Text Chunk type.""" content: str metadata: ChunkMetadata token_count: int = 0 # Keep service responsibilities explicit so downstream behavior stays predictable. class RegulationChunker: """Represent the Regulation Chunker type.""" # Keep service responsibilities explicit so downstream behavior stays predictable. CHAPTER_PATTERN = re.compile(r'^第[一二三四五六七八九十百]+章\s+[^\n]+') SECTION_PATTERN = re.compile(r'^第[一二三四五六七八九十百]+节\s+[^\n]+') CLAUSE_PATTERN = re.compile(r'^第[一二三四五六七八九十百]+条\s') # Keep service responsibilities explicit so downstream behavior stays predictable. SUB_ITEM_PATTERN = re.compile(r'^[\(（][一二三四五六七八九十]+[\)）]\s') NUMBER_ITEM_PATTERN = re.compile(r'^[\d]+[\.、]\s') def __init__( self, chunk_size: int = 512, chunk_overlap: int = 50, max_chunk_size: int = 2048, min_chunk_size: int = 100 ): """Initialize the Regulation Chunker instance.""" self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap self.max_chunk_size = max_chunk_size self.min_chunk_size = min_chunk_size def chunk_document( self, markdown_text: str, doc_id: str = "", doc_name: str = "", regulation_type: str = "", version: str = "" ) -> List[TextChunk]: """Handle chunk document for the Regulation Chunker instance.""" logger.info(f"开始分块文档: {doc_name}") # Keep service responsibilities explicit so downstream behavior stays predictable. sections = self._split_by_sections(markdown_text) # Keep service responsibilities explicit so downstream behavior stays predictable. chunks = [] global_position = 0 for section_num, section_title, section_content, section_start in sections: # Keep service responsibilities explicit so downstream behavior stays predictable. clause_chunks = self._split_by_clauses( section_content, section_num, section_title, section_start + global_position ) for chunk_content, clause_num, clause_title, start_pos, end_pos in clause_chunks: # Keep service responsibilities explicit so downstream behavior stays predictable. if len(chunk_content) > self.max_chunk_size: sub_chunks = self._split_long_clause( chunk_content, clause_num, clause_title ) for sub_content, sub_start, sub_end in sub_chunks: chunk = self._create_chunk( sub_content, doc_id, doc_name, section_num, section_title, clause_num, sub_start + start_pos, sub_end + start_pos, regulation_type, version ) chunks.append(chunk) else: chunk = self._create_chunk( chunk_content, doc_id, doc_name, section_num, section_title, clause_num, start_pos, end_pos, regulation_type, version ) chunks.append(chunk) logger.success(f"分块完成，共{len(chunks)}个chunk") return chunks def _split_by_sections(self, markdown_text: str) -> List[Tuple[str, str, str, int]]: """Handle split by sections for this module for the Regulation Chunker instance.""" sections = [] lines = markdown_text.split('\n') current_section_num = "" current_section_title = "" current_section_content = [] current_section_start = 0 for i, line in enumerate(lines): # Keep service responsibilities explicit so downstream behavior stays predictable. chapter_match = self.CHAPTER_PATTERN.match(line.strip()) section_match = self.SECTION_PATTERN.match(line.strip()) if chapter_match or section_match: # Keep service responsibilities explicit so downstream behavior stays predictable. if current_section_content: content = '\n'.join(current_section_content) sections.append(( current_section_num, current_section_title, content, current_section_start )) # Keep service responsibilities explicit so downstream behavior stays predictable. current_section_start = sum(len(l) + 1 for l in lines[:i]) current_section_content = [] if chapter_match: current_section_num = line.strip() current_section_title = self._extract_title(line.strip()) else: current_section_num = line.strip() current_section_title = self._extract_title(line.strip()) current_section_content.append(line) # Keep service responsibilities explicit so downstream behavior stays predictable. if current_section_content: content = '\n'.join(current_section_content) sections.append(( current_section_num, current_section_title, content, current_section_start )) # Keep service responsibilities explicit so downstream behavior stays predictable. if not sections: sections.append(( "", "全文", markdown_text, 0 )) return sections def _split_by_clauses( self, section_content: str, section_num: str, section_title: str, section_start: int ) -> List[Tuple[str, str, str, int, int]]: """Handle split by clauses for this module for the Regulation Chunker instance.""" clauses = [] lines = section_content.split('\n') current_clause_num = "" current_clause_title = "" current_clause_content = [] current_clause_start = section_start for i, line in enumerate(lines): # Keep service responsibilities explicit so downstream behavior stays predictable. clause_match = self.CLAUSE_PATTERN.match(line.strip()) if clause_match: # Keep service responsibilities explicit so downstream behavior stays predictable. if current_clause_content: content = '\n'.join(current_clause_content) end_pos = current_clause_start + len(content) clauses.append(( content, current_clause_num, current_clause_title, current_clause_start, end_pos )) # Keep service responsibilities explicit so downstream behavior stays predictable. current_clause_start = section_start + sum(len(l) + 1 for l in lines[:i]) current_clause_content = [] current_clause_num = self._extract_clause_number(line.strip()) current_clause_title = line.strip() current_clause_content.append(line) # Keep service responsibilities explicit so downstream behavior stays predictable. if current_clause_content: content = '\n'.join(current_clause_content) end_pos = current_clause_start + len(content) clauses.append(( content, current_clause_num, current_clause_title, current_clause_start, end_pos )) # Keep service responsibilities explicit so downstream behavior stays predictable. if not clauses: clauses.append(( section_content, "", section_title, section_start, section_start + len(section_content) )) return clauses def _split_long_clause( self, content: str, clause_num: str, clause_title: str ) -> List[Tuple[str, int, int]]: """Handle split long clause for this module for the Regulation Chunker instance.""" sub_chunks = [] lines = content.split('\n') # Keep service responsibilities explicit so downstream behavior stays predictable. has_sub_items = any( self.SUB_ITEM_PATTERN.match(line.strip()) or self.NUMBER_ITEM_PATTERN.match(line.strip()) for line in lines ) if has_sub_items: # Keep service responsibilities explicit so downstream behavior stays predictable. current_sub_content = [] current_sub_start = 0 for i, line in enumerate(lines): is_sub_item = ( self.SUB_ITEM_PATTERN.match(line.strip()) or self.NUMBER_ITEM_PATTERN.match(line.strip()) ) if is_sub_item and current_sub_content: sub_content = '\n'.join(current_sub_content) sub_end = current_sub_start + len(sub_content) if len(sub_content) >= self.min_chunk_size: sub_chunks.append((sub_content, current_sub_start, sub_end)) current_sub_content = [] current_sub_start = sum(len(l) + 1 for l in lines[:i]) current_sub_content.append(line) # Keep service responsibilities explicit so downstream behavior stays predictable. if current_sub_content: sub_content = '\n'.join(current_sub_content) sub_end = current_sub_start + len(sub_content) sub_chunks.append((sub_content, current_sub_start, sub_end)) else: # Keep service responsibilities explicit so downstream behavior stays predictable. paragraphs = [] current_para = [] for line in lines: if line.strip(): current_para.append(line) else: if current_para: paragraphs.append('\n'.join(current_para)) current_para = [] if current_para: paragraphs.append('\n'.join(current_para)) # Keep service responsibilities explicit so downstream behavior stays predictable. current_chunk = [] current_length = 0 chunk_start = 0 for para in paragraphs: if current_length + len(para) > self.chunk_size and current_chunk: chunk_content = '\n'.join(current_chunk) chunk_end = chunk_start + len(chunk_content) sub_chunks.append((chunk_content, chunk_start, chunk_end)) current_chunk = [] current_length = 0 chunk_start = chunk_end current_chunk.append(para) current_length += len(para) # Keep service responsibilities explicit so downstream behavior stays predictable. if current_chunk: chunk_content = '\n'.join(current_chunk) chunk_end = chunk_start + len(chunk_content) sub_chunks.append((chunk_content, chunk_start, chunk_end)) return sub_chunks def _extract_title(self, header_line: str) -> str: """Handle extract title for this module for the Regulation Chunker instance.""" # Keep service responsibilities explicit so downstream behavior stays predictable. title = re.sub(r'^第[一二三四五六七八九十百]+[章节]\s+', '', header_line) return title.strip() def _extract_clause_number(self, clause_line: str) -> str: """Handle extract clause number for this module for the Regulation Chunker instance.""" match = self.CLAUSE_PATTERN.match(clause_line) if match: return match.group(0).strip() return "" def _create_chunk( self, content: str, doc_id: str, doc_name: str, section_num: str, section_title: str, clause_num: str, start_pos: int, end_pos: int, regulation_type: str, version: str ) -> TextChunk: """Handle create chunk for this module for the Regulation Chunker instance.""" # Keep service responsibilities explicit so downstream behavior stays predictable. content = content.strip() # Keep service responsibilities explicit so downstream behavior stays predictable. token_count = int(len(content) * 0.7) # Keep service responsibilities explicit so downstream behavior stays predictable. # Keep service responsibilities explicit so downstream behavior stays predictable. chunk_id = f"{doc_id}_{section_num}_{clause_num}_{start_pos}" metadata = ChunkMetadata( doc_id=doc_id, doc_name=doc_name, chunk_id=chunk_id, section_number=section_num, section_title=section_title, clause_number=clause_num, start_position=start_pos, end_position=end_pos, regulation_type=regulation_type, version=version ) return TextChunk( content=content, metadata=metadata, token_count=token_count ) def chunk_regulation_document( markdown_text: str, doc_id: str = "", doc_name: str = "", regulation_type: str = "", version: str = "", chunk_size: int = 512 ) -> List[TextChunk]: """Handle chunk regulation document.""" chunker = RegulationChunker(chunk_size=chunk_size) return chunker.chunk_document( markdown_text, doc_id, doc_name, regulation_type, version )