alita-sdk 0.3.379__py3-none-any.whl → 0.3.627__py3-none-any.whl

alita_sdk/runtime/clients/mcp_discovery.py

@@ -0,0 +1,342 @@
+"""
+Dynamic MCP Server Discovery Client.
+Implements the MCP protocol for discovering tools from remote servers.
+"""
+
+import asyncio
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Any, Set
+from urllib.parse import urlparse
+import aiohttp
+from datetime import datetime, timedelta
+
+from ..models.mcp_models import McpConnectionConfig, McpToolMetadata
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class McpServerInfo:
+    """Information about an MCP server."""
+    name: str
+    url: str
+    headers: Optional[Dict[str, str]] = None
+    last_discovery: Optional[datetime] = None
+    tools: List[McpToolMetadata] = field(default_factory=list)
+    status: str = "unknown"  # unknown, online, offline, error
+    error: Optional[str] = None
+
+
+class McpDiscoveryClient:
+    """
+    Client for dynamically discovering tools from MCP servers.
+    Implements the MCP protocol for tool discovery.
+    """
+
+    def __init__(
+        self,
+        discovery_interval: int = 300,  # 5 minutes
+        request_timeout: int = 30,
+        max_retries: int = 3,
+        cache_ttl: int = 600  # 10 minutes
+    ):
+        self.discovery_interval = discovery_interval
+        self.request_timeout = request_timeout
+        self.max_retries = max_retries
+        self.cache_ttl = cache_ttl
+
+        # Server registry
+        self.servers: Dict[str, McpServerInfo] = {}
+        self.session: Optional[aiohttp.ClientSession] = None
+
+        # Discovery state
+        self._discovery_task: Optional[asyncio.Task] = None
+        self._running = False
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.stop()
+
+    async def start(self):
+        """Start the discovery client."""
+        if self._running:
+            return
+
+        self.session = aiohttp.ClientSession(
+            timeout=aiohttp.ClientTimeout(total=self.request_timeout)
+        )
+        self._running = True
+
+        # Start background discovery task
+        self._discovery_task = asyncio.create_task(self._discovery_loop())
+        logger.info("MCP Discovery Client started")
+
+    async def stop(self):
+        """Stop the discovery client."""
+        if not self._running:
+            return
+
+        self._running = False
+
+        if self._discovery_task:
+            self._discovery_task.cancel()
+            try:
+                await self._discovery_task
+            except asyncio.CancelledError:
+                pass
+
+        if self.session:
+            await self.session.close()
+
+        logger.info("MCP Discovery Client stopped")
+
+    def add_server(self, server_name: str, connection_config: McpConnectionConfig):
+        """Add an MCP server to discovery."""
+        server_info = McpServerInfo(
+            name=server_name,
+            url=connection_config.url,
+            headers=connection_config.headers
+        )
+        self.servers[server_name] = server_info
+        logger.info(f"Added MCP server for discovery: {server_name} at {connection_config.url}")
+
+    def remove_server(self, server_name: str):
+        """Remove an MCP server from discovery."""
+        if server_name in self.servers:
+            del self.servers[server_name]
+            logger.info(f"Removed MCP server from discovery: {server_name}")
+
+    async def discover_server_tools(self, server_name: str, force: bool = False) -> List[McpToolMetadata]:
+        """Discover tools from a specific MCP server."""
+        if server_name not in self.servers:
+            raise ValueError(f"Server {server_name} not registered for discovery")
+
+        server_info = self.servers[server_name]
+
+        # Check cache unless force refresh
+        if not force and self._is_cache_valid(server_info):
+            logger.debug(f"Using cached tools for server {server_name}")
+            return server_info.tools
+
+        try:
+            tools = await self._fetch_server_tools(server_info)
+            server_info.tools = tools
+            server_info.last_discovery = datetime.now()
+            server_info.status = "online"
+            server_info.error = None
+
+            logger.info(f"Discovered {len(tools)} tools from server {server_name}")
+            return tools
+
+        except Exception as e:
+            error_msg = f"Failed to discover tools from {server_name}: {e}"
+            logger.error(error_msg)
+            server_info.status = "error"
+            server_info.error = str(e)
+            return []
+
+    async def get_all_tools(self) -> Dict[str, List[McpToolMetadata]]:
+        """Get all discovered tools from all servers."""
+        all_tools = {}
+
+        for server_name in self.servers:
+            tools = await self.discover_server_tools(server_name)
+            all_tools[server_name] = tools
+
+        return all_tools
+
+    def get_server_status(self, server_name: str) -> Optional[McpServerInfo]:
+        """Get status information for a server."""
+        return self.servers.get(server_name)
+
+    def get_all_server_status(self) -> Dict[str, McpServerInfo]:
+        """Get status information for all servers."""
+        return self.servers.copy()
+
+    async def _discovery_loop(self):
+        """Background task for periodic tool discovery."""
+        while self._running:
+            try:
+                await self._perform_discovery()
+                await asyncio.sleep(self.discovery_interval)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error in discovery loop: {e}")
+                await asyncio.sleep(60)  # Wait before retrying
+
+    async def _perform_discovery(self):
+        """Perform discovery on all registered servers."""
+        if not self.servers:
+            return
+
+        discovery_tasks = [
+            self.discover_server_tools(server_name)
+            for server_name in self.servers
+        ]
+
+        results = await asyncio.gather(*discovery_tasks, return_exceptions=True)
+
+        # Log any errors
+        for i, result in enumerate(results):
+            if isinstance(result, Exception):
+                server_name = list(self.servers.keys())[i]
+                logger.error(f"Discovery failed for server {server_name}: {result}")
+
+    async def _fetch_server_tools(self, server_info: McpServerInfo) -> List[McpToolMetadata]:
+        """Fetch tools from an MCP server using HTTP requests."""
+        if not self.session:
+            raise RuntimeError("Discovery client not started")
+
+        # MCP protocol: list_tools request
+        mcp_request = {
+            "jsonrpc": "2.0",
+            "id": f"discover_{int(time.time())}",
+            "method": "tools/list",
+            "params": {}
+        }
+
+        headers = {"Content-Type": "application/json"}
+        if server_info.headers:
+            headers.update(server_info.headers)
+
+        async with self.session.post(
+            server_info.url,
+            json=mcp_request,
+            headers=headers
+        ) as response:
+
+            if response.status != 200:
+                raise Exception(f"HTTP {response.status}: {await response.text()}")
+
+            data = await response.json()
+
+            if "error" in data:
+                raise Exception(f"MCP Error: {data['error']}")
+
+            # Parse MCP response
+            tools_data = data.get("result", {}).get("tools", [])
+            tools = []
+
+            for tool_data in tools_data:
+                try:
+                    tool_metadata = McpToolMetadata(
+                        name=tool_data.get("name", ""),
+                        description=tool_data.get("description", ""),
+                        server=server_info.name,
+                        input_schema=tool_data.get("inputSchema", {}),
+                        enabled=True
+                    )
+                    tools.append(tool_metadata)
+                except Exception as e:
+                    logger.warning(f"Failed to parse tool from {server_info.name}: {e}")
+
+            return tools
+
+    def _is_cache_valid(self, server_info: McpServerInfo) -> bool:
+        """Check if cached tools are still valid."""
+        if not server_info.last_discovery:
+            return False
+
+        cache_age = datetime.now() - server_info.last_discovery
+        return cache_age.total_seconds() < self.cache_ttl
+
+
+class McpDiscoveryService:
+    """
+    High-level service for managing MCP server discovery.
+    Integrates with the existing toolkit system.
+    """
+
+    def __init__(self, discovery_client: Optional[McpDiscoveryClient] = None):
+        self.client = discovery_client or McpDiscoveryClient()
+        self._started = False
+
+    async def start(self):
+        """Start the discovery service."""
+        if not self._started:
+            await self.client.start()
+            self._started = True
+
+    async def stop(self):
+        """Stop the discovery service."""
+        if self._started:
+            await self.client.stop()
+            self._started = False
+
+    async def register_server(self, server_name: str, connection_config: McpConnectionConfig):
+        """Register an MCP server for discovery."""
+        self.client.add_server(server_name, connection_config)
+
+        # Perform immediate discovery
+        await self.client.discover_server_tools(server_name, force=True)
+
+    def unregister_server(self, server_name: str):
+        """Unregister an MCP server."""
+        self.client.remove_server(server_name)
+
+    async def get_server_tools(self, server_name: str) -> List[McpToolMetadata]:
+        """Get tools from a specific server."""
+        return await self.client.discover_server_tools(server_name)
+
+    async def get_all_available_tools(self) -> Dict[str, List[McpToolMetadata]]:
+        """Get all available tools from all registered servers."""
+        return await self.client.get_all_tools()
+
+    def get_server_health(self) -> Dict[str, Dict[str, Any]]:
+        """Get health status of all servers."""
+        status_info = {}
+
+        for name, server_info in self.client.get_all_server_status().items():
+            status_info[name] = {
+                "status": server_info.status,
+                "url": server_info.url,
+                "last_discovery": server_info.last_discovery.isoformat() if server_info.last_discovery else None,
+                "tool_count": len(server_info.tools),
+                "error": server_info.error
+            }
+
+        return status_info
+
+    async def refresh_server(self, server_name: str):
+        """Force refresh tools from a specific server."""
+        await self.client.discover_server_tools(server_name, force=True)
+
+    async def refresh_all_servers(self):
+        """Force refresh tools from all servers."""
+        for server_name in self.client.servers:
+            await self.client.discover_server_tools(server_name, force=True)
+
+
+# Global discovery service instance
+_discovery_service: Optional[McpDiscoveryService] = None
+
+
+def get_discovery_service() -> McpDiscoveryService:
+    """Get the global MCP discovery service instance."""
+    global _discovery_service
+    if _discovery_service is None:
+        _discovery_service = McpDiscoveryService()
+    return _discovery_service
+
+
+async def init_discovery_service():
+    """Initialize the global discovery service."""
+    service = get_discovery_service()
+    await service.start()
+
+
+async def shutdown_discovery_service():
+    """Shutdown the global discovery service."""
+    global _discovery_service
+    if _discovery_service:
+        await _discovery_service.stop()
+        _discovery_service = None

alita_sdk/runtime/clients/mcp_manager.py

@@ -0,0 +1,262 @@
+"""
+MCP Manager - Unified interface for both static and dynamic MCP tool discovery.
+Provides a single API that can work with both registry-based and live discovery.
+"""
+
+import asyncio
+import logging
+from typing import Dict, List, Optional, Any, Union
+from enum import Enum
+
+from ..models.mcp_models import McpConnectionConfig, McpToolMetadata
+from .mcp_discovery import McpDiscoveryService, get_discovery_service
+
+logger = logging.getLogger(__name__)
+
+
+class DiscoveryMode(Enum):
+    """MCP discovery modes."""
+    STATIC = "static"      # Use alita.get_mcp_toolkits() registry
+    DYNAMIC = "dynamic"    # Live discovery from MCP servers
+    HYBRID = "hybrid"      # Try dynamic first, fallback to static
+
+
+class McpManager:
+    """
+    Unified manager for MCP tool discovery supporting multiple modes.
+    """
+
+    def __init__(
+        self,
+        default_mode: DiscoveryMode = DiscoveryMode.DYNAMIC,
+        discovery_service: Optional[McpDiscoveryService] = None
+    ):
+        self.default_mode = default_mode
+        self.discovery_service = discovery_service or get_discovery_service()
+        self._static_fallback_enabled = True
+
+    async def discover_server_tools(
+        self,
+        server_name: str,
+        connection_config: Optional[McpConnectionConfig] = None,
+        alita_client=None,
+        mode: Optional[DiscoveryMode] = None,
+        **kwargs
+    ) -> List[McpToolMetadata]:
+        """
+        Discover tools from an MCP server using the specified mode.
+
+        Args:
+            server_name: Name of the MCP server
+            connection_config: Connection configuration (required for dynamic mode)
+            alita_client: Alita client (required for static mode)
+            mode: Discovery mode to use (defaults to manager's default)
+            **kwargs: Additional options
+
+        Returns:
+            List of discovered tool metadata
+        """
+        discovery_mode = mode or self.default_mode
+
+        if discovery_mode == DiscoveryMode.DYNAMIC:
+            return await self._discover_dynamic(server_name, connection_config)
+
+        elif discovery_mode == DiscoveryMode.STATIC:
+            return await self._discover_static(server_name, alita_client)
+
+        elif discovery_mode == DiscoveryMode.HYBRID:
+            return await self._discover_hybrid(server_name, connection_config, alita_client)
+
+        else:
+            raise ValueError(f"Unknown discovery mode: {discovery_mode}")
+
+    async def _discover_dynamic(
+        self,
+        server_name: str,
+        connection_config: Optional[McpConnectionConfig]
+    ) -> List[McpToolMetadata]:
+        """Discover tools using dynamic MCP protocol."""
+        if not connection_config:
+            raise ValueError("Connection configuration required for dynamic discovery")
+
+        try:
+            # Ensure discovery service is started
+            await self.discovery_service.start()
+
+            # Register and discover
+            await self.discovery_service.register_server(server_name, connection_config)
+            tools = await self.discovery_service.get_server_tools(server_name)
+
+            logger.info(f"Dynamic discovery found {len(tools)} tools from {server_name}")
+            return tools
+
+        except Exception as e:
+            logger.error(f"Dynamic discovery failed for {server_name}: {e}")
+            raise
+
+    async def _discover_static(
+        self,
+        server_name: str,
+        alita_client
+    ) -> List[McpToolMetadata]:
+        """Discover tools using static registry."""
+        if not alita_client or not hasattr(alita_client, 'get_mcp_toolkits'):
+            raise ValueError("Alita client with get_mcp_toolkits() required for static discovery")
+
+        try:
+            # Use existing registry approach
+            all_toolkits = alita_client.get_mcp_toolkits()
+            server_toolkit = next((tk for tk in all_toolkits if tk.get('name') == server_name), None)
+
+            if not server_toolkit:
+                logger.warning(f"Static registry: Server {server_name} not found")
+                return []
+
+            # Convert to metadata format
+            tools = []
+            for tool_info in server_toolkit.get('tools', []):
+                metadata = McpToolMetadata(
+                    name=tool_info.get('name', ''),
+                    description=tool_info.get('description', ''),
+                    server=server_name,
+                    input_schema=tool_info.get('inputSchema', {}),
+                    enabled=True
+                )
+                tools.append(metadata)
+
+            logger.info(f"Static discovery found {len(tools)} tools from {server_name}")
+            return tools
+
+        except Exception as e:
+            logger.error(f"Static discovery failed for {server_name}: {e}")
+            raise
+
+    async def _discover_hybrid(
+        self,
+        server_name: str,
+        connection_config: Optional[McpConnectionConfig],
+        alita_client
+    ) -> List[McpToolMetadata]:
+        """Discover tools using hybrid approach (dynamic first, static fallback)."""
+
+        # Try dynamic discovery first
+        if connection_config:
+            try:
+                return await self._discover_dynamic(server_name, connection_config)
+            except Exception as e:
+                logger.warning(f"Dynamic discovery failed for {server_name}, trying static: {e}")
+
+        # Fallback to static discovery
+        if self._static_fallback_enabled and alita_client:
+            try:
+                return await self._discover_static(server_name, alita_client)
+            except Exception as e:
+                logger.error(f"Static fallback also failed for {server_name}: {e}")
+
+        logger.error(f"All discovery methods failed for {server_name}")
+        return []
+
+    async def get_server_health(
+        self,
+        server_name: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Get health information for servers."""
+        try:
+            if server_name:
+                # Get specific server health from discovery service
+                all_health = self.discovery_service.get_server_health()
+                return all_health.get(server_name, {"status": "unknown"})
+            else:
+                # Get all server health
+                return self.discovery_service.get_server_health()
+        except Exception as e:
+            logger.error(f"Failed to get server health: {e}")
+            return {"status": "error", "error": str(e)}
+
+    async def refresh_server(self, server_name: str):
+        """Force refresh a specific server's tools."""
+        try:
+            await self.discovery_service.refresh_server(server_name)
+        except Exception as e:
+            logger.error(f"Failed to refresh server {server_name}: {e}")
+
+    async def start(self):
+        """Start the MCP manager."""
+        await self.discovery_service.start()
+
+    async def stop(self):
+        """Stop the MCP manager."""
+        await self.discovery_service.stop()
+
+    def set_static_fallback(self, enabled: bool):
+        """Enable or disable static fallback in hybrid mode."""
+        self._static_fallback_enabled = enabled
+
+
+# Global manager instance
+_mcp_manager: Optional[McpManager] = None
+
+
+def get_mcp_manager(mode: DiscoveryMode = DiscoveryMode.HYBRID) -> McpManager:
+    """Get the global MCP manager instance."""
+    global _mcp_manager
+    if _mcp_manager is None:
+        _mcp_manager = McpManager(default_mode=mode)
+    return _mcp_manager
+
+
+async def discover_mcp_tools(
+    server_name: str,
+    connection_config: Optional[McpConnectionConfig] = None,
+    alita_client=None,
+    mode: Optional[DiscoveryMode] = None
+) -> List[McpToolMetadata]:
+    """
+    Convenience function for discovering MCP tools.
+
+    Args:
+        server_name: Name of the MCP server
+        connection_config: Connection config (for dynamic discovery)
+        alita_client: Alita client (for static discovery)
+        mode: Discovery mode (defaults to HYBRID)
+
+    Returns:
+        List of discovered tool metadata
+    """
+    manager = get_mcp_manager()
+    return await manager.discover_server_tools(
+        server_name=server_name,
+        connection_config=connection_config,
+        alita_client=alita_client,
+        mode=mode or DiscoveryMode.HYBRID
+    )
+
+
+async def init_mcp_manager(mode: DiscoveryMode = DiscoveryMode.HYBRID):
+    """Initialize the global MCP manager."""
+    manager = get_mcp_manager(mode)
+    await manager.start()
+
+
+async def shutdown_mcp_manager():
+    """Shutdown the global MCP manager."""
+    global _mcp_manager
+    if _mcp_manager:
+        await _mcp_manager.stop()
+        _mcp_manager = None
+
+
+# Configuration helpers
+def create_discovery_config(
+    mode: str = "hybrid",
+    discovery_interval: int = 300,
+    enable_static_fallback: bool = True,
+    **kwargs
+) -> Dict[str, Any]:
+    """Create a discovery configuration dictionary."""
+    return {
+        "discovery_mode": mode,
+        "discovery_interval": discovery_interval,
+        "enable_static_fallback": enable_static_fallback,
+        **kwargs
+    }

alita_sdk/runtime/clients/sandbox_client.py

@@ -48,27 +48,6 @@ class SandboxArtifact:
             return f'{data['error']}. {data['content'] if data['content'] else ''}'
         detected = chardet.detect(data)
         return data
-        # TODO: add proper handling for binary files (images, pdf, etc.) for sandbox
-        # if detected['encoding'] is not None:
-        #     try:
-        #         return data.decode(detected['encoding'])
-        #     except Exception:
-        #         logger.error('Error while default encoding')
-        #         return parse_file_content(file_name=artifact_name,
-        #                                   file_content=data,
-        #                                   is_capture_image=is_capture_image,
-        #                                   page_number=page_number,
-        #                                   sheet_name=sheet_name,
-        #                                   excel_by_sheets=excel_by_sheets,
-        #                                   llm=llm)
-        # else:
-        #     return parse_file_content(file_name=artifact_name,
-        #                               file_content=data,
-        #                               is_capture_image=is_capture_image,
-        #                               page_number=page_number,
-        #                               sheet_name=sheet_name,
-        #                               excel_by_sheets=excel_by_sheets,
-        #                               llm=llm)
 
     def delete(self, artifact_name: str, bucket_name=None):
         if not bucket_name:
@@ -143,6 +122,7 @@ class SandboxClient:
         self.configurations_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=configurations&unsecret=true'
         self.ai_section_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=ai'
         self.image_generation_url = f'{self.base_url}{self.llm_path}/images/generations'
+        self.auth_user_url = f'{self.base_url}{self.api_path}/auth/user'
         self.configurations: list = configurations or []
         self.model_timeout = kwargs.get('model_timeout', 120)
         self.model_image_generation = kwargs.get('model_image_generation')
@@ -363,3 +343,10 @@ class SandboxClient:
         url = f'{self.artifact_url}/{bucket_name}'
         data = requests.delete(url, headers=self.headers, verify=False, params={'filename': quote(artifact_name)})
         return self._process_requst(data)
+
+    def get_user_data(self) -> Dict[str, Any]:
+        resp = requests.get(self.auth_user_url, headers=self.headers, verify=False)
+        if resp.ok:
+            return resp.json()
+        logger.error(f'Failed to fetch user data: {resp.status_code} - {resp.text}')
+        raise ApiDetailsRequestError(f'Failed to fetch user data with status code {resp.status_code}.')