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

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

@@ -143,6 +143,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 +364,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}.')

alita_sdk/runtime/langchain/assistant.py

@@ -17,6 +17,7 @@ from .constants import REACT_ADDON, REACT_VARS, XML_ADDON
 from .chat_message_template import Jinja2TemplatedChatMessagesTemplate
 from ..tools.echo import EchoTool
 from langchain_core.tools import BaseTool, ToolException
+from jinja2 import Environment, DebugUndefined
 
 logger = logging.getLogger(__name__)
 
@@ -29,7 +30,9 @@ class Assistant:
                  app_type: str = "openai",
                  tools: Optional[list] = [],
                  memory: Optional[Any] = None,
-                 store: Optional[BaseStore] = None):
+                 store: Optional[BaseStore] = None,
+                 debug_mode: Optional[bool] = False,
+                 mcp_tokens: Optional[dict] = None):
 
         self.app_type = app_type
         self.memory = memory
@@ -87,7 +90,14 @@ class Assistant:
             for internal_tool_name in meta.get("internal_tools"):
                 version_tools.append({"type": "internal_tool", "name": internal_tool_name})
 
-        self.tools = get_tools(version_tools, alita_client=alita, llm=self.client, memory_store=self.store)
+        self.tools = get_tools(
+            version_tools,
+            alita_client=alita,
+            llm=self.client,
+            memory_store=self.store,
+            debug_mode=debug_mode,
+            mcp_tokens=mcp_tokens
+        )
         if tools:
             self.tools += tools
         # Handle prompt setup
@@ -184,17 +194,6 @@ class Assistant:
         """
         # Exclude compiled graph runnables from simple tool agents
         simple_tools = [t for t in self.tools if isinstance(t, (BaseTool, CompiledStateGraph))]
-        # Add image generation tool if model is configured
-        if self.alita_client.model_image_generation is not None:
-            try:
-                from ..tools.image_generation import (
-                    create_image_generation_tool
-                )
-                image_tool = create_image_generation_tool(self.alita_client)
-                simple_tools.append(image_tool)
-                logger.info("Added ImageGenerationTool to react agent")
-            except Exception as e:
-                logger.error(f"Error adding ImageGenerationTool: {e}")
         
         # Set up memory/checkpointer if available
         checkpointer = None
@@ -274,6 +273,9 @@ class Assistant:
         schema_dict = {
             'name': 'react_agent',
             'state': {
+                'input': {
+                    'type': 'str'
+                },
                 'messages': state_messages_config
             },
             'nodes': [{
@@ -282,6 +284,21 @@ class Assistant:
                 'prompt': {
                     'template': escaped_prompt
                 },
+                'input_mapping': {
+                    'system': {
+                        'type': 'fixed',
+                        'value': escaped_prompt
+                    },
+                    'task': {
+                        'type': 'variable',
+                        'value': 'input'
+                    },
+                    'chat_history': {
+                        'type': 'variable',
+                        'value': 'messages'
+                    }
+                },
+                'step_limit': self.max_iterations,
                 'input': ['messages'],
                 'output': ['messages'],
                 'transition': 'END'
@@ -307,7 +324,8 @@ class Assistant:
             store=self.store,
             debug=False,
             for_subgraph=False,
-            alita_client=self.alita_client
+            alita_client=self.alita_client,
+            steps_limit=self.max_iterations
         )
         
         return agent
@@ -322,7 +340,8 @@ class Assistant:
         agent = create_graph(
             client=self.client, tools=self.tools,
             yaml_schema=self.prompt, memory=memory,
-            alita_client=self.alita_client
+            alita_client=self.alita_client,
+            steps_limit=self.max_iterations
         )
         #
         return agent
@@ -342,5 +361,7 @@ class Assistant:
                     continue
                 # take only the content of the system message from the openai prompt
                 if isinstance(message, SystemMessage):
-                    return message.content
+                    environment = Environment(undefined=DebugUndefined)
+                    template = environment.from_string(message.content)
+                    return template.render(self.prompt.partial_variables)
         return None

alita_sdk/runtime/langchain/constants.py

@@ -27,7 +27,7 @@ Use this if you want to respond directly to the human. Markdown code snippet for
 ```json
 {
     "action": "Final Answer",
-    "action_input": string \ You should put what you want to return to use here
+    "action_input": string // You should put what you want to return to use here
 }
 ```
 
@@ -80,3 +80,8 @@ DEFAULT_MULTIMODAL_PROMPT = """
 - Maintain a structured and logical flow in the output to enhance understanding and usability.
 - Avoid presenting the entire prompt for user.
 """
+
+ELITEA_RS = "elitea_response"
+PRINTER = "printer"
+PRINTER_NODE_RS = "printer_output"
+PRINTER_COMPLETED_STATE = "PRINTER_COMPLETED"