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

unique_toolkit/agentic/tools/config.py

@@ -0,0 +1,167 @@
+import json
+from enum import StrEnum
+from typing import Annotated, Any, Dict
+
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    Field,
+    ValidationInfo,
+    model_validator,
+)
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+
+
+class ToolIcon(StrEnum):
+    ANALYTICS = "IconAnalytics"
+    BOOK = "IconBook"
+    FOLDERDATA = "IconFolderData"
+    INTEGRATION = "IconIntegration"
+    TEXT_COMPARE = "IconTextCompare"
+    WORLD = "IconWorld"
+    QUICK_REPLY = "IconQuickReply"
+    CHAT_PLUS = "IconChatPlus"
+    TELESCOPE = "IconTelescope"
+
+
+class ToolSelectionPolicy(StrEnum):
+    """Determine the usage policy of tools."""
+
+    FORCED_BY_DEFAULT = "ForcedByDefault"
+    ON_BY_DEFAULT = "OnByDefault"
+    BY_USER = "ByUser"
+
+
+def handle_undefined_icon(value: Any) -> ToolIcon:
+    try:
+        if isinstance(value, str):
+            return ToolIcon(value)
+        else:
+            return ToolIcon.BOOK
+    except ValueError:
+        return ToolIcon.BOOK
+
+
+class ToolBuildConfig(BaseModel):
+    model_config = get_configuration_dict()
+    """Main tool configuration"""
+
+    name: str
+    configuration: BaseToolConfig
+    display_name: str = ""
+    icon: Annotated[ToolIcon, BeforeValidator(handle_undefined_icon)] = Field(
+        default=ToolIcon.BOOK,
+        description="The icon name that will be used to display the tool in the user interface.",
+    )
+    selection_policy: ToolSelectionPolicy = Field(
+        default=ToolSelectionPolicy.BY_USER,
+    )
+    is_exclusive: bool = Field(
+        default=False,
+        description="This tool must be chosen by the user and no other tools are used for this iteration.",
+    )
+    is_sub_agent: bool = False
+
+    is_enabled: bool = Field(default=True)
+
+    @model_validator(mode="before")
+    def initialize_config_based_on_tool_name(
+        cls,
+        value: Any,
+        info: ValidationInfo,
+    ) -> Any:
+        """Check the given values for."""
+        if not isinstance(value, dict):
+            return value
+
+        is_mcp_tool = value.get("mcp_source_id", "") != ""
+        mcp_configuration = value.get("configuration", {})
+
+        # Import at runtime to avoid circular imports
+        from unique_toolkit.agentic.tools.mcp.models import MCPToolConfig
+
+        if (
+            isinstance(mcp_configuration, MCPToolConfig)
+            and mcp_configuration.mcp_source_id
+        ):
+            return value
+        if is_mcp_tool:
+            # For MCP tools, skip ToolFactory validation
+            # Configuration can remain as a dict
+            return value
+
+        is_sub_agent_tool = (
+            value.get("is_sub_agent") or value.get("isSubAgent") or False
+        )
+
+        configuration = value.get("configuration", {})
+
+        if is_sub_agent_tool:
+            from unique_toolkit.agentic.tools.a2a import ExtendedSubAgentToolConfig
+
+            config = ExtendedSubAgentToolConfig.model_validate(configuration)
+        elif isinstance(configuration, dict):
+            # Local import to avoid circular import at module import time
+            from unique_toolkit.agentic.tools.factory import ToolFactory
+
+            config = ToolFactory.build_tool_config(
+                value["name"],
+                **configuration,
+            )
+        else:
+            # Check that the type of config matches the tool name
+            from unique_toolkit.agentic.tools.factory import ToolFactory
+
+            assert isinstance(
+                configuration,
+                ToolFactory.tool_config_map[value["name"]],  # type: ignore
+            )
+            config = configuration
+        value["configuration"] = config
+        return value
+
+    def model_dump(self) -> Dict[str, Any]:
+        """
+        Returns a dict representation of the tool config that preserves
+        subclass fields from `configuration` by delegating to its own
+        model_dump. This prevents `{}` when `configuration` is typed
+        as `BaseToolConfig` but holds a subclass instance.
+        """
+        data: Dict[str, Any] = {
+            "name": self.name,
+            "configuration": self.configuration.model_dump()
+            if self.configuration
+            else None,
+            "display_name": self.display_name,
+            "icon": self.icon,
+            "selection_policy": self.selection_policy,
+            "is_exclusive": self.is_exclusive,
+            "is_sub_agent": self.is_sub_agent,
+            "is_enabled": self.is_enabled,
+        }
+        return data
+
+    def model_dump_json(self) -> str:
+        """
+        Returns a JSON string representation of the tool config.
+        Ensures `configuration` is fully serialized by using the
+        subclass's `model_dump_json()` when available.
+        """
+        config_json = (
+            self.configuration.model_dump_json() if self.configuration else None
+        )
+        config = json.loads(config_json) if config_json else None
+
+        data: Dict[str, Any] = {
+            "name": self.name,
+            "configuration": config,
+            "display_name": self.display_name,
+            "icon": self.icon,
+            "selection_policy": self.selection_policy,
+            "is_exclusive": self.is_exclusive,
+            "is_sub_agent": self.is_sub_agent,
+            "is_enabled": self.is_enabled,
+        }
+        return json.dumps(data)

unique_toolkit/agentic/tools/factory.py

@@ -0,0 +1,44 @@
+from typing import TYPE_CHECKING, Callable
+
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+from unique_toolkit.agentic.tools.tool import Tool
+
+if TYPE_CHECKING:
+    from unique_toolkit.agentic.tools.config import ToolBuildConfig
+
+
+class ToolFactory:
+    tool_map: dict[str, type[Tool]] = {}
+    tool_config_map: dict[str, Callable] = {}
+
+    @classmethod
+    def register_tool_config(cls, tool_name: str, tool_config: type[BaseToolConfig]):
+        cls.tool_config_map[tool_name] = tool_config
+
+    @classmethod
+    def register_tool(
+        cls,
+        tool: type[Tool],
+        tool_config: type[BaseToolConfig],
+    ):
+        cls.tool_map[tool.name] = tool
+        cls.tool_config_map[tool.name] = tool_config
+
+    @classmethod
+    def build_tool(cls, tool_name: str, *args, **kwargs) -> Tool[BaseToolConfig]:
+        tool = cls.tool_map[tool_name](*args, **kwargs)
+        return tool
+
+    @classmethod
+    def build_tool_with_settings(
+        cls, tool_name: str, settings: "ToolBuildConfig", *args, **kwargs
+    ) -> Tool[BaseToolConfig]:
+        tool = cls.tool_map[tool_name](*args, **kwargs)
+        tool.settings = settings
+        return tool
+
+    @classmethod
+    def build_tool_config(cls, tool_name: str, **kwargs) -> BaseToolConfig:
+        if tool_name not in cls.tool_config_map:
+            raise ValueError(f"Tool {tool_name} not found")
+        return cls.tool_config_map[tool_name](**kwargs)

unique_toolkit/agentic/tools/mcp/__init__.py

@@ -0,0 +1,4 @@
+from .models import MCPToolConfig
+from .tool_wrapper import MCPToolWrapper
+
+__all__ = ["MCPToolWrapper", "MCPToolConfig"]

unique_toolkit/agentic/tools/mcp/manager.py

@@ -0,0 +1,71 @@
+import logging
+
+from unique_toolkit.agentic.tools.config import (
+    ToolBuildConfig,
+    ToolIcon,
+    ToolSelectionPolicy,
+)
+from unique_toolkit.agentic.tools.mcp.models import MCPToolConfig
+from unique_toolkit.agentic.tools.mcp.tool_wrapper import MCPToolWrapper
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+from unique_toolkit.agentic.tools.tool import Tool
+from unique_toolkit.agentic.tools.tool_progress_reporter import ToolProgressReporter
+from unique_toolkit.app.schemas import ChatEvent, McpServer
+
+
+class MCPManager:
+    def __init__(
+        self,
+        mcp_servers: list[McpServer],
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter,
+    ):
+        self._mcp_servers = mcp_servers
+        self._event = event
+        self._tool_progress_reporter = tool_progress_reporter
+
+    def get_mcp_servers(self):
+        return self._mcp_servers
+
+    def get_mcp_server_by_id(self, id: str):
+        return next((server for server in self._mcp_servers if server.id == id), None)
+
+    def get_all_mcp_tools(self) -> list[Tool[BaseToolConfig]]:
+        selected_tools = []
+        for server in self._mcp_servers:
+            if not hasattr(server, "tools"):
+                continue
+            if not server.tools:
+                continue
+
+            for tool in server.tools:
+                try:
+                    config = MCPToolConfig(
+                        server_id=server.id,
+                        server_name=server.name,
+                        server_system_prompt=server.system_prompt,
+                        server_user_prompt=server.user_prompt,
+                        mcp_source_id=server.id,
+                    )
+                    wrapper = MCPToolWrapper(
+                        mcp_server=server,
+                        mcp_tool=tool,
+                        config=config,
+                        event=self._event,
+                        tool_progress_reporter=self._tool_progress_reporter,
+                    )
+                    wrapper.settings = ToolBuildConfig(  # TODO: this must be refactored to behave like the other tools.
+                        name=tool.name,
+                        configuration=config,
+                        display_name=tool.title or tool.name,
+                        is_exclusive=False,
+                        is_enabled=True,
+                        icon=ToolIcon.BOOK,
+                        selection_policy=ToolSelectionPolicy.BY_USER,
+                    )
+                    selected_tools.append(wrapper)
+                except Exception as e:
+                    logging.error(
+                        f"Error creating MCP tool wrapper for {tool.name}: {e}"
+                    )
+        return selected_tools

unique_toolkit/agentic/tools/mcp/models.py

@@ -0,0 +1,28 @@
+from typing import Any, Dict, Optional
+
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+
+
+class MCPTool:
+    """Protocol defining the expected structure of an MCP tool."""
+
+    name: str
+    description: Optional[str]
+    input_schema: Dict[str, Any]
+    output_schema: Optional[Dict[str, Any]]
+    annotations: Optional[Dict[str, Any]]
+    title: Optional[str]
+    icon: Optional[str]
+    system_prompt: Optional[str]
+    user_prompt: Optional[str]
+    is_connected: bool
+
+
+class MCPToolConfig(BaseToolConfig):
+    """Configuration for MCP tools"""
+
+    server_id: str
+    server_name: str
+    server_system_prompt: Optional[str] = None
+    server_user_prompt: Optional[str] = None
+    mcp_source_id: str

unique_toolkit/agentic/tools/mcp/tool_wrapper.py

@@ -0,0 +1,234 @@
+import json
+import logging
+from typing import Any, Dict
+
+import unique_sdk
+
+from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
+from unique_toolkit.agentic.tools.mcp.models import MCPToolConfig
+from unique_toolkit.agentic.tools.schemas import ToolCallResponse
+from unique_toolkit.agentic.tools.tool import Tool
+from unique_toolkit.agentic.tools.tool_progress_reporter import (
+    ProgressState,
+    ToolProgressReporter,
+)
+from unique_toolkit.app.schemas import ChatEvent, McpServer, McpTool
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+    LanguageModelToolDescription,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MCPToolWrapper(Tool[MCPToolConfig]):
+    """Wrapper class for MCP tools that implements the Tool interface"""
+
+    def __init__(
+        self,
+        mcp_server: McpServer,
+        mcp_tool: McpTool,
+        config: MCPToolConfig,
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter | None = None,
+    ):
+        self.name = mcp_tool.name
+        super().__init__(config, event, tool_progress_reporter)
+        self._mcp_tool = mcp_tool
+        self._mcp_server = mcp_server
+
+    def tool_description(self) -> LanguageModelToolDescription:
+        """Convert MCP tool schema to LanguageModelToolDescription"""
+        # Create a Pydantic model from the MCP tool's input schema
+        logger.info(
+            "MCP tool %s schema %s", self._mcp_tool.name, self._mcp_tool.input_schema
+        )
+
+        return LanguageModelToolDescription(
+            name=self.name,
+            description=self._mcp_tool.description or "",
+            parameters=self._mcp_tool.input_schema,
+        )
+
+    def _json_schema_to_python_type(self, schema: Dict[str, Any]) -> type:
+        """Convert JSON schema type to Python type"""
+        json_type = schema.get("type", "string")
+
+        type_mapping = {
+            "string": str,
+            "integer": int,
+            "number": float,
+            "boolean": bool,
+            "array": list,
+            "object": dict,
+        }
+
+        return type_mapping.get(json_type, str)
+
+    def tool_description_for_system_prompt(self) -> str:
+        """Return tool description for system prompt"""
+        # Not using jinja here to keep it simple and not import new packages.
+        description = (
+            f"**MCP Server**: {self._mcp_server.name}\n"
+            f"**Tool Name**: {self.name}\n"
+            f"{self._mcp_tool.system_prompt}"
+        )
+
+        return description
+
+    def tool_description_for_user_prompt(self) -> str:
+        return self._mcp_tool.user_prompt or ""
+
+    def tool_format_information_for_user_prompt(self) -> str:
+        return ""
+
+    def tool_format_information_for_system_prompt(self) -> str:
+        """Return formatting information for system prompt"""
+        return ""  # this is empty for now as it requires to add this to the MCP model of the backend.
+
+    def evaluation_check_list(self) -> list[EvaluationMetricName]:
+        """Return evaluation check list - empty for MCP tools for now"""
+        # TODO: this is empty for now as it requires a setting in the backend for choosing a suitable validator.
+        return []
+
+    def get_evaluation_checks_based_on_tool_response(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> list[EvaluationMetricName]:
+        """Return evaluation checks based on tool response"""
+        return []
+
+    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
+        """Execute the MCP tool using SDK to call public API"""
+        self.logger.info(f"Running MCP tool: {self.name}")
+
+        # Notify progress if reporter is available
+        if self._tool_progress_reporter:
+            await self._tool_progress_reporter.notify_from_tool_call(
+                tool_call=tool_call,
+                name=f"**{self.display_name()}**",
+                message=f"Executing MCP tool: {self.display_name()}",
+                state=ProgressState.RUNNING,
+            )
+
+        try:
+            # Robust argument extraction and validation
+            arguments = self._extract_and_validate_arguments(tool_call)
+
+            # Use SDK to call the public API
+            result = await self._call_mcp_tool_via_sdk(arguments)
+
+            # Create successful response
+            tool_response = ToolCallResponse(  # TODO: Why result here not applied directly to the body of the tool_response? like so how does it know the results in the history?
+                id=tool_call.id or "",
+                name=self.name,
+                debug_info={
+                    "mcp_tool": self.name,
+                    "arguments": arguments,
+                },
+                error_message="",
+                content=json.dumps(result),
+            )
+
+            # Notify completion
+            if self._tool_progress_reporter:
+                await self._tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=f"**{self.display_name()}**",
+                    message=f"MCP tool completed: {self.display_name()}",
+                    state=ProgressState.FINISHED,
+                )
+
+            return tool_response
+
+        except Exception as e:
+            self.logger.error(f"Error executing MCP tool {self.name}: {e}")
+
+            # Notify failure
+            if self._tool_progress_reporter:
+                await self._tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=f"**{self.display_name()}**",
+                    message=f"MCP tool failed: {str(e)}",
+                    state=ProgressState.FAILED,
+                )
+
+            return ToolCallResponse(
+                id=tool_call.id or "",
+                name=self.name,
+                debug_info={
+                    "mcp_tool": self.name,
+                    "error": str(e),
+                    "original_arguments": getattr(tool_call, "arguments", None),
+                },
+                error_message=str(e),
+            )
+
+    def _extract_and_validate_arguments(
+        self, tool_call: LanguageModelFunction
+    ) -> Dict[str, Any]:
+        """
+        Extract and validate arguments from tool call, handling various formats robustly.
+
+        The arguments field can come in different formats:
+        1. As a JSON string (expected format from OpenAI API)
+        2. As a dictionary (from internal processing)
+        3. As None or empty (edge cases)
+        """
+        raw_arguments = tool_call.arguments
+
+        # Handle None or empty arguments
+        if not raw_arguments:
+            self.logger.warning(f"MCP tool {self.name} called with empty arguments")
+            return {}
+
+        # Handle string arguments (JSON format)
+        if isinstance(raw_arguments, str):
+            try:
+                parsed_arguments = json.loads(raw_arguments)
+                if not isinstance(parsed_arguments, dict):
+                    self.logger.warning(
+                        f"MCP tool {self.name}: arguments JSON parsed to non-dict: {type(parsed_arguments)}"
+                    )
+                    return {}
+                return parsed_arguments
+            except json.JSONDecodeError as e:
+                self.logger.error(
+                    f"MCP tool {self.name}: failed to parse arguments JSON '{raw_arguments}': {e}"
+                )
+                raise ValueError(
+                    f"Invalid JSON arguments for MCP tool {self.name}: {e}"
+                )
+
+        # Handle dictionary arguments (already parsed)
+        if isinstance(raw_arguments, dict):
+            self.logger.debug(f"MCP tool {self.name}: arguments already in dict format")
+            return raw_arguments
+
+        # Handle unexpected argument types
+        self.logger.error(
+            f"MCP tool {self.name}: unexpected arguments type {type(raw_arguments)}: {raw_arguments}"
+        )
+        raise ValueError(
+            f"Unexpected arguments type for MCP tool {self.name}: {type(raw_arguments)}"
+        )
+
+    async def _call_mcp_tool_via_sdk(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Call MCP tool via SDK to public API"""
+        try:
+            result = unique_sdk.MCP.call_tool(
+                user_id=self.event.user_id,
+                company_id=self.event.company_id,
+                name=self.name,
+                arguments=arguments,
+            )
+
+            self.logger.info(
+                f"Calling MCP tool {self.name} with arguments: {arguments}"
+            )
+            self.logger.debug(f"Result: {result}")
+
+            return result
+        except Exception as e:
+            self.logger.error(f"SDK call failed for MCP tool {self.name}: {e}")
+            raise

unique_toolkit/agentic/tools/openai_builtin/__init__.py

@@ -0,0 +1,11 @@
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter import (
+    OpenAICodeInterpreterConfig,
+    OpenAICodeInterpreterTool,
+)
+from unique_toolkit.agentic.tools.openai_builtin.manager import OpenAIBuiltInToolManager
+
+__all__ = [
+    "OpenAIBuiltInToolManager",
+    "OpenAICodeInterpreterTool",
+    "OpenAICodeInterpreterConfig",
+]

unique_toolkit/agentic/tools/openai_builtin/base.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from enum import StrEnum
+from typing import Generic, TypeVar
+
+from openai.types.responses.tool_param import CodeInterpreter
+
+from unique_toolkit.agentic.tools.schemas import ToolPrompts
+
+
+class OpenAIBuiltInToolName(StrEnum):
+    CODE_INTERPRETER = "code_interpreter"
+
+
+BuiltInToolType = CodeInterpreter  # Add other tool types when needed
+ToolType = TypeVar("ToolType", bound=BuiltInToolType)
+
+
+class OpenAIBuiltInTool(ABC, Generic[ToolType]):
+    @property
+    @abstractmethod
+    def name(self) -> OpenAIBuiltInToolName:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def tool_description(self) -> BuiltInToolType:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_tool_prompts(self) -> ToolPrompts:
+        raise NotImplementedError()

unique_toolkit/agentic/tools/openai_builtin/code_interpreter/__init__.py

@@ -0,0 +1,8 @@
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.config import (
+    OpenAICodeInterpreterConfig,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.service import (
+    OpenAICodeInterpreterTool,
+)
+
+__all__ = ["OpenAICodeInterpreterConfig", "OpenAICodeInterpreterTool"]

unique_toolkit/agentic/tools/openai_builtin/code_interpreter/config.py

@@ -0,0 +1,57 @@
+from pydantic import Field
+
+from unique_toolkit.agentic.tools.factory import ToolFactory
+from unique_toolkit.agentic.tools.openai_builtin.base import (
+    OpenAIBuiltInToolName,
+)
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+
+DEFAULT_TOOL_DESCRIPTION = "Use this tool to run python code, e.g to generate plots, process excel files, perform calculations, etc."
+
+DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = """
+Use this tool to run python code, e.g to generate plots, process excel files, perform calculations, etc.
+Instructions:
+- All files uploaded to the chat are available in the code interpreter under the path `/mnt/data/<filename>
+- All files generated through code should be saved in the `/mnt/data` folder
+
+Instructions for displaying images and files in the chat:
+Once files are generated in the `/mnt/data` folder you MUST reference them in the chat using markdown syntax in order to display them in the chat.
+
+- If you want to display an image, use the following syntax: `![Image Name](sandbox:/mnt/data/<filename>)`
+    - Images will be converted and shown in the chat.
+    - Do NOT display an extra download link for images a part from the markdown above.
+    - Not using markdown syntax will FAIL to show images to the user. 
+    - YOU MUST use the syntax above to display images, otherwise the image will not be displayed in the chat.
+    - For displaying a link to a file, use the following syntax: `[filename](sandbox:/mnt/data/<filename>)`
+        - Files are converted to references the user can click on to download the file
+
+You MUST always use this syntax, otherwise the files will not be displayed in the chat.
+""".strip()
+
+DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT = ""
+
+DEFAULT_TOOL_FORMAT_INFORMATION_FOR_USER_PROMPT = ""
+
+DEFAULT_TOOL_DESCRIPTION_FOR_USER_PROMPT = ""
+
+
+class OpenAICodeInterpreterConfig(BaseToolConfig):
+    upload_files_in_chat: bool = Field(default=True)
+
+    tool_description: str = DEFAULT_TOOL_DESCRIPTION
+    tool_description_for_system_prompt: str = DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT
+    tool_format_information_for_system_prompt: str = (
+        DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT
+    )
+    tool_description_for_user_prompt: str = DEFAULT_TOOL_DESCRIPTION_FOR_USER_PROMPT
+    tool_format_information_for_user_prompt: str = (
+        DEFAULT_TOOL_FORMAT_INFORMATION_FOR_USER_PROMPT
+    )
+
+    expires_after_minutes: int = 20
+    use_auto_container: bool = False
+
+
+ToolFactory.register_tool_config(
+    OpenAIBuiltInToolName.CODE_INTERPRETER, OpenAICodeInterpreterConfig
+)