rag-db/generate_db.py

#!/usr/bin/env python3
"""
RAG Database Generation Script for gprMax Documentation
Generates a ChromaDB vector database from gprMax documentation
"""

import os
import sys
import shutil
import argparse
import logging
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any
import json
import hashlib

import chromadb
import git
from tqdm import tqdm

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


class GprMaxDocumentProcessor:
    """Process gprMax documentation files for vectorization"""
    
    SUPPORTED_EXTENSIONS = {'.rst', '.md', '.txt'}
    CHUNK_SIZE = 1000  # Characters per chunk
    CHUNK_OVERLAP = 200  # Overlap between chunks
    
    def __init__(self, repo_path: Path):
        self.repo_path = repo_path
        self.doc_path = repo_path / "docs"
        
    def extract_documents(self) -> List[Dict[str, Any]]:
        """Extract and chunk all documentation files"""
        documents = []
        
        if not self.doc_path.exists():
            logger.warning(f"Documentation path {self.doc_path} does not exist")
            return documents
            
        for file_path in self._find_doc_files():
            try:
                chunks = self._process_file(file_path)
                documents.extend(chunks)
            except Exception as e:
                logger.error(f"Error processing {file_path}: {e}")
                
        logger.info(f"Extracted {len(documents)} document chunks")
        return documents
    
    def _find_doc_files(self) -> List[Path]:
        """Find all documentation files"""
        doc_files = []
        for ext in self.SUPPORTED_EXTENSIONS:
            doc_files.extend(self.doc_path.rglob(f"*{ext}"))
        return doc_files
    
    def _process_file(self, file_path: Path) -> List[Dict[str, Any]]:
        """Process a single file into chunks"""
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
            
        # Calculate relative path for metadata
        rel_path = file_path.relative_to(self.repo_path)
        
        # Create chunks with overlap
        chunks = []
        for i in range(0, len(content), self.CHUNK_SIZE - self.CHUNK_OVERLAP):
            chunk_text = content[i:i + self.CHUNK_SIZE]
            
            # Skip empty or very small chunks
            if len(chunk_text.strip()) < 50:
                continue
                
            # Generate unique ID for chunk
            chunk_id = hashlib.md5(f"{rel_path}_{i}_{chunk_text[:50]}".encode()).hexdigest()
            
            chunks.append({
                "id": chunk_id,
                "text": chunk_text,
                "metadata": {
                    "source": str(rel_path),
                    "file_type": file_path.suffix,
                    "chunk_index": len(chunks),
                    "char_start": i,
                    "char_end": min(i + self.CHUNK_SIZE, len(content))
                }
            })
            
        return chunks


# Removed custom embedding model - using ChromaDB's default


class ChromaRAGDatabase:
    """ChromaDB-based RAG database"""
    
    def __init__(self, db_path: Path):
        self.db_path = db_path
        
        # Initialize ChromaDB with persistent storage
        self.client = chromadb.PersistentClient(path=str(db_path))
        
        # Collection name with version for easy updates
        self.collection_name = "gprmax_docs_v1"
        
    def create_collection(self, recreate: bool = False):
        """Create or get the document collection"""
        if recreate:
            try:
                self.client.delete_collection(self.collection_name)
                logger.info(f"Deleted existing collection: {self.collection_name}")
            except:
                pass
                
        # Let ChromaDB use its default embedding function
        self.collection = self.client.create_collection(
            name=self.collection_name,
            metadata={"created_at": datetime.now().isoformat()}
        )
        logger.info(f"Created collection: {self.collection_name}")
        
    def add_documents(self, documents: List[Dict[str, Any]]):
        """Add documents to the collection"""
        if not documents:
            logger.warning("No documents to add")
            return
            
        # Prepare data for ChromaDB
        ids = [doc["id"] for doc in documents]
        texts = [doc["text"] for doc in documents]
        metadatas = [doc["metadata"] for doc in documents]
        
        # Add to collection in batches (ChromaDB will generate embeddings automatically)
        batch_size = 100
        logger.info(f"Adding {len(documents)} documents to database...")
        for i in tqdm(range(0, len(ids), batch_size), desc="Adding to database"):
            end_idx = min(i + batch_size, len(ids))
            self.collection.add(
                ids=ids[i:end_idx],
                documents=texts[i:end_idx],
                metadatas=metadatas[i:end_idx]
                # No embeddings parameter - ChromaDB will generate them
            )
            
        logger.info(f"Added {len(documents)} documents to database")
        
        # Verify documents were added
        actual_count = self.collection.count()
        logger.info(f"Verified collection now contains {actual_count} documents")
        
    def save_metadata(self):
        """Save database metadata for reference"""
        # Get fresh count
        doc_count = self.collection.count()
        
        metadata = {
            "created_at": datetime.now().isoformat(),
            "embedding_model": "ChromaDB Default (all-MiniLM-L6-v2)",
            "collection_name": self.collection_name,
            "chunk_size": GprMaxDocumentProcessor.CHUNK_SIZE,
            "chunk_overlap": GprMaxDocumentProcessor.CHUNK_OVERLAP,
            "total_documents": doc_count
        }
        
        metadata_path = self.db_path / "metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)
            
        logger.info(f"Saved metadata to {metadata_path}")


def clone_gprmax_repo(target_dir: Path) -> Path:
    """Clone or update gprMax repository"""
    repo_path = target_dir / "gprMax"
    
    if repo_path.exists():
        logger.info(f"Updating existing repository at {repo_path}")
        repo = git.Repo(repo_path)
        repo.remotes.origin.pull()
    else:
        logger.info(f"Cloning gprMax repository to {repo_path}")
        git.Repo.clone_from(
            "https://github.com/gprMax/gprMax.git",
            repo_path,
            depth=1  # Shallow clone for faster download
        )
        
    return repo_path


def main():
    parser = argparse.ArgumentParser(description="Generate RAG database from gprMax documentation")
    parser.add_argument(
        "--db-path",
        type=Path,
        default=Path(__file__).parent / "chroma_db",
        help="Path to store the ChromaDB database"
    )
    parser.add_argument(
        "--temp-dir",
        type=Path,
        default=Path(__file__).parent / "temp",
        help="Temporary directory for cloning repository"
    )
    parser.add_argument(
        "--recreate",
        action="store_true",
        help="Recreate database from scratch (delete existing)"
    )
    
    args = parser.parse_args()
    
    try:
        # Step 1: Clone/update gprMax repository
        logger.info("Step 1: Fetching gprMax repository...")
        repo_path = clone_gprmax_repo(args.temp_dir)
        
        # Step 2: Process documentation
        logger.info("Step 2: Processing documentation files...")
        processor = GprMaxDocumentProcessor(repo_path)
        documents = processor.extract_documents()
        
        if not documents:
            logger.error("No documents found to process")
            return 1
            
        # Step 3: Create database
        logger.info("Step 3: Creating vector database...")
        db = ChromaRAGDatabase(args.db_path)
        db.create_collection(recreate=args.recreate)
        
        # Step 4: Add documents
        logger.info("Step 4: Adding documents to database...")
        db.add_documents(documents)
        
        # Step 5: Save metadata
        db.save_metadata()
        
        logger.info(f"✅ Database successfully created at {args.db_path}")
        logger.info(f"Total documents: {len(documents)}")
        
        # Cleanup temp files if needed
        if args.temp_dir.exists() and args.temp_dir != args.db_path.parent:
            logger.info("Cleaning up temporary files...")
            shutil.rmtree(args.temp_dir, ignore_errors=True)
            
        return 0
        
    except Exception as e:
        logger.error(f"Failed to generate database: {e}")
        return 1


if __name__ == "__main__":
    sys.exit(main())