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

sdg_hub/core/blocks/filtering/column_value_filter.py

@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Filter by value block for dataset filtering operations.
+
+This module provides a block for filtering datasets based on column values
+using various operations with optional data type conversion.
+"""
+
+# Standard
+from typing import Any, Optional, Union
+import operator
+
+# Third Party
+from datasets import Dataset
+from pydantic import Field, field_validator
+
+# Local
+from ...utils.logger_config import setup_logger
+from ..base import BaseBlock
+from ..registry import BlockRegistry
+
+logger = setup_logger(__name__)
+
+# Supported operations mapping
+OPERATION_MAP = {
+    "eq": operator.eq,
+    "ne": operator.ne,
+    "lt": operator.lt,
+    "le": operator.le,
+    "gt": operator.gt,
+    "ge": operator.ge,
+    "contains": operator.contains,
+    "in": lambda x, y: x in y,  # Reverse contains for "x in y" semantics
+}
+
+# Supported data types mapping
+DTYPE_MAP = {
+    "float": float,
+    "int": int,
+}
+
+
+@BlockRegistry.register(
+    "ColumnValueFilterBlock",
+    "filtering",
+    "Filters datasets based on column values using various comparison operations",
+)
+class ColumnValueFilterBlock(BaseBlock):
+    """A block for filtering datasets based on column values.
+
+    This block allows filtering of datasets using various operations (e.g., equals, contains)
+    on specified column values, with optional data type conversion.
+
+    Attributes
+    ----------
+    block_name : str
+        Name of the block.
+    input_cols : Union[str, List[str]]
+        Input column name(s). The first column will be used for filtering.
+    filter_value : Union[Any, List[Any]]
+        The value(s) to filter by.
+    operation : str
+        A string representing the binary operation to perform (e.g., "eq", "contains", "gt").
+        Supported operations: "eq", "ne", "lt", "le", "gt", "ge", "contains", "in".
+    convert_dtype : Optional[str], optional
+        String representation of type to convert the filter column to. Can be "float" or "int".
+        If None, no conversion is performed.
+    """
+
+    filter_value: Union[Any, list[Any]] = Field(
+        ..., description="The value(s) to filter by"
+    )
+    operation: str = Field(
+        ...,
+        description="String name of binary operator for comparison (e.g., 'eq', 'contains')",
+    )
+    convert_dtype: Optional[str] = Field(
+        None,
+        description="String name of type to convert filter column to ('float' or 'int')",
+    )
+
+    @field_validator("operation")
+    @classmethod
+    def validate_operation(cls, v):
+        """Validate that operation is a supported operation string."""
+        if v not in OPERATION_MAP:
+            raise ValueError(
+                f"Unsupported operation '{v}'. Supported operations: {list(OPERATION_MAP.keys())}"
+            )
+        return v
+
+    @field_validator("convert_dtype")
+    @classmethod
+    def validate_convert_dtype(cls, v):
+        """Validate that convert_dtype is a supported type string."""
+        if v is not None and v not in DTYPE_MAP:
+            raise ValueError(
+                f"Unsupported dtype '{v}'. Supported dtypes: {list(DTYPE_MAP.keys())}"
+            )
+        return v
+
+    @field_validator("input_cols", mode="after")
+    @classmethod
+    def validate_input_cols_not_empty(cls, v):
+        """Validate that we have at least one input column."""
+        if not v or len(v) == 0:
+            raise ValueError(
+                "ColumnValueFilterBlock requires at least one input column"
+            )
+        return v
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize derived attributes after Pydantic validation."""
+        super().model_post_init(__context) if hasattr(
+            super(), "model_post_init"
+        ) else None
+
+        # Ensure output_cols is empty list for filtering operations (doesn't create new columns)
+        if self.output_cols is None:
+            self.output_cols = []
+
+        # Set derived attributes
+        self.value = (
+            self.filter_value
+            if isinstance(self.filter_value, list)
+            else [self.filter_value]
+        )
+        self.column_name = self.input_cols[0]  # Use first input column for filtering
+
+        # Convert string operation to actual callable
+        self._operation_func = OPERATION_MAP[self.operation]
+
+        # Convert string dtype to actual type if specified
+        self._convert_dtype_func = (
+            DTYPE_MAP[self.convert_dtype] if self.convert_dtype else None
+        )
+
+    def _convert_dtype(self, sample: dict[str, Any]) -> dict[str, Any]:
+        """Convert the data type of the filter column.
+
+        Parameters
+        ----------
+        sample : Dict[str, Any]
+            The sample dictionary containing the column to convert.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The sample with converted column value.
+        """
+        try:
+            sample[self.column_name] = self._convert_dtype_func(
+                sample[self.column_name]
+            )
+        except ValueError as e:
+            logger.error(
+                "Error converting dtype: %s, filling with None to be filtered later", e
+            )
+            sample[self.column_name] = None
+        return sample
+
+    def generate(self, samples: Dataset, **_kwargs: Any) -> Dataset:
+        """Generate filtered dataset based on specified conditions.
+
+        Parameters
+        ----------
+        samples : Dataset
+            The input dataset to filter.
+
+        Returns
+        -------
+        Dataset
+            The filtered dataset.
+        """
+        if self._convert_dtype_func:
+            samples = samples.map(self._convert_dtype)
+
+        samples = samples.filter(
+            lambda x: x[self.column_name] is not None,
+        )
+
+        # Apply filter operation
+        samples = samples.filter(
+            lambda x: any(
+                self._operation_func(x[self.column_name], value) for value in self.value
+            )
+        )
+
+        return samples

sdg_hub/core/blocks/llm/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LLM blocks for provider-agnostic text generation.
+
+This module provides blocks for interacting with language models through
+LiteLLM, supporting 100+ providers including OpenAI, Anthropic, Google,
+local models (vLLM, Ollama), and more.
+"""
+
+# Local
+from .client_manager import LLMClientManager
+from .config import LLMConfig
+from .error_handler import ErrorCategory, LLMErrorHandler
+from .llm_chat_block import LLMChatBlock
+from .prompt_builder_block import PromptBuilderBlock
+from .text_parser_block import TextParserBlock
+
+__all__ = [
+    "LLMConfig",
+    "LLMClientManager",
+    "LLMErrorHandler",
+    "ErrorCategory",
+    "LLMChatBlock",
+    "PromptBuilderBlock",
+    "TextParserBlock",
+]

sdg_hub/core/blocks/llm/client_manager.py

@@ -0,0 +1,398 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Client manager for LLM operations supporting all providers via LiteLLM."""
+
+# Standard
+from typing import Any, Optional, Union
+import asyncio
+
+# Third Party
+from litellm import acompletion, completion
+import litellm
+
+# Local
+from ...utils.logger_config import setup_logger
+from .config import LLMConfig
+from .error_handler import LLMErrorHandler
+
+logger = setup_logger(__name__)
+
+
+class LLMClientManager:
+    """Client manager for LLM operations using LiteLLM.
+
+    This class provides a unified interface for calling any LLM provider
+    supported by LiteLLM, with robust error handling and retry logic.
+
+    Parameters
+    ----------
+    config : LLMConfig
+        Configuration for the LLM client.
+    error_handler : Optional[LLMErrorHandler], optional
+        Custom error handler. If None, a default one will be created.
+    """
+
+    def __init__(
+        self, config: LLMConfig, error_handler: Optional[LLMErrorHandler] = None
+    ) -> None:
+        self.config = config
+        self.error_handler = error_handler or LLMErrorHandler(
+            max_retries=config.max_retries
+        )
+        self._is_loaded = False
+
+    def load(self) -> None:
+        """Load and configure the LLM client.
+
+        This method sets up LiteLLM configuration and validates the setup.
+        """
+        if self._is_loaded:
+            return
+
+        # Configure LiteLLM
+        self._configure_litellm()
+
+        # Test the configuration
+        self._validate_setup()
+
+        self._is_loaded = True
+
+        # Only log when model is actually configured
+        if self.config.model:
+            logger.info(
+                f"Loaded LLM client for model '{self.config.model}'",
+                extra={
+                    "model": self.config.model,
+                    "provider": self.config.get_provider(),
+                    "is_local": self.config.is_local_model(),
+                    "api_base": self.config.api_base,
+                },
+            )
+
+    def unload(self) -> None:
+        """Unload the client and clean up resources."""
+        self._is_loaded = False
+        try:
+            logger.info(f"Unloaded LLM client for model '{self.config.model}'")
+        except Exception:
+            # Ignore logging errors during cleanup to prevent issues during shutdown
+            pass
+
+    def _configure_litellm(self) -> None:
+        """Configure LiteLLM settings."""
+        # Set global timeout for LiteLLM
+        litellm.request_timeout = self.config.timeout
+
+        # Note: API keys are now passed directly in completion calls
+        # instead of modifying environment variables for thread-safety
+
+    def _validate_setup(self) -> None:
+        """Validate that the LLM setup is working."""
+        try:
+            # For testing/development, skip validation if using dummy API key
+            if self.config.api_key == "test-key":
+                logger.debug(
+                    f"Skipping validation for model '{self.config.model}' (test mode)"
+                )
+                return
+
+            # TODO: Skip validation for now to avoid API calls during initialization
+            # we might want to make a minimal test call
+            logger.debug(
+                f"Setup configured for model '{self.config.model}'. "
+                f"Validation will occur on first actual call."
+            )
+
+        except Exception as e:
+            logger.warning(
+                f"Could not validate setup for model '{self.config.model}': {e}"
+            )
+
+    def create_completion(
+        self, messages: list[dict[str, Any]], **overrides: Any
+    ) -> Union[str, list[str]]:
+        """Create a completion using LiteLLM.
+
+        Parameters
+        ----------
+        messages : List[Dict[str, Any]]
+            Messages in OpenAI format.
+        **overrides : Any
+            Runtime parameter overrides.
+
+        Returns
+        -------
+        Union[str, List[str]]
+            The completion text(s). Returns a single string when n=1 or n is None,
+            returns a list of strings when n>1.
+
+        Raises
+        ------
+        Exception
+            If the completion fails after all retries.
+        """
+        if not self._is_loaded:
+            self.load()
+
+        # Merge configuration with overrides
+        final_config = self.config.merge_overrides(**overrides)
+        kwargs = self._build_completion_kwargs(messages, final_config)
+
+        # Create retry wrapper
+        context = {
+            "model": final_config.model,
+            "provider": final_config.get_provider(),
+            "message_count": len(messages),
+        }
+
+        completion_func = self.error_handler.wrap_completion(
+            self._call_litellm_completion, context=context
+        )
+
+        # Make the completion call
+        response = completion_func(kwargs)
+
+        # Extract content from response
+        # Check if n > 1 to determine return type
+        n_value = final_config.n or 1
+        if n_value > 1:
+            return [choice.message.content for choice in response.choices]
+        else:
+            return response.choices[0].message.content
+
+    async def acreate_completion(
+        self, messages: list[dict[str, Any]], **overrides: Any
+    ) -> Union[str, list[str]]:
+        """Create an async completion using LiteLLM.
+
+        Parameters
+        ----------
+        messages : List[Dict[str, Any]]
+            Messages in OpenAI format.
+        **overrides : Any
+            Runtime parameter overrides.
+
+        Returns
+        -------
+        Union[str, List[str]]
+            The completion text(s). Returns a single string when n=1 or n is None,
+            returns a list of strings when n>1.
+
+        Raises
+        ------
+        Exception
+            If the completion fails after all retries.
+        """
+        if not self._is_loaded:
+            self.load()
+
+        # Merge configuration with overrides
+        final_config = self.config.merge_overrides(**overrides)
+        kwargs = self._build_completion_kwargs(messages, final_config)
+
+        # Create retry wrapper for async
+        context = {
+            "model": final_config.model,
+            "provider": final_config.get_provider(),
+            "message_count": len(messages),
+        }
+
+        completion_func = self.error_handler.wrap_completion(
+            self._call_litellm_acompletion, context=context
+        )
+
+        # Make the async completion call
+        response = await completion_func(kwargs)
+
+        # Extract content from response
+        # Check if n > 1 to determine return type
+        n_value = final_config.n or 1
+        if n_value > 1:
+            return [choice.message.content for choice in response.choices]
+        else:
+            return response.choices[0].message.content
+
+    def create_completions_batch(
+        self, messages_list: list[list[dict[str, Any]]], **overrides: Any
+    ) -> list[Union[str, list[str]]]:
+        """Create multiple completions in batch.
+
+        Parameters
+        ----------
+        messages_list : List[List[Dict[str, Any]]]
+            List of message lists to process.
+        **overrides : Any
+            Runtime parameter overrides.
+
+        Returns
+        -------
+        List[Union[str, List[str]]]
+            List of completion texts. Each element is a single string when n=1 or n is None,
+            or a list of strings when n>1.
+        """
+        results = []
+        for messages in messages_list:
+            result = self.create_completion(messages, **overrides)
+            results.append(result)
+        return results
+
+    async def acreate_completions_batch(
+        self, messages_list: list[list[dict[str, Any]]], **overrides: Any
+    ) -> list[Union[str, list[str]]]:
+        """Create multiple completions in batch asynchronously.
+
+        Parameters
+        ----------
+        messages_list : List[List[Dict[str, Any]]]
+            List of message lists to process.
+        **overrides : Any
+            Runtime parameter overrides.
+
+        Returns
+        -------
+        List[Union[str, List[str]]]
+            List of completion texts. Each element is a single string when n=1 or n is None,
+            or a list of strings when n>1.
+        """
+        tasks = [
+            self.acreate_completion(messages, **overrides) for messages in messages_list
+        ]
+        return await asyncio.gather(*tasks)
+
+    def _build_completion_kwargs(
+        self, messages: list[dict[str, Any]], config: LLMConfig
+    ) -> dict[str, Any]:
+        """Build kwargs for LiteLLM completion call.
+
+        Parameters
+        ----------
+        messages : List[Dict[str, Any]]
+            Messages in OpenAI format.
+        config : LLMConfig
+            Final configuration after merging overrides.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Kwargs for litellm.completion().
+        """
+        kwargs = {
+            "model": config.model,
+            "messages": messages,
+        }
+
+        # Add API configuration
+        if config.api_key:
+            kwargs["api_key"] = config.api_key
+
+        if config.api_base:
+            kwargs["api_base"] = config.api_base
+
+        # Add generation parameters
+        generation_kwargs = config.get_generation_kwargs()
+        kwargs.update(generation_kwargs)
+
+        return kwargs
+
+    def _call_litellm_completion(self, kwargs: dict[str, Any]) -> Any:
+        """Call LiteLLM completion with error handling.
+
+        Parameters
+        ----------
+        kwargs : Dict[str, Any]
+            Arguments for litellm.completion().
+
+        Returns
+        -------
+        Any
+            LiteLLM completion response.
+        """
+        logger.debug(
+            f"Calling LiteLLM completion for model '{kwargs['model']}'",
+            extra={
+                "model": kwargs["model"],
+                "message_count": len(kwargs["messages"]),
+                "generation_params": {
+                    k: v
+                    for k, v in kwargs.items()
+                    if k in ["temperature", "max_tokens", "top_p", "n"]
+                },
+            },
+        )
+
+        response = completion(**kwargs)
+
+        logger.debug(
+            f"LiteLLM completion successful for model '{kwargs['model']}'",
+            extra={
+                "model": kwargs["model"],
+                "choices_count": len(response.choices),
+            },
+        )
+
+        return response
+
+    async def _call_litellm_acompletion(self, kwargs: dict[str, Any]) -> Any:
+        """Call LiteLLM async completion with error handling.
+
+        Parameters
+        ----------
+        kwargs : Dict[str, Any]
+            Arguments for litellm.acompletion().
+
+        Returns
+        -------
+        Any
+            LiteLLM completion response.
+        """
+        logger.debug(
+            f"Calling LiteLLM async completion for model '{kwargs['model']}'",
+            extra={
+                "model": kwargs["model"],
+                "message_count": len(kwargs["messages"]),
+            },
+        )
+
+        response = await acompletion(**kwargs)
+
+        logger.debug(
+            f"LiteLLM async completion successful for model '{kwargs['model']}'",
+            extra={
+                "model": kwargs["model"],
+                "choices_count": len(response.choices),
+            },
+        )
+
+        return response
+
+    def get_model_info(self) -> dict[str, Any]:
+        """Get information about the configured model.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Model information.
+        """
+        return {
+            "model": self.config.model,
+            "provider": self.config.get_provider(),
+            "model_name": self.config.get_model_name(),
+            "is_local": self.config.is_local_model(),
+            "api_base": self.config.api_base,
+            "is_loaded": self._is_loaded,
+        }
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.load()
+        return self
+
+    def __exit__(self, _exc_type, _exc_val, _exc_tb):
+        """Context manager exit."""
+        self.unload()
+
+    def __repr__(self) -> str:
+        """String representation."""
+        return (
+            f"LLMClientManager(model='{self.config.model}', "
+            f"provider='{self.config.get_provider()}', loaded={self._is_loaded})"
+        )