unique_toolkit 0.7.7__py3-none-any.whl → 1.23.0__py3-none-any.whl

unique_toolkit/agentic/tools/tool_progress_reporter.py

@@ -0,0 +1,285 @@
+import re
+from datetime import datetime
+from enum import StrEnum
+from functools import wraps
+from typing import Protocol, TypedDict
+
+from pydantic import BaseModel, Field
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+    LanguageModelStreamResponse,
+)
+
+ARROW = "→ "
+DUMMY_REFERENCE_PLACEHOLDER = "<sup></sup>"
+
+
+class ProgressState(StrEnum):
+    STARTED = "started"
+    RUNNING = "running"
+    FAILED = "failed"
+    FINISHED = "finished"
+
+
+class ToolExecutionStatus(BaseModel):
+    name: str
+    message: str
+    state: ProgressState
+    references: list[ContentReference] = []
+    timestamp: datetime = Field(default_factory=datetime.now)
+
+
+class StateToDisplayTemplate(TypedDict):
+    started: str
+    running: str
+    failed: str
+    finished: str
+
+
+_DEFAULT_STATE_TO_DISPLAY_TEMPLATE: StateToDisplayTemplate = {
+    "started": "{arrow}**{{tool_name}}** ⚪: {{message}}".format(arrow=ARROW),
+    "running": "{arrow}**{{tool_name}}** 🟡: {{message}}".format(arrow=ARROW),
+    "finished": "{arrow}**{{tool_name}}** 🟢: {{message}}".format(arrow=ARROW),
+    "failed": "{arrow}**{{tool_name}}** 🔴: {{message}}".format(arrow=ARROW),
+}
+
+
+state_to_display_template_description = """
+Display templates for the different progress states.
+The template is a string that will be used to display the progress status.
+It can contain the following placeholders:
+- `{tool_name}`: The name of the tool
+- `{message}`: The message to display (sent by the tool)
+""".strip()
+
+
+class ToolProgressReporterConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    state_to_display_template: StateToDisplayTemplate = Field(
+        default=_DEFAULT_STATE_TO_DISPLAY_TEMPLATE,
+        description=state_to_display_template_description,
+        title="Display Templates",
+    )
+
+
+class ToolProgressReporter:
+    def __init__(
+        self,
+        chat_service: ChatService,
+        config: ToolProgressReporterConfig | None = None,
+    ):
+        self.chat_service = chat_service
+        self.tool_statuses: dict[str, ToolExecutionStatus] = {}
+        self._progress_start_text = ""
+        self._requires_new_assistant_message = False
+        self._config = config or ToolProgressReporterConfig()
+
+    @property
+    def requires_new_assistant_message(self):
+        return self._requires_new_assistant_message
+
+    @requires_new_assistant_message.setter
+    def requires_new_assistant_message(self, value: bool):
+        self._requires_new_assistant_message = value
+
+    @property
+    def tool_statuses_is_empty(self):
+        return len(self.tool_statuses) == 0
+
+    def empty_tool_statuses_if_stream_has_text(
+        self, stream_response: LanguageModelStreamResponse
+    ):
+        if stream_response.message.text:
+            self.tool_statuses = {}
+
+    async def notify_from_tool_call(
+        self,
+        tool_call: LanguageModelFunction,
+        name: str,
+        message: str,
+        state: ProgressState,
+        references: list[ContentReference] = [],
+        requires_new_assistant_message: bool = False,
+    ):
+        """
+        Notifies about a tool call execution status and updates the assistant message.
+
+        Args:
+            tool_call (LanguageModelFunction): The tool call being executed
+            name (str): Name of the tool being executed
+            message (str): Status message to display
+            state (ProgressState): Current execution state of the tool
+            references (list[ContentReference], optional): List of content references. Defaults to [].
+            requires_new_assistant_message (bool, optional): Whether a new assistant message is needed when tool call is finished.
+            Defaults to False. If yes, the agentic steps will remain in chat history and will be overwritten by the stream response.
+        """
+        self.tool_statuses[tool_call.id] = ToolExecutionStatus(
+            name=name,
+            message=message,
+            state=state,
+            references=references,
+            timestamp=self._get_timestamp_for_tool_call(tool_call),
+        )
+        self.requires_new_assistant_message = (
+            self.requires_new_assistant_message or requires_new_assistant_message
+        )
+        await self.publish()
+
+    async def publish(self):
+        messages = []
+        all_references = []
+        for item in sorted(self.tool_statuses.values(), key=lambda x: x.timestamp):
+            references = item.references
+            start_number = len(all_references) + 1
+            message = self._replace_placeholders(item.message, start_number)
+            references = self._correct_reference_sequence(references, start_number)
+            all_references.extend(references)
+
+            display_message = self._get_tool_status_display_message(
+                name=item.name, message=message, state=item.state
+            )
+            if display_message is not None:
+                messages.append(display_message)
+
+        await self.chat_service.modify_assistant_message_async(
+            content=self._progress_start_text + "\n\n" + "\n\n".join(messages),
+            references=all_references,
+        )
+
+    @staticmethod
+    def _replace_placeholders(message: str, start_number: int = 1) -> str:
+        counter = start_number
+
+        def replace_match(match):
+            nonlocal counter
+            result = f"<sup>{counter}</sup>"
+            counter += 1
+            return result
+
+        return re.sub(r"<sup></sup>", replace_match, message)
+
+    @staticmethod
+    def _correct_reference_sequence(
+        references: list[ContentReference], start_number: int = 1
+    ) -> list[ContentReference]:
+        for i, reference in enumerate(references, start_number):
+            reference.sequence_number = i
+        return references
+
+    def _get_timestamp_for_tool_call(
+        self, tool_call: LanguageModelFunction
+    ) -> datetime:
+        """
+        Keep the same timestamp if the tool call is already in the statuses.
+        This ensures the display order stays consistent.
+        """
+        if tool_call.id in self.tool_statuses:
+            return self.tool_statuses[tool_call.id].timestamp
+
+        return datetime.now()
+
+    def _get_tool_status_display_message(
+        self, name: str, message: str, state: ProgressState
+    ) -> str | None:
+        display_message = self._config.state_to_display_template[state.value].format(
+            tool_name=name,
+            message=message,
+        )
+        # Don't display empty messages
+        if display_message.strip() == "":
+            return None
+
+        return display_message
+
+
+class ToolWithToolProgressReporter(Protocol):
+    tool_progress_reporter: ToolProgressReporter
+
+
+def track_tool_progress(
+    message: str,
+    on_start_state: ProgressState = ProgressState.RUNNING,
+    on_success_state: ProgressState = ProgressState.RUNNING,
+    on_success_message: str | None = None,
+    on_error_message: str = "Unexpected error occurred",
+    requires_new_assistant_message: bool = False,
+):
+    """
+    Decorator to add progress reporting and status tracking steps to tool functions. Can be used with async and sync functions.
+
+    Args:
+        name (str): Display name for the tool progress status
+        message (str): Message to show during tool execution
+        on_error_message (str, optional): Message to show if tool execution fails. Defaults to empty string.
+        on_success_state (ProgressState, optional): State to set after successful execution. Defaults to RUNNING.
+        requires_new_assistant_message (bool, optional): Whether to create a new assistant message. Defaults to False.
+
+    The decorator will:
+    1. Show a RUNNING status when the tool starts executing
+    2. Update the status to on_success_state if execution succeeds
+    3. Update the status to FAILED if execution fails
+    4. Include any references from the tool result in the status update if the result has a 'references' attribute or item.
+    5. Create a new assistant message if requires_new_assistant_message is True
+
+    The decorated function must be a method of a class that implements ToolWithToolProgressReporter.
+    """
+
+    def decorator(func):
+        @wraps(func)  # Preserve the original function's metadata
+        async def async_wrapper(
+            self: ToolWithToolProgressReporter,
+            tool_call: LanguageModelFunction,
+            notification_tool_name: str,
+            *args,
+            **kwargs,
+        ):
+            try:
+                # Start status
+                await self.tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=notification_tool_name,
+                    message=message,
+                    state=on_start_state,
+                )
+
+                # Execute the tool function
+                result = await func(
+                    self, tool_call, notification_tool_name, *args, **kwargs
+                )
+
+                # Success status
+                await self.tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=notification_tool_name,
+                    message=on_success_message or message,
+                    state=on_success_state,
+                    references=_get_references_from_results(result),
+                    requires_new_assistant_message=requires_new_assistant_message,
+                )
+                return result
+
+            except Exception as e:
+                # Failure status
+                await self.tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=notification_tool_name,
+                    message=on_error_message,
+                    state=ProgressState.FAILED,
+                    requires_new_assistant_message=requires_new_assistant_message,
+                )
+                raise e
+
+        return async_wrapper
+
+    return decorator
+
+
+def _get_references_from_results(result):
+    if isinstance(result, dict):
+        return result.get("references", [])
+    return getattr(result, "references", [])

unique_toolkit/agentic/tools/utils/__init__.py

@@ -0,0 +1,19 @@
+"""Utilities for tools."""
+
+from unique_toolkit.agentic.tools.utils.execution.execution import (
+    Result,
+    SafeTaskExecutor,
+    failsafe,
+    failsafe_async,
+    safe_execute,
+    safe_execute_async,
+)
+
+__all__ = [
+    "failsafe",
+    "failsafe_async",
+    "safe_execute",
+    "safe_execute_async",
+    "SafeTaskExecutor",
+    "Result",
+]

unique_toolkit/agentic/tools/utils/execution/__init__.py

@@ -0,0 +1 @@
+"""Execution utilities for tools."""

unique_toolkit/agentic/tools/utils/execution/execution.py

@@ -0,0 +1,286 @@
+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 or _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:
+                self._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:
+                self._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/agentic/tools/utils/source_handling/schema.py

@@ -0,0 +1,21 @@
+# 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._common.pydantic_helpers 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"\|")