ai-agent-book/chapter3/contextual-retrieval/chunking.py

"""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()