unique_toolkit 0.8.22__py3-none-any.whl → 0.8.24__py3-none-any.whl

unique_toolkit/evals/evaluation_manager.py

@@ -78,11 +78,9 @@ class EvaluationManager:
         self,
         logger: Logger,
         chat_service: ChatService,
-        assistant_message_id: str,
     ):
         self._logger = logger
         self._chat_service = chat_service
-        self._assistant_message_id = assistant_message_id
         self._evaluations: dict[EvaluationMetricName, Evaluation] = {}
         self._evaluation_passed: bool = True
 
@@ -96,6 +94,7 @@ class EvaluationManager:
         self,
         selected_evaluation_names: list[EvaluationMetricName],
         loop_response: LanguageModelStreamResponse,
+        assistant_message_id: str,
     ) -> list[EvaluationMetricResult]:
         task_executor = SafeTaskExecutor(
             logger=self._logger,
@@ -106,6 +105,7 @@ class EvaluationManager:
                 self.execute_evaluation_call,
                 loop_response=loop_response,
                 evaluation_name=evaluation_name,
+                assistant_message_id=assistant_message_id,
             )
             for evaluation_name in selected_evaluation_names
         ]
@@ -126,6 +126,7 @@ class EvaluationManager:
         self,
         evaluation_name: EvaluationMetricName,
         loop_response: LanguageModelStreamResponse,
+        assistant_message_id: str,
     ) -> EvaluationMetricResult:
         self._logger.info(f"Processing tool call: {evaluation_name}")
 
@@ -133,13 +134,15 @@ class EvaluationManager:
 
         if evaluation_instance:
             # Execute the evaluation
-            await self._create_assistant_message(evaluation_instance)
+            await self._create_assistant_message(
+                evaluation_instance, assistant_message_id
+            )
             evaluation_metric_result: EvaluationMetricResult = (
                 await evaluation_instance.run(loop_response)
             )
             # show results to the user
             await self._show_message_assessment(
-                evaluation_instance, evaluation_metric_result
+                evaluation_instance, evaluation_metric_result, assistant_message_id
             )
 
             return evaluation_metric_result
@@ -182,6 +185,7 @@ class EvaluationManager:
         self,
         evaluation_instance: Evaluation,
         evaluation_metric_result: EvaluationMetricResult,
+        assistant_message_id: str,
     ) -> None:
         evaluation_assessment_message = (
             await evaluation_instance.evaluation_metric_to_assessment(
@@ -189,7 +193,7 @@ class EvaluationManager:
             )
         )
         await self._chat_service.modify_message_assessment_async(
-            assistant_message_id=self._assistant_message_id,
+            assistant_message_id=assistant_message_id,
             status=evaluation_assessment_message.status,
             title=evaluation_assessment_message.title,
             explanation=evaluation_assessment_message.explanation,
@@ -197,9 +201,11 @@ class EvaluationManager:
             type=evaluation_assessment_message.type,
         )
 
-    async def _create_assistant_message(self, evaluation_instance: Evaluation):
+    async def _create_assistant_message(
+        self, evaluation_instance: Evaluation, assistant_message_id: str
+    ):
         await self._chat_service.create_message_assessment_async(
-            assistant_message_id=self._assistant_message_id,
+            assistant_message_id=assistant_message_id,
             status=ChatMessageAssessmentStatus.PENDING,
             type=evaluation_instance.get_assessment_type(),
         )

unique_toolkit/tools/config.py

@@ -85,6 +85,22 @@ class ToolBuildConfig(BaseModel):
         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.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
+
         configuration = value.get("configuration", {})
         if isinstance(configuration, dict):
             # Local import to avoid circular import at module import time
@@ -105,3 +121,11 @@ class ToolBuildConfig(BaseModel):
             config = configuration
         value["configuration"] = config
         return value
+
+
+def _rebuild_config_model():
+    """Rebuild the ToolBuildConfig model to resolve forward references."""
+    # Import here to avoid circular imports
+    from unique_toolkit.tools.schemas import BaseToolConfig  # noqa: F401
+
+    ToolBuildConfig.model_rebuild()

unique_toolkit/tools/mcp/__init__.py

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

unique_toolkit/tools/mcp/manager.py

@@ -0,0 +1,82 @@
+import logging
+
+from unique_toolkit.app.schemas import ChatEvent, McpServer, McpTool
+from unique_toolkit.tools.mcp.models import EnrichedMCPTool, MCPToolConfig
+from unique_toolkit.tools.mcp.tool_wrapper import MCPToolWrapper
+from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
+
+
+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 _enrich_tool_with_mcp_info(
+        self, mcp_tool: McpTool, server: McpServer
+    ) -> EnrichedMCPTool:
+        enriched_tool = type("EnrichedMcpTool", (), {})()
+
+        # Copy all attributes from the original tool
+        for attr in dir(mcp_tool):
+            if not attr.startswith("_"):
+                setattr(enriched_tool, attr, getattr(mcp_tool, attr))
+
+        # Add server-specific attributes
+        enriched_tool.server_id = server.id
+        enriched_tool.server_name = server.name
+        enriched_tool.server_system_prompt = getattr(server, "system_prompt", None)
+        enriched_tool.server_user_prompt = getattr(server, "user_prompt", None)
+        enriched_tool.mcp_source_id = server.id
+
+        return enriched_tool
+
+    def create_mcp_tool_wrapper(
+        self, mcp_tool: EnrichedMCPTool, tool_progress_reporter: ToolProgressReporter
+    ) -> MCPToolWrapper:
+        """Create MCP tool wrapper that behave like internal tools"""
+        try:
+            config = MCPToolConfig(
+                server_id=mcp_tool.server_id,
+                server_name=mcp_tool.server_name,
+                server_system_prompt=mcp_tool.server_system_prompt,
+                server_user_prompt=mcp_tool.server_user_prompt,
+                mcp_source_id=mcp_tool.mcp_source_id,
+            )
+            wrapper = MCPToolWrapper(
+                mcp_tool=mcp_tool,
+                config=config,
+                event=self._event,
+                tool_progress_reporter=tool_progress_reporter,
+            )
+            return wrapper
+        except Exception as e:
+            import traceback
+
+            logging.error(f"Error creating MCP tool wrapper for {mcp_tool.name}: {e}")
+            logging.error(f"Full traceback: {traceback.format_exc()}")
+            return None
+
+    def get_all_mcp_tools(self, selected_by_user: list[str]) -> list[MCPToolWrapper]:
+        selected_tools = []
+        for server in self._mcp_servers:
+            if hasattr(server, "tools") and server.tools:
+                for tool in server.tools:
+                    enriched_tool = self._enrich_tool_with_mcp_info(tool, server)
+                    wrapper = self.create_mcp_tool_wrapper(
+                        enriched_tool, self._tool_progress_reporter
+                    )
+                    if wrapper is not None:
+                        selected_tools.append(wrapper)
+        return selected_tools

unique_toolkit/tools/mcp/models.py

@@ -0,0 +1,38 @@
+from typing import Any, Dict, Optional, Protocol
+
+from unique_toolkit.tools.schemas import BaseToolConfig
+
+
+class MCPTool(Protocol):
+    """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 EnrichedMCPTool(MCPTool, Protocol):
+    """Protocol for MCP tools enriched with server information."""
+
+    server_id: str
+    server_name: str
+    server_system_prompt: Optional[str]
+    server_user_prompt: Optional[str]
+    mcp_source_id: str
+
+
+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/tools/mcp/tool_wrapper.py

@@ -0,0 +1,278 @@
+import json
+from typing import Any, Dict
+
+import unique_sdk
+from pydantic import BaseModel, Field, create_model
+
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.evals.schemas import EvaluationMetricName
+from unique_toolkit.language_model import LanguageModelMessage
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+    LanguageModelToolDescription,
+    LanguageModelToolMessage,
+)
+from unique_toolkit.tools.mcp.models import EnrichedMCPTool, MCPToolConfig
+from unique_toolkit.tools.schemas import ToolCallResponse
+from unique_toolkit.tools.tool import Tool
+from unique_toolkit.tools.tool_progress_reporter import (
+    ProgressState,
+    ToolProgressReporter,
+)
+
+
+class MCPToolWrapper(Tool[MCPToolConfig]):
+    """Wrapper class for MCP tools that implements the Tool interface"""
+
+    def __init__(
+        self,
+        mcp_tool: EnrichedMCPTool,
+        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._tool_description = mcp_tool.description or ""
+        self._parameters_schema = mcp_tool.input_schema
+
+        # Set the display name for user-facing messages
+        # Priority: title > annotations.title > name
+        self.display_name = (
+            getattr(mcp_tool, "title", None)
+            or (getattr(mcp_tool, "annotations", {}) or {}).get("title")
+            or mcp_tool.name
+        )
+
+    def tool_description(self) -> LanguageModelToolDescription:
+        """Convert MCP tool schema to LanguageModelToolDescription"""
+        # Create a Pydantic model from the MCP tool's input schema
+        parameters_model = self._create_parameters_model()
+
+        return LanguageModelToolDescription(
+            name=self.name,
+            description=self._tool_description,
+            parameters=parameters_model,
+        )
+
+    def _create_parameters_model(self) -> type[BaseModel]:
+        """Create a Pydantic model from MCP tool's input schema"""
+        properties = self._parameters_schema.get("properties", {})
+        required_fields = self._parameters_schema.get("required", [])
+
+        # Convert JSON schema properties to Pydantic fields
+        fields = {}
+        for prop_name, prop_schema in properties.items():
+            field_type = self._json_schema_to_python_type(prop_schema)
+            field_description = prop_schema.get("description", "")
+
+            if prop_name in required_fields:
+                fields[prop_name] = (
+                    field_type,
+                    Field(description=field_description),
+                )
+            else:
+                fields[prop_name] = (
+                    field_type,
+                    Field(default=None, description=field_description),
+                )
+
+        # Create dynamic model
+        return create_model(f"{self.name}Parameters", **fields)
+
+    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 display_name(self) -> str:
+        """The display name of the tool."""
+        return self._display_name
+
+    def tool_description_for_system_prompt(self) -> str:
+        """Return tool description for system prompt"""
+        return self._tool_description
+
+    def tool_format_information_for_system_prompt(self) -> str:
+        """Return formatting information for system prompt"""
+        return f"Use this MCP tool to {self._tool_description.lower()}"
+
+    def evaluation_check_list(self) -> list[EvaluationMetricName]:
+        """Return evaluation check list - empty for MCP tools for now"""
+        return []
+
+    def get_evaluation_checks_based_on_tool_response(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> list[EvaluationMetricName]:
+        """Return evaluation checks based on tool response"""
+        return []
+
+    def get_tool_call_result_for_loop_history(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> LanguageModelMessage:
+        """Convert tool response to message for loop history"""
+        # Convert the tool response to a message for the conversation history
+        content = (
+            tool_response.error_message
+            if tool_response.error_message
+            else "Tool executed successfully"
+        )
+
+        if hasattr(tool_response, "content") and tool_response.content:
+            content = str(tool_response.content)
+        elif tool_response.debug_info:
+            content = json.dumps(tool_response.debug_info)
+
+        return LanguageModelToolMessage(
+            content=content,
+            tool_call_id=tool_response.id,
+            name=tool_response.name,
+        )
+
+    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(
+                id=tool_call.id or "",
+                name=self.name,
+                debug_info={
+                    "mcp_tool": self.name,
+                    "arguments": arguments,
+                    "result": result,
+                },
+                error_message="",
+            )
+
+            # 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/tools/test/test_mcp_manager.py

@@ -0,0 +1,399 @@
+import logging
+from unittest.mock import Mock
+
+import pytest
+from pydantic import BaseModel
+
+from tests.test_obj_factory import get_event_obj
+from unique_toolkit.app.schemas import McpServer, McpTool
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.tools.config import ToolBuildConfig, ToolIcon, ToolSelectionPolicy
+from unique_toolkit.tools.factory import ToolFactory
+from unique_toolkit.tools.mcp.manager import MCPManager
+from unique_toolkit.tools.schemas import BaseToolConfig
+from unique_toolkit.tools.tool import Tool
+from unique_toolkit.tools.tool_manager import ToolManager, ToolManagerConfig
+from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
+
+
+class MockParameters(BaseModel):
+    pass
+
+
+class MockInternalSearchTool(Tool[BaseToolConfig]):
+    """Mock internal search tool for testing"""
+
+    name = "internal_search"
+
+    def __init__(self, config, event, tool_progress_reporter=None):
+        super().__init__(config, event, tool_progress_reporter)
+
+    def tool_description(self):
+        from unique_toolkit.language_model.schemas import LanguageModelToolDescription
+
+        return LanguageModelToolDescription(
+            name="internal_search",
+            description="Internal search tool for testing",
+            parameters=MockParameters,
+        )
+
+    def tool_description_for_system_prompt(self) -> str:
+        return "Internal search tool for searching content"
+
+    def tool_format_information_for_system_prompt(self) -> str:
+        return "Use this tool to search for content"
+
+    def get_tool_call_result_for_loop_history(self, tool_response):
+        from unique_toolkit.language_model.schemas import LanguageModelMessage
+
+        return LanguageModelMessage(role="tool", content="Mock search result")
+
+    def evaluation_check_list(self):
+        return []
+
+    def get_evaluation_checks_based_on_tool_response(self, tool_response):
+        return []
+
+    def get_tool_prompts(self):
+        from unique_toolkit.tools.schemas import ToolPrompts
+
+        return ToolPrompts()
+
+    async def run(self, tool_call):
+        from unique_toolkit.tools.schemas import ToolCallResponse
+
+        return ToolCallResponse(id=tool_call.id, name=tool_call.name, content_chunks=[])
+
+
+class MockInternalSearchConfig(BaseToolConfig):
+    """Mock configuration for internal search tool"""
+
+    pass
+
+
+class TestMCPManager:
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """Setup test environment"""
+        self.logger = logging.getLogger(__name__)
+
+        # Register mock internal tool
+        ToolFactory.register_tool(MockInternalSearchTool, MockInternalSearchConfig)
+
+        self.event = get_event_obj(
+            user_id="test_user",
+            company_id="test_company",
+            assistant_id="test_assistant",
+            chat_id="test_chat",
+        )
+
+        # Set tool choices to include both internal and MCP tools
+        self.event.payload.tool_choices = ["internal_search", "mcp_test_tool"]
+        self.event.payload.disabled_tools = []
+
+    @pytest.fixture
+    def mock_chat_service(self):
+        """Create mock chat service for tool progress reporter"""
+        return Mock(spec=ChatService)
+
+    @pytest.fixture
+    def tool_progress_reporter(self, mock_chat_service):
+        """Create tool progress reporter fixture"""
+        return ToolProgressReporter(mock_chat_service)
+
+    @pytest.fixture
+    def mcp_tools(self):
+        """Create mock MCP tools fixture"""
+        mcp_tool = McpTool(
+            id="mcp_test_tool_id",
+            name="mcp_test_tool",
+            description="Test MCP tool",
+            input_schema={
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string", "description": "Search query"}
+                },
+                "required": ["query"],
+            },
+            output_schema=None,
+            annotations=None,
+            title="Test MCP Tool",
+            icon=None,
+            system_prompt=None,
+            user_prompt=None,
+            is_connected=True,
+        )
+        return [mcp_tool]
+
+    @pytest.fixture
+    def mcp_servers(self, mcp_tools):
+        """Create mock MCP servers fixture"""
+        server = McpServer(
+            id="test_server_id",
+            name="test_server",
+            description="Test MCP server",
+            tools=mcp_tools,
+            system_prompt="Test system prompt",
+            user_prompt="Test user prompt",
+            is_connected=True,
+        )
+        return [server]
+
+    @pytest.fixture
+    def internal_tools(self):
+        """Create internal tools fixture"""
+        internal_tool_config = ToolBuildConfig(
+            name="internal_search",
+            configuration=MockInternalSearchConfig(),
+            display_name="Internal Search",
+            icon=ToolIcon.BOOK,
+            selection_policy=ToolSelectionPolicy.BY_USER,
+            is_exclusive=False,
+            is_enabled=True,
+        )
+        return [internal_tool_config]
+
+    @pytest.fixture
+    def mcp_manager(self, mcp_servers, tool_progress_reporter):
+        """Create MCP manager fixture"""
+        return MCPManager(
+            mcp_servers=mcp_servers,
+            event=self.event,
+            tool_progress_reporter=tool_progress_reporter,
+        )
+
+    @pytest.fixture
+    def tool_manager_config(self, internal_tools):
+        """Create tool manager configuration fixture"""
+        return ToolManagerConfig(tools=internal_tools, max_tool_calls=10)
+
+    @pytest.fixture
+    def tool_manager(self, tool_manager_config, mcp_manager, tool_progress_reporter):
+        """Create tool manager fixture"""
+        return ToolManager(
+            logger=self.logger,
+            config=tool_manager_config,
+            event=self.event,
+            tool_progress_reporter=tool_progress_reporter,
+            mcp_manager=mcp_manager,
+        )
+
+    def test_tool_manager_initialization(self, tool_manager):
+        """Test tool manager is initialized correctly"""
+        assert tool_manager is not None
+        assert (
+            len(tool_manager.get_tools()) >= 2
+        )  # Should have both internal and MCP tools
+
+    def test_tool_manager_has_both_tool_types(self, tool_manager):
+        """Test that tool manager contains both MCP and internal tools"""
+        tools = tool_manager.get_tools()
+        tool_names = [tool.name for tool in tools]
+
+        # Should contain internal search tool
+        assert "internal_search" in tool_names
+
+        # Should contain MCP tool (wrapped)
+        assert "mcp_test_tool" in tool_names
+
+        # Should have at least 2 tools total
+        assert len(tools) >= 2
+
+    def test_tool_manager_can_get_tools_by_name(self, tool_manager):
+        """Test that tool manager can retrieve tools by name"""
+        # Test getting internal tool
+        internal_tool = tool_manager.get_tool_by_name("internal_search")
+        assert internal_tool is not None
+        assert internal_tool.name == "internal_search"
+
+        # Test getting MCP tool
+        mcp_tool = tool_manager.get_tool_by_name("mcp_test_tool")
+        assert mcp_tool is not None
+        assert mcp_tool.name == "mcp_test_tool"
+
+    def test_tool_manager_tools_property_contains_both_types(self, tool_manager):
+        """Test that the _tools property contains both internal and MCP tools"""
+        # Access the private _tools attribute directly to verify integration
+        tools = tool_manager._tools
+        tool_names = [tool.name for tool in tools]
+
+        # Verify both tool types are present
+        assert "internal_search" in tool_names, (
+            f"Internal tool missing. Available tools: {tool_names}"
+        )
+        assert "mcp_test_tool" in tool_names, (
+            f"MCP tool missing. Available tools: {tool_names}"
+        )
+
+        # Verify we have the expected number of tools
+        assert len(tools) == 2, f"Expected 2 tools, got {len(tools)}: {tool_names}"
+
+    def test_tool_manager_logs_loaded_tools(self, tool_manager, caplog):
+        """Test that tool manager logs the loaded tools correctly"""
+        with caplog.at_level(logging.INFO):
+            tool_manager.log_loaded_tools()
+
+        # Check that both tools are mentioned in the logs
+        log_output = caplog.text
+        assert "internal_search" in log_output
+        assert "mcp_test_tool" in log_output
+
+    def test_tool_manager_gets_tool_definitions(self, tool_manager):
+        """Test that tool manager returns tool definitions for both tool types"""
+        definitions = tool_manager.get_tool_definitions()
+
+        # Should have definitions for both tools
+        assert len(definitions) == 2
+
+        definition_names = [defn.name for defn in definitions]
+        assert "internal_search" in definition_names
+        assert "mcp_test_tool" in definition_names
+
+    def test_init_tools_method(
+        self, tool_manager_config, mcp_manager, tool_progress_reporter
+    ):
+        """Test the _init__tools method behavior with different scenarios"""
+
+        # Test 1: Normal initialization with both tool types
+        tool_manager = ToolManager(
+            logger=self.logger,
+            config=tool_manager_config,
+            event=self.event,
+            tool_progress_reporter=tool_progress_reporter,
+            mcp_manager=mcp_manager,
+        )
+
+        # Verify both tools are loaded
+        tools = tool_manager.get_tools()
+        tool_names = [tool.name for tool in tools]
+        assert "internal_search" in tool_names
+        assert "mcp_test_tool" in tool_names
+        assert len(tools) == 2
+
+    def test_init_tools_with_disabled_tools(
+        self, tool_manager_config, mcp_manager, tool_progress_reporter
+    ):
+        """Test _init__tools method when some tools are disabled"""
+
+        # Modify event to disable the internal search tool
+        event_with_disabled = get_event_obj(
+            user_id="test_user",
+            company_id="test_company",
+            assistant_id="test_assistant",
+            chat_id="test_chat",
+        )
+        event_with_disabled.payload.tool_choices = ["internal_search", "mcp_test_tool"]
+        event_with_disabled.payload.disabled_tools = ["internal_search"]
+
+        tool_manager = ToolManager(
+            logger=self.logger,
+            config=tool_manager_config,
+            event=event_with_disabled,
+            tool_progress_reporter=tool_progress_reporter,
+            mcp_manager=mcp_manager,
+        )
+
+        # Should only have MCP tool, internal tool should be filtered out
+        tools = tool_manager.get_tools()
+        tool_names = [tool.name for tool in tools]
+        assert "internal_search" not in tool_names
+        assert "mcp_test_tool" in tool_names
+        assert len(tools) == 1
+
+    def test_init_tools_with_limited_tool_choices(
+        self, tool_manager_config, mcp_manager, tool_progress_reporter
+    ):
+        """Test _init__tools method when only specific tools are chosen"""
+
+        # Modify event to only choose internal search tool
+        event_with_limited_choices = get_event_obj(
+            user_id="test_user",
+            company_id="test_company",
+            assistant_id="test_assistant",
+            chat_id="test_chat",
+        )
+        event_with_limited_choices.payload.tool_choices = ["internal_search"]
+        event_with_limited_choices.payload.disabled_tools = []
+
+        tool_manager = ToolManager(
+            logger=self.logger,
+            config=tool_manager_config,
+            event=event_with_limited_choices,
+            tool_progress_reporter=tool_progress_reporter,
+            mcp_manager=mcp_manager,
+        )
+
+        # Should only have internal search tool
+        tools = tool_manager.get_tools()
+        tool_names = [tool.name for tool in tools]
+        assert "internal_search" in tool_names
+        assert "mcp_test_tool" not in tool_names
+        assert len(tools) == 1
+
+    def test_init_tools_with_exclusive_tool(self, mcp_manager, tool_progress_reporter):
+        """Test _init__tools method when an exclusive tool is present"""
+
+        # Create an exclusive tool configuration
+        exclusive_tool_config = ToolBuildConfig(
+            name="internal_search",
+            configuration=MockInternalSearchConfig(),
+            display_name="Internal Search",
+            icon=ToolIcon.BOOK,
+            selection_policy=ToolSelectionPolicy.BY_USER,
+            is_exclusive=True,  # Make it exclusive
+            is_enabled=True,
+        )
+
+        config_with_exclusive = ToolManagerConfig(
+            tools=[exclusive_tool_config], max_tool_calls=10
+        )
+
+        tool_manager = ToolManager(
+            logger=self.logger,
+            config=config_with_exclusive,
+            event=self.event,
+            tool_progress_reporter=tool_progress_reporter,
+            mcp_manager=mcp_manager,
+        )
+
+        # Should only have the exclusive tool, MCP tools should be ignored
+        tools = tool_manager.get_tools()
+        tool_names = [tool.name for tool in tools]
+        assert "internal_search" in tool_names
+        assert "mcp_test_tool" not in tool_names
+        assert len(tools) == 1
+
+    def test_init_tools_with_disabled_tool_config(
+        self, mcp_manager, tool_progress_reporter
+    ):
+        """Test _init__tools method when a tool is disabled in its configuration"""
+
+        # Create a disabled tool configuration
+        disabled_tool_config = ToolBuildConfig(
+            name="internal_search",
+            configuration=MockInternalSearchConfig(),
+            display_name="Internal Search",
+            icon=ToolIcon.BOOK,
+            selection_policy=ToolSelectionPolicy.BY_USER,
+            is_exclusive=False,
+            is_enabled=False,  # Disable the tool
+        )
+
+        config_with_disabled = ToolManagerConfig(
+            tools=[disabled_tool_config], max_tool_calls=10
+        )
+
+        tool_manager = ToolManager(
+            logger=self.logger,
+            config=config_with_disabled,
+            event=self.event,
+            tool_progress_reporter=tool_progress_reporter,
+            mcp_manager=mcp_manager,
+        )
+
+        # Should only have MCP tool, disabled internal tool should be filtered out
+        tools = tool_manager.get_tools()
+        tool_names = [tool.name for tool in tools]
+        assert "internal_search" not in tool_names
+        assert "mcp_test_tool" in tool_names
+        assert len(tools) == 1

unique_toolkit/tools/tool_manager.py

@@ -11,13 +11,24 @@ from unique_toolkit.language_model.schemas import (
     LanguageModelTool,
     LanguageModelToolDescription,
 )
-from unique_toolkit.tools.config import ToolBuildConfig
+from unique_toolkit.tools.config import ToolBuildConfig, _rebuild_config_model
 from unique_toolkit.tools.factory import ToolFactory
+from unique_toolkit.tools.mcp.manager import MCPManager
 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
 
+# Rebuild the config model now that all imports are resolved
+_rebuild_config_model()
+
+
+class ForcedToolOption:
+    type: str = "function"
+
+    def __init__(self, name: str):
+        self.name = name
+
 
 class ToolManagerConfig(BaseModel):
     tools: list[ToolBuildConfig] = Field(
@@ -59,6 +70,7 @@ class ToolManager:
         config: ToolManagerConfig,
         event: ChatEvent,
         tool_progress_reporter: ToolProgressReporter,
+        mcp_manager: MCPManager,
     ):
         self._logger = logger
         self._config = config
@@ -68,6 +80,7 @@ class ToolManager:
         self._disabled_tools = event.payload.disabled_tools
         # this needs to be a set of strings to avoid duplicates
         self._tool_evaluation_check_list: set[EvaluationMetricName] = set()
+        self._mcp_manager = mcp_manager
         self._init__tools(event)
 
     def _init__tools(self, event: ChatEvent) -> None:
@@ -77,7 +90,8 @@ class ToolManager:
         self._logger.info(f"Tool choices: {tool_choices}")
         self._logger.info(f"Tool configs: {tool_configs}")
 
-        self.available_tools = [
+        # Build internal tools from configurations
+        internal_tools = [
             ToolFactory.build_tool_with_settings(
                 t.name,
                 t,
@@ -88,6 +102,11 @@ class ToolManager:
             for t in tool_configs
         ]
 
+        # Get MCP tools (these are already properly instantiated)
+        mcp_tools = self._mcp_manager.get_all_mcp_tools(tool_choices)
+        # Combine both types of tools
+        self.available_tools = internal_tools + mcp_tools
+
         for t in self.available_tools:
             if t.is_exclusive():
                 self._tools = [t]

{unique_toolkit-0.8.22.dist-info → unique_toolkit-0.8.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.8.22
+Version: 0.8.24
 Summary: 
 License: Proprietary
 Author: Martin Fadler
@@ -114,6 +114,12 @@ All notable changes to this project will be documented in this file.
 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).
 
+## [0.8.24] - 2025-08-25
+- Optimized hallucination manager
+
+## [0.8.23] - 2025-08-27
+- Add MCP manager that handles MCP related logic
+
 
 ## [0.8.22] - 2025-08-25
 - Add DeepSeek-R1, DeepSeek-V3.1, Qwen3-235B-A22B and Qwen3-235B-A22B-Thinking-2507 to supported model list

{unique_toolkit-0.8.22.dist-info → unique_toolkit-0.8.24.dist-info}/RECORD

@@ -40,7 +40,7 @@ unique_toolkit/evals/config.py,sha256=ywHIrJs5SFdKr1WXfrofWuFfzb0iPQw8iZDpq5oEug
 unique_toolkit/evals/context_relevancy/prompts.py,sha256=EdHFUOB581yVxcOL8482KUv_LzaRjuiem71EF8udYMc,1331
 unique_toolkit/evals/context_relevancy/schema.py,sha256=lm9x0jExOinUk9itqC8ZpgReC7yj1VDwEMppxlZGqpY,2923
 unique_toolkit/evals/context_relevancy/service.py,sha256=txTWIhV65QGFhxG1jCb5TTqZc_c7K9I8pi2HVDTAfm8,8384
-unique_toolkit/evals/evaluation_manager.py,sha256=luavQrMTrxEi7hRhpUSlGLYpEza7eQCg5Nk46ogbabw,7721
+unique_toolkit/evals/evaluation_manager.py,sha256=g-8qa_6_p53C9Okx8iNkuoIXYSJrf-6sQ-xku7bo9kI,7895
 unique_toolkit/evals/exception.py,sha256=7lcVbCyoN4Md1chNJDFxpUYyWbVrcr9dcc3TxWykJTc,115
 unique_toolkit/evals/hallucination/constants.py,sha256=FLcXl5XU07jCvS8YPX9l6UjTaqyQ8YvnSKpx4Z6wZ2Y,1997
 unique_toolkit/evals/hallucination/hallucination_evaluation.py,sha256=TuZ88jeVn0tVr9d0GhWyJSxKNA16nhvr2xRPo-yK8OM,3063
@@ -96,18 +96,23 @@ unique_toolkit/short_term_memory/service.py,sha256=5PeVBu1ZCAfyDb2HLVvlmqSbyzBBu
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=cxWjb2dxEI2HGsakKdVCkSNi7VK9mr08w5sDcFCQyWI,9553
 unique_toolkit/thinking_manager/thinking_manager.py,sha256=AJfmrTXTr-DxBnJ2_zYYpYo57kr5deqT0LiZb8UdaDQ,4175
-unique_toolkit/tools/config.py,sha256=CHelTOEd2uqktK-bSYIyaqb9M1hKz8fy_DAR2mVzXpc,2985
+unique_toolkit/tools/config.py,sha256=E4lFQA4gCi3_j4hcdnd5YUHhbaWTyIGWrgD18QardQw,3796
 unique_toolkit/tools/factory.py,sha256=w3uNHuYBIJ330Xi8PTdAkr8G3OMbQH2cBgvk5UT16oE,1253
+unique_toolkit/tools/mcp/__init__.py,sha256=_2_im3F4IgAwKyHysndhQSV1GbfoRviNOZlAS5vkF3Q,158
+unique_toolkit/tools/mcp/manager.py,sha256=7J9oE_4WxYcPDXx8RcWP67pt5oVMCUvxk_FDsu2J3YY,3218
+unique_toolkit/tools/mcp/models.py,sha256=4SFE12Y8xoJ0GBaXdz8oylhUBMHrrm66BRurCm3tCAI,987
+unique_toolkit/tools/mcp/tool_wrapper.py,sha256=lclzLZ9q4q6dLg7wUJq_aYjV0sTlQwk5abB9DeXJsr8,10287
 unique_toolkit/tools/schemas.py,sha256=PJBWN9NcEBSUmhrAMznbnw06ordrf4lfbhcxYZFmAOc,4720
+unique_toolkit/tools/test/test_mcp_manager.py,sha256=jnVKxqJyeJBDb6LvI8bM9dd8XPxrbXV0VCOHMhCdlug,14398
 unique_toolkit/tools/test/test_tool_progress_reporter.py,sha256=GTtmBqOUo0-4fh_q0lRgxDhwKeankc3FHFD5ULZAm4Y,6299
 unique_toolkit/tools/tool.py,sha256=zreY6fGBAfOIPNujWkl4cV5fBojjzGzVyvYuNLDrbzQ,5626
-unique_toolkit/tools/tool_manager.py,sha256=vYGmNF9IvQoAG3N20TIbCpzmsJfoTFn_SVaENTJIk1w,9358
+unique_toolkit/tools/tool_manager.py,sha256=-DTfX608iPSHo2_UlC_lVG8DGJfWhuzjAk_Wshc3YTQ,9988
 unique_toolkit/tools/tool_progress_reporter.py,sha256=ixud9VoHey1vlU1t86cW0-WTvyTwMxNSWBon8I11SUk,7955
 unique_toolkit/tools/utils/execution/execution.py,sha256=vjG2Y6awsGNtlvyQAGCTthQ5thWHYnn-vzZXaYLb3QE,7922
 unique_toolkit/tools/utils/source_handling/schema.py,sha256=vzAyf6ZWNexjMO0OrnB8y2glGkvAilmGGQXd6zcDaKw,870
 unique_toolkit/tools/utils/source_handling/source_formatting.py,sha256=C7uayNbdkNVJdEARA5CENnHtNY1SU6etlaqbgHNyxaQ,9152
 unique_toolkit/tools/utils/source_handling/tests/test_source_formatting.py,sha256=oM5ZxEgzROrnX1229KViCAFjRxl9wCTzWZoinYSHleM,6979
-unique_toolkit-0.8.22.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-0.8.22.dist-info/METADATA,sha256=yFz8RDjtGih0zLFKHQNfgyWM1P0W1URIBGu2EjK3ZWc,28453
-unique_toolkit-0.8.22.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-0.8.22.dist-info/RECORD,,
+unique_toolkit-0.8.24.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-0.8.24.dist-info/METADATA,sha256=0vvS5x-HsApDZlOrB7vnLQPal3u7Ucqxv7hZa35ebOo,28588
+unique_toolkit-0.8.24.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-0.8.24.dist-info/RECORD,,