unique_toolkit 0.8.22__tar.gz → 0.8.24__tar.gz

{unique_toolkit-0.8.22 → unique_toolkit-0.8.24}/CHANGELOG.md

@@ -5,6 +5,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 → unique_toolkit-0.8.24}/PKG-INFO

@@ -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 → unique_toolkit-0.8.24}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "0.8.22"
+version = "0.8.24"
 description = ""
 authors = [
     "Martin Fadler <martin.fadler@unique.ch>",

{unique_toolkit-0.8.22 → unique_toolkit-0.8.24}/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-0.8.22 → unique_toolkit-0.8.24}/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-0.8.24/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-0.8.24/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-0.8.24/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-0.8.24/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