RAG MCP Tool

import os
import time
import uuid
from typing import List, Dict, Set
from .config import AppConfig
from .storage import RAGStorage
from .utils import is_text_file, read_file_content, chunk_text

from .logger import logger

class Indexer:
    def __init__(self, target_dir: str, config: AppConfig):
        self.target_dir = os.path.abspath(target_dir)
        self.config = config
        self.storage = RAGStorage(self.target_dir, config)

    def index(self):
        logger.info(f"Indexing directory: {self.target_dir}")
        self.storage.initialize()
        
        # Get existing files in DB
        existing_data = self.storage.collection.get(include=['metadatas'])
        
        existing_files: Dict[str, float] = {} # path -> mtime
        file_ids: Dict[str, List[str]] = {} # path -> list of chunk ids
        
        if existing_data and existing_data['ids']:
            for i, meta in enumerate(existing_data['metadatas']):
                if meta and 'file_path' in meta:
                    fpath = meta['file_path']
                    fmtime = meta.get('mtime', 0)
                    existing_files[fpath] = fmtime
                    if fpath not in file_ids:
                        file_ids[fpath] = []
                    file_ids[fpath].append(existing_data['ids'][i])

        # Walk directory
        current_files: Set[str] = set()
        files_to_process = []
        
        for root, dirs, files in os.walk(self.target_dir):
            # Skip hidden directories
            dirs[:] = [d for d in dirs if not d.startswith('.')]
                
            for file in files:
                file_path = os.path.join(root, file)
                if not is_text_file(file_path):
                    continue
                    
                current_files.add(file_path)
                mtime = os.path.getmtime(file_path)
                
                # Check if needs update
                # Use a small epsilon for float comparison
                if file_path not in existing_files or abs(existing_files[file_path] - mtime) > 1e-6:
                    files_to_process.append(file_path)

        # Identify deleted files
        files_to_delete = []
        for fpath in existing_files:
            if fpath not in current_files:
                files_to_delete.append(fpath)

        # Delete removed files
        if files_to_delete:
            logger.info(f"Removing {len(files_to_delete)} deleted files from index...")
            all_ids_to_delete = []
            for fpath in files_to_delete:
                if fpath in file_ids:
                    all_ids_to_delete.extend(file_ids[fpath])
            if all_ids_to_delete:
                self.storage.collection.delete(ids=all_ids_to_delete)

        # Process new/modified files
        # First, delete old chunks for modified files
        ids_to_remove_for_update = []
        for fpath in files_to_process:
            if fpath in file_ids:
                ids_to_remove_for_update.extend(file_ids[fpath])
        
        if ids_to_remove_for_update:
             self.storage.collection.delete(ids=ids_to_remove_for_update)

        logger.info(f"Processing {len(files_to_process)} files...")
        
        for i, file_path in enumerate(files_to_process):
            try:
                content = read_file_content(file_path)
                if not content:
                    continue
                    
                chunks = chunk_text(content, self.config.processing.chunk_count)
                
                if not chunks:
                    continue
                    
                mtime = os.path.getmtime(file_path)
                file_name = os.path.basename(file_path)
                
                doc_ids = [str(uuid.uuid4()) for _ in chunks]
                metadatas = []
                for chunk in chunks:
                    metadatas.append({
                        "file_path": file_path,
                        "file_name": file_name,
                        "mtime": mtime,
                        "chunk_index": chunks.index(chunk),
                        "total_chunks": len(chunks)
                    })
                
                self.storage.add_documents(chunks, metadatas, doc_ids)
                
                if (i + 1) % 10 == 0:
                    logger.info(f"Processed {i + 1}/{len(files_to_process)} files")
                    
            except Exception as e:
                logger.error(f"Error processing {file_path}: {e}")
                
        logger.info("Indexing complete.")

    def clean(self):
        self.storage.clear()