alita-sdk 0.3.457__py3-none-any.whl → 0.3.486__py3-none-any.whl

alita_sdk/runtime/toolkits/mcp.py

@@ -6,12 +6,11 @@ Following MCP specification: https://modelcontextprotocol.io/specification/2025-
 
 import logging
 import re
-import requests
 import asyncio
 from typing import List, Optional, Any, Dict, Literal, ClassVar, Union
 
 from langchain_core.tools import BaseToolkit, BaseTool
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, SecretStr
 
 from ..tools.mcp_server_tool import McpServerTool
 from ..tools.mcp_remote_tool import McpRemoteTool
@@ -208,28 +207,32 @@ class McpToolkit(BaseToolkit):
                     }
                 )
             ),
-            timeout=(
-                Union[int, str], # TODO: remove one I will figure out why UI sends str
+            client_id=(
+                Optional[str],
                 Field(
-                    default=300,
-                    description="Request timeout in seconds (1-3600)"
+                    default=None,
+                    description="OAuth Client ID (if applicable)"
                 )
             ),
-            discovery_mode=(
-                Literal['static', 'dynamic', 'hybrid'],
+            client_secret=(
+                Optional[SecretStr],
                 Field(
-                    default="dynamic",
-                    description="Discovery mode",
-                    json_schema_extra={
-                        'tooltip': 'static: use registry, dynamic: live discovery, hybrid: try dynamic first'
-                    }
+                    default=None,
+                    description="OAuth Client Secret (if applicable)"
                 )
             ),
-            discovery_interval=(
-                Union[int, str],
+            scopes=(
+                Optional[List[str]],
+                Field(
+                    default=None,
+                    description="OAuth Scopes (if applicable)"
+                )
+            ),
+            timeout=(
+                Union[int, str], # TODO: remove one I will figure out why UI sends str
                 Field(
                     default=300,
-                    description="Discovery interval in seconds (60-3600, for periodic discovery)"
+                    description="Request timeout in seconds (1-3600)"
                 )
             ),
             selected_tools=(
@@ -259,7 +262,7 @@ class McpToolkit(BaseToolkit):
             __config__=ConfigDict(
                 json_schema_extra={
                     'metadata': {
-                        "label": "Remove MCP",
+                        "label": "Remote MCP",
                         "icon_url": None,
                         "categories": ["other"],
                         "extra_categories": ["remote tools", "sse", "http"],
@@ -275,8 +278,6 @@ class McpToolkit(BaseToolkit):
         url: str,
         headers: Optional[Dict[str, str]] = None,
         timeout: int = 60,
-        discovery_mode: str = "hybrid",
-        discovery_interval: int = 300,
         selected_tools: List[str] = None,
         enable_caching: bool = True,
         cache_ttl: int = 300,
@@ -297,8 +298,6 @@ class McpToolkit(BaseToolkit):
             url: MCP server HTTP URL
             headers: HTTP headers for authentication
             timeout: Request timeout in seconds
-            discovery_mode: Discovery mode ('static', 'dynamic', 'hybrid')
-            discovery_interval: Discovery interval in seconds (for periodic discovery)
             selected_tools: List of specific tools to enable (empty = all tools)
             enable_caching: Whether to enable caching
             cache_ttl: Cache TTL in seconds
@@ -317,7 +316,6 @@ class McpToolkit(BaseToolkit):
 
         # Convert numeric parameters that may come as strings from UI
         timeout = safe_int(timeout, 60)
-        discovery_interval = safe_int(discovery_interval, 300)
         cache_ttl = safe_int(cache_ttl, 300)
 
         logger.info(f"Creating MCP toolkit: {toolkit_name}")
@@ -363,8 +361,7 @@ class McpToolkit(BaseToolkit):
             connection_config=connection_config,
             timeout=timeout,
             selected_tools=selected_tools,
-            client=client,
-            discovery_mode=discovery_mode
+            client=client
         )
 
         return toolkit
@@ -376,8 +373,7 @@ class McpToolkit(BaseToolkit):
         connection_config: McpConnectionConfig,
         timeout: int,
         selected_tools: List[str],
-        client,
-        discovery_mode: str = "dynamic"
+        client
     ) -> List[BaseTool]:
         """
         Create tools from a single MCP server. Always performs live discovery when connection config is provided.
@@ -423,19 +419,28 @@ class McpToolkit(BaseToolkit):
 
             logger.info(f"Successfully created {len(tools)} MCP tools from toolkit '{toolkit_name}' via direct discovery")
 
+        except McpAuthorizationRequired:
+            # Authorization is required; surface upstream so the caller can prompt the user
+            logger.info(f"MCP toolkit '{toolkit_name}' requires OAuth authorization")
+            raise
         except Exception as e:
             logger.error(f"Direct discovery failed for MCP toolkit '{toolkit_name}': {e}", exc_info=True)
             logger.error(f"Discovery error details - URL: {connection_config.url}, Timeout: {timeout}s")
 
-            # Fallback to static mode if available and not already static
-            if isinstance(e, McpAuthorizationRequired):
-                # Authorization is required; surface upstream so the caller can prompt the user
+            # Check if the exception wraps McpAuthorizationRequired (can happen with asyncio)
+            if hasattr(e, '__cause__') and isinstance(e.__cause__, McpAuthorizationRequired):
+                logger.info(f"Found wrapped McpAuthorizationRequired, re-raising")
+                raise e.__cause__
+            
+            # For new MCP toolkits (no client), don't silently return empty - surface the error
+            # This helps users understand why tool discovery failed
+            if not client:
+                logger.warning(f"No fallback available for toolkit '{toolkit_name}' - re-raising discovery error")
                 raise
-            if client and discovery_mode != "static":
-                logger.info(f"Falling back to static discovery for toolkit '{toolkit_name}'")
-                tools = cls._create_tools_static(toolkit_name, selected_tools, timeout, client)
-            else:
-                logger.warning(f"No fallback available for toolkit '{toolkit_name}' - returning empty tools list")
+            
+            # Only fall back to static discovery for existing toolkits with a client
+            logger.info(f"Falling back to static discovery for toolkit '{toolkit_name}'")
+            tools = cls._create_tools_static(toolkit_name, selected_tools, timeout, client)
 
         # Don't add inspection tool to agent - it's only for internal use by toolkit
         # inspection_tool = cls._create_inspection_tool(
@@ -481,8 +486,15 @@ class McpToolkit(BaseToolkit):
                 )
             )
             return all_tools, session_id
+        except McpAuthorizationRequired:
+            # Re-raise auth required exceptions directly
+            logger.info(f"[MCP SSE] Authorization required for '{toolkit_name}'")
+            raise
         except Exception as e:
             logger.error(f"[MCP SSE] Discovery failed for '{toolkit_name}': {e}")
+            # Check if the exception wraps McpAuthorizationRequired
+            if hasattr(e, '__cause__') and isinstance(e.__cause__, McpAuthorizationRequired):
+                raise e.__cause__
             raise
     
     @classmethod
@@ -858,8 +870,6 @@ def get_tools(tool_config: dict, alita_client, llm=None, memory_store=None) -> L
         url=url,
         headers=headers,
         timeout=safe_int(settings.get('timeout'), 60),
-        discovery_mode=settings.get('discovery_mode', 'dynamic'),
-        discovery_interval=safe_int(settings.get('discovery_interval'), 300),
         selected_tools=settings.get('selected_tools', []),
         enable_caching=settings.get('enable_caching', True),
         cache_ttl=safe_int(settings.get('cache_ttl'), 300),

alita_sdk/runtime/toolkits/planning.py

@@ -0,0 +1,171 @@
+"""
+PlanningToolkit - Runtime toolkit for agent plan management.
+
+Provides tools for creating, tracking, and completing multi-step execution plans.
+Supports two storage backends:
+1. PostgreSQL - when pgvector_configuration with connection_string is provided
+2. Filesystem - when no connection string (local CLI usage)
+"""
+
+from typing import ClassVar, List, Any, Literal, Optional, Callable
+
+from langchain_community.agent_toolkits.base import BaseToolkit
+from langchain_core.tools import BaseTool
+from pydantic import create_model, BaseModel, ConfigDict, Field
+from pydantic.fields import FieldInfo
+
+from ..tools.planning import PlanningWrapper
+from ...tools.base.tool import BaseAction
+from ...tools.utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+
+
+class PlanningToolkit(BaseToolkit):
+    """
+    Toolkit for agent plan management.
+    
+    Provides tools for creating, updating, and tracking execution plans.
+    Supports PostgreSQL (production) and filesystem (local) storage backends.
+    Plans are scoped by conversation_id.
+    """
+    tools: List[BaseTool] = []
+    _toolkit_max_length: ClassVar[int] = 50  # Use ClassVar to avoid Pydantic treating it as field
+    
+    @staticmethod
+    def toolkit_config_schema() -> BaseModel:
+        """
+        Returns the configuration schema for the Planning toolkit.
+        
+        Used by the UI to generate the toolkit configuration form.
+        """
+        # Define available tools
+        selected_tools = {
+            'update_plan': {
+                'title': 'UpdatePlanInput',
+                'type': 'object',
+                'properties': {
+                    'title': {'type': 'string', 'description': "Title for the plan"},
+                    'steps': {'type': 'array', 'items': {'type': 'string'}, 'description': "List of step descriptions"},
+                    'conversation_id': {'type': 'string', 'description': "Conversation ID (auto-injected)"}
+                },
+                'required': ['title', 'steps', 'conversation_id']
+            },
+            'complete_step': {
+                'title': 'CompleteStepInput', 
+                'type': 'object',
+                'properties': {
+                    'step_number': {'type': 'integer', 'description': "Step number to complete (1-indexed)"},
+                    'conversation_id': {'type': 'string', 'description': "Conversation ID (auto-injected)"}
+                },
+                'required': ['step_number', 'conversation_id']
+            },
+            'get_plan_status': {
+                'title': 'GetPlanStatusInput',
+                'type': 'object', 
+                'properties': {
+                    'conversation_id': {'type': 'string', 'description': "Conversation ID (auto-injected)"}
+                },
+                'required': ['conversation_id']
+            },
+            'delete_plan': {
+                'title': 'DeletePlanInput',
+                'type': 'object',
+                'properties': {
+                    'conversation_id': {'type': 'string', 'description': "Conversation ID (auto-injected)"}
+                },
+                'required': ['conversation_id']
+            }
+        }
+        
+        PlanningToolkit._toolkit_max_length = get_max_toolkit_length(selected_tools)
+        
+        return create_model(
+            "planning",
+            # Tool selection
+            selected_tools=(
+                List[Literal[tuple(selected_tools)]], 
+                Field(
+                    default=list(selected_tools.keys()),
+                    json_schema_extra={'args_schemas': selected_tools}
+                )
+            ),
+            __config__=ConfigDict(
+                json_schema_extra={
+                    'metadata': {
+                        "label": "Planning",
+                        "description": "Tools for managing multi-step execution plans with progress tracking. Uses PostgreSQL when configured, filesystem otherwise.",
+                        "icon_url": None,
+                        "max_length": PlanningToolkit._toolkit_max_length
+                    }
+                }
+            )
+        )
+    
+    @classmethod
+    def get_toolkit(
+        cls,
+        toolkit_name: Optional[str] = None,
+        selected_tools: Optional[List[str]] = None,
+        pgvector_configuration: Optional[dict] = None,
+        storage_dir: Optional[str] = None,
+        plan_callback: Optional[Any] = None,
+        conversation_id: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Create a PlanningToolkit instance with configured tools.
+        
+        Args:
+            toolkit_name: Optional name prefix for tools
+            selected_tools: List of tool names to include (default: all)
+            pgvector_configuration: PostgreSQL configuration dict with connection_string.
+                                   If not provided, uses filesystem storage.
+            storage_dir: Directory for filesystem storage (when no pgvector_configuration)
+            plan_callback: Optional callback function called when plan changes (for CLI UI)
+            conversation_id: Conversation ID for scoping plans.
+                            For server: from elitea_core payload. For CLI: session_id.
+            **kwargs: Additional configuration options
+            
+        Returns:
+            PlanningToolkit instance with configured tools
+        """
+        if selected_tools is None:
+            selected_tools = ['update_plan', 'complete_step', 'get_plan_status', 'delete_plan']
+        
+        tools = []
+        
+        # Extract connection string from pgvector configuration (if provided)
+        connection_string = None
+        if pgvector_configuration:
+            connection_string = pgvector_configuration.get('connection_string', '')
+            if hasattr(connection_string, 'get_secret_value'):
+                connection_string = connection_string.get_secret_value()
+        
+        # Create wrapper - it will auto-select storage backend
+        wrapper = PlanningWrapper(
+            connection_string=connection_string if connection_string else None,
+            conversation_id=conversation_id,
+            storage_dir=storage_dir,
+            plan_callback=plan_callback,
+        )
+        
+        # Build tool name prefix
+        prefix = clean_string(toolkit_name, cls._toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
+        
+        # Create tools from wrapper
+        available_tools = wrapper.get_available_tools()
+        for tool in available_tools:
+            if tool["name"] not in selected_tools:
+                continue
+            
+            tools.append(BaseAction(
+                api_wrapper=wrapper,
+                name=prefix + tool["name"],
+                description=tool["description"],
+                args_schema=tool["args_schema"]
+            ))
+        
+        return cls(tools=tools)
+    
+    def get_tools(self) -> List[BaseTool]:
+        """Return the list of configured tools."""
+        return self.tools

alita_sdk/runtime/toolkits/tools.py

@@ -9,6 +9,7 @@ from alita_sdk.tools import get_tools as alita_tools
 from .application import ApplicationToolkit
 from .artifact import ArtifactToolkit
 from .datasource import DatasourcesToolkit
+from .planning import PlanningToolkit
 from .prompt import PromptToolkit
 from .subgraph import SubgraphToolkit
 from .vectorstore import VectorStoreToolkit
@@ -21,6 +22,7 @@ from ...community import get_toolkits as community_toolkits, get_tools as commun
 from ...tools.memory import MemoryToolkit
 from ..utils.mcp_oauth import canonical_resource, McpAuthorizationRequired
 from ...tools.utils import TOOLKIT_SPLITTER
+from alita_sdk.tools import _inject_toolkit_id
 
 logger = logging.getLogger(__name__)
 
@@ -29,6 +31,7 @@ def get_toolkits():
     core_toolkits = [
         ArtifactToolkit.toolkit_config_schema(),
         MemoryToolkit.toolkit_config_schema(),
+        PlanningToolkit.toolkit_config_schema(),
         VectorStoreToolkit.toolkit_config_schema(),
         SandboxToolkit.toolkit_config_schema(),
         ImageGenerationToolkit.toolkit_config_schema(),
@@ -38,7 +41,7 @@ def get_toolkits():
     return core_toolkits + community_toolkits() + alita_toolkits()
 
 
-def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None, debug_mode: Optional[bool] = False, mcp_tokens: Optional[dict] = None) -> list:
+def get_tools(tools_list: list, alita_client=None, llm=None, memory_store: BaseStore = None, debug_mode: Optional[bool] = False, mcp_tokens: Optional[dict] = None, conversation_id: Optional[str] = None) -> list:
     prompts = []
     tools = []
 
@@ -92,7 +95,7 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                         logger.warning("Image generation internal tool requested "
                                        "but no image generation model configured")
             elif tool['type'] == 'artifact':
-                tools.extend(ArtifactToolkit.get_toolkit(
+                toolkit_tools = ArtifactToolkit.get_toolkit(
                     client=alita_client,
                     bucket=tool['settings']['bucket'],
                     toolkit_name=tool.get('toolkit_name', ''),
@@ -102,13 +105,36 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                     pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
                     embedding_model=tool['settings'].get('embedding_model'),
                     collection_name=f"{tool.get('toolkit_name')}",
-                    collection_schema = str(tool['id'])
-                ).get_tools())
+                    collection_schema=str(tool['id']),
+                ).get_tools()
+                # Inject toolkit_id for artifact tools as well
+                _inject_toolkit_id(tool, toolkit_tools)
+                tools.extend(toolkit_tools)
+
             elif tool['type'] == 'vectorstore':
                 tools.extend(VectorStoreToolkit.get_toolkit(
                     llm=llm,
                     toolkit_name=tool.get('toolkit_name', ''),
                     **tool['settings']).get_tools())
+            elif tool['type'] == 'planning':
+                # Planning toolkit for multi-step task tracking
+                # Try to fetch pgvector_project_connstr from project secrets
+                pgvector_connstr = None
+                if alita_client:
+                    try:
+                        pgvector_connstr = alita_client.unsecret('pgvector_project_connstr')
+                        if pgvector_connstr:
+                            logger.info("Using pgvector_project_connstr for planning toolkit")
+                    except Exception as e:
+                        logger.debug(f"pgvector_project_connstr not available: {e}")
+                
+                pgvector_config = {'connection_string': pgvector_connstr} if pgvector_connstr else {}
+                tools.extend(PlanningToolkit.get_toolkit(
+                    toolkit_name=tool.get('toolkit_name', ''),
+                    selected_tools=tool['settings'].get('selected_tools', []),
+                    pgvector_configuration=pgvector_config,
+                    conversation_id=conversation_id or tool['settings'].get('conversation_id'),
+                ).get_tools())
             elif tool['type'] == 'mcp':
                 # remote mcp tool initialization with token injection
                 settings = dict(tool['settings'])
@@ -153,8 +179,15 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                     toolkit_name=tool.get('toolkit_name', ''),
                     client=alita_client,
                     **settings).get_tools())
+        except McpAuthorizationRequired:
+            # Re-raise auth required exceptions directly
+            raise
         except Exception as e:
-            if isinstance(e, McpAuthorizationRequired):
+            # Check for wrapped McpAuthorizationRequired
+            if hasattr(e, '__cause__') and isinstance(e.__cause__, McpAuthorizationRequired):
+                raise e.__cause__
+            # Check exception class name as fallback
+            if e.__class__.__name__ == 'McpAuthorizationRequired':
                 raise
             logger.error(f"Error initializing toolkit for tool '{tool.get('name', 'unknown')}': {e}", exc_info=True)
             if debug_mode:
@@ -271,7 +304,7 @@ def _init_single_mcp_tool(server_toolkit_name, toolkit_name, available_tool, ali
         tool_name = f'{toolkit_name}{TOOLKIT_SPLITTER}{available_tool["name"]}'
         return McpServerTool(
             name=tool_name,
-            description=f"MCP for a tool '{tool_name}': {available_tool.get("description", "")}",
+            description=f"MCP for a tool '{tool_name}': {available_tool.get('description', '')}",
             args_schema=McpServerTool.create_pydantic_model_from_schema(
                 available_tool.get("inputSchema", {})
             ),

alita_sdk/runtime/tools/function.py

@@ -16,6 +16,18 @@ from ..langchain.utils import propagate_the_input_mapping
 logger = logging.getLogger(__name__)
 
 
+def replace_escaped_newlines(data):
+    """
+        Replace \\n with \n in all string values recursively.
+        Required for sanitization of state variables in code node
+    """
+    if isinstance(data, dict):
+        return {key: replace_escaped_newlines(value) for key, value in data.items()}
+    elif isinstance(data, str):
+        return data.replace('\\n', '\n')
+    else:
+        return data
+
 class FunctionTool(BaseTool):
     name: str = 'FunctionalTool'
     description: str = 'This is direct call node for tools'
@@ -30,11 +42,13 @@ class FunctionTool(BaseTool):
     def _prepare_pyodide_input(self, state: Union[str, dict, ToolCall]) -> str:
         """Prepare input for PyodideSandboxTool by injecting state into the code block."""
         # add state into the code block here since it might be changed during the execution of the code
-        state_copy = deepcopy(state)
+        state_copy = replace_escaped_newlines(deepcopy(state))
 
         del state_copy['messages']  # remove messages to avoid issues with pickling without langchain-core
         # inject state into the code block as alita_state variable
-        pyodide_predata = f"#state dict\nalita_state = {state_copy}\n"
+        state_json = json.dumps(state_copy, ensure_ascii=False)
+        pyodide_predata = f'#state dict\nimport json\nalita_state = json.loads({json.dumps(state_json)})\n'
+
         return pyodide_predata
 
     def _handle_pyodide_output(self, tool_result: Any) -> dict:
@@ -94,9 +108,7 @@ class FunctionTool(BaseTool):
         # special handler for PyodideSandboxTool
         if self._is_pyodide_tool():
             code = func_args['code']
-            func_args['code'] = (f"{self._prepare_pyodide_input(state)}\n{code}"
-                                # handle new lines in the code properly
-                                 .replace('\\n','\\\\n'))
+            func_args['code'] = f"{self._prepare_pyodide_input(state)}\n{code}"
         try:
             tool_result = self.tool.invoke(func_args, config, **kwargs)
             dispatch_custom_event(