unique_toolkit 1.20.1__tar.gz → 1.21.0__tar.gz

{unique_toolkit-1.20.1 → unique_toolkit-1.21.0}/CHANGELOG.md

@@ -6,6 +6,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [1.21.0] - 2025-10-30
+- Add option to customize the display of tool progress statuses.
+
 ## [1.20.1] - 2025-10-30
 - Fix typing issues in `LanguageModelFunction`.
 

{unique_toolkit-1.20.1 → unique_toolkit-1.21.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.20.1
+Version: 1.21.0
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -119,6 +119,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [1.21.0] - 2025-10-30
+- Add option to customize the display of tool progress statuses.
+
 ## [1.20.1] - 2025-10-30
 - Fix typing issues in `LanguageModelFunction`.
 

{unique_toolkit-1.20.1 → unique_toolkit-1.21.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "1.20.1"
+version = "1.21.0"
 description = ""
 authors = [
     "Cedric Klinkert <cedric.klinkert@unique.ch>",

unique_toolkit-1.21.0/unique_toolkit/agentic/tools/test/test_tool_progress_reporter.py

@@ -0,0 +1,429 @@
+from unittest.mock import AsyncMock
+
+import pytest
+
+from unique_toolkit.agentic.tools.tool_progress_reporter import (
+    DUMMY_REFERENCE_PLACEHOLDER,
+    ProgressState,
+    ToolExecutionStatus,
+    ToolProgressReporter,
+    ToolProgressReporterConfig,
+    ToolWithToolProgressReporter,
+    track_tool_progress,
+)
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.language_model.schemas import LanguageModelFunction
+
+
+@pytest.fixture
+def chat_service():
+    return AsyncMock(spec=ChatService)
+
+
+@pytest.fixture
+def tool_progress_reporter(chat_service):
+    return ToolProgressReporter(chat_service)
+
+
+@pytest.fixture
+def tool_call():
+    return LanguageModelFunction(id="test_id", name="test_tool")
+
+
+class TestToolProgressReporter:
+    @pytest.mark.asyncio
+    async def test_notify_from_tool_call(self, tool_progress_reporter, tool_call):
+        # Arrange
+        name = "Test Tool"
+        message = "Processing..."
+        state = ProgressState.RUNNING
+        references = [
+            ContentReference(
+                sequence_number=1,
+                id="1",
+                message_id="1",
+                name="1",
+                source="1",
+                source_id="1",
+                url="1",
+            )
+        ]
+
+        # Act
+        await tool_progress_reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name=name,
+            message=message,
+            state=state,
+            references=references,
+            requires_new_assistant_message=True,
+        )
+
+        # Assert
+        assert tool_call.id in tool_progress_reporter.tool_statuses
+        status = tool_progress_reporter.tool_statuses[tool_call.id]
+        assert status.name == name
+        assert status.message == message
+        assert status.state == state
+        assert status.references == references
+        assert tool_progress_reporter.requires_new_assistant_message is True
+
+    def test_replace_placeholders(self, tool_progress_reporter):
+        # Arrange
+        message = (
+            f"Test{DUMMY_REFERENCE_PLACEHOLDER}message{DUMMY_REFERENCE_PLACEHOLDER}"
+        )
+
+        # Act
+        result = tool_progress_reporter._replace_placeholders(message, start_number=1)
+
+        # Assert
+        assert result == "Test<sup>1</sup>message<sup>2</sup>"
+
+    def test_correct_reference_sequence(self, tool_progress_reporter):
+        # Arrange
+        references = [
+            ContentReference(
+                sequence_number=0,
+                id="1",
+                message_id="1",
+                name="1",
+                source="1",
+                source_id="1",
+                url="1",
+            ),
+            ContentReference(
+                sequence_number=0,
+                id="2",
+                message_id="2",
+                name="2",
+                source="2",
+                source_id="2",
+                url="2",
+            ),
+        ]
+
+        # Act
+        result = tool_progress_reporter._correct_reference_sequence(
+            references, start_number=1
+        )
+
+        # Assert
+        assert len(result) == 2
+        assert result[0].sequence_number == 1
+        assert result[1].sequence_number == 2
+
+    @pytest.mark.asyncio
+    async def test_publish_updates_chat_service(
+        self, tool_progress_reporter, tool_call
+    ):
+        # Arrange
+        status = ToolExecutionStatus(
+            name="Test Tool",
+            message="Test message",
+            state=ProgressState.FINISHED,
+            references=[
+                ContentReference(
+                    sequence_number=1,
+                    id="1",
+                    message_id="1",
+                    name="1",
+                    source="1",
+                    source_id="1",
+                    url="1",
+                )
+            ],
+        )
+        tool_progress_reporter.tool_statuses[tool_call.id] = status
+
+        # Act
+        await tool_progress_reporter.publish()
+
+        # Assert
+        tool_progress_reporter.chat_service.modify_assistant_message_async.assert_called_once()
+
+
+class TestToolProgressDecorator:
+    class DummyTool(ToolWithToolProgressReporter):
+        def __init__(self, tool_progress_reporter):
+            self.tool_progress_reporter = tool_progress_reporter
+
+        @track_tool_progress(
+            message="Processing",
+            on_start_state=ProgressState.STARTED,
+            on_success_state=ProgressState.FINISHED,
+            on_success_message="Completed",
+            requires_new_assistant_message=True,
+        )
+        async def execute(self, tool_call, notification_tool_name):
+            return {
+                "references": [
+                    ContentReference(
+                        sequence_number=1,
+                        id="1",
+                        message_id="1",
+                        name="1",
+                        source="1",
+                        source_id="1",
+                        url="1",
+                    )
+                ]
+            }
+
+    @pytest.mark.asyncio
+    async def test_decorator_success_flow(self, tool_progress_reporter, tool_call):
+        # Arrange
+        tool = self.DummyTool(tool_progress_reporter)
+
+        # Act
+        await tool.execute(tool_call, "Test Tool")
+
+        # Assert
+        assert len(tool_progress_reporter.tool_statuses) == 1
+        status = tool_progress_reporter.tool_statuses[tool_call.id]
+        assert status.state == ProgressState.FINISHED
+        assert status.message == "Completed"
+
+    @pytest.mark.asyncio
+    async def test_decorator_error_flow(self, tool_progress_reporter, tool_call):
+        # Arrange
+        class ErrorTool(ToolWithToolProgressReporter):
+            def __init__(self, tool_progress_reporter):
+                self.tool_progress_reporter = tool_progress_reporter
+
+            @track_tool_progress(message="Processing")
+            async def execute(self, tool_call, notification_tool_name):
+                raise ValueError("Test error")
+
+        tool = ErrorTool(tool_progress_reporter)
+
+        # Act & Assert
+        with pytest.raises(ValueError):
+            await tool.execute(tool_call, "Test Tool")
+
+        status = tool_progress_reporter.tool_statuses[tool_call.id]
+        assert status.state == ProgressState.FAILED
+
+
+class TestToolProgressReporterConfig:
+    """Tests for ToolProgressReporterConfig and custom display configuration."""
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_config__uses_default_templates__when_no_config_provided(
+        self, chat_service, tool_call
+    ) -> None:
+        """
+        Purpose: Verify that default state-to-display templates are used when no config is provided.
+        Why this matters: Ensures backward compatibility and default behavior.
+        Setup summary: Create reporter without config, add status, verify default template is used.
+        """
+        # Arrange
+        reporter = ToolProgressReporter(chat_service)
+
+        # Act
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name="Test Tool",
+            message="Processing data",
+            state=ProgressState.RUNNING,
+        )
+
+        # Assert
+        assert tool_call.id in reporter.tool_statuses
+        chat_service.modify_assistant_message_async.assert_called()
+        call_args = chat_service.modify_assistant_message_async.call_args
+        content = call_args.kwargs["content"]
+        assert "Test Tool" in content
+        assert "🟡" in content  # Default emoji for RUNNING state
+        assert "Processing data" in content
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_config__uses_custom_templates__when_config_provided(
+        self, chat_service, tool_call
+    ) -> None:
+        """
+        Purpose: Verify that custom templates are used when provided via config.
+        Why this matters: Enables customization of progress display format.
+        Setup summary: Create reporter with custom template, verify custom format is used.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(
+            state_to_display_template={
+                ProgressState.RUNNING: "⏳ {tool_name}: {message}",
+                ProgressState.FINISHED: "✅ {tool_name}: {message}",
+            }
+        )
+        reporter = ToolProgressReporter(chat_service, config=custom_config)
+
+        # Act
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name="My Tool",
+            message="Working on it",
+            state=ProgressState.RUNNING,
+        )
+
+        # Assert
+        chat_service.modify_assistant_message_async.assert_called()
+        call_args = chat_service.modify_assistant_message_async.call_args
+        content = call_args.kwargs["content"]
+        assert "⏳ My Tool: Working on it" in content
+        assert "🟡" not in content  # Default emoji should not appear
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_config__skips_states_not_in_template__when_state_excluded(
+        self, chat_service, tool_call
+    ) -> None:
+        """
+        Purpose: Verify that states not in the config mapping are not displayed.
+        Why this matters: Allows selective display of only certain states (e.g., hide STARTED).
+        Setup summary: Create config excluding RUNNING state, verify message is not displayed.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(
+            state_to_display_template={
+                ProgressState.FINISHED: "✅ {tool_name}: {message}",
+                # RUNNING state is intentionally excluded
+            }
+        )
+        reporter = ToolProgressReporter(chat_service, config=custom_config)
+
+        # Act
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name="Test Tool",
+            message="Processing",
+            state=ProgressState.RUNNING,
+        )
+
+        # Assert
+        chat_service.modify_assistant_message_async.assert_called()
+        call_args = chat_service.modify_assistant_message_async.call_args
+        content = call_args.kwargs["content"]
+        # Content should not contain the message since RUNNING is excluded
+        assert "Processing" not in content
+        assert "Test Tool" not in content
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_config__formats_placeholders_correctly__with_multiple_tools(
+        self, chat_service
+    ) -> None:
+        """
+        Purpose: Verify that {tool_name} and {message} placeholders are replaced correctly for multiple tools.
+        Why this matters: Ensures template formatting works correctly in multi-tool scenarios.
+        Setup summary: Add multiple tool statuses with different names/messages, verify formatting.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(
+            state_to_display_template={
+                ProgressState.RUNNING: "▶️ {tool_name} - {message}",
+                ProgressState.FINISHED: "✓ {tool_name} - {message}",
+            }
+        )
+        reporter = ToolProgressReporter(chat_service, config=custom_config)
+        tool_call_1 = LanguageModelFunction(id="tool_1", name="search")
+        tool_call_2 = LanguageModelFunction(id="tool_2", name="analyze")
+
+        # Act
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call_1,
+            name="Search Tool",
+            message="Searching database",
+            state=ProgressState.RUNNING,
+        )
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call_2,
+            name="Analysis Tool",
+            message="Analyzing results",
+            state=ProgressState.FINISHED,
+        )
+
+        # Assert
+        call_args = chat_service.modify_assistant_message_async.call_args
+        content = call_args.kwargs["content"]
+        assert "▶️ Search Tool - Searching database" in content
+        assert "✓ Analysis Tool - Analyzing results" in content
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_config__shows_only_finished_state__when_only_finished_configured(
+        self, chat_service, tool_call
+    ) -> None:
+        """
+        Purpose: Verify selective state display shows only FINISHED when it's the only state configured.
+        Why this matters: Use case where user only wants final results, not intermediate steps.
+        Setup summary: Configure only FINISHED state, send STARTED and FINISHED, verify only FINISHED appears.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(
+            state_to_display_template={
+                ProgressState.FINISHED: "Done: {tool_name} - {message}",
+            }
+        )
+        reporter = ToolProgressReporter(chat_service, config=custom_config)
+
+        # Act - Send STARTED state (should not appear)
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name="Test Tool",
+            message="Starting",
+            state=ProgressState.STARTED,
+        )
+
+        # Get first call content
+        first_call_args = chat_service.modify_assistant_message_async.call_args
+        first_content = first_call_args.kwargs["content"]
+
+        # Act - Update to FINISHED state (should appear)
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name="Test Tool",
+            message="Completed successfully",
+            state=ProgressState.FINISHED,
+        )
+
+        # Assert
+        final_call_args = chat_service.modify_assistant_message_async.call_args
+        final_content = final_call_args.kwargs["content"]
+
+        # STARTED state should not appear in first call
+        assert "Starting" not in first_content
+
+        # FINISHED state should appear in final call
+        assert "Done: Test Tool - Completed successfully" in final_content
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_config__handles_empty_template_dict__when_all_states_excluded(
+        self, chat_service, tool_call
+    ) -> None:
+        """
+        Purpose: Verify that an empty template dict results in no messages being displayed.
+        Why this matters: Edge case handling and allows disabling all progress display.
+        Setup summary: Create config with empty dict, verify no tool messages appear.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(state_to_display_template={})
+        reporter = ToolProgressReporter(chat_service, config=custom_config)
+
+        # Act
+        await reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name="Test Tool",
+            message="Processing",
+            state=ProgressState.RUNNING,
+        )
+
+        # Assert
+        chat_service.modify_assistant_message_async.assert_called()
+        call_args = chat_service.modify_assistant_message_async.call_args
+        content = call_args.kwargs["content"]
+
+        # Should only have the progress start text and newlines, no actual messages
+        assert "Test Tool" not in content
+        assert "Processing" not in content

{unique_toolkit-1.20.1 → unique_toolkit-1.21.0}/unique_toolkit/agentic/tools/tool_progress_reporter.py

@@ -1,11 +1,12 @@
 import re
 from datetime import datetime
-from enum import Enum
+from enum import StrEnum
 from functools import wraps
 from typing import Protocol
 
-from pydantic import BaseModel
+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 (
@@ -17,11 +18,11 @@ ARROW = "→ "
 DUMMY_REFERENCE_PLACEHOLDER = "<sup></sup>"
 
 
-class ProgressState(Enum):
-    STARTED = "⚪"
-    RUNNING = "🟡"
-    FAILED = "🔴"
-    FINISHED = "🟢"
+class ProgressState(StrEnum):
+    STARTED = "started"
+    RUNNING = "running"
+    FAILED = "failed"
+    FINISHED = "finished"
 
 
 class ToolExecutionStatus(BaseModel):
@@ -29,15 +30,56 @@ class ToolExecutionStatus(BaseModel):
     message: str
     state: ProgressState
     references: list[ContentReference] = []
-    timestamp: datetime = datetime.now()
+    timestamp: datetime = Field(default_factory=datetime.now)
+
+
+_DEFAULT_STATE_TO_DISPLAY_TEMPLATE = {
+    ProgressState.FINISHED: "{arrow}**{{tool_name}}** 🟢: {{message}}".format(
+        arrow=ARROW
+    ),
+    ProgressState.RUNNING: "{arrow}**{{tool_name}}** 🟡: {{message}}".format(
+        arrow=ARROW
+    ),
+    ProgressState.FAILED: "{arrow}**{{tool_name}}** 🔴: {{message}}".format(
+        arrow=ARROW
+    ),
+    ProgressState.STARTED: "{arrow}**{{tool_name}}** ⚪: {{message}}".format(
+        arrow=ARROW
+    ),
+}
+
+
+state_to_display_template_description = """
+A mapping progress states to display templates.
+The display template is a string that will be used to display the progress status.
+The template can contain the following placeholders:
+- `{tool_name}`: The name of the tool
+- `{message}`: The message to display (sent by the tool)
+
+If a state is not present in the mapping, then updates for that state will not be displayed.
+""".strip()
+
+
+class ToolProgressReporterConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    state_to_display_template: dict[ProgressState, str] = Field(
+        default=_DEFAULT_STATE_TO_DISPLAY_TEMPLATE,
+        description=state_to_display_template_description,
+    )
 
 
 class ToolProgressReporter:
-    def __init__(self, chat_service: ChatService):
+    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):
@@ -77,16 +119,13 @@ class ToolProgressReporter:
             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.
-
-        Raises:
-            AssertionError: If tool_call.id is None
         """
-        assert tool_call.id is not None
         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
@@ -103,7 +142,11 @@ class ToolProgressReporter:
             references = self._correct_reference_sequence(references, start_number)
             all_references.extend(references)
 
-            messages.append(f"{ARROW}**{item.name} {item.state.value}**: {message}")
+            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),
@@ -130,6 +173,28 @@ class ToolProgressReporter:
             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:
+        if state in self._config.state_to_display_template:
+            return self._config.state_to_display_template[state].format(
+                tool_name=name,
+                message=message,
+            )
+        return None
+
 
 class ToolWithToolProgressReporter(Protocol):
     tool_progress_reporter: ToolProgressReporter