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

"""Script to chunk and index local legal documents using Contextual Retrieval

This script:
1. Cleans up existing indexes
2. Reads legal documents from local laws directory  
3. Chunks them with paragraph-aware boundaries (soft limit 1024, hard limit 2048)
4. Generates contextual descriptions for each chunk
5. Indexes them with contextual enhancement
"""

import os
import sys
import json
import logging
from typing import List, Dict, Any, Optional
from pathlib import Path
from dataclasses import dataclass
import time
import requests

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent))

from config import Config, ChunkingConfig, LLMConfig
from contextual_chunking import ContextualChunker, ContextualChunk

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Configuration
RETRIEVAL_PIPELINE_URL = "http://localhost:4242"  # Default retrieval pipeline URL
LAWS_DIR = Path("laws")  # Local laws directory

# Custom chunking configuration for legal documents
LEGAL_CHUNKING_CONFIG = ChunkingConfig(
    chunk_size=1024,  # Soft limit
    max_chunk_size=2048,  # Hard limit
    chunk_overlap=200,
    respect_paragraph_boundary=True,
    min_chunk_size=500
)


class ContextualLegalIndexer:
    """Handles contextual chunking and indexing of local legal documents"""
    
    def __init__(self, 
                 laws_dir: Path = LAWS_DIR, 
                 pipeline_url: str = RETRIEVAL_PIPELINE_URL,
                 use_contextual: bool = True,
                 llm_config: Optional[LLMConfig] = None):
        self.laws_dir = laws_dir
        self.pipeline_url = pipeline_url
        self.use_contextual = use_contextual
        
        # Initialize contextual chunker
        self.chunker = ContextualChunker(
            chunking_config=LEGAL_CHUNKING_CONFIG,
            llm_config=llm_config or LLMConfig(),
            use_contextual=use_contextual
        )
        
        self.stats = {
            "documents_processed": 0,
            "chunks_created": 0,
            "contextual_chunks": 0,
            "chunks_indexed": 0,
            "total_context_tokens": 0,
            "errors": 0,
            "categories_processed": set()
        }
        
        # Document store for tracking (must match tools.py expectation)
        self.doc_store_path = Path("document_store.json")
        
        logger.info(f"Initialized contextual indexer for local laws in {laws_dir}")
        logger.info(f"Pipeline URL: {pipeline_url}")
        logger.info(f"Contextual mode: {use_contextual}")
    
    def cleanup_existing_index(self):
        """Clean up existing indexes and document store"""
        logger.info("Cleaning up existing indexes...")
        
        # Clean local document store
        if self.doc_store_path.exists():
            try:
                # Load existing store to get document IDs
                with open(self.doc_store_path, 'r', encoding='utf-8') as f:
                    existing_docs = json.load(f)
                    
                logger.info(f"Found {len(existing_docs)} existing documents in store")
                
                # Clear the store
                self.doc_store_path.unlink()
                logger.info("Cleared local document store")
                
            except Exception as e:
                logger.error(f"Error cleaning document store: {e}")
        
        # Try to clear the retrieval pipeline
        try:
            response = requests.delete(f"{self.pipeline_url}/clear", timeout=30)
            if response.status_code == 200:
                logger.info("Cleared retrieval pipeline index")
            else:
                logger.warning(f"Failed to clear pipeline: {response.status_code}")
        except Exception as e:
            logger.warning(f"Could not clear retrieval pipeline: {e}")
        
        logger.info("Cleanup complete")
    
    def get_all_legal_documents(self) -> List[Dict[str, Any]]:
        """Get all legal documents from local laws directory"""
        documents = []
        
        if not self.laws_dir.exists():
            logger.error(f"Laws directory not found: {self.laws_dir}")
            return documents
        
        # Iterate through category directories
        for category_dir in sorted(self.laws_dir.iterdir()):
            if not category_dir.is_dir():
                continue
            
            category_name = category_dir.name
            logger.info(f"Found category: {category_name}")
            
            # Find all .md files in this category
            for md_file in category_dir.glob("*.md"):
                doc_info = {
                    "path": md_file,
                    "name": md_file.stem,  # filename without extension
                    "category": category_name,
                    "full_name": md_file.name
                }
                documents.append(doc_info)
        
        logger.info(f"Found {len(documents)} legal documents")
        return documents
    
    def read_document(self, doc_info: Dict[str, Any]) -> Optional[str]:
        """Read a legal document from disk"""
        try:
            doc_path = doc_info["path"]
            content = doc_path.read_text(encoding='utf-8')
            logger.debug(f"Read {doc_info['name']} ({len(content)} chars)")
            return content
        except Exception as e:
            logger.error(f"Error reading {doc_info['name']}: {e}")
            self.stats["errors"] += 1
            return None
    
    def generate_document_id(self, doc_info: Dict[str, Any]) -> str:
        """
        Generate a semantically meaningful document ID from file name.
        
        Examples:
        - "宪法.md" -> "宪法"
        - "劳动法（2018-12-29）.md" -> "劳动法_2018-12-29"
        - "民法典/第一编_总则.md" -> "民法典_第一编_总则"
        """
        # Use the file stem (name without extension)
        base_name = doc_info["name"]
        
        # Clean up the name: remove problematic characters but keep Chinese and alphanumeric
        import re
        # Replace parentheses with underscore
        clean_name = base_name.replace('（', '_').replace('）', '')
        clean_name = clean_name.replace('(', '_').replace(')', '')
        # Replace spaces and other separators with underscore
        clean_name = re.sub(r'[\s\-]+', '_', clean_name)
        # Remove trailing underscores
        clean_name = clean_name.strip('_')
        
        # Don't add category prefix - just use the clean file name
        # This makes the IDs cleaner and more readable
        doc_id = clean_name
        
        # Ensure the ID is not too long
        if len(doc_id) > 100:
            # Truncate but keep it meaningful
            doc_id = doc_id[:97] + "..."
        
        logger.debug(f"Generated document ID: {doc_info['name']} -> {doc_id}")
        return doc_id
    
    def process_document(self, doc_info: Dict[str, Any], content: str, 
                        index_immediately: bool = True) -> List[ContextualChunk]:
        """Process a document using contextual chunking with optional immediate indexing"""
        # Generate semantically meaningful document ID
        doc_id = self.generate_document_id(doc_info)
        
        # Prepare document metadata
        doc_metadata = {
            "category": doc_info["category"],
            "file_name": doc_info["full_name"],
            "document_type": "legal",
            "language": "zh-CN"
        }
        
        # Create callback for immediate indexing if enabled
        indexed_chunks = []
        
        def on_chunk_ready(chunk: ContextualChunk):
            """Callback to index chunk immediately when ready"""
            # Update metadata
            chunk.metadata["category"] = doc_info["category"]
            chunk.metadata["doc_title"] = doc_info["name"]
            
            if index_immediately:
                # Index the chunk immediately
                success = self.index_chunk(chunk)
                if success:
                    logger.info(f"  → Indexed chunk {chunk.chunk_index + 1} immediately")
                    indexed_chunks.append(chunk)
                else:
                    logger.warning(f"  → Failed to index chunk {chunk.chunk_index + 1}")
        
        # Use the contextual chunker with callback
        logger.info(f"Chunking {doc_info['name']} with contextual enhancement...")
        chunks = self.chunker.chunk_document(
            text=content,
            doc_id=doc_id,
            doc_metadata=doc_metadata,
            on_chunk_ready=on_chunk_ready if index_immediately else None
        )
        
        # If not indexing immediately, update metadata for all chunks
        if not index_immediately:
            for chunk in chunks:
                chunk.metadata["category"] = doc_info["category"]
                chunk.metadata["doc_title"] = doc_info["name"]
        
        return chunks
    
    def index_chunk(self, chunk: ContextualChunk) -> bool:
        """Index a contextual chunk in the retrieval pipeline"""
        try:
            # Prepare the indexing request
            # Use contextualized text for better retrieval
            index_data = {
                "text": chunk.contextualized_text if self.use_contextual else chunk.text,
                "doc_id": chunk.chunk_id,
                "metadata": {
                    **chunk.metadata,
                    "original_text": chunk.text,
                    "context": chunk.context,
                    "chunk_index": chunk.chunk_index,
                    "char_count": chunk.char_count,
                    "contextual": self.use_contextual
                }
            }
            
            # Send to retrieval pipeline
            response = requests.post(
                f"{self.pipeline_url}/index",
                json=index_data,
                headers={"Content-Type": "application/json"}, timeout=30
            )
            
            if response.status_code == 200:
                self.stats["chunks_indexed"] += 1
                if chunk.context:
                    self.stats["contextual_chunks"] += 1
                
                # IMPORTANT: Save to document store immediately after successful indexing
                self.save_chunk_to_document_store(chunk)
                
                return True
            else:
                logger.warning(f"Failed to index chunk {chunk.chunk_id}: {response.status_code}")
                return False
                
        except Exception as e:
            logger.error(f"Error indexing chunk {chunk.chunk_id}: {e}")
            self.stats["errors"] += 1
            return False
    
    def save_chunk_to_document_store(self, chunk: ContextualChunk):
        """Save individual chunk to document store immediately (compatible with tools.py)"""
        try:
            # Load existing store or create new
            if self.doc_store_path.exists():
                with open(self.doc_store_path, 'r', encoding='utf-8') as f:
                    doc_store = json.load(f)
            else:
                doc_store = {}
            
            # Save chunk in format expected by tools.py get_document method
            doc_store[chunk.chunk_id] = {
                "doc_id": chunk.chunk_id,
                "content": chunk.contextualized_text if self.use_contextual else chunk.text,
                "metadata": {
                    **chunk.metadata,
                    "original_text": chunk.text,
                    "context": chunk.context,
                    "chunk_index": chunk.chunk_index,
                    "char_count": chunk.char_count,
                    "contextual": self.use_contextual
                }
            }
            
            # Write immediately - this is critical per user requirement
            with open(self.doc_store_path, 'w', encoding='utf-8') as f:
                json.dump(doc_store, f, ensure_ascii=False, indent=2)
            
            logger.debug(f"Saved chunk {chunk.chunk_id} to document store")
            
        except Exception as e:
            logger.error(f"Error saving chunk {chunk.chunk_id} to document store: {e}")
    
    def save_document_info(self, doc_info: Dict[str, Any], chunks: List[ContextualChunk]):
        """Save document summary information to local store"""
        # Load existing store or create new
        if self.doc_store_path.exists():
            with open(self.doc_store_path, 'r', encoding='utf-8') as f:
                doc_store = json.load(f)
        else:
            doc_store = {}
        
        # Add document info with meaningful ID
        doc_id = self.generate_document_id(doc_info)
        
        # Calculate statistics
        total_context_tokens = sum(c.context_tokens for c in chunks)
        avg_context_tokens = total_context_tokens / len(chunks) if chunks else 0
        
        # Store document-level summary (compatible with tools.py format)
        doc_store[doc_id] = {
            "doc_id": doc_id,
            "content": f"Document: {doc_info['name']}\nCategory: {doc_info['category']}\n\nThis is a summary entry for the complete document. Individual chunks are stored separately.",
            "metadata": {
                "title": doc_info["name"],
                "category": doc_info["category"],
                "file": str(doc_info["path"]),
                "chunks": len(chunks),
                "contextual_chunks": sum(1 for c in chunks if c.context),
                "total_chars": sum(c.char_count for c in chunks),
                "total_context_tokens": total_context_tokens,
                "avg_context_tokens": avg_context_tokens,
                "indexed_at": time.strftime("%Y-%m-%d %H:%M:%S"),
                "chunk_ids": [c.chunk_id for c in chunks],
                "is_summary": True
            }
        }
        
        # Save store
        with open(self.doc_store_path, 'w', encoding='utf-8') as f:
            json.dump(doc_store, f, ensure_ascii=False, indent=2)
    
    def process_all_documents(self, 
                            max_docs: Optional[int] = None, 
                            categories: Optional[List[str]] = None,
                            batch_size: int = 10):
        """Process all legal documents with contextual chunking"""
        start_time = time.time()
        
        # Clean up first
        self.cleanup_existing_index()
        
        # Get all documents
        all_documents = self.get_all_legal_documents()
        
        # Filter by categories if specified
        if categories:
            all_documents = [d for d in all_documents if any(cat in d["category"] for cat in categories)]
        
        # Limit documents if specified
        if max_docs:
            all_documents = all_documents[:max_docs]
        
        logger.info(f"Processing {len(all_documents)} documents...")
        
        for i, doc_info in enumerate(all_documents):
            logger.info(f"\n[{i+1}/{len(all_documents)}] Processing: {doc_info['name']}")
            logger.info(f"  Category: {doc_info['category']}")
            
            # Track category
            self.stats["categories_processed"].add(doc_info["category"])
            
            # Read document
            content = self.read_document(doc_info)
            if not content:
                continue
            
            # Process with contextual chunking and immediate indexing
            try:
                # Process and index chunks immediately as they're generated
                chunks = self.process_document(doc_info, content, index_immediately=True)
                self.stats["chunks_created"] += len(chunks)
                
                # Log contextual stats
                if self.use_contextual:
                    contextual_count = sum(1 for c in chunks if c.context)
                    logger.info(f"  ✓ Created and indexed {len(chunks)} chunks ({contextual_count} with context)")
                else:
                    logger.info(f"  ✓ Created and indexed {len(chunks)} chunks (no context)")
                
                # Save document info immediately after processing
                self.save_document_info(doc_info, chunks)
                
                self.stats["documents_processed"] += 1
                
                # Update token statistics
                if self.use_contextual:
                    self.stats["total_context_tokens"] += sum(c.context_tokens for c in chunks)
                
            except Exception as e:
                logger.error(f"Error processing {doc_info['name']}: {e}")
                self.stats["errors"] += 1
        
        elapsed = time.time() - start_time
        
        # Get final chunker statistics
        chunker_stats = self.chunker.get_statistics()
        
        # Print statistics
        self._print_statistics(elapsed, chunker_stats)
    
    def _print_statistics(self, elapsed_time: float, chunker_stats: Dict[str, Any]):
        """Print processing statistics"""
        print("\n" + "="*60)
        print("CONTEXTUAL INDEXING COMPLETE")
        print("="*60)
        print(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
        print(f"Time elapsed: {elapsed_time:.2f} seconds")
        print(f"Categories processed: {len(self.stats['categories_processed'])}")
        
        # Show categories
        categories_list = sorted(self.stats['categories_processed'])
        for cat in categories_list:
            print(f"  • {cat}")
        
        print(f"\nDocuments processed: {self.stats['documents_processed']}")
        print(f"Chunks created: {self.stats['chunks_created']}")
        
        if self.use_contextual:
            print(f"Contextual chunks: {self.stats['contextual_chunks']}")
            print(f"Chunks indexed: {self.stats['chunks_indexed']}")
            
            # Token usage
            print(f"\nContext Generation Statistics:")
            print(f"  Total context tokens: {chunker_stats.get('total_context_tokens', 0):,}")
            print(f"  Average tokens per chunk: {chunker_stats.get('avg_context_tokens', 0):.1f}")
            print(f"  Total generation time: {chunker_stats.get('total_generation_time', 0):.2f}s")
            print(f"  Average time per chunk: {chunker_stats.get('avg_generation_time', 0):.2f}s")
            
            # Cache statistics
            print(f"\nCache Performance:")
            print(f"  Cache hits: {chunker_stats.get('cache_hits', 0)}")
            print(f"  Cache misses: {chunker_stats.get('cache_misses', 0)}")
            print(f"  Hit rate: {chunker_stats.get('cache_hit_rate', 0):.1%}")
            
            # Cost estimation
            print(f"\nCost Estimation:")
            print(f"  Estimated cost: ${chunker_stats.get('estimated_cost', 0):.2f}")
        else:
            print(f"Chunks indexed: {self.stats['chunks_indexed']}")
        
        print(f"\nErrors: {self.stats['errors']}")
        
        if self.stats['chunks_created'] > 0:
            avg_chunks = self.stats['chunks_created'] / max(1, self.stats['documents_processed'])
            print(f"Average chunks per document: {avg_chunks:.1f}")
            
        if elapsed_time > 0 and self.stats['documents_processed'] > 0:
            docs_per_min = (self.stats['documents_processed'] / elapsed_time) * 60
            print(f"Processing speed: {docs_per_min:.1f} docs/minute")
        
        print("="*60 + "\n")
    
    def compare_retrieval_methods(self, test_queries: Optional[List[str]] = None):
        """Compare contextual vs non-contextual retrieval"""
        if not test_queries:
            test_queries = [
                "什么是合同的成立条件",
                "劳动者的基本权利有哪些",
                "如何处理交通事故责任",
                "什么是正当防卫",
                "公司股东的权利和义务"
            ]
        
        print("\n" + "="*60)
        print("RETRIEVAL COMPARISON TEST")
        print("="*60)
        print(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
        
        for query in test_queries:
            try:
                response = requests.post(
                    f"{self.pipeline_url}/search",
                    json={
                        "query": query,
                        "mode": "hybrid",
                        "top_k": 10,
                        "rerank_top_k": 5
                    }, timeout=30
                )
                
                if response.status_code == 200:
                    results = response.json()
                    print(f"\n✓ Query: '{query}'")
                    
                    if "results" in results:
                        print(f"  Found {len(results['results'])} results")
                        
                        # Show top 3 results
                        for i, result in enumerate(results['results'][:3]):
                            score = result.get('score', result.get('rerank_score', 'N/A'))
                            metadata = result.get('metadata', {})
                            
                            print(f"\n  Result {i+1}:")
                            print(f"    Score: {score}")
                            print(f"    Category: {metadata.get('category', 'Unknown')}")
                            print(f"    Document: {metadata.get('doc_title', 'Unknown')}")
                            
                            # Show context if available
                            if metadata.get('context'):
                                print(f"    Context: {metadata['context'][:100]}...")
                            
                            # Show text preview
                            text = result.get('text', '')
                            # If contextual, try to extract original text
                            if metadata.get('contextual') and metadata.get('original_text'):
                                text = metadata['original_text']
                            
                            print(f"    Preview: {text[:150]}...")
                    else:
                        print(f"  No results found")
                else:
                    print(f"✗ Query '{query}' failed: {response.status_code}")
                    
            except Exception as e:
                print(f"✗ Error testing '{query}': {e}")
        
        print("\n" + "="*60 + "\n")


def main():
    """Main function"""
    import argparse
    
    parser = argparse.ArgumentParser(description="Index local legal documents with contextual retrieval")
    parser.add_argument("--pipeline-url", default=RETRIEVAL_PIPELINE_URL, help="Retrieval pipeline URL")
    parser.add_argument("--max-docs", type=int, help="Maximum number of documents to process")
    parser.add_argument("--categories", nargs="+", help="Specific categories to process")
    parser.add_argument("--no-contextual", action="store_true", help="Disable contextual enhancement")
    parser.add_argument("--no-cleanup", action="store_true", help="Don't clean existing indexes")
    parser.add_argument("--compare", action="store_true", help="Run comparison test after indexing")
    parser.add_argument("--batch-size", type=int, default=10, help="Batch size for indexing")
    parser.add_argument("--llm-provider", default="kimi", help="LLM provider for context generation")
    parser.add_argument("--llm-model", help="Specific LLM model to use")
    
    args = parser.parse_args()
    
    # Configure LLM if specified
    llm_config = None
    if args.llm_provider or args.llm_model:
        llm_config = LLMConfig(
            provider=args.llm_provider,
            model=args.llm_model
        )
    
    # Create indexer
    indexer = ContextualLegalIndexer(
        pipeline_url=args.pipeline_url,
        use_contextual=not args.no_contextual,
        llm_config=llm_config
    )
    
    # Skip cleanup if requested
    if args.no_cleanup:
        indexer.cleanup_existing_index = lambda: logger.info("Skipping cleanup (--no-cleanup flag)")
    
    # Process documents
    indexer.process_all_documents(
        max_docs=args.max_docs,
        categories=args.categories,
        batch_size=args.batch_size
    )
    
    # Run comparison test if requested
    if args.compare:
        indexer.compare_retrieval_methods()


if __name__ == "__main__":
    main()