kreuzberg 3.6.2__py3-none-any.whl → 3.8.0__py3-none-any.whl

kreuzberg/_types.py

@@ -100,6 +100,96 @@ class Metadata(TypedDict, total=False):
     width: NotRequired[int]
     """Width of the document page/slide/image, if applicable."""
 
+    # Email-specific fields
+    email_from: NotRequired[str]
+    """Email sender (from field)."""
+    email_to: NotRequired[str]
+    """Email recipient (to field)."""
+    email_cc: NotRequired[str]
+    """Email carbon copy recipients."""
+    email_bcc: NotRequired[str]
+    """Email blind carbon copy recipients."""
+    date: NotRequired[str]
+    """Email date or document date."""
+    attachments: NotRequired[list[str]]
+    """List of attachment names."""
+
+    # Additional metadata fields for various extractors
+    content: NotRequired[str]
+    """Content metadata field."""
+    parse_error: NotRequired[str]
+    """Parse error information."""
+    warning: NotRequired[str]
+    """Warning messages."""
+
+    # Table extraction metadata
+    table_count: NotRequired[int]
+    """Number of tables extracted from the document."""
+    tables_summary: NotRequired[str]
+    """Summary of table extraction results."""
+    quality_score: NotRequired[float]
+    """Quality score for extracted content (0.0-1.0)."""
+
+
+# Cache valid metadata keys at module level for performance
+_VALID_METADATA_KEYS = {
+    "authors",
+    "categories",
+    "citations",
+    "comments",
+    "content",
+    "copyright",
+    "created_at",
+    "created_by",
+    "description",
+    "fonts",
+    "height",
+    "identifier",
+    "keywords",
+    "languages",
+    "license",
+    "modified_at",
+    "modified_by",
+    "organization",
+    "parse_error",
+    "publisher",
+    "references",
+    "status",
+    "subject",
+    "subtitle",
+    "summary",
+    "title",
+    "version",
+    "warning",
+    "width",
+    "email_from",
+    "email_to",
+    "email_cc",
+    "email_bcc",
+    "date",
+    "attachments",
+    "table_count",
+    "tables_summary",
+    "quality_score",
+}
+
+
+def normalize_metadata(data: dict[str, Any] | None) -> Metadata:
+    """Normalize any dict to proper Metadata TypedDict.
+
+    Filters out invalid keys and ensures type safety.
+    """
+    if not data:
+        return {}
+
+    # Filter and return only valid metadata
+    normalized: Metadata = {}
+    for key, value in data.items():
+        if key in _VALID_METADATA_KEYS and value is not None:
+            normalized[key] = value  # type: ignore[literal-required]
+
+    return normalized
+
 
 @dataclass(frozen=True)
 class Entity:
@@ -140,6 +230,45 @@ class ExtractionResult:
         """Converts the ExtractionResult to a dictionary."""
         return asdict(self)
 
+    def export_tables_to_csv(self) -> list[str]:
+        """Export all tables to CSV format.
+
+        Returns:
+            List of CSV strings, one per table
+        """
+        if not self.tables:
+            return []
+
+        from kreuzberg._utils._table import export_table_to_csv
+
+        return [export_table_to_csv(table) for table in self.tables]
+
+    def export_tables_to_tsv(self) -> list[str]:
+        """Export all tables to TSV format.
+
+        Returns:
+            List of TSV strings, one per table
+        """
+        if not self.tables:
+            return []
+
+        from kreuzberg._utils._table import export_table_to_tsv
+
+        return [export_table_to_tsv(table) for table in self.tables]
+
+    def get_table_summaries(self) -> list[dict[str, Any]]:
+        """Get structural information for all tables.
+
+        Returns:
+            List of table structure dictionaries
+        """
+        if not self.tables:
+            return []
+
+        from kreuzberg._utils._table import extract_table_structure_info
+
+        return [extract_table_structure_info(table) for table in self.tables]
+
 
 PostProcessingHook = Callable[[ExtractionResult], ExtractionResult | Awaitable[ExtractionResult]]
 ValidationHook = Callable[[ExtractionResult], None | Awaitable[None]]
@@ -193,6 +322,8 @@ class ExtractionConfig:
     """Configuration for language detection. If None, uses default settings."""
     spacy_entity_extraction_config: SpacyEntityExtractionConfig | None = None
     """Configuration for spaCy entity extraction. If None, uses default settings."""
+    enable_quality_processing: bool = True
+    """Whether to apply quality post-processing to improve extraction results."""
 
     def __post_init__(self) -> None:
         if self.custom_entity_patterns is not None and isinstance(self.custom_entity_patterns, dict):

kreuzberg/_utils/_cache.py

@@ -64,8 +64,23 @@ class KreuzbergCache(Generic[T]):
         Returns:
             Unique cache key string
         """
-        # Sort for consistent hashing  # ~keep
-        cache_str = str(sorted(kwargs.items()))
+        # Use more efficient string building for cache key
+        if not kwargs:
+            return "empty"
+
+        # Build key string efficiently
+        parts = []
+        for key in sorted(kwargs):
+            value = kwargs[key]
+            # Convert common types efficiently
+            if isinstance(value, (str, int, float, bool)):
+                parts.append(f"{key}={value}")
+            elif isinstance(value, bytes):
+                parts.append(f"{key}=bytes:{len(value)}")
+            else:
+                parts.append(f"{key}={type(value).__name__}:{value!s}")
+
+        cache_str = "&".join(parts)
         return hashlib.sha256(cache_str.encode()).hexdigest()[:16]
 
     def _get_cache_path(self, cache_key: str) -> Path:

kreuzberg/_utils/_process_pool.py

@@ -7,7 +7,12 @@ from concurrent.futures import ProcessPoolExecutor
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, TypeVar
 
+import anyio
+import psutil
+from typing_extensions import Self
+
 if TYPE_CHECKING:
+    import types
     from collections.abc import Callable, Generator
 
 T = TypeVar("T")
@@ -62,7 +67,7 @@ def _extract_pdf_text_worker(pdf_path: str) -> tuple[str, str]:
         text_parts = []
         for page in pdf:
             text_page = page.get_textpage()
-            text = text_page.get_text_range()
+            text = text_page.get_text_bounded()
             text_parts.append(text)
             text_page.close()
             page.close()
@@ -98,3 +103,175 @@ def _extract_pdf_images_worker(pdf_path: str, scale: float = 4.25) -> tuple[str,
     finally:
         if pdf:
             pdf.close()
+
+
+class ProcessPoolManager:
+    """Resource-aware process pool manager for CPU-intensive tasks."""
+
+    def __init__(
+        self,
+        max_processes: int | None = None,
+        memory_limit_gb: float | None = None,
+    ) -> None:
+        """Initialize the process pool manager.
+
+        Args:
+            max_processes: Maximum number of processes. Defaults to CPU count.
+            memory_limit_gb: Memory limit in GB. Defaults to 75% of available memory.
+        """
+        self.max_processes = max_processes or mp.cpu_count()
+
+        if memory_limit_gb is None:
+            available_memory = psutil.virtual_memory().available
+            self.memory_limit_bytes = int(available_memory * 0.75)  # Use 75% of available  # ~keep
+        else:
+            self.memory_limit_bytes = int(memory_limit_gb * 1024**3)
+
+        self._executor: ProcessPoolExecutor | None = None
+        self._active_tasks = 0
+
+    def get_optimal_workers(self, task_memory_mb: float = 100) -> int:
+        """Calculate optimal number of workers based on memory constraints.
+
+        Args:
+            task_memory_mb: Estimated memory usage per task in MB.
+
+        Returns:
+            Optimal number of workers.
+        """
+        task_memory_bytes = task_memory_mb * 1024**2
+        memory_based_limit = max(1, int(self.memory_limit_bytes / task_memory_bytes))
+
+        return min(self.max_processes, memory_based_limit)
+
+    def _ensure_executor(self, max_workers: int | None = None) -> ProcessPoolExecutor:
+        """Ensure process pool executor is initialized."""
+        if self._executor is None or getattr(self._executor, "_max_workers", None) != max_workers:
+            if self._executor is not None:
+                self._executor.shutdown(wait=False)
+
+            workers = max_workers or self.max_processes
+            self._executor = ProcessPoolExecutor(max_workers=workers)
+
+        return self._executor
+
+    async def submit_task(
+        self,
+        func: Callable[..., T],
+        *args: Any,
+        task_memory_mb: float = 100,
+    ) -> T:
+        """Submit a task to the process pool.
+
+        Args:
+            func: Function to execute.
+            *args: Positional arguments for the function.
+            task_memory_mb: Estimated memory usage in MB.
+
+        Returns:
+            Result of the function execution.
+        """
+        workers = self.get_optimal_workers(task_memory_mb)
+        self._ensure_executor(workers)
+
+        self._active_tasks += 1
+
+        try:
+            return await anyio.to_thread.run_sync(func, *args)
+        finally:
+            self._active_tasks -= 1
+
+    async def submit_batch(
+        self,
+        func: Callable[..., T],
+        arg_batches: list[tuple[Any, ...]],
+        task_memory_mb: float = 100,
+        max_concurrent: int | None = None,
+    ) -> list[T]:
+        """Submit a batch of tasks to the process pool.
+
+        Args:
+            func: Function to execute.
+            arg_batches: List of argument tuples for each task.
+            task_memory_mb: Estimated memory usage per task in MB.
+            max_concurrent: Maximum concurrent tasks. Defaults to optimal workers.
+
+        Returns:
+            List of results in the same order as input.
+        """
+        if not arg_batches:
+            return []
+
+        workers = self.get_optimal_workers(task_memory_mb)
+        max_concurrent = max_concurrent or workers
+
+        self._ensure_executor(workers)
+
+        semaphore = anyio.CapacityLimiter(max_concurrent)
+
+        async def submit_single(args: tuple[Any, ...]) -> T:
+            async with semaphore:
+                self._active_tasks += 1
+                try:
+                    return await anyio.to_thread.run_sync(func, *args)
+                finally:
+                    self._active_tasks -= 1
+
+        async with anyio.create_task_group() as tg:
+            results: list[T] = [None] * len(arg_batches)  # type: ignore[list-item]
+
+            async def run_task(idx: int, args: tuple[Any, ...]) -> None:
+                results[idx] = await submit_single(args)
+
+            for idx, args in enumerate(arg_batches):
+                tg.start_soon(run_task, idx, args)
+
+        return results
+
+    def get_system_info(self) -> dict[str, Any]:
+        """Get current system resource information."""
+        memory = psutil.virtual_memory()
+        cpu_percent = psutil.cpu_percent(interval=1)
+
+        return {
+            "cpu_count": mp.cpu_count(),
+            "cpu_percent": cpu_percent,
+            "memory_total": memory.total,
+            "memory_available": memory.available,
+            "memory_percent": memory.percent,
+            "active_tasks": self._active_tasks,
+            "max_processes": self.max_processes,
+            "memory_limit": self.memory_limit_bytes,
+        }
+
+    def shutdown(self, wait: bool = True) -> None:
+        """Shutdown the process pool."""
+        if self._executor is not None:
+            self._executor.shutdown(wait=wait)
+            self._executor = None
+
+    def __enter__(self) -> Self:
+        """Context manager entry."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Context manager exit."""
+        self.shutdown()
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Async context manager exit."""
+        self.shutdown()

kreuzberg/_utils/_quality.py

@@ -0,0 +1,237 @@
+"""Quality post-processing utilities for extracted text."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+# Pre-compiled patterns for performance
+_OCR_ARTIFACTS = {
+    # Common OCR misreads
+    "scattered_chars": re.compile(r"\b[a-zA-Z]\s{2,}[a-zA-Z]\s{2,}[a-zA-Z]\b"),
+    "repeated_punctuation": re.compile(r"[.]{3,}|[-]{3,}|[_]{3,}"),
+    "isolated_punctuation": re.compile(r"\s[.,;:!?]\s"),
+    "malformed_words": re.compile(r"\b[a-zA-Z]+[0-9]+[a-zA-Z]+[a-zA-Z0-9]*\b"),
+    "excessive_whitespace": re.compile(r"\s{3,}"),
+    "broken_sentences": re.compile(r"[a-z]\s{3,}[A-Z][a-z]"),
+}
+
+# Combined pattern for faster OCR penalty calculation
+_COMBINED_OCR_PATTERN = re.compile(
+    r"(?P<scattered>\b[a-zA-Z]\s{2,}[a-zA-Z]\s{2,}[a-zA-Z]\b)|"
+    r"(?P<repeated>[.]{3,}|[-]{3,}|[_]{3,})|"
+    r"(?P<isolated>\s[.,;:!?]\s)|"
+    r"(?P<malformed>\b[a-zA-Z]+[0-9]+[a-zA-Z]+[a-zA-Z0-9]*\b)|"
+    r"(?P<whitespace>\s{3,})|"
+    r"(?P<broken>[a-z]\s{3,}[A-Z][a-z])"
+)
+
+# Pre-compiled patterns for text normalization
+_WHITESPACE_NORMALIZE = re.compile(r"[ \t\f\v\r\xa0\u2000-\u200b\u2028\u2029\u3000]+")
+_NEWLINE_NORMALIZE = re.compile(r"\n\s*\n\s*\n+")
+_SENTENCE_DETECT = re.compile(r"[.!?]\s+[A-Z]")
+_PUNCTUATION_DETECT = re.compile(r"[.!?]")
+
+_SCRIPT_PATTERNS = {
+    # JavaScript and CSS content
+    "js_functions": re.compile(r"function\s+\w+\s*\([^)]*\)\s*\{[^}]*\}", re.IGNORECASE),
+    "css_rules": re.compile(r"\.[a-zA-Z][\w-]*\s*\{[^}]*\}", re.IGNORECASE),
+    "script_tags": re.compile(r"<script[^>]*>.*?</script>", re.DOTALL | re.IGNORECASE),
+    "style_tags": re.compile(r"<style[^>]*>.*?</style>", re.DOTALL | re.IGNORECASE),
+}
+
+_NAVIGATION_PATTERNS = {
+    "nav_words": re.compile(r"\b(?:Skip to main content|Back to top|Main navigation|Site navigation)\b", re.IGNORECASE),
+    "breadcrumbs": re.compile(r"(?:Home\s*[>»]\s*|[>»]\s*){2,}"),
+    "pagination": re.compile(
+        r"\b(?:Page \d+ of \d+|First page|Last page|Previous page|Next page|^\d+ of \d+$)\b", re.IGNORECASE
+    ),
+}
+
+
+def calculate_quality_score(text: str, metadata: dict[str, Any] | None = None) -> float:
+    """Calculate overall quality score for extracted text.
+
+    Args:
+        text: The extracted text content
+        metadata: Optional metadata for additional scoring
+
+    Returns:
+        Quality score between 0.0 and 1.0
+    """
+    if not text or not text.strip():
+        return 0.0
+
+    # Initialize score
+    score = 1.0
+    total_chars = len(text)
+
+    # Penalize OCR artifacts
+    ocr_penalty = _calculate_ocr_penalty(text, total_chars)
+    score -= ocr_penalty * 0.3
+
+    # Penalize script/style content
+    script_penalty = _calculate_script_penalty(text, total_chars)
+    score -= script_penalty * 0.2
+
+    # Penalize navigation content
+    nav_penalty = _calculate_navigation_penalty(text, total_chars)
+    score -= nav_penalty * 0.1
+
+    # Bonus for structure (sentences, paragraphs)
+    structure_bonus = _calculate_structure_bonus(text)
+    score += structure_bonus * 0.2
+
+    # Bonus for metadata richness
+    if metadata:
+        metadata_bonus = _calculate_metadata_bonus(metadata)
+        score += metadata_bonus * 0.1
+
+    return max(0.0, min(1.0, score))
+
+
+def clean_extracted_text(text: str) -> str:
+    """Clean extracted text by removing artifacts and improving quality.
+
+    Args:
+        text: The raw extracted text
+
+    Returns:
+        Cleaned text with artifacts removed
+    """
+    if not text:
+        return text
+
+    # Remove script and style content
+    for pattern in _SCRIPT_PATTERNS.values():
+        text = pattern.sub(" ", text)
+
+    # Clean OCR artifacts
+    text = _clean_ocr_artifacts(text)
+
+    # Clean navigation elements
+    text = _clean_navigation_elements(text)
+
+    # Normalize whitespace using pre-compiled patterns
+    text = _WHITESPACE_NORMALIZE.sub(" ", text)
+    text = _NEWLINE_NORMALIZE.sub("\n\n", text)
+
+    return text.strip()
+
+
+def _calculate_ocr_penalty(text: str, total_chars: int) -> float:
+    """Calculate penalty for OCR artifacts."""
+    if total_chars == 0:
+        return 0.0
+
+    # Use combined pattern for single-pass processing
+    artifact_chars = sum(len(match.group()) for match in _COMBINED_OCR_PATTERN.finditer(text))
+    return min(1.0, artifact_chars / total_chars)
+
+
+def _calculate_script_penalty(text: str, total_chars: int) -> float:
+    """Calculate penalty for script/style content."""
+    if total_chars == 0:
+        return 0.0
+
+    script_chars = 0
+    for pattern in _SCRIPT_PATTERNS.values():
+        matches = pattern.findall(text)
+        script_chars += sum(len(match) for match in matches)
+
+    return min(1.0, script_chars / total_chars)
+
+
+def _calculate_navigation_penalty(text: str, total_chars: int) -> float:
+    """Calculate penalty for navigation content."""
+    if total_chars == 0:
+        return 0.0
+
+    nav_chars = 0
+    for pattern in _NAVIGATION_PATTERNS.values():
+        matches = pattern.findall(text)
+        nav_chars += sum(len(match) for match in matches)
+
+    return min(1.0, nav_chars / total_chars)
+
+
+def _calculate_structure_bonus(text: str) -> float:
+    """Calculate bonus for proper text structure."""
+    if not text:
+        return 0.0
+
+    # Count sentences (rough heuristic)
+    sentence_count = len(_SENTENCE_DETECT.findall(text))
+
+    # Count paragraphs
+    paragraph_count = len(text.split("\n\n"))
+
+    # Calculate structure score
+    words = len(text.split())
+    if words == 0:
+        return 0.0
+
+    # Good structure: reasonable sentence and paragraph distribution
+    avg_words_per_sentence = words / max(1, sentence_count)
+    avg_words_per_paragraph = words / max(1, paragraph_count)
+
+    structure_score = 0.0
+
+    # Bonus for reasonable sentence length (10-30 words)
+    if 10 <= avg_words_per_sentence <= 30:
+        structure_score += 0.3
+
+    # Bonus for reasonable paragraph length (50-300 words)
+    if 50 <= avg_words_per_paragraph <= 300:
+        structure_score += 0.3
+
+    # Bonus for having multiple paragraphs
+    if paragraph_count > 1:
+        structure_score += 0.2
+
+    # Bonus for having punctuation
+    if _PUNCTUATION_DETECT.search(text):
+        structure_score += 0.2
+
+    return min(1.0, structure_score)
+
+
+def _calculate_metadata_bonus(metadata: dict[str, Any]) -> float:
+    """Calculate bonus for rich metadata."""
+    if not metadata:
+        return 0.0
+
+    important_fields = {"title", "author", "subject", "description", "keywords"}
+    present_fields = sum(1 for field in important_fields if metadata.get(field))
+
+    return present_fields / len(important_fields)
+
+
+def _clean_ocr_artifacts(text: str) -> str:
+    """Remove common OCR artifacts from text."""
+    # Fix scattered characters (likely OCR errors)
+    text = _OCR_ARTIFACTS["scattered_chars"].sub(lambda m: m.group().replace(" ", ""), text)
+
+    # Clean repeated punctuation
+    text = _OCR_ARTIFACTS["repeated_punctuation"].sub("...", text)
+
+    # Fix isolated punctuation
+    text = _OCR_ARTIFACTS["isolated_punctuation"].sub(" ", text)
+
+    # Remove malformed words with numbers mixed in
+    text = _OCR_ARTIFACTS["malformed_words"].sub(" ", text)
+
+    # Normalize excessive whitespace
+    return _OCR_ARTIFACTS["excessive_whitespace"].sub(" ", text)
+
+
+def _clean_navigation_elements(text: str) -> str:
+    """Remove navigation elements from text."""
+    # Remove navigation words
+    text = _NAVIGATION_PATTERNS["nav_words"].sub(" ", text)
+
+    # Remove breadcrumbs
+    text = _NAVIGATION_PATTERNS["breadcrumbs"].sub(" ", text)
+
+    # Remove pagination
+    return _NAVIGATION_PATTERNS["pagination"].sub(" ", text)

kreuzberg/_utils/_serialization.py

@@ -29,8 +29,10 @@ def encode_hook(obj: Any) -> Any:
         "to_list",
         "tolist",
     ):
-        if hasattr(obj, key) and callable(getattr(obj, key)):
-            return getattr(obj, key)()
+        if hasattr(obj, key):
+            method = getattr(obj, key)  # Cache the attribute lookup
+            if callable(method):
+                return method()
 
     if is_dataclass(obj) and not isinstance(obj, type):
         return {k: v if not isinstance(v, Enum) else v.value for (k, v) in asdict(obj).items()}