ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/contextual-retrieval/chunking.py

@@ -0,0 +1,423 @@
+"""Document chunking and indexing script"""
+
+import os
+import json
+import hashlib
+import logging
+import requests
+from typing import List, Dict, Any, Optional, Tuple
+from pathlib import Path
+from datetime import datetime
+from config import ChunkingConfig, KnowledgeBaseConfig, KnowledgeBaseType
+
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+class DocumentChunker:
+    """Document chunking with configurable strategies"""
+    
+    def __init__(self, config: Optional[ChunkingConfig] = None):
+        self.config = config or ChunkingConfig()
+        
+    def chunk_text(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
+        """
+        Chunk text into smaller segments.
+        
+        Args:
+            text: Document text to chunk
+            doc_id: Document identifier
+            
+        Returns:
+            List of chunks with metadata
+        """
+        chunks = []
+        
+        if self.config.respect_paragraph_boundary:
+            chunks = self._chunk_by_paragraphs(text, doc_id)
+        else:
+            chunks = self._chunk_by_size(text, doc_id)
+        
+        logger.info(f"Created {len(chunks)} chunks for document {doc_id}")
+        return chunks
+    
+    def _chunk_by_paragraphs(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
+        """Chunk text respecting paragraph boundaries"""
+        paragraphs = text.split('\n\n')
+        chunks = []
+        current_chunk = []
+        current_size = 0
+        
+        for para in paragraphs:
+            para = para.strip()
+            if not para:
+                continue
+            
+            para_size = len(para)
+            
+            # If single paragraph exceeds max size, split it
+            if para_size > self.config.max_chunk_size:
+                # Save current chunk if exists
+                if current_chunk:
+                    chunk_text = '\n\n'.join(current_chunk)
+                    chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
+                    current_chunk = []
+                    current_size = 0
+                
+                # Split large paragraph
+                sentences = self._split_into_sentences(para)
+                for sent in sentences:
+                    if len(sent) > self.config.max_chunk_size:
+                        # Force split very long sentences
+                        for i in range(0, len(sent), self.config.chunk_size):
+                            sub_chunk = sent[i:i + self.config.chunk_size]
+                            chunks.append(self._create_chunk(sub_chunk, doc_id, len(chunks)))
+                    else:
+                        chunks.append(self._create_chunk(sent, doc_id, len(chunks)))
+                continue
+            
+            # Check if adding this paragraph exceeds chunk size
+            if current_size + para_size > self.config.chunk_size and current_chunk:
+                # Save current chunk
+                chunk_text = '\n\n'.join(current_chunk)
+                chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
+                
+                # Start new chunk with overlap
+                if self.config.chunk_overlap > 0 and current_chunk:
+                    # Keep last paragraph for overlap
+                    current_chunk = [current_chunk[-1], para]
+                    current_size = len(current_chunk[0]) + para_size
+                else:
+                    current_chunk = [para]
+                    current_size = para_size
+            else:
+                current_chunk.append(para)
+                current_size += para_size
+        
+        # Save final chunk
+        if current_chunk:
+            chunk_text = '\n\n'.join(current_chunk)
+            if len(chunk_text) >= self.config.min_chunk_size:
+                chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
+        
+        return chunks
+    
+    def _chunk_by_size(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
+        """Simple size-based chunking"""
+        chunks = []
+        
+        for i in range(0, len(text), self.config.chunk_size - self.config.chunk_overlap):
+            chunk_text = text[i:i + self.config.chunk_size]
+            
+            if len(chunk_text) >= self.config.min_chunk_size:
+                chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
+        
+        return chunks
+    
+    def _split_into_sentences(self, text: str) -> List[str]:
+        """Split text into sentences (simple implementation)"""
+        # Simple sentence splitting for Chinese and English
+        import re
+        
+        # Split on common sentence endings
+        sentences = re.split(r'([。！？\.!?]+)', text)
+        
+        # Reconstruct sentences with their endings
+        result = []
+        # Step to the end of the list: re.split with a capturing group yields
+        # [text, delim, text, delim, ..., trailing_text], so stopping at
+        # len(sentences) - 1 dropped the trailing fragment whenever the text
+        # did not end in terminal punctuation (and returned [] for text with
+        # none at all). The strip-and-filter below still discards the empty
+        # tail that re.split produces when the text does end in punctuation.
+        for i in range(0, len(sentences), 2):
+            if i + 1 < len(sentences):
+                result.append(sentences[i] + sentences[i + 1])
+            else:
+                result.append(sentences[i])
+        
+        return [s.strip() for s in result if s.strip()]
+    
+    def _create_chunk(self, text: str, doc_id: str, chunk_index: int) -> Dict[str, Any]:
+        """Create a chunk with metadata"""
+        chunk_id = f"{doc_id}_chunk_{chunk_index}"
+        
+        return {
+            "chunk_id": chunk_id,
+            "doc_id": doc_id,
+            "text": text,
+            "chunk_index": chunk_index,
+            "char_count": len(text),
+            "hash": hashlib.md5(text.encode()).hexdigest()
+        }
+
+
+class DocumentIndexer:
+    """Index documents to knowledge base"""
+    
+    def __init__(self, 
+                 kb_config: Optional[KnowledgeBaseConfig] = None,
+                 chunking_config: Optional[ChunkingConfig] = None):
+        self.kb_config = kb_config or KnowledgeBaseConfig()
+        self.chunker = DocumentChunker(chunking_config)
+        self.indexed_docs = {}
+        
+    def index_file(self, file_path: str, doc_id: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Index a single file.
+        
+        Args:
+            file_path: Path to the file
+            doc_id: Optional document ID
+            
+        Returns:
+            Indexing result
+        """
+        file_path = Path(file_path)
+        if not file_path.exists():
+            return {"error": f"File not found: {file_path}"}
+        
+        # Generate doc_id if not provided
+        if not doc_id:
+            doc_id = file_path.stem
+        
+        # Read file content
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+        except Exception as e:
+            return {"error": f"Error reading file: {e}"}
+        
+        # Chunk the document
+        chunks = self.chunker.chunk_text(content, doc_id)
+        
+        # Index chunks
+        result = self._index_chunks(chunks, doc_id, content)
+        
+        # Store full document
+        self._store_document(doc_id, content, {"source_file": str(file_path)})
+        
+        return result
+    
+    def index_directory(self, dir_path: str, extensions: List[str] = None) -> Dict[str, Any]:
+        """
+        Index all files in a directory.
+        
+        Args:
+            dir_path: Directory path
+            extensions: File extensions to include (e.g., ['.txt', '.md'])
+            
+        Returns:
+            Indexing results
+        """
+        dir_path = Path(dir_path)
+        if not dir_path.exists():
+            return {"error": f"Directory not found: {dir_path}"}
+        
+        extensions = extensions or ['.txt', '.md', '.json']
+        results = {"indexed": [], "errors": []}
+        
+        for file_path in dir_path.rglob('*'):
+            if file_path.is_file() and file_path.suffix in extensions:
+                doc_id = f"{file_path.parent.name}/{file_path.stem}"
+                result = self.index_file(str(file_path), doc_id)
+                
+                if "error" in result:
+                    results["errors"].append({
+                        "file": str(file_path),
+                        "error": result["error"]
+                    })
+                else:
+                    results["indexed"].append({
+                        "file": str(file_path),
+                        "doc_id": doc_id,
+                        "chunks": result.get("chunks_indexed", 0)
+                    })
+        
+        logger.info(f"Indexed {len(results['indexed'])} files, {len(results['errors'])} errors")
+        return results
+    
+    def _index_chunks(self, chunks: List[Dict[str, Any]], doc_id: str, full_content: str) -> Dict[str, Any]:
+        """Index chunks to the knowledge base"""
+        if self.kb_config.type == KnowledgeBaseType.LOCAL:
+            return self._index_to_local(chunks, doc_id)
+        elif self.kb_config.type == KnowledgeBaseType.DIFY:
+            return self._index_to_dify(chunks, doc_id, full_content)
+        else:
+            return {"error": f"Unsupported KB type: {self.kb_config.type}"}
+    
+    def _index_to_local(self, chunks: List[Dict[str, Any]], doc_id: str) -> Dict[str, Any]:
+        """Index to local retrieval pipeline"""
+        indexed_count = 0
+        errors = []
+        
+        for chunk in chunks:
+            try:
+                # Index each chunk
+                response = requests.post(
+                    f"{self.kb_config.local_base_url}/index",
+                    json={
+                        "text": chunk["text"],
+                        "doc_id": chunk["doc_id"],
+                        "metadata": {
+                            "chunk_id": chunk["chunk_id"],
+                            "chunk_index": chunk["chunk_index"],
+                            "char_count": chunk["char_count"]
+                        }
+                    }, timeout=30
+                )
+                response.raise_for_status()
+                indexed_count += 1
+                
+            except Exception as e:
+                errors.append(f"Error indexing chunk {chunk['chunk_id']}: {e}")
+        
+        result = {
+            "doc_id": doc_id,
+            "chunks_indexed": indexed_count,
+            "total_chunks": len(chunks)
+        }
+        
+        if errors:
+            result["errors"] = errors
+        
+        return result
+    
+    def _index_to_dify(self, chunks: List[Dict[str, Any]], doc_id: str, full_content: str) -> Dict[str, Any]:
+        """Index to Dify knowledge base"""
+        if not self.kb_config.dify_api_key:
+            return {"error": "Dify API key not configured"}
+        
+        try:
+            headers = {
+                "Authorization": f"Bearer {self.kb_config.dify_api_key}",
+                "Content-Type": "application/json"
+            }
+            
+            # Dify expects documents, not individual chunks
+            # So we'll create segments from our chunks
+            segments = []
+            for chunk in chunks:
+                segments.append({
+                    "content": chunk["text"],
+                    "keywords": [],  # Can add keywords if needed
+                    "enabled": True
+                })
+            
+            payload = {
+                "name": doc_id,
+                "text": full_content,
+                "indexing_technique": "high_quality",  # or "economy"
+                "process_rule": {
+                    "mode": "custom",
+                    "rules": {
+                        "pre_processing_rules": [],
+                        "segmentation": {
+                            "separator": "\n\n",
+                            "max_tokens": self.chunker.config.chunk_size // 4  # Rough token estimate
+                        }
+                    }
+                }
+            }
+            
+            if self.kb_config.dify_dataset_id:
+                # Add to existing dataset
+                response = requests.post(
+                    f"{self.kb_config.dify_base_url}/datasets/{self.kb_config.dify_dataset_id}/documents",
+                    headers=headers,
+                    json=payload, timeout=30
+                )
+            else:
+                # Create new document
+                response = requests.post(
+                    f"{self.kb_config.dify_base_url}/documents",
+                    headers=headers,
+                    json=payload, timeout=30
+                )
+            
+            response.raise_for_status()
+            
+            return {
+                "doc_id": doc_id,
+                "chunks_indexed": len(chunks),
+                "total_chunks": len(chunks),
+                "dify_response": response.json()
+            }
+            
+        except Exception as e:
+            return {"error": f"Error indexing to Dify: {e}"}
+    
+    def _store_document(self, doc_id: str, content: str, metadata: Dict[str, Any]):
+        """Store full document locally"""
+        # Store in local file for retrieval
+        store_path = self.kb_config.document_store_path
+        
+        try:
+            # Load existing store
+            if os.path.exists(store_path):
+                with open(store_path, 'r', encoding='utf-8') as f:
+                    store = json.load(f)
+            else:
+                store = {}
+            
+            # Add document
+            store[doc_id] = {
+                "doc_id": doc_id,
+                "content": content,
+                "metadata": metadata,
+                "indexed_at": datetime.now().isoformat()
+            }
+            
+            # Save store
+            with open(store_path, 'w', encoding='utf-8') as f:
+                json.dump(store, f, ensure_ascii=False, indent=2)
+            
+            self.indexed_docs[doc_id] = True
+            logger.info(f"Stored document {doc_id}")
+            
+        except Exception as e:
+            logger.error(f"Error storing document: {e}")
+
+
+def main():
+    """Main function for standalone chunking and indexing"""
+    import argparse
+    from config import Config
+    
+    parser = argparse.ArgumentParser(description="Chunk and index documents")
+    parser.add_argument("path", help="File or directory path to index")
+    parser.add_argument("--chunk-size", type=int, default=2048, help="Chunk size in characters")
+    parser.add_argument("--max-chunk-size", type=int, default=1024, help="Max chunk size")
+    parser.add_argument("--overlap", type=int, default=200, help="Chunk overlap")
+    parser.add_argument("--kb-type", choices=["local", "dify"], default="local", help="Knowledge base type")
+    parser.add_argument("--extensions", nargs="+", default=[".txt", ".md"], help="File extensions to index")
+    
+    args = parser.parse_args()
+    
+    # Create config
+    config = Config.from_env()
+    config.chunking.chunk_size = args.chunk_size
+    config.chunking.max_chunk_size = args.max_chunk_size
+    config.chunking.chunk_overlap = args.overlap
+    config.knowledge_base.type = KnowledgeBaseType(args.kb_type)
+    
+    # Create indexer
+    indexer = DocumentIndexer(config.knowledge_base, config.chunking)
+    
+    # Index path
+    path = Path(args.path)
+    if path.is_file():
+        result = indexer.index_file(str(path))
+    elif path.is_dir():
+        result = indexer.index_directory(str(path), args.extensions)
+    else:
+        print(f"Path not found: {path}")
+        return
+    
+    # Print results
+    print(json.dumps(result, indent=2, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()