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

unique_toolkit/tools/test/test_tool_progress_reporter.py

@@ -0,0 +1,204 @@
+from unittest.mock import AsyncMock
+
+import pytest
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.language_model.schemas import LanguageModelFunction
+from unique_toolkit.tools.tool_progress_reporter import (
+    DUMMY_REFERENCE_PLACEHOLDER,
+    ProgressState,
+    ToolExecutionStatus,
+    ToolProgressReporter,
+    ToolWithToolProgressReporter,
+    track_tool_progress,
+)
+
+
+@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

unique_toolkit/tools/tool.py

@@ -0,0 +1,168 @@
+from abc import ABC, abstractmethod
+from enum import StrEnum
+from logging import getLogger
+from typing import Generic, TypeVar
+from typing import Any, cast
+
+from pydantic import Field
+from typing_extensions import deprecated
+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,
+    LanguageModelMessage,
+)
+from unique_toolkit.language_model.service import LanguageModelService
+
+
+from unique_toolkit.evaluators.schemas import EvaluationMetricName
+from unique_toolkit.tools.agent_chunks_handler import AgentChunksHandler
+from unique_toolkit.tools.config import ToolBuildConfig, ToolSelectionPolicy
+from unique_toolkit.tools.schemas import BaseToolConfig, ToolCallResponse, ToolPrompts
+from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
+
+ConfigType = TypeVar("ConfigType", bound=BaseToolConfig)
+
+
+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
+
+    @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)
+
+    @abstractmethod
+    def tool_description_for_system_prompt(self) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def tool_format_information_for_system_prompt(self) -> str:
+        raise NotImplementedError
+
+    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 as is bound to loop agent only")
+    @abstractmethod
+    def get_tool_call_result_for_loop_history(
+        self,
+        tool_response: ToolCallResponse,
+        agent_chunks_handler: AgentChunksHandler,
+    ) -> LanguageModelMessage:
+        raise NotImplementedError
+
+    @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_format_information_for_system_prompt=self.tool_format_information_for_system_prompt(),
+            input_model=self.tool_description_as_json(),
+        )
+
+    # 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,  # type: ignore
+            # the ToolBuildConfig has a wrong type in it to be fixed later.
+        )
+
+        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)

unique_toolkit/tools/tool_manager.py

@@ -0,0 +1,242 @@
+import asyncio
+from logging import Logger, getLogger
+from pydantic import BaseModel, Field
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+    LanguageModelToolDescription,
+)
+from unique_toolkit.tools.config import ToolBuildConfig
+from unique_toolkit.tools.factory import ToolFactory
+from unique_toolkit.tools.schemas import ToolCallResponse, ToolPrompts
+from unique_toolkit.tools.tool import Tool
+from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
+from unique_toolkit.tools.utils.execution.execution import Result, SafeTaskExecutor
+
+
+class ForcedToolOption:
+    type: str = "function"
+
+    def __init__(self, name: str):
+        self.name = name
+
+
+class ToolManagerConfig(BaseModel):
+    tools: list[ToolBuildConfig] = Field(
+        default=[],
+        description="List of tools that the agent can use.",
+    )
+
+    max_tool_calls: int = Field(
+        default=10,
+        ge=1,
+        description="Maximum number of tool calls that can be executed in one iteration.",
+    )
+
+    def __init__(self, tools: list[ToolBuildConfig], max_tool_calls: int = 10):
+        self.tools = tools
+        self.max_tool_calls = max_tool_calls
+
+
+class ToolManager:
+    """
+    Manages the tools available to the agent and executes tool calls.
+
+    This class is responsible for:
+    - Initializing tools based on the provided configuration and runtime events.
+    - Filtering tools based on availability, exclusivity, and user-defined constraints.
+    - Managing the lifecycle of tools, including retrieval, execution, and logging.
+    - Executing tool calls in parallel when possible to optimize performance.
+    - Enforcing limits on the number of tool calls and handling duplicate requests.
+
+    Key Features:
+    - Dynamic Tool Initialization: Tools are dynamically selected and initialized
+      based on runtime events and user preferences.
+    - Parallel Execution: Supports asynchronous execution of tools for efficiency.
+    - Error Handling: Provides detailed error messages and logs for failed tool calls.
+    - Scalability: Designed to handle a large number of tools and tool calls efficiently.
+
+    Only the ToolManager is allowed to interact with the tools directly.
+    """
+
+    def __init__(
+        self,
+        logger: Logger,
+        config: ToolManagerConfig,
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter,
+    ):
+        self._logger = logger
+        self._config  = config
+        self._tool_progress_reporter = tool_progress_reporter
+        self._tools = []
+        self._tool_choices = event.payload.tool_choices
+        self._disabled_tools = event.payload.disabled_tools
+        self._init__tools(event)
+
+    def _init__tools(self,  event: ChatEvent) -> None:
+        tool_choices = self._tool_choices
+        tool_configs = self._config .tools
+        self._logger.info("Initializing tool definitions...")
+        self._logger.info(f"Tool choices: {tool_choices}")
+        self._logger.info(f"Tool configs: {tool_configs}")
+
+        self.available_tools = [
+            ToolFactory.build_tool_with_settings(
+                t.name,
+                t,
+                t.configuration,
+                event,
+                tool_progress_reporter=self._tool_progress_reporter,
+            )
+            for t in tool_configs
+        ]
+
+        for t in self.available_tools:
+            if t.is_exclusive():
+                self._tools = [t]
+                return
+            if not t.is_enabled():
+                continue
+            if t.name in self._disabled_tools:
+                continue
+            if len(tool_choices) > 0 and t.name not in tool_choices:
+                continue
+
+            self._tools.append(t)
+
+    def log_loaded_tools(self):
+        self._logger.info(f"Loaded tools: {[tool.name for tool in self._tools]}")
+
+    def get_tools(self) -> list[Tool]:
+        return self._tools
+
+    def get_tool_by_name(self, name: str) -> Tool | None:
+        for tool in self._tools:
+            if tool.name == name:
+                return tool
+        return None
+
+    def get_forced_tools(self) -> list[ForcedToolOption]:
+        return [ForcedToolOption(t.name) for t in self._tools if t.name in self._tool_choices]
+
+    def get_tool_definitions(self) -> list[LanguageModelToolDescription]:
+        return [tool.tool_description() for tool in self._tools]
+
+    def get_tool_prompts(self) -> list[ToolPrompts]:
+        return [tool.get_tool_prompts() for tool in self._tools]
+
+    async def execute_selected_tools(
+        self,
+        tool_calls: list[LanguageModelFunction],
+    ) -> list[ToolCallResponse]:
+        tool_calls = tool_calls
+
+        tool_calls = self.filter_duplicate_tool_calls(
+            tool_calls=tool_calls,
+        )
+        num_tool_calls = len(tool_calls)
+
+        if num_tool_calls > self._config .max_tool_calls:
+            self._logger.warning(
+                (
+                    "Number of tool calls %s exceeds the allowed maximum of %s."
+                    "The tool calls will be reduced to the first %s."
+                ),
+                num_tool_calls,
+                self._config .max_tool_calls,
+                self._config .max_tool_calls,
+            )
+            tool_calls = tool_calls[: self._config .max_tool_calls]
+
+        tool_call_responses = await self._execute_parallelized(tool_calls=tool_calls)
+        return tool_call_responses
+
+    async def _execute_parallelized(
+        self,
+        tool_calls: list[LanguageModelFunction],
+    ) -> list[ToolCallResponse]:
+        self._logger.info("Execute tool calls")
+
+        task_executor = SafeTaskExecutor(
+            logger=self._logger,
+        )
+
+        # Create tasks for each tool call
+        tasks = [
+            task_executor.execute_async(
+                self.execute_tool_call,
+                tool_call=tool_call,
+            )
+            for tool_call in tool_calls
+        ]
+
+        # Wait until all tasks are finished
+        tool_call_results = await asyncio.gather(*tasks)
+        tool_call_results_unpacked: list[ToolCallResponse] = []
+        for i, result in enumerate(tool_call_results):
+            unpacked_tool_call_result = self._create_tool_call_response(
+                result, tool_calls[i]
+            )
+            tool_call_results_unpacked.append(unpacked_tool_call_result)
+
+        return tool_call_results_unpacked
+
+    async def execute_tool_call(
+        self, tool_call: LanguageModelFunction
+    ) -> ToolCallResponse:
+        self._logger.info(f"Processing tool call: {tool_call.name}")
+
+        tool_instance = self.get_tool_by_name(tool_call.name)
+
+        if tool_instance:
+            # Execute the tool
+            tool_response: ToolCallResponse = await tool_instance.run(
+                tool_call=tool_call
+            )
+            return tool_response
+
+        return ToolCallResponse(
+            id=tool_call.id,  # type: ignore
+            name=tool_call.name,
+            error_message=f"Tool of name {tool_call.name} not found",
+        )
+
+    def _create_tool_call_response(
+        self, result: Result[ToolCallResponse], tool_call: LanguageModelFunction
+    ) -> ToolCallResponse:
+        if not result.success:
+            return ToolCallResponse(
+                id=tool_call.id or "unknown_id",
+                name=tool_call.name,
+                error_message=str(result.exception),
+            )
+        unpacked = result.unpack()
+        if not isinstance(unpacked, ToolCallResponse):
+            return ToolCallResponse(
+                id=tool_call.id or "unknown_id",
+                name=tool_call.name,
+                error_message="Tool call response is not of type ToolCallResponse",
+            )
+        return unpacked
+
+    def filter_duplicate_tool_calls(
+        self,
+        tool_calls: list[LanguageModelFunction],
+    ) -> list[LanguageModelFunction]:
+        """
+        Filter out duplicate tool calls based on name and arguments.
+        """
+
+        unique_tool_calls = []
+
+        for call in tool_calls:
+            if all(not call == other_call for other_call in unique_tool_calls):
+                unique_tool_calls.append(call)
+
+        if len(tool_calls) != len(unique_tool_calls):
+            self._logger = getLogger(__name__)
+            self._logger.warning(
+                f"Filtered out {len(tool_calls) - len(unique_tool_calls)} duplicate tool calls."
+            )
+        return unique_tool_calls

unique_toolkit/tools/tool_progress_reporter.py

@@ -88,28 +88,21 @@ class ToolProgressReporter:
             references=references,
         )
         self.requires_new_assistant_message = (
-            self.requires_new_assistant_message
-            or 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
-        ):
+        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
-            )
+            references = self._correct_reference_sequence(references, start_number)
             all_references.extend(references)
 
-            messages.append(
-                f"{ARROW}**{item.name} {item.state.value}**: {message}"
-            )
+            messages.append(f"{ARROW}**{item.name} {item.state.value}**: {message}")
 
         await self.chat_service.modify_assistant_message_async(
             content=self._progress_start_text + "\n\n" + "\n\n".join(messages),