Build lightweight AI agent admin

2026-06-08 18:14:59 +08:00
commit e164840f43
2530 changed files with 435693 additions and 0 deletions
@@ -0,0 +1,172 @@
+"""
+Markdown 结构化分块策略
+
+按 Markdown 标题层级分割文档，保留文档结构信息
+"""
+import re
+import logging
+from typing import Dict, Any, List, Optional
+
+from .base import BaseChunker, ChunkResult
+
+logger = logging.getLogger(__name__)
+
+
+class MarkdownChunker(BaseChunker):
+    """
+    Markdown 分块器
+    
+    按标题层级分割 Markdown 文档，每个标题下的内容作为一个分块
+    如果单个标题下的内容超过 chunk_size，则使用递归分割
+    """
+
+    def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
+        """按 Markdown 标题分块"""
+        if not text or not text.strip():
+            return []
+
+        text = self._clean_text(text)
+        metadata = metadata or {}
+
+        # 按标题分割
+        sections = self._split_by_headers(text)
+
+        # 处理每个 section
+        raw_chunks = []
+        for section in sections:
+            header = section.get('header', '')
+            content = section.get('content', '')
+            level = section.get('level', 0)
+
+            if not content.strip():
+                continue
+
+            # 组合标题和内容
+            full_text = f"{header}\n{content}" if header else content
+
+            if len(full_text) <= self.chunk_size:
+                raw_chunks.append({
+                    'content': full_text.strip(),
+                    'metadata': {
+                        **metadata,
+                        'header': header,
+                        'header_level': level,
+                    }
+                })
+            else:
+                # 内容超长，递归分割
+                sub_chunks = self._split_long_section(content, header)
+                for i, sub in enumerate(sub_chunks):
+                    raw_chunks.append({
+                        'content': sub.strip(),
+                        'metadata': {
+                            **metadata,
+                            'header': header,
+                            'header_level': level,
+                            'sub_chunk': i,
+                        }
+                    })
+
+        # 合并过小的分块
+        merged = self._merge_small_section_chunks(raw_chunks)
+
+        # 构建结果
+        results = []
+        for i, item in enumerate(merged):
+            if item['content'].strip():
+                results.append(ChunkResult(
+                    content=item['content'],
+                    position=i,
+                    metadata=item.get('metadata', {}),
+                ))
+
+        return results
+
+    def _split_by_headers(self, text: str) -> List[Dict[str, Any]]:
+        """按 Markdown 标题分割"""
+        # 匹配 Markdown 标题: # Title, ## Title, ### Title 等
+        header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
+
+        sections = []
+        last_end = 0
+        last_header = ''
+        last_level = 0
+
+        for match in header_pattern.finditer(text):
+            # 保存前一个 section 的内容
+            if last_end > 0 or match.start() > 0:
+                content = text[last_end:match.start()]
+                if content.strip() or last_header:
+                    sections.append({
+                        'header': last_header,
+                        'content': content.strip(),
+                        'level': last_level,
+                    })
+
+            last_header = match.group(0)
+            last_level = len(match.group(1))
+            last_end = match.end()
+
+        # 最后一个 section
+        remaining = text[last_end:]
+        if remaining.strip() or last_header:
+            sections.append({
+                'header': last_header,
+                'content': remaining.strip(),
+                'level': last_level,
+            })
+
+        # 如果没有找到任何标题，整个文本作为一个 section
+        if not sections:
+            sections.append({
+                'header': '',
+                'content': text.strip(),
+                'level': 0,
+            })
+
+        return sections
+
+    def _split_long_section(self, content: str, header: str = '') -> List[str]:
+        """分割超长的 section 内容"""
+        from .recursive import RecursiveChunker
+
+        chunker = RecursiveChunker(
+            chunk_size=self.chunk_size,
+            chunk_overlap=self.chunk_overlap,
+        )
+        results = chunker.chunk(content)
+
+        chunks = []
+        for i, result in enumerate(results):
+            # 第一个分块带上标题
+            if i == 0 and header:
+                chunks.append(f"{header}\n{result.content}")
+            else:
+                chunks.append(result.content)
+
+        return chunks if chunks else [content]
+
+    def _merge_small_section_chunks(self, chunks: List[Dict], min_size: int = 80) -> List[Dict]:
+        """合并过小的 section 分块"""
+        if not chunks:
+            return []
+
+        merged = []
+        buffer = None
+
+        for chunk in chunks:
+            if buffer is None:
+                buffer = chunk
+            elif len(buffer['content']) < min_size and len(buffer['content']) + len(chunk['content']) <= self.chunk_size:
+                buffer['content'] = buffer['content'] + "\n\n" + chunk['content']
+            else:
+                merged.append(buffer)
+                buffer = chunk
+
+        if buffer:
+            if merged and len(buffer['content']) < min_size:
+                merged[-1]['content'] = merged[-1]['content'] + "\n\n" + buffer['content']
+            else:
+                merged.append(buffer)
+
+        return merged