sdg-hub 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

sdg_hub/core/blocks/llm/text_parser_block.py

@@ -0,0 +1,357 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Text parser block for parsing and post-processing LLM outputs.
+
+This module provides the TextParserBlock for handling output parsing using
+start/end tags, custom regex patterns, and cleanup operations.
+"""
+
+# Standard
+from typing import Any, Optional
+import re
+
+# Third Party
+from datasets import Dataset
+from pydantic import Field, field_validator, model_validator
+
+# Local
+from ...utils.logger_config import setup_logger
+from ..base import BaseBlock
+from ..registry import BlockRegistry
+
+logger = setup_logger(__name__)
+
+
+@BlockRegistry.register(
+    "TextParserBlock",
+    "llm",
+    "Parses and post-processes LLM outputs using tags or regex patterns",
+)
+class TextParserBlock(BaseBlock):
+    """Block for parsing and post-processing LLM outputs.
+
+    This block handles output parsing using start/end tags, custom regex patterns,
+    and cleanup operations. It expects exactly one input column containing raw LLM output.
+
+    Attributes
+    ----------
+    block_name : str
+        Unique identifier for this block instance.
+    input_cols : Union[str, List[str], Dict[str, Any], None]
+        Input column name(s) containing raw LLM output. Must specify exactly one column.
+    output_cols : Union[str, List[str], Dict[str, Any], None]
+        Output column name(s) for parsed results.
+    start_tags : List[str]
+        List of start tags for tag-based parsing.
+    end_tags : List[str]
+        List of end tags for tag-based parsing.
+    parsing_pattern : Optional[str]
+        Regex pattern for custom parsing.
+    parser_cleanup_tags : Optional[List[str]]
+        List of tags to clean from parsed output.
+    expand_lists : bool
+        Whether to expand list inputs into individual rows (True) or preserve lists (False).
+        Default is True for backward compatibility.
+    """
+
+    start_tags: list[str] = Field(
+        default_factory=list, description="List of start tags for tag-based parsing"
+    )
+    end_tags: list[str] = Field(
+        default_factory=list, description="List of end tags for tag-based parsing"
+    )
+    parsing_pattern: Optional[str] = Field(
+        default=None, description="Regex pattern for custom parsing"
+    )
+    parser_cleanup_tags: Optional[list[str]] = Field(
+        default=None, description="List of tags to clean from parsed output"
+    )
+    expand_lists: bool = Field(
+        default=True,
+        description="Whether to expand list inputs into individual rows (True) or preserve lists (False). ",
+    )
+
+    @field_validator("start_tags", "end_tags", mode="before")
+    @classmethod
+    def normalize_tags(cls, v):
+        """Normalize tag lists to ensure they are always lists."""
+        if v is None:
+            return []
+        if isinstance(v, str):
+            return [v]
+        if isinstance(v, list):
+            return v
+        raise ValueError(f"Tags must be a string, list, or None, got {type(v)}")
+
+    @field_validator("parser_cleanup_tags", mode="before")
+    @classmethod
+    def normalize_cleanup_tags(cls, v):
+        """Normalize cleanup tags to ensure they are always lists when not None."""
+        if v is None:
+            return None
+        if isinstance(v, str):
+            return [v]
+        if isinstance(v, list):
+            return v
+        raise ValueError(f"Cleanup tags must be a string, list, or None, got {type(v)}")
+
+    @model_validator(mode="after")
+    def validate_parsing_configuration(self):
+        """Validate that parsing configuration is consistent."""
+        # Validate that at least one parsing method is configured
+        has_regex = self.parsing_pattern is not None
+        has_tags = bool(self.start_tags) or bool(self.end_tags)
+
+        if not has_regex and not has_tags:
+            raise ValueError(
+                "TextParserBlock requires at least one parsing method: "
+                "either 'parsing_pattern' (regex) or 'start_tags'/'end_tags' (tag-based parsing)"
+            )
+
+        # Validate tag parsing configuration
+        if has_tags:
+            if len(self.start_tags) != len(self.end_tags):
+                raise ValueError(
+                    f"start_tags and end_tags must have the same length. "
+                    f"Got {len(self.start_tags)} start_tags and {len(self.end_tags)} end_tags"
+                )
+
+            # We can't validate against output_cols here since they might not be normalized yet
+            # This validation will be moved to _validate_custom
+
+        return self
+
+    def _validate_custom(self, dataset: Dataset) -> None:
+        """Validate TextParserBlock specific requirements.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset to validate.
+
+        Raises
+        ------
+        ValueError
+            If TextParserBlock requirements are not met.
+        """
+        # Validate that we have exactly one input column
+        if len(self.input_cols) == 0:
+            raise ValueError("TextParserBlock expects at least one input column")
+        if len(self.input_cols) > 1:
+            logger.warning(
+                f"TextParserBlock expects exactly one input column, but got {len(self.input_cols)}. "
+                f"Using the first column: {self.input_cols[0]}"
+            )
+
+        # Validate tag parsing against output columns (can only be done after model creation)
+        has_tags = bool(self.start_tags) or bool(self.end_tags)
+        if has_tags and len(self.start_tags) != len(self.output_cols):
+            raise ValueError(
+                f"When using tag-based parsing, the number of tag pairs must match output_cols. "
+                f"Got {len(self.start_tags)} tag pairs and {len(self.output_cols)} output columns"
+            )
+
+    def _extract_matches(
+        self, text: str, start_tag: Optional[str], end_tag: Optional[str]
+    ) -> list[str]:
+        if not text:
+            return []
+        if not start_tag and not end_tag:
+            return [text.strip()]
+
+        pattern = ""
+        if start_tag:
+            pattern += re.escape(start_tag)
+        pattern += r"(.*?)"
+        if end_tag:
+            pattern += re.escape(end_tag)
+        elif start_tag:
+            pattern += "$"
+
+        return [match.strip() for match in re.findall(pattern, text, re.DOTALL)]
+
+    def _parse(self, generated_string: str) -> dict[str, list[str]]:
+        if self.parsing_pattern is not None:
+            return self._parse_with_regex(generated_string)
+        return self._parse_with_tags(generated_string)
+
+    def _parse_with_regex(self, generated_string: str) -> dict[str, list[str]]:
+        """Parse using regex pattern."""
+        if self.parsing_pattern is None:
+            raise ValueError("parsing_pattern is required for regex parsing")
+        pattern = re.compile(self.parsing_pattern, re.DOTALL)
+        all_matches = pattern.findall(generated_string)
+        matches: dict[str, list[str]] = {
+            column_name: [] for column_name in self.output_cols
+        }
+
+        logger.debug(
+            f"Regex parsing found {len(all_matches)} matches with pattern: {self.parsing_pattern}"
+        )
+
+        if all_matches and isinstance(all_matches[0], tuple):
+            return self._process_tuple_matches(all_matches, matches)
+        return self._process_single_matches(all_matches, matches)
+
+    def _parse_with_tags(self, generated_string: str) -> dict[str, list[str]]:
+        """Parse using start/end tags."""
+        matches: dict[str, list[str]] = {
+            column_name: [] for column_name in self.output_cols
+        }
+
+        for start_tag, end_tag, output_col in zip(
+            self.start_tags, self.end_tags, self.output_cols
+        ):
+            extracted = self._extract_matches(generated_string, start_tag, end_tag)
+            matches[output_col] = extracted
+            logger.debug(
+                f"Tag parsing for '{output_col}' with tags '{start_tag}'/'{end_tag}' found {len(extracted)} matches"
+            )
+
+        return matches
+
+    def _process_tuple_matches(
+        self, all_matches: list, matches: dict[str, list[str]]
+    ) -> dict[str, list[str]]:
+        """Process regex matches that are tuples."""
+        for match in all_matches:
+            for column_name, value in zip(self.output_cols, match):
+                value = self._clean_value(value.strip())
+                matches[column_name].append(value)
+        return matches
+
+    def _process_single_matches(
+        self, all_matches: list, matches: dict[str, list[str]]
+    ) -> dict[str, list[str]]:
+        """Process regex matches that are single values."""
+        cleaned_matches = [self._clean_value(match.strip()) for match in all_matches]
+        matches[self.output_cols[0]] = cleaned_matches
+        return matches
+
+    def _clean_value(self, value: str) -> str:
+        """Clean value by removing cleanup tags."""
+        if self.parser_cleanup_tags:
+            for clean_tag in self.parser_cleanup_tags:
+                value = value.replace(clean_tag, "")
+        return value
+
+    def _generate(self, sample: dict) -> list[dict]:
+        input_column = self.input_cols[0]
+        raw_output = sample[input_column]
+
+        # Handle list inputs (e.g., from LLMChatBlock with n > 1)
+        if isinstance(raw_output, list):
+            if not raw_output:
+                logger.warning(f"Input column '{input_column}' contains empty list")
+                return []
+
+            if not self.expand_lists:
+                # When expand_lists=False, preserve the list structure
+                # Parse each response in the list and collect results as lists
+                all_parsed_outputs = {col: [] for col in self.output_cols}
+                valid_responses = 0
+
+                for i, response in enumerate(raw_output):
+                    if not response or not isinstance(response, str):
+                        logger.warning(
+                            f"List item {i} in column '{input_column}' contains invalid data "
+                            f"(empty or non-string): {type(response)}"
+                        )
+                        continue
+
+                    parsed_outputs = self._parse(response)
+
+                    if not parsed_outputs or not any(
+                        len(value) > 0 for value in parsed_outputs.values()
+                    ):
+                        logger.warning(
+                            f"Failed to parse content from list item {i}. Raw output length: {len(response)}, "
+                            f"parsing method: {'regex' if self.parsing_pattern else 'tags'}"
+                        )
+                        continue
+
+                    valid_responses += 1
+                    # Collect all parsed values for each column as lists
+                    for col in self.output_cols:
+                        all_parsed_outputs[col].extend(parsed_outputs.get(col, []))
+
+                if valid_responses == 0:
+                    return []
+
+                # Return single row with lists as values
+                # TODO: This breaks retry counting in LLMChatWithParsingRetryBlock until LLMChatWithParsingRetryBlock is re-based
+                # which expects one row per successful parse for counting
+                return [{**sample, **all_parsed_outputs}]
+
+            else:
+                # When expand_lists=True, use existing expanding behavior
+                all_results = []
+                for i, response in enumerate(raw_output):
+                    if not response or not isinstance(response, str):
+                        logger.warning(
+                            f"List item {i} in column '{input_column}' contains invalid data "
+                            f"(empty or non-string): {type(response)}"
+                        )
+                        continue
+
+                    parsed_outputs = self._parse(response)
+
+                    if not parsed_outputs or not any(
+                        len(value) > 0 for value in parsed_outputs.values()
+                    ):
+                        logger.warning(
+                            f"Failed to parse content from list item {i}. Raw output length: {len(response)}, "
+                            f"parsing method: {'regex' if self.parsing_pattern else 'tags'}"
+                        )
+                        continue
+
+                    # Create output rows for this response
+                    max_length = max(len(value) for value in parsed_outputs.values())
+                    for values in zip(
+                        *(lst[:max_length] for lst in parsed_outputs.values())
+                    ):
+                        all_results.append(
+                            {**sample, **dict(zip(parsed_outputs.keys(), values))}
+                        )
+
+                return all_results
+
+        # Handle string inputs (existing logic)
+        elif isinstance(raw_output, str):
+            if not raw_output:
+                logger.warning(f"Input column '{input_column}' contains empty string")
+                return []
+
+            parsed_outputs = self._parse(raw_output)
+
+            if not parsed_outputs or not any(
+                len(value) > 0 for value in parsed_outputs.values()
+            ):
+                logger.warning(
+                    f"Failed to parse any content from input. Raw output length: {len(raw_output)}, "
+                    f"parsing method: {'regex' if self.parsing_pattern else 'tags'}"
+                )
+                return []
+
+            result = []
+            max_length = max(len(value) for value in parsed_outputs.values())
+            for values in zip(*(lst[:max_length] for lst in parsed_outputs.values())):
+                result.append({**sample, **dict(zip(parsed_outputs.keys(), values))})
+            return result
+
+        else:
+            logger.warning(
+                f"Input column '{input_column}' contains invalid data type: {type(raw_output)}. "
+                f"Expected str or List[str]"
+            )
+            return []
+
+    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+        logger.debug(f"Parsing outputs for {len(samples)} samples")
+        if len(samples) == 0:
+            logger.warning("No samples to parse, returning empty dataset")
+            return Dataset.from_list([])
+
+        new_data = []
+        for sample in samples:
+            new_data.extend(self._generate(sample))
+        return Dataset.from_list(new_data)

sdg_hub/core/blocks/registry.py

@@ -0,0 +1,331 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Enhanced BlockRegistry with metadata and better error handling.
+
+This module provides a clean registry system for blocks with metadata,
+categorization, and improved error handling.
+"""
+
+# Standard
+from dataclasses import dataclass
+from difflib import get_close_matches
+from typing import Optional
+import inspect
+
+# Third Party
+from rich.console import Console
+from rich.table import Table
+
+# Local
+from ..utils.logger_config import setup_logger
+
+logger = setup_logger(__name__)
+console = Console()
+
+
+@dataclass
+class BlockMetadata:
+    """Metadata for registered blocks.
+
+    Parameters
+    ----------
+    name : str
+        The registered name of the block.
+    block_class : Type
+        The actual block class.
+    category : str
+        Category for organization (e.g., 'llm', 'utility', 'filtering').
+    description : str, optional
+        Human-readable description of what the block does.
+    deprecated : bool, optional
+        Whether this block is deprecated.
+    replacement : str, optional
+        Suggested replacement if deprecated.
+    """
+
+    name: str
+    block_class: type
+    category: str
+    description: str = ""
+    deprecated: bool = False
+    replacement: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        """Validate metadata after initialization."""
+        if not self.name:
+            raise ValueError("Block name cannot be empty")
+        if not inspect.isclass(self.block_class):
+            raise ValueError("block_class must be a class")
+
+
+class BlockRegistry:
+    """Registry for block classes with metadata and enhanced error handling."""
+
+    _metadata: dict[str, BlockMetadata] = {}
+    _categories: dict[str, set[str]] = {}
+
+    @classmethod
+    def register(
+        cls,
+        block_name: str,
+        category: str,
+        description: str = "",
+        deprecated: bool = False,
+        replacement: Optional[str] = None,
+    ):
+        """Register a block class with metadata.
+
+        Parameters
+        ----------
+        block_name : str
+            Name under which to register the block.
+        category : str
+            Category for organization.
+        description : str, optional
+            Human-readable description of the block.
+        deprecated : bool, optional
+            Whether this block is deprecated.
+        replacement : str, optional
+            Suggested replacement if deprecated.
+
+        Returns
+        -------
+        callable
+            Decorator function.
+        """
+
+        def decorator(block_class: type) -> type:
+            # Validate the class
+            cls._validate_block_class(block_class)
+
+            # Create metadata
+            metadata = BlockMetadata(
+                name=block_name,
+                block_class=block_class,
+                category=category,
+                description=description,
+                deprecated=deprecated,
+                replacement=replacement,
+            )
+
+            # Register the metadata
+            cls._metadata[block_name] = metadata
+
+            # Update category index
+            if category not in cls._categories:
+                cls._categories[category] = set()
+            cls._categories[category].add(block_name)
+
+            logger.debug(
+                f"Registered block '{block_name}' "
+                f"({block_class.__name__}) in category '{category}'"
+            )
+
+            if deprecated:
+                warning_msg = f"Block '{block_name}' is deprecated."
+                if replacement:
+                    warning_msg += f" Use '{replacement}' instead."
+                logger.warning(warning_msg)
+
+            return block_class
+
+        return decorator
+
+    @classmethod
+    def _validate_block_class(cls, block_class: type) -> None:
+        """Validate that a class is a proper block class.
+
+        Parameters
+        ----------
+        block_class : Type
+            The class to validate.
+
+        Raises
+        ------
+        ValueError
+            If the class is not a valid block class.
+        """
+        if not inspect.isclass(block_class):
+            raise ValueError(f"Expected a class, got {type(block_class)}")
+
+        # Validate BaseBlock inheritance
+        try:
+            # Local
+            from .base import BaseBlock
+
+            if not issubclass(block_class, BaseBlock):
+                raise ValueError(
+                    f"Block class '{block_class.__name__}' must inherit from BaseBlock"
+                )
+        except ImportError as exc:
+            # BaseBlock not available, check for generate method
+            if not hasattr(block_class, "generate"):
+                raise ValueError(
+                    f"Block class '{block_class.__name__}' must implement 'generate' method"
+                ) from exc
+
+    @classmethod
+    def get(cls, block_name: str) -> type:
+        """Get a block class with enhanced error handling.
+
+        Parameters
+        ----------
+        block_name : str
+            Name of the block to retrieve.
+
+        Returns
+        -------
+        Type
+            The block class.
+
+        Raises
+        ------
+        KeyError
+            If the block is not found, with helpful suggestions.
+        """
+        if block_name not in cls._metadata:
+            available_blocks = list(cls._metadata.keys())
+            suggestions = get_close_matches(
+                block_name, available_blocks, n=3, cutoff=0.6
+            )
+
+            error_msg = f"Block '{block_name}' not found in registry."
+
+            if suggestions:
+                error_msg += f" Did you mean: {', '.join(suggestions)}?"
+
+            if available_blocks:
+                error_msg += (
+                    f"\nAvailable blocks: {', '.join(sorted(available_blocks))}"
+                )
+
+            if cls._categories:
+                error_msg += (
+                    f"\nCategories: {', '.join(sorted(cls._categories.keys()))}"
+                )
+
+            logger.error(error_msg)
+            raise KeyError(error_msg)
+
+        metadata = cls._metadata[block_name]
+
+        if metadata.deprecated:
+            warning_msg = f"Block '{block_name}' is deprecated."
+            if metadata.replacement:
+                warning_msg += f" Use '{metadata.replacement}' instead."
+            logger.warning(warning_msg)
+
+        return metadata.block_class
+
+    @classmethod
+    def info(cls, block_name: str) -> BlockMetadata:
+        """Get metadata for a specific block.
+
+        Parameters
+        ----------
+        block_name : str
+            Name of the block.
+
+        Returns
+        -------
+        BlockMetadata
+            The block's metadata.
+
+        Raises
+        ------
+        KeyError
+            If the block is not found.
+        """
+        if block_name not in cls._metadata:
+            raise KeyError(f"Block '{block_name}' not found in registry.")
+        return cls._metadata[block_name]
+
+    @classmethod
+    def categories(cls) -> list[str]:
+        """Get all available categories.
+
+        Returns
+        -------
+        List[str]
+            Sorted list of categories.
+        """
+        return sorted(cls._categories.keys())
+
+    @classmethod
+    def category(cls, category: str) -> list[str]:
+        """Get all blocks in a specific category.
+
+        Parameters
+        ----------
+        category : str
+            The category to filter by.
+
+        Returns
+        -------
+        List[str]
+            List of block names in the category.
+
+        Raises
+        ------
+        KeyError
+            If the category doesn't exist.
+        """
+        if category not in cls._categories:
+            available_categories = sorted(cls._categories.keys())
+            raise KeyError(
+                f"Category '{category}' not found. "
+                f"Available categories: {', '.join(available_categories)}"
+            )
+        return sorted(cls._categories[category])
+
+    @classmethod
+    def all(cls) -> dict[str, list[str]]:
+        """List all blocks organized by category.
+
+        Returns
+        -------
+        Dict[str, List[str]]
+            Dictionary mapping categories to lists of block names.
+        """
+        return {
+            category: sorted(blocks) for category, blocks in cls._categories.items()
+        }
+
+    @classmethod
+    def discover_blocks(cls) -> None:
+        """Print a Rich-formatted table of all available blocks."""
+        if not cls._metadata:
+            console.print("[yellow]No blocks registered yet.[/yellow]")
+            return
+
+        table = Table(
+            title="Available Blocks", show_header=True, header_style="bold magenta"
+        )
+        table.add_column("Block Name", style="cyan", no_wrap=True)
+        table.add_column("Category", style="green")
+        table.add_column("Description", style="white")
+
+        # Sort blocks by category, then by name
+        sorted_blocks = sorted(
+            cls._metadata.items(), key=lambda x: (x[1].category, x[0])
+        )
+
+        for name, metadata in sorted_blocks:
+            description = metadata.description or "No description"
+
+            # Show deprecated blocks with a warning indicator in the name
+            block_name = f"⚠️ {name}" if metadata.deprecated else name
+
+            table.add_row(block_name, metadata.category, description)
+
+        console.print(table)
+
+        # Show summary
+        total_blocks = len(cls._metadata)
+        total_categories = len(cls._categories)
+        deprecated_count = sum(1 for m in cls._metadata.values() if m.deprecated)
+
+        console.print(
+            f"\n[bold]Summary:[/bold] {total_blocks} blocks across {total_categories} categories"
+        )
+        if deprecated_count > 0:
+            console.print(f"[yellow]⚠️  {deprecated_count} deprecated blocks[/yellow]")