lexiredact 0.1.0__py3-none-any.whl

lexiredact/__init__.py

@@ -0,0 +1,139 @@
+"""
+LexiRedact - Privacy-Preserving RAG Middleware
+
+A Python SDK for protecting PII in vector databases while maintaining
+semantic search quality through intelligent embedding and redaction.
+
+Key Features:
+- Automatic PII detection and redaction using Microsoft Presidio
+- Embedding generation from original text (Shadow Mode architecture)
+- Only sanitized text stored in vector databases
+- Redis caching for performance optimization
+- Pluggable architecture via dependency injection
+- Comprehensive metrics and tracking
+
+Basic Usage:
+    >>> import lexiredact as vs
+    >>> 
+    >>> # Create pipeline with defaults
+    >>> pipeline = vs.IngestionPipeline()
+    >>> await pipeline.initialize()
+    >>> 
+    >>> # Process documents
+    >>> doc = vs.Document(id="1", text="Contact John at john@example.com")
+    >>> result = await pipeline.process_document(doc)
+    >>> 
+    >>> print(result.clean_text)  # "Contact <PERSON> at <EMAIL_ADDRESS>"
+    >>> print(result.pii_entities)  # ["PERSON", "EMAIL_ADDRESS"]
+    >>> 
+    >>> await pipeline.shutdown()
+
+Custom Configuration:
+    >>> from lexiredact import IngestionPipeline, load_config
+    >>> 
+    >>> config = load_config(config_dict={
+    ...     "embedding_model": "BAAI/bge-base-en-v1.5",
+    ...     "cache_backend": "redis",
+    ...     "redis_host": "localhost"
+    ... })
+    >>> 
+    >>> pipeline = IngestionPipeline(config=config)
+
+Custom Components:
+    >>> from lexiredact import IngestionPipeline
+    >>> from lexiredact.interfaces import Embedder
+    >>> 
+    >>> class MyEmbedder(Embedder):
+    ...     # Custom implementation
+    ...     pass
+    >>> 
+    >>> pipeline = IngestionPipeline(embedder=MyEmbedder())
+"""
+
+__version__ = "0.1.0"
+
+# Core pipeline
+from .pipeline import IngestionPipeline, Document, ProcessedDocument
+
+# Configuration
+from .config import load_config, get_default_config, save_config_to_yaml
+
+# Privacy components
+from .privacy import PIIDetector, PIIRedactor, PIIPolicy
+
+# Interfaces (for custom implementations)
+from .interfaces import CacheBackend, Embedder, VectorStore, Tracker
+
+# Default implementations
+from .implementations import (
+    MemoryCache,  
+    RedisCache,
+    GenericCache,
+    FastEmbedEmbedder,
+    GenericEmbedder,  
+    ChromaVectorStore,
+    GenericVectorStore, 
+    MLflowTracker,
+)
+
+# Metrics
+from .metrics import (
+    MetricsCollector,
+    AggregateStats,
+    RetrievalAggregateStats,
+    RetrievalMetricsEvaluator,
+    RetrievalQueryMetrics,
+)
+
+# Utilities
+from .utils import hash_text, generate_cache_key, Timer
+
+__all__ = [
+    # Version
+    "__version__",
+    
+    # Core
+    "IngestionPipeline",
+    "Document",
+    "ProcessedDocument",
+    
+    # Configuration
+    "load_config",
+    "get_default_config",
+    "save_config_to_yaml",
+    
+    # Privacy
+    "PIIDetector",
+    "PIIRedactor",
+    "PIIPolicy",
+    
+    # Interfaces
+    "CacheBackend",
+    "Embedder",
+    "VectorStore",
+    "Tracker",
+    
+    #custom models
+    "GenericCache",
+    "GenericEmbedder", 
+    "GenericVectorStore",
+    
+    # Implementations
+    "MemoryCache",
+    "RedisCache",
+    "FastEmbedEmbedder",
+    "ChromaVectorStore",
+    "MLflowTracker",
+    
+    # Metrics
+    "MetricsCollector",
+    "AggregateStats",
+    "RetrievalAggregateStats",
+    "RetrievalMetricsEvaluator",
+    "RetrievalQueryMetrics",
+    
+    # Utils
+    "hash_text",
+    "generate_cache_key",
+    "Timer",
+]

lexiredact/chunking/__init__.py

@@ -0,0 +1,18 @@
+"""
+Document chunking module for LexiRedact.
+
+Converts PDFs and large text documents into manageable chunks
+suitable for embedding and PII detection.
+"""
+
+from .chunker import Chunk, DocumentChunker, ChunkingStrategy
+from .json_exporter import JSONExporter
+from .pdf_loader import PDFLoader
+
+__all__ = [
+    "Chunk",
+    "DocumentChunker",
+    "ChunkingStrategy",
+    "PDFLoader",
+    "JSONExporter",
+]

lexiredact/chunking/chunker.py

@@ -0,0 +1,321 @@
+"""
+Core document chunking logic.
+Splits large documents into smaller chunks with overlap.
+"""
+
+from typing import List, Dict, Any, Optional, Literal
+from dataclasses import dataclass
+import uuid
+from enum import Enum
+import re
+
+
+class ChunkingStrategy(str, Enum):
+    """Chunking strategies."""
+    FIXED_SIZE = "fixed_size"          # Fixed token/char chunks
+    SENTENCE = "sentence"               # Split by sentences
+    PARAGRAPH = "paragraph"             # Split by paragraphs
+    HYBRID = "hybrid"                   # Sentences grouped into chunks
+
+
+@dataclass
+class Chunk:
+    """Single chunk of text."""
+    id: str
+    text: str
+    chunk_index: int
+    start_char: int
+    end_char: int
+    metadata: Dict[str, Any]
+
+
+class DocumentChunker:
+    """
+    Convert large documents into LexiRedact-compatible chunks.
+
+    Support for:
+    - Fixed-size chunking (tokens or characters)
+    - Sentence-based chunking
+    - Paragraph-based chunking
+    - Overlap between chunks
+    """
+
+    def __init__(
+        self,
+        chunk_size: int = 512,              # Max characters per chunk
+        overlap: int = 100,                 # Overlap between chunks (chars)
+        strategy: ChunkingStrategy = ChunkingStrategy.FIXED_SIZE,
+        preserve_sentences: bool = True,    # Don't split mid-sentence
+    ):
+        """
+        Initialize chunker.
+
+        Args:
+            chunk_size: Target chunk size in characters
+            overlap: Overlap between chunks (to preserve context)
+            strategy: Chunking strategy to use
+            preserve_sentences: Don't split in middle of sentence
+        """
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be greater than 0")
+        if overlap < 0:
+            raise ValueError("overlap must be greater than or equal to 0")
+        if overlap >= chunk_size:
+            raise ValueError("overlap must be smaller than chunk_size")
+
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+        self.strategy = strategy
+        self.preserve_sentences = preserve_sentences
+
+    def chunk_text(
+        self,
+        text: str,
+        doc_id: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[Chunk]:
+        """
+        Chunk a document.
+
+        Args:
+            text: Document text to chunk
+            doc_id: Original document ID (source)
+            metadata: Optional metadata to attach to chunks
+
+        Returns:
+            List of Chunk objects
+        """
+        metadata = metadata or {}
+
+        if self.strategy == ChunkingStrategy.FIXED_SIZE:
+            return self._chunk_fixed_size(text, doc_id, metadata)
+        elif self.strategy == ChunkingStrategy.SENTENCE:
+            return self._chunk_by_sentence(text, doc_id, metadata)
+        elif self.strategy == ChunkingStrategy.PARAGRAPH:
+            return self._chunk_by_paragraph(text, doc_id, metadata)
+        elif self.strategy == ChunkingStrategy.HYBRID:
+            return self._chunk_hybrid(text, doc_id, metadata)
+        else:
+            raise ValueError(f"Unknown strategy: {self.strategy}")
+
+    def _chunk_fixed_size(
+        self,
+        text: str,
+        doc_id: str,
+        metadata: Dict[str, Any]
+    ) -> List[Chunk]:
+        """Split into fixed-size chunks with overlap."""
+        chunks = []
+        chunk_index = 0
+        start = 0
+
+        while start < len(text):
+            end = min(start + self.chunk_size, len(text))
+
+            # If preserve_sentences, adjust end to not split mid-sentence
+            if self.preserve_sentences and end < len(text):
+                window = text[start:end]
+                boundaries = list(re.finditer(r"[.!?](?=\s|$)|\n", window))
+                if boundaries:
+                    end = start + boundaries[-1].end()
+
+            chunk_text = text[start:end].strip()
+
+            if chunk_text:  # Skip empty chunks
+                chunk = Chunk(
+                    id=f"{doc_id}_chunk_{chunk_index}",
+                    text=chunk_text,
+                    chunk_index=chunk_index,
+                    start_char=start,
+                    end_char=end,
+                    metadata={
+                        **metadata,
+                        "source_doc_id": doc_id,
+                        "chunk_number": chunk_index,
+                        "strategy": self.strategy.value,
+                    }
+                )
+                chunks.append(chunk)
+                chunk_index += 1
+
+            if end >= len(text):
+                break
+
+            # Move start position (with overlap)
+            start = end - self.overlap
+
+        return chunks
+
+    def _chunk_by_sentence(
+        self,
+        text: str,
+        doc_id: str,
+        metadata: Dict[str, Any]
+    ) -> List[Chunk]:
+        """Split by sentences while preserving punctuation-heavy tokens."""
+        sentences = self._split_sentences(text)
+        chunks = []
+        chunk_index = 0
+
+        current_chunk = []
+        current_size = 0
+        start_char = 0
+
+        for sentence in sentences:
+            sentence_size = len(sentence)
+
+            # If adding this sentence exceeds chunk_size, save current chunk
+            if current_size + sentence_size > self.chunk_size and current_chunk:
+                chunk_text = ' '.join(current_chunk).strip()
+                if chunk_text:
+                    chunk = Chunk(
+                        id=f"{doc_id}_chunk_{chunk_index}",
+                        text=chunk_text,
+                        chunk_index=chunk_index,
+                        start_char=start_char,
+                        end_char=start_char + len(chunk_text),
+                        metadata={
+                            **metadata,
+                            "source_doc_id": doc_id,
+                            "chunk_number": chunk_index,
+                            "strategy": self.strategy.value,
+                        }
+                    )
+                    chunks.append(chunk)
+                    chunk_index += 1
+
+                start_char += len(chunk_text) + 1
+                current_chunk = [sentence]
+                current_size = sentence_size
+            else:
+                current_chunk.append(sentence)
+                current_size += sentence_size
+
+        # Add remaining chunk
+        if current_chunk:
+            chunk_text = ' '.join(current_chunk).strip()
+            chunk = Chunk(
+                id=f"{doc_id}_chunk_{chunk_index}",
+                text=chunk_text,
+                chunk_index=chunk_index,
+                start_char=start_char,
+                end_char=start_char + len(chunk_text),
+                metadata={
+                    **metadata,
+                    "source_doc_id": doc_id,
+                    "chunk_number": chunk_index,
+                    "strategy": self.strategy.value,
+                }
+            )
+            chunks.append(chunk)
+
+        return chunks
+
+    def _chunk_by_paragraph(
+        self,
+        text: str,
+        doc_id: str,
+        metadata: Dict[str, Any]
+    ) -> List[Chunk]:
+        """Split by paragraphs (double newline)."""
+        paragraphs = text.split('\n\n')
+        chunks = []
+        chunk_index = 0
+        start_char = 0
+
+        for para in paragraphs:
+            para = para.strip()
+            if para:
+                chunk = Chunk(
+                    id=f"{doc_id}_chunk_{chunk_index}",
+                    text=para,
+                    chunk_index=chunk_index,
+                    start_char=start_char,
+                    end_char=start_char + len(para),
+                    metadata={
+                        **metadata,
+                        "source_doc_id": doc_id,
+                        "chunk_number": chunk_index,
+                        "strategy": self.strategy.value,
+                    }
+                )
+                chunks.append(chunk)
+                chunk_index += 1
+                start_char += len(para) + 2  # +2 for '\n\n'
+
+        return chunks
+
+    def _chunk_hybrid(
+        self,
+        text: str,
+        doc_id: str,
+        metadata: Dict[str, Any]
+    ) -> List[Chunk]:
+        """
+        Hybrid: Group sentences into chunks of target size.
+        Better than fixed_size because sentences stay together.
+        """
+        sentences = self._split_sentences(text)
+        chunks = []
+        chunk_index = 0
+        start_char = 0
+
+        current_chunk = []
+        current_size = 0
+
+        for sentence in sentences:
+            # If adding sentence exceeds size and we have content, save chunk
+            if current_size + len(sentence) > self.chunk_size and current_chunk:
+                chunk_text = ' '.join(current_chunk).strip()
+                chunk = Chunk(
+                    id=f"{doc_id}_chunk_{chunk_index}",
+                    text=chunk_text,
+                    chunk_index=chunk_index,
+                    start_char=start_char,
+                    end_char=start_char + len(chunk_text),
+                    metadata={
+                        **metadata,
+                        "source_doc_id": doc_id,
+                        "chunk_number": chunk_index,
+                        "strategy": self.strategy.value,
+                    }
+                )
+                chunks.append(chunk)
+                chunk_index += 1
+                start_char += len(chunk_text) + 1
+                current_chunk = [sentence]
+                current_size = len(sentence)
+            else:
+                current_chunk.append(sentence)
+                current_size += len(sentence)
+
+        # Add final chunk
+        if current_chunk:
+            chunk_text = ' '.join(current_chunk).strip()
+            chunk = Chunk(
+                id=f"{doc_id}_chunk_{chunk_index}",
+                text=chunk_text,
+                chunk_index=chunk_index,
+                start_char=start_char,
+                end_char=start_char + len(chunk_text),
+                metadata={
+                    **metadata,
+                    "source_doc_id": doc_id,
+                    "chunk_number": chunk_index,
+                    "strategy": self.strategy.value,
+                }
+            )
+            chunks.append(chunk)
+
+        return chunks
+
+    def _split_sentences(self, text: str) -> List[str]:
+        """
+        Split text on sentence boundaries without breaking emails or domains.
+        """
+        normalized = text.strip()
+        if not normalized:
+            return []
+
+        parts = re.split(r"(?<=[.!?])\s+(?=[A-Z0-9])", normalized)
+        return [part.strip() for part in parts if part.strip()]

lexiredact/chunking/json_exporter.py

@@ -0,0 +1,103 @@
+"""
+Export chunks as JSON in LexiRedact format.
+"""
+
+import json
+from typing import List, Dict, Any
+from pathlib import Path
+from .chunker import Chunk
+
+
+class JSONExporter:
+    """Export chunks to JSON format compatible with LexiRedact."""
+
+    @staticmethod
+    def to_lexiredact_format(
+        chunks: List[Chunk],
+        output_path: str | Path | None = None,
+        wrap_documents: bool = False,
+    ) -> str | List[Dict[str, Any]] | Dict[str, List[Dict[str, Any]]]:
+        """
+        Convert chunks to LexiRedact JSON format.
+
+        Args:
+            chunks: List of Chunk objects
+            output_path: Optional path to save JSON file
+            wrap_documents: Export CLI-ready payload as
+                {"documents": [...]} when True
+
+        Returns:
+            JSON string (if output_path) or payload object (if None)
+
+        Example output format:
+            [
+              {
+                "id": "doc1_chunk_0",
+                "text": "First chunk of text...",
+                "metadata": {
+                  "source_doc_id": "doc1",
+                  "chunk_number": 0,
+                  "chunk_size": 512
+                }
+              },
+              ...
+            ]
+        """
+        chunk_dicts = [
+            {
+                "id": chunk.id,
+                "text": chunk.text,
+                "metadata": chunk.metadata
+            }
+            for chunk in chunks
+        ]
+        payload: List[Dict[str, Any]] | Dict[str, List[Dict[str, Any]]]
+        payload = {"documents": chunk_dicts} if wrap_documents else chunk_dicts
+
+        if output_path:
+            output_path = Path(output_path)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+
+            with open(output_path, 'w', encoding='utf-8') as f:
+                json.dump(payload, f, indent=2, ensure_ascii=False)
+
+            return json.dumps(payload, indent=2, ensure_ascii=False)
+
+        return payload
+
+    @staticmethod
+    def to_cli_input(
+        chunks: List[Chunk],
+        output_path: str | Path | None = None,
+    ) -> str | Dict[str, List[Dict[str, Any]]]:
+        """Export chunks in the JSON shape expected by `lexiredact process`."""
+        return JSONExporter.to_lexiredact_format(
+            chunks,
+            output_path=output_path,
+            wrap_documents=True,
+        )
+
+    @staticmethod
+    def to_jsonl(
+        chunks: List[Chunk],
+        output_path: str | Path,
+    ) -> None:
+        """
+        Export chunks as JSONL (one JSON per line).
+        Useful for streaming/large datasets.
+
+        Args:
+            chunks: List of Chunk objects
+            output_path: Path to save JSONL file
+        """
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(output_path, 'w', encoding='utf-8') as f:
+            for chunk in chunks:
+                line = {
+                    "id": chunk.id,
+                    "text": chunk.text,
+                    "metadata": chunk.metadata
+                }
+                f.write(json.dumps(line, ensure_ascii=False) + '\n')

lexiredact/chunking/pdf_loader.py

@@ -0,0 +1,86 @@
+"""
+PDF extraction for document chunking.
+"""
+
+from typing import Dict, Any
+from pathlib import Path
+
+
+class PDFLoader:
+    """Load and extract text from PDF files."""
+
+    def __init__(self, use_ocr: bool = False):
+        """
+        Initialize PDF loader.
+
+        Args:
+            use_ocr: Use OCR for scanned PDFs (requires pytesseract)
+        """
+        self.use_ocr = use_ocr
+
+    @staticmethod
+    def extract_text(pdf_path: str | Path) -> str:
+        """
+        Extract text from PDF.
+
+        Args:
+            pdf_path: Path to PDF file
+
+        Returns:
+            Extracted text
+
+        Raises:
+            ImportError: If pypdf not installed
+            FileNotFoundError: If PDF not found
+        """
+        try:
+            from pypdf import PdfReader
+        except ImportError:
+            raise ImportError(
+                "PDF support requires pypdf. "
+                "Install with: pip install pypdf"
+            )
+
+        pdf_path = Path(pdf_path)
+        if not pdf_path.exists():
+            raise FileNotFoundError(f"PDF not found: {pdf_path}")
+
+        text = []
+        with open(pdf_path, 'rb') as file:
+            reader = PdfReader(file)
+            for page in reader.pages:
+                text.append(page.extract_text())
+
+        return '\n\n'.join(text)
+
+    @staticmethod
+    def extract_metadata(pdf_path: str | Path) -> Dict[str, Any]:
+        """
+        Extract metadata from PDF.
+
+        Args:
+            pdf_path: Path to PDF file
+
+        Returns:
+            Metadata dictionary
+        """
+        try:
+            from pypdf import PdfReader
+        except ImportError:
+            raise ImportError("PDF support requires pypdf")
+
+        pdf_path = Path(pdf_path)
+        metadata = {}
+
+        with open(pdf_path, 'rb') as file:
+            reader = PdfReader(file)
+            if reader.metadata:
+                metadata = {
+                    "title": reader.metadata.get("/Title", ""),
+                    "author": reader.metadata.get("/Author", ""),
+                    "subject": reader.metadata.get("/Subject", ""),
+                    "pages": len(reader.pages),
+                    "pdf_file": str(pdf_path.name),
+                }
+
+        return metadata