tooluniverse 1.0.9__py3-none-any.whl → 1.0.10__py3-none-any.whl

tooluniverse/humanbase_tool.py

@@ -35,19 +35,19 @@ class HumanBaseTool(BaseTool):
         """
         Retrieve the official gene symbol (same as EnrichrTool method)
 
-        Parameters:
+        Parameters
             gene_name (str): The gene name or synonym to query.
 
-        Returns:
+        Returns
             str: The official gene symbol.
         """
         """
         Retrieve the official gene symbol for a given gene name or synonym using the MyGene.info API.
 
-        Parameters:
+        Parameters
             gene_name (str): The gene name or synonym to query.
 
-        Returns:
+        Returns
             str: The official gene symbol if found; otherwise, raises an Exception.
         """
         # URL-encode the gene_name to handle special characters
@@ -96,10 +96,10 @@ class HumanBaseTool(BaseTool):
         """
         Convert gene names to Entrez IDs using NCBI Entrez API.
 
-        Parameters:
+        Parameters
             gene_names (list): List of gene names to convert.
 
-        Returns:
+        Returns
             list: List of Entrez IDs corresponding to the gene names.
         """
         # Define the NCBI Entrez API URL for querying gene information
@@ -149,13 +149,13 @@ class HumanBaseTool(BaseTool):
         """
         Retrieve protein-protein interactions and biological processes from HumanBase.
 
-        Parameters:
+        Parameters
             genes (list): List of gene names to analyze.
             tissue (str): Tissue type for tissue-specific interactions.
             max_node (int): Maximum number of nodes to retrieve.
             interaction (str): Specific interaction type to filter by.
 
-        Returns:
+        Returns
             tuple: (NetworkX Graph of interactions, list of biological processes)
         """
         genes = self.get_entrez_ids(genes)
@@ -243,13 +243,13 @@ class HumanBaseTool(BaseTool):
         """
         Convert NetworkX graph and biological processes to string representation.
 
-        Parameters:
+        Parameters
             graph (networkx.Graph): The network graph.
             bp_collection (list): List of biological processes.
             original_genes (list): Original gene list provided by user.
             tissue (str): Tissue type used for analysis.
 
-        Returns:
+        Returns
             str: Comprehensive string representation of the network data.
         """
         output = []

tooluniverse/logging_config.py

@@ -232,7 +232,7 @@ def get_logger(name: Optional[str] = None) -> logging.Logger:
     Args:
         name (str, optional): Logger name (usually __name__)
 
-    Returns:
+    Returns
         logging.Logger: Logger instance
     """
     return _logger_manager.get_logger(name)
@@ -291,7 +291,7 @@ def get_hook_logger(name: str = "HookManager") -> logging.Logger:
     Args:
         name (str): Name of the logger. Defaults to 'HookManager'
 
-    Returns:
+    Returns
         logging.Logger: Configured logger for hook operations
     """
     return get_logger(name)

tooluniverse/mcp_client_tool.py

@@ -8,11 +8,11 @@ supporting all MCP functionality including tools, resources, and prompts.
 import json
 import asyncio
 import websockets
-import aiohttp
-import uuid
 from typing import Dict, List, Any, Optional
 from urllib.parse import urljoin
 import warnings
+from mcp.client.session import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
 from .base_tool import BaseTool
 from .tool_registry import register_tool
 import os
@@ -35,8 +35,6 @@ class BaseMCPClient:
         self.transport = normalized_transport
         self.timeout = timeout
         self.session = None
-        self.mcp_session_id = None
-        self._initialized = False
 
         # Validate transport (accept 'stdio' via normalization above)
         supported_transports = ["http", "websocket"]
@@ -44,24 +42,9 @@ class BaseMCPClient:
             # Keep message concise to satisfy line length rules
             raise ValueError("Invalid transport")
 
-    async def _ensure_session(self):
-        """Ensure HTTP session is available for HTTP transport"""
-        if self.transport == "http" and self.session is None:
-            connector = aiohttp.TCPConnector()
-            timeout = aiohttp.ClientTimeout(total=self.timeout)
-            self.session = aiohttp.ClientSession(connector=connector, timeout=timeout)
-
     async def _close_session(self):
-        """Close HTTP session if exists"""
-        if self.session:
-            try:
-                await self.session.close()
-            except Exception:
-                pass  # Ignore errors during cleanup
-            finally:
-                self.session = None
-                self.mcp_session_id = None
-                self._initialized = False
+        """Placeholder for compatibility; HTTP client calls are scoped per request."""
+        return
 
     def _get_mcp_endpoint(self, path: str) -> str:
         """Get the full MCP endpoint URL"""
@@ -72,115 +55,55 @@ class BaseMCPClient:
             return urljoin(base_url + "/", path)
         return self.server_url
 
-    async def _initialize_mcp_session(self):
-        """Initialize MCP session if needed (for compatibility with different MCP servers)"""
-        if self._initialized:
-            return
-
-        await self._ensure_session()
-
-        # Try to get session ID from server
-        try:
-            url = f"{self.server_url.rstrip('/')}/mcp"
-            test_payload = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
-
-            headers = {
-                "Content-Type": "application/json",
-                "Accept": "application/json, text/event-stream",
-            }
-
-            async with self.session.post(
-                url, json=test_payload, headers=headers
-            ) as response:
-                session_id = response.headers.get("mcp-session-id")
-                if session_id:
-                    self.mcp_session_id = session_id
-
-                if response.status in [200, 400, 406, 500]:
-                    self._initialized = True
-                    return
-
-        except Exception:
-            pass
-
-        # Fallback: generate session ID
-        if not self.mcp_session_id:
-            self.mcp_session_id = str(uuid.uuid4()).replace("-", "")
-
-        self._initialized = True
-
     async def _make_mcp_request(
         self, method: str, params: Optional[Dict] = None
     ) -> Dict[str, Any]:
         """Make an MCP JSON-RPC request"""
-        request_id = "1"
-
-        payload = {"jsonrpc": "2.0", "id": request_id, "method": method}
-
-        if params:
-            payload["params"] = params
-
         if self.transport == "http":
-            await self._ensure_session()
-            await self._initialize_mcp_session()  # Ensure session is initialized
-
-            headers = {
-                "Content-Type": "application/json",
-                "Accept": "application/json, text/event-stream",
-            }
-
-            # Add session ID if available
-            if self.mcp_session_id:
-                headers["mcp-session-id"] = self.mcp_session_id
-
             endpoint = self._get_mcp_endpoint("")
-
-            async with self.session.post(
-                endpoint, json=payload, headers=headers
-            ) as response:
-                if response.status != 200:
-                    raise Exception(
-                        f"MCP request failed with status {response.status}: {await response.text()}"
-                    )
-
-                content_type = response.headers.get("content-type", "").lower()
-
-                if "text/event-stream" in content_type:
-                    # Handle Server-Sent Events format
-                    response_text = await response.text()
-
-                    for line in response_text.split("\n"):
-                        line = line.strip()
-                        if line.startswith("data: "):
-                            json_data = line[6:]
-                            try:
-                                result = json.loads(json_data)
-                                break
-                            except json.JSONDecodeError:
-                                continue
-                    else:
-                        raise Exception(
-                            f"Failed to parse SSE response: {response_text}"
+            async with streamablehttp_client(endpoint, timeout=self.timeout) as (
+                read_stream,
+                write_stream,
+                _,
+            ):
+                async with ClientSession(read_stream, write_stream) as session:
+                    await session.initialize()
+
+                    if method == "tools/list":
+                        result = await session.list_tools()
+                    elif method == "tools/call":
+                        if not params or "name" not in params:
+                            raise ValueError("Missing tool name for tools/call")
+                        result = await session.call_tool(
+                            params["name"], params.get("arguments") or {}
                         )
-
-                elif "application/json" in content_type:
-                    result = await response.json()
-                else:
-                    try:
-                        result = await response.json()
-                    except Exception:
-                        response_text = await response.text()
-                        raise Exception(
-                            f"Unexpected content type {content_type}. Response: {response_text}"
+                    elif method == "resources/list":
+                        result = await session.list_resources()
+                    elif method == "resources/read":
+                        if not params or "uri" not in params:
+                            raise ValueError("Missing uri for resources/read")
+                        result = await session.read_resource(params["uri"])
+                    elif method == "prompts/list":
+                        result = await session.list_prompts()
+                    elif method == "prompts/get":
+                        if not params or "name" not in params:
+                            raise ValueError("Missing prompt name for prompts/get")
+                        result = await session.get_prompt(
+                            params["name"], params.get("arguments")
                         )
+                    else:
+                        raise ValueError(f"Unsupported MCP method: {method}")
 
-                if "error" in result:
-                    raise Exception(f"MCP error: {result['error']}")
-
-                return result.get("result", {})
+                    if hasattr(result, "model_dump"):
+                        return result.model_dump(mode="json")
+                    return result
 
         elif self.transport == "websocket":
             async with websockets.connect(self.server_url) as websocket:
+                request_id = "1"
+                payload = {"jsonrpc": "2.0", "id": request_id, "method": method}
+                if params:
+                    payload["params"] = params
                 await websocket.send(json.dumps(payload))
                 response = await websocket.recv()
                 result = json.loads(response)
@@ -575,7 +498,6 @@ class MCPAutoLoaderTool(BaseTool, BaseMCPClient):
     async def discover_tools(self) -> Dict[str, Any]:
         """Discover all available tools from the MCP server"""
         try:
-            await self._initialize_mcp_session()
             tools_response = await self._make_mcp_request("tools/list")
             tools = tools_response.get("tools", [])
 
@@ -631,22 +553,28 @@ class MCPAutoLoaderTool(BaseTool, BaseMCPClient):
         return configs
 
     def register_tools_in_engine(self, engine):
-        """Register discovered tools directly in the ToolUniverse engine"""
+        """Register discovered tools using ToolUniverse public API"""
         try:
             configs = self.generate_proxy_tool_configs()
 
             for config in configs:
-                # Add configuration to engine's all_tools list for validation
-                engine.all_tools.append(config)
+                config["type"]
+                tool_name = config["name"]
 
-                # Create MCPProxyTool instance for execution
-                proxy_tool = MCPProxyTool(config)
+                # Use public API to register tool
+                engine.register_custom_tool(
+                    tool_class=MCPProxyTool,
+                    tool_name=tool_name,
+                    tool_config=config,
+                    instantiate=True,  # Immediately instantiate and cache
+                )
 
-                # Register both config (for validation) and tool instance (for execution)
-                tool_name = config["name"]
-                engine.all_tool_dict[tool_name] = config  # For validation
-                engine.callable_functions[tool_name] = proxy_tool  # For execution
-                self._registered_tools[tool_name] = proxy_tool
+                # Keep reference for tracking
+                # Use the actual key that register_custom_tool uses
+                actual_key = config.get("name", tool_name)
+                self._registered_tools[tool_name] = engine.callable_functions[
+                    actual_key
+                ]
 
             return len(configs)
         except Exception as e:

tooluniverse/mcp_integration.py

@@ -16,8 +16,8 @@ def load_mcp_tools(self, server_urls: List[str] = None, **kwargs):
     This method automatically discovers tools from MCP servers and registers them
     as ToolUniverse tools, enabling seamless usage of remote capabilities.
 
-    Parameters:
-    ===========
+    Parameters
+    ----------
     server_urls : list of str, optional
         List of MCP server URLs to load tools from. Examples:
 
@@ -35,31 +35,34 @@ def load_mcp_tools(self, server_urls: List[str] = None, **kwargs):
         - selected_tools (list): Specific tools to load from each server
         - categories (list): Tool categories to filter by
 
-    Returns:
-    ========
+    Returns
+    -------
     dict
         Summary of loaded tools with counts and any errors encountered.
 
-    Examples:
-    =========
+    Examples
+    --------
 
     Load from specific servers:
-    ```python
-    tu = ToolUniverse()
 
-    # Load tools from multiple MCP servers
-    result = tu.load_mcp_tools([
-        "http://localhost:8001",  # Local analysis server
-        "http://ml-server:8002",  # Remote ML server
-        "ws://realtime:9000"      # WebSocket server
-    ])
+    .. code-block:: python
 
-    print(f"Loaded {result['total_tools']} tools from {result['servers_connected']} servers")
-    ```
+        tu = ToolUniverse()
+
+        # Load tools from multiple MCP servers
+        result = tu.load_mcp_tools([
+            "http://localhost:8001",  # Local analysis server
+            "http://ml-server:8002",  # Remote ML server
+            "ws://realtime:9000"      # WebSocket server
+        ])
+
+        print(f"Loaded {result['total_tools']} tools from {result['servers_connected']} servers")
 
     Load with custom configuration:
-    ```python
-    tu.load_mcp_tools(
+
+    .. code-block:: python
+
+        tu.load_mcp_tools(
         server_urls=["http://localhost:8001"],
         tool_prefix="analysis\\_",
         timeout=60,
@@ -223,8 +226,8 @@ def discover_mcp_tools(self, server_urls: List[str] = None, **kwargs) -> Dict[st
     without actually registering them in ToolUniverse. Useful for exploration
     and selective tool loading.
 
-    Parameters:
-    ===========
+    Parameters
+    ----------
     server_urls : list of str, optional
         List of MCP server URLs to discover from
     **kwargs
@@ -232,28 +235,28 @@ def discover_mcp_tools(self, server_urls: List[str] = None, **kwargs) -> Dict[st
         - timeout (int): Connection timeout (default: 30)
         - include_schemas (bool): Include tool parameter schemas (default: True)
 
-    Returns:
-    ========
+    Returns
+    -------
     dict
         Discovery results with tools organized by server
 
-    Examples:
-    =========
-    ```python
-    tu = ToolUniverse()
-
-    # Discover what's available
-    discovery = tu.discover_mcp_tools([
-        "http://localhost:8001",
-        "http://ml-server:8002"
-    ])
-
-    # Show available tools
-    for server, info in discovery["servers"].items():
-        print(f"\\n{server}:")
-        for tool in info.get("tools", []):
-            print(f"  - {tool['name']}: {tool['description']}")
-    ```
+    Examples
+    --------
+    .. code-block:: python
+
+        tu = ToolUniverse()
+
+        # Discover what's available
+        discovery = tu.discover_mcp_tools([
+            "http://localhost:8001",
+            "http://ml-server:8002"
+        ])
+
+        # Show available tools
+        for server, info in discovery["servers"].items():
+            print(f"\\n{server}:")
+            for tool in info.get("tools", []):
+                print(f"  - {tool['name']}: {tool['description']}")
     """
     if server_urls is None:
         try:
@@ -334,20 +337,20 @@ def list_mcp_connections(self) -> Dict[str, Any]:
     """
     List all active MCP connections and loaded tools.
 
-    Returns:
-    ========
+    Returns
+    -------
     dict
         Information about MCP connections, auto-loaders, and loaded tools
 
-    Examples:
-    =========
-    ```python
-    tu = ToolUniverse()
-    tu.load_mcp_tools(["http://localhost:8001"])
+    Examples
+    --------
+    .. code-block:: python
 
-    connections = tu.list_mcp_connections()
-    print(f"Active MCP connections: {len(connections['connections'])}")
-    ```
+        tu = ToolUniverse()
+        tu.load_mcp_tools(["http://localhost:8001"])
+
+        connections = tu.list_mcp_connections()
+        print(f"Active MCP connections: {len(connections['connections'])}")
     """
     mcp_tools = {}
     mcp_loaders = {}