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

unique_toolkit/agentic/tools/test/test_tool_progress_reporter.py

@@ -0,0 +1,445 @@
+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={
+                "started": "⚪ {tool_name}: {message}",
+                "running": "⏳ {tool_name}: {message}",
+                "finished": "✅ {tool_name}: {message}",
+                "failed": "❌ {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_with_empty_template__when_state_hidden(
+        self, chat_service, tool_call
+    ) -> None:
+        """
+        Purpose: Verify that states with empty string templates are not displayed.
+        Why this matters: Allows selective display of only certain states (e.g., hide STARTED).
+        Setup summary: Create config with empty string for RUNNING state, verify message is not displayed.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(
+            state_to_display_template={
+                "started": "",
+                "running": "",  # Empty string hides RUNNING state
+                "finished": "✅ {tool_name}: {message}",
+                "failed": "❌ {tool_name}: {message}",
+            }
+        )
+        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 template is empty
+        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={
+                "started": "○ {tool_name} - {message}",
+                "running": "▶️ {tool_name} - {message}",
+                "finished": "✓ {tool_name} - {message}",
+                "failed": "✗ {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 other states use empty templates.
+        Why this matters: Use case where user only wants final results, not intermediate steps.
+        Setup summary: Configure only FINISHED with content, send STARTED and FINISHED, verify only FINISHED appears.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(
+            state_to_display_template={
+                "started": "",  # Empty template hides STARTED
+                "running": "",  # Empty template hides RUNNING
+                "finished": "Done: {tool_name} - {message}",
+                "failed": "Failed: {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_all_empty_templates__when_all_states_hidden(
+        self, chat_service, tool_call
+    ) -> None:
+        """
+        Purpose: Verify that all empty string templates result in no messages being displayed.
+        Why this matters: Edge case handling and allows disabling all progress display.
+        Setup summary: Create config with all empty templates, verify no tool messages appear.
+        """
+        # Arrange
+        custom_config = ToolProgressReporterConfig(
+            state_to_display_template={
+                "started": "",
+                "running": "",
+                "finished": "",
+                "failed": "",
+            }
+        )
+        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/agentic/tools/tool.py

@@ -0,0 +1,183 @@
+from abc import ABC, abstractmethod
+from logging import getLogger
+from typing import Any, Generic, TypeVar, cast
+
+from typing_extensions import deprecated
+
+from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
+from unique_toolkit.agentic.tools.config import ToolBuildConfig, ToolSelectionPolicy
+from unique_toolkit.agentic.tools.schemas import (
+    BaseToolConfig,
+    ToolCallResponse,
+    ToolPrompts,
+)
+from unique_toolkit.agentic.tools.tool_progress_reporter import ToolProgressReporter
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.chat.service import (
+    ChatService,
+)
+from unique_toolkit.language_model import LanguageModelToolDescription
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+)
+from unique_toolkit.language_model.service import LanguageModelService
+
+ConfigType = TypeVar("ConfigType", bound=BaseToolConfig)
+
+ToolBuildConfig.model_rebuild()
+
+
+class Tool(ABC, Generic[ConfigType]):
+    name: str
+    settings: ToolBuildConfig
+
+    def display_name(self) -> str:
+        """The display name of the tool."""
+        return self.settings.display_name
+
+    def icon(self) -> str:
+        """The icon of the tool."""
+        return self.settings.icon
+
+    def selection_policy(self) -> ToolSelectionPolicy:
+        """The selection policy of the tool."""
+        return self.settings.selection_policy
+
+    def is_exclusive(self) -> bool:
+        """Whether the tool is exclusive or not."""
+        return self.settings.is_exclusive
+
+    def is_enabled(self) -> bool:
+        """Whether the tool is enabled or not."""
+        return self.settings.is_enabled
+
+    def takes_control(self):
+        """
+        Some tools require to take control of the conversation with the user and do not want the orchestrator to intervene.
+        this function indicates whether the tool takes control or not. It yanks the control away from the orchestrator.
+        A typical use-case is deep-research.
+        """
+        return False
+
+    @property
+    def configuration(self) -> BaseToolConfig:
+        """The configuration of the tool."""
+        return self.settings.configuration
+
+    @abstractmethod
+    def tool_description(self) -> LanguageModelToolDescription:
+        raise NotImplementedError
+
+    def tool_description_as_json(self) -> dict[str, Any]:
+        parameters = self.tool_description().parameters
+        if not isinstance(parameters, dict):
+            return parameters.model_json_schema()
+        else:
+            return cast("dict[str, Any]", parameters)
+
+    # TODO: This method should be a property
+    def tool_description_for_system_prompt(self) -> str:
+        return ""
+
+    # TODO: This method should be a property
+    def tool_format_information_for_system_prompt(self) -> str:
+        return ""
+
+    # TODO: This method should be a property
+    def tool_description_for_user_prompt(self) -> str:
+        return ""
+
+    # TODO: This method should be a property
+    def tool_format_information_for_user_prompt(self) -> str:
+        return ""
+
+    # TODO: This method should be a property
+    def tool_format_reminder_for_user_prompt(self) -> str:
+        """A short reminder for the user prompt for formatting rules for the tool.
+        You can use this if the LLM fails to follow the formatting rules.
+        """
+        return ""
+
+    @deprecated(
+        "Do not use. The tool should not determine how"
+        "it is checked. This should be defined by the user"
+        "of the tool."
+    )
+    @abstractmethod
+    def evaluation_check_list(self) -> list[EvaluationMetricName]:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
+        raise NotImplementedError
+
+    @deprecated(
+        "Do not use as the evaluation checks should not be determined by\n"
+        "the tool. The decision on what check should be done is up to the\n"
+        "user of the tool or the dev.",
+    )
+    @abstractmethod
+    def get_evaluation_checks_based_on_tool_response(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> list[EvaluationMetricName]:
+        raise NotImplementedError
+
+    def get_tool_prompts(self) -> ToolPrompts:
+        return ToolPrompts(
+            name=self.name,
+            display_name=self.display_name(),
+            tool_description=self.tool_description().description,
+            tool_system_prompt=self.tool_description_for_system_prompt(),
+            tool_format_information_for_system_prompt=self.tool_format_information_for_system_prompt(),
+            input_model=self.tool_description_as_json(),
+            tool_user_prompt=self.tool_description_for_user_prompt(),
+            tool_format_information_for_user_prompt=self.tool_format_information_for_user_prompt(),
+        )
+
+    # Properties that we should soon deprecate
+
+    @property
+    @deprecated("Never reuse event. Dangerous")
+    def event(self) -> ChatEvent:
+        return self._event
+
+    @property
+    @deprecated("Do not use this property as directly tied to chat frontend")
+    def chat_service(self) -> ChatService:
+        return self._chat_service
+
+    @property
+    @deprecated("Do not use this property as directly tied to chat frontend")
+    def language_model_service(self) -> LanguageModelService:
+        return self._language_model_service
+
+    @property
+    @deprecated("Do not use this as directly tied to chat frontend")
+    def tool_progress_reporter(self) -> ToolProgressReporter | None:
+        return self._tool_progress_reporter
+
+    def __init__(
+        self,
+        config: ConfigType,
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter | None = None,
+    ):
+        self.settings = ToolBuildConfig(
+            name=self.name,
+            configuration=config,
+        )
+
+        self.config = config
+        module_name = "default overwrite for module name"
+        self.logger = getLogger(f"{module_name}.{__name__}")
+        self.debug_info: dict = {}
+
+        # TODO: Remove these properties as soon as possible
+        self._event: ChatEvent = event
+        self._tool_progress_reporter: ToolProgressReporter | None = (
+            tool_progress_reporter
+        )
+
+        self._chat_service = ChatService(event)
+        self._language_model_service = LanguageModelService(event)