unique_toolkit 0.8.4__py3-none-any.whl → 0.8.6__py3-none-any.whl

unique_toolkit/tools/utils/execution/execution.py

@@ -0,0 +1,282 @@
+import functools
+import logging
+from typing import (
+    Awaitable,
+    Callable,
+    Generic,
+    Iterable,
+    ParamSpec,
+    Type,
+    TypeVar,
+    cast,
+)
+
+# Function types
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+logger = logging.getLogger(__name__)
+
+
+class Result(Generic[R]):
+    def __init__(
+        self,
+        success: bool,
+        result: R | None = None,
+        exception: Exception | None = None,
+    ) -> None:
+        self._success = success
+        self._result = result
+        self._exception = exception
+
+    @property
+    def exception(self) -> Exception | None:
+        return self._exception
+
+    @property
+    def success(self) -> bool:
+        return self._success
+
+    def unpack(self, default: R | None = None) -> R:
+        return cast(R, self._result) if self.success else cast(R, default)
+
+    def __str__(self) -> str:
+        return (
+            f"Success: {str(self._result)}"
+            if self.success
+            else f"Failure: {str(self._exception)}"
+        )
+
+
+class SafeTaskExecutor:
+    """
+    Execute function calls "safely": exceptions are caught and logged,
+    and the function result is returned as a `Result` object.
+
+    Several parameters are available to customize the behavior of the executor:
+    - `exceptions`: a list of exceptions that should be caught and logged
+    - `ignored_exceptions`: a list of exceptions that should be passed through
+    - `log_exceptions`: whether to log exceptions
+    - `log_exc_info`: whether to log exception info
+    - `logger`: a logger to use for logging
+
+
+    Usage:
+    ```python
+    executor = SafeTaskExecutor(
+        exceptions=(ValueError,),
+        ignored_exceptions=(KeyError,),
+    )
+
+    executor.execute(failing_function, "test")
+
+    executor.execute_async(async_failing_function, "test")
+    ```
+    """
+
+    def __init__(
+        self,
+        exceptions: Iterable[Type[Exception]] = (Exception,),
+        ignored_exceptions: Iterable[Type[Exception]] = (),
+        log_exceptions: bool = True,
+        log_exc_info: bool = True,
+        logger: logging.Logger | None = None,
+    ) -> None:
+        self._exceptions = tuple(exceptions)
+        self._ignored_exceptions = tuple(ignored_exceptions)
+        self._log_exceptions = log_exceptions
+        self._log_exc_info = log_exc_info
+        self._logger = logger
+
+    def execute(
+        self, f: Callable[P, R], *args: P.args, **kwargs: P.kwargs
+    ) -> Result[R]:
+        try:
+            return Result(True, f(*args, **kwargs))
+        except self._exceptions as e:
+            if isinstance(e, self._ignored_exceptions):
+                raise e
+            if self._log_exceptions:
+                logger.error(f"Error in {f.__name__}: {e}", exc_info=self._log_exc_info)
+            return Result(False, exception=e)
+
+    async def execute_async(
+        self, f: Callable[P, Awaitable[R]], *args: P.args, **kwargs: P.kwargs
+    ) -> Result[R]:
+        try:
+            return Result(True, await f(*args, **kwargs))
+        except self._exceptions as e:
+            if isinstance(e, self._ignored_exceptions):
+                raise e
+            if self._log_exceptions:
+                logger.error(f"Error in {f.__name__}: {e}", exc_info=self._log_exc_info)
+            return Result(False, exception=e)
+
+
+def safe_execute(f: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> Result[R]:
+    """
+    Execute a function call "safely": exceptions are caught and logged,
+    and the function result is returned as a `Result` object.
+
+    Usage:
+    ```python
+    def failing_function(a : str) -> int:
+        raise ValueError(a)
+
+    result = safe_execute(failing_function, "test")
+    print(result)
+    >> Failure: ValueError('test')
+
+    result.success
+    >> False
+
+    result.unpack()
+    >> None
+
+    result.exception
+    >> ValueError('test')
+
+    result.unpack(default=1)
+    >> 1
+    ```
+
+    ```python
+    def succeeding_function(a : str):
+        return a
+
+
+    result = safe_execute(succeeding_function, "test")
+
+    print(result)
+    >> Success: test
+
+    result.success
+    >> True
+
+    result.unpack()
+    >> 'test'
+
+    result.exception
+    >> None
+    ```
+    """
+    return SafeTaskExecutor().execute(f, *args, **kwargs)
+
+
+async def safe_execute_async(
+    f: Callable[P, Awaitable[R]], *args: P.args, **kwargs: P.kwargs
+) -> Result[R]:
+    """
+    Equivalent to `safe_execute` for async functions.
+    """
+    return await SafeTaskExecutor().execute_async(f, *args, **kwargs)
+
+
+FailureReturnType = TypeVar("FailureReturnType")
+
+
+def failsafe(
+    failure_return_value: FailureReturnType,
+    exceptions: Iterable[Type[Exception]] = (Exception,),
+    ignored_exceptions: Iterable[Type[Exception]] = (),
+    log_exceptions: bool = True,
+    log_exc_info: bool = True,
+    logger: logging.Logger | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R | FailureReturnType]]:
+    """
+    Decorator that executes sync functions with failsafe behavior: exceptions are caught and logged,
+    and a fallback return value is returned on failure instead of raising the exception.
+
+    Parameters are the same as SafeTaskExecutor plus:
+    - `failure_return_value`: value to return when an exception occurs
+
+    Usage:
+    ```python
+    @failsafe(
+        failure_return_value="default",
+        exceptions=(ValueError,),
+        ignored_exceptions=(KeyError,),
+    )
+    def failing_function(a: str) -> str:
+        raise ValueError(a)
+
+
+    result = failing_function("test")
+    # Returns "default" instead of raising ValueError
+    ```
+    """
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R | FailureReturnType]:
+        executor = SafeTaskExecutor(
+            exceptions=exceptions,
+            ignored_exceptions=ignored_exceptions,
+            log_exceptions=log_exceptions,
+            log_exc_info=log_exc_info,
+            logger=logger,
+        )
+
+        @functools.wraps(func)
+        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R | FailureReturnType:
+            result = executor.execute(func, *args, **kwargs)
+            return result.unpack(default=cast(R, failure_return_value))
+
+        return sync_wrapper
+
+    return decorator
+
+
+def failsafe_async(
+    failure_return_value: FailureReturnType,
+    exceptions: Iterable[Type[Exception]] = (Exception,),
+    ignored_exceptions: Iterable[Type[Exception]] = (),
+    log_exceptions: bool = True,
+    log_exc_info: bool = True,
+    logger: logging.Logger | None = None,
+) -> Callable[
+    [Callable[P, Awaitable[R]]], Callable[P, Awaitable[R | FailureReturnType]]
+]:
+    """
+    Decorator that executes async functions with failsafe behavior: exceptions are caught and logged,
+    and a fallback return value is returned on failure instead of raising the exception.
+
+    Parameters are the same as SafeTaskExecutor plus:
+    - `failure_return_value`: value to return when an exception occurs
+
+    Usage:
+    ```python
+    @failsafe_async(
+        failure_return_value=[],
+        exceptions=(ValueError,),
+        ignored_exceptions=(KeyError,),
+    )
+    async def async_failing_function(a: str) -> list:
+        raise ValueError(a)
+
+
+    result = await async_failing_function("test")
+    # Returns [] instead of raising ValueError
+    ```
+    """
+
+    def decorator(
+        func: Callable[P, Awaitable[R]],
+    ) -> Callable[P, Awaitable[R | FailureReturnType]]:
+        executor = SafeTaskExecutor(
+            exceptions=exceptions,
+            ignored_exceptions=ignored_exceptions,
+            log_exceptions=log_exceptions,
+            log_exc_info=log_exc_info,
+            logger=logger,
+        )
+
+        @functools.wraps(func)
+        async def async_wrapper(
+            *args: P.args, **kwargs: P.kwargs
+        ) -> R | FailureReturnType:
+            result = await executor.execute_async(func, *args, **kwargs)
+            return result.unpack(default=cast(R, failure_return_value))
+
+        return async_wrapper
+
+    return decorator

unique_toolkit/tools/utils/source_handling/schema.py

@@ -0,0 +1,22 @@
+# default schema follows logic in node-ingestion-worker: https://github.com/Unique-AG/monorepo/blob/76b4923611199a80abf9304639b3aa0538ec41ed/node/apps/node-ingestion-worker/src/ingestors/lib/text-manipulations.ts#L181C17-L181C28
+from pydantic import BaseModel
+
+from unique_toolkit.tools.config import get_configuration_dict
+
+
+SOURCE_TEMPLATE = "<source${index}>${document}${info}${text}</source${index}>"
+SECTIONS = {
+    "document": "<|document|>{}<|/document|>\n",
+    "info": "<|info|>{}<|/info|>\n",
+}
+
+
+class SourceFormatConfig(BaseModel):
+    model_config = get_configuration_dict()
+    source_template: str = SOURCE_TEMPLATE
+    sections: dict[str, str] = SECTIONS
+
+    @staticmethod
+    def template_to_pattern(template: str) -> str:
+        """Convert a template string into a regex pattern."""
+        return template.replace("{}", "(.*?)").replace("|", r"\|")

unique_toolkit/tools/utils/source_handling/source_formatting.py

@@ -0,0 +1,207 @@
+import re
+from string import Template
+
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.tools.utils.source_handling.schema import SourceFormatConfig
+
+
+def _format_page_range(chunk: ContentChunk) -> str:
+    """Format page range string from chunk metadata."""
+    if not (
+        chunk.start_page
+        and chunk.end_page
+        and chunk.start_page > 0
+        and chunk.end_page > 0
+    ):
+        return ""
+    return (
+        str(chunk.start_page)
+        if chunk.start_page == chunk.end_page
+        else f"{chunk.start_page} - {chunk.end_page}"
+    )
+
+
+def _parse_chunk(
+    chunk: ContentChunk, section_templates: dict[str, str]
+) -> dict[str, str]:
+    """Extract sections from chunk text using regex patterns."""
+    text = chunk.text
+    result = dict()
+
+    for section, template in section_templates.items():
+        # Document and info are the only sections that are included in the text
+        if section in [
+            "document",
+            "info",
+        ]:  # Skip page as it's derived from metadata
+            pattern = SourceFormatConfig.template_to_pattern(template)
+            match = re.search(pattern, text, re.DOTALL)
+            result[section] = match.group(1) if match else ""
+            text = text.replace(match.group(0), "") if match else text
+
+    result["text"] = text.strip()
+    return result
+
+
+def format_chunk(index: int, chunk: ContentChunk, config: SourceFormatConfig) -> str:
+    """
+    This function formats a content chunk based on a given configuration template and its sections. Each chunk in the database includes a document section, an optional info section, and a text section, with the text section being the primary content. Typically, chunks are added to sources in search modules without any changes. However, certain scenarios necessitate extra formatting, such as incorporating page numbers or other metadata. This function enables the custom formatting of chunks when they are appended as sources.
+
+    Args:
+        index (int): The source index number to be used in the template.
+        chunk (ContentChunk): A ContentChunk object containing:
+            - text (str): The main content text
+            - start_page (int, optional): Starting page number
+            - end_page (int, optional): Ending page number
+            - metadata (dict, optional): Additional metadata key-value pairs
+        config (SourceFormatConfig): Configuration object containing:
+            - source_template (str): The overall template for the output
+            - sections (dict): Mapping of section names to their format templates
+
+    Returns:
+        str: Formatted string according to the template
+
+    Examples:
+        Using XML-style config without page numbers (default):
+        >>> config = SourceFormatConfig(
+        ...     source_template="<source${index}>${document}${info}${text}</source${index}>",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "info": "<|info|>{}<|/info|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content"
+        ... )
+        >>> format_chunk(1, chunk, config)
+        '<source1><|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content</source1>'
+
+        Using XML-style config with page numbers:
+        >>> config = SourceFormatConfig(
+        ...     source_template="<source${index}>${document}${page}${info}${text}</source${index}>",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "info": "<|info|>{}<|/info|>\n",
+        ...         "page": "<|page|>{}<|/page|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content",
+        ...     start_page=1,
+        ...     end_page=3,
+        ... )
+        >>> format_chunk(1, chunk, config)
+        '<source1><|document|>Sample Doc.pdf</|document|>\n<|page|>1 - 3</|page|>\n<|info|>Important info</|info|>\nMain content</source1>'
+
+        Using XML-style config with metadata:
+        >>> config = SourceFormatConfig(
+        ...     source_template="<source${index}>${document}${date}${text}</source${index}>",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "date": "<|DateFromMetaData|>{}<|/DateFromMetaData|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\nMain content",
+        ...     metadata={
+        ...         "key": "metadata-key",
+        ...         "mimeType": "text/plain",
+        ...         "date": "12.03.2025",
+        ...     },
+        ... )
+        >>> format_chunk(1, chunk, config)
+        '<source1><|document|>Sample Doc.pdf</|document|>\n<|DateFromMetaData|>12.03.2025</|DateFromMetaData|>\nMain content</source1>'
+
+        Using JSON-style config:
+        >>> config = SourceFormatConfig(
+        ...     source_template="{'source_number': ${index}, 'content': '${document}${page}${info}${text}'}",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "info": "<|info|>{}<|/info|>\n",
+        ...         "page": "<|page|>{}<|/page|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content",
+        ...     start_page=5,
+        ...     end_page=5,
+        ... )
+        >>> format_chunk(1, chunk, config)
+        "{'source_number': 1, 'content': '<|document|>Sample Doc.pdf</|document|>\n<|page|>5</|page|>\n<|info|>Important info</|info|>\nMain content'}"
+
+    Notes:
+        - The function extracts document and info sections from the chunk text using regex patterns
+        - Page numbers are formatted as single numbers when start_page equals end_page
+        - Page numbers are formatted as ranges (e.g., "1 - 3") when start_page differs from end_page
+        - If page numbers are not available (None or 0), the page section will be empty
+        - Metadata keys that match section names (except 'document' and 'text') will be included in the output
+        - Metadata is processed by the _process_metadata function to update the parsed dictionary
+        - When using custom metadata tags like '<|DateFromMetaData|>', the key in chunk.metadata must match
+          the key in the sections dictionary (e.g., 'date' in the example above), not the tag name
+    """
+    sections = config.sections
+    source_template = config.source_template
+
+    parsed = _parse_chunk(chunk, sections)
+    parsed["page"] = _format_page_range(chunk)
+
+    # Update parsed with metadata values
+    parsed = _process_metadata(chunk, parsed, sections)
+
+    # Create a new dictionary to hold the formatted sections
+    formatted_sections = {}
+
+    # Process each section
+    for section, template in sections.items():
+        if parsed.get(section):
+            formatted_sections[section] = template.format(parsed.get(section, ""))
+        else:
+            formatted_sections[section] = ""
+
+    # Add the text section
+    formatted_sections["text"] = parsed["text"]
+
+    return Template(source_template).substitute(index=index, **formatted_sections)
+
+
+def _process_metadata(
+    chunk: ContentChunk, parsed: dict[str, str], sections: dict[str, str]
+) -> dict[str, str]:
+    """
+    Process metadata from chunk and update the parsed dictionary.
+
+    This function extracts metadata from a ContentChunk object and updates the parsed
+    dictionary with values whose keys match section names defined in SourceFormatConfig.
+
+    Args:
+        chunk (ContentChunk): The content chunk containing metadata
+        parsed (dict): The dictionary of already parsed sections to update
+
+    Returns:
+        dict: The updated parsed dictionary with metadata values added
+
+    Notes:
+        - Keys 'document' and 'text' are explicitly excluded from metadata processing
+        - Only metadata keys that match section names in SourceFormatConfig will be processed
+        - If chunk.metadata is None or not iterable, the parsed dict is returned unchanged
+        - Metadata values are added directly to the parsed dictionary using their original keys
+    """
+    # Return unchanged parsed dict if metadata is None
+    if not hasattr(chunk, "metadata") or chunk.metadata is None:
+        return parsed
+
+    # Ensure metadata is a dictionary
+    metadata_dict = dict(chunk.metadata) if hasattr(chunk.metadata, "__iter__") else {}
+
+    # Define keys that should not be treated as metadata keys
+    excluded_keys = {"document", "info"}
+
+    # Get the keys from SourceFormatConfig.sections
+    valid_section_keys = set(sections.keys()) - excluded_keys
+
+    # Update parsed with valid metadata entries
+    for key, value in metadata_dict.items():
+        if key in valid_section_keys:
+            parsed[key] = value
+
+    return parsed