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

alita_sdk/runtime/toolkits/tools.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Optional
 
 from langchain_core.tools import ToolException
 from langgraph.store.base import BaseStore
@@ -8,15 +9,23 @@ 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
+from .mcp import McpToolkit
+from .mcp_config import McpConfigToolkit, get_mcp_config_toolkit_schemas
+from .skill_router import SkillRouterToolkit
 from ..tools.mcp_server_tool import McpServerTool
 from ..tools.sandbox import SandboxToolkit
+from ..tools.image_generation import ImageGenerationToolkit
+from ..tools.data_analysis import DataAnalysisToolkit
 # Import community tools
 from ...community import get_toolkits as community_toolkits, get_tools as community_tools
 from ...tools.memory import MemoryToolkit
-from ...tools.utils import TOOLKIT_SPLITTER
+from ..utils.mcp_oauth import canonical_resource, McpAuthorizationRequired
+from ...tools.utils import clean_string
+from alita_sdk.tools import _inject_toolkit_id
 
 logger = logging.getLogger(__name__)
 
@@ -25,82 +34,344 @@ 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()
+        SandboxToolkit.toolkit_config_schema(),
+        ImageGenerationToolkit.toolkit_config_schema(),
+        DataAnalysisToolkit.toolkit_config_schema(),
+        McpToolkit.toolkit_config_schema(),
+        McpConfigToolkit.toolkit_config_schema(),
+        SkillRouterToolkit.toolkit_config_schema()
     ]
 
-    return core_toolkits + community_toolkits() + alita_toolkits()
+    # Add configured MCP servers (stdio and http) as available toolkits
+    mcp_config_toolkits = get_mcp_config_toolkit_schemas()
 
+    return core_toolkits + mcp_config_toolkits + community_toolkits() + alita_toolkits()
+
+
+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, ignored_mcp_servers: Optional[list] = None) -> list:
+    # Sanitize tools_list to handle corrupted tool configurations
+    sanitized_tools = []
+    for tool in tools_list:
+        if isinstance(tool, dict):
+            # Check for corrupted structure where 'type' and 'name' contain the full tool config
+            if 'type' in tool and isinstance(tool['type'], dict):
+                # This is a corrupted tool - use the inner dict instead
+                logger.warning(f"Detected corrupted tool configuration (type=dict), fixing: {tool}")
+                actual_tool = tool['type']  # or tool['name'], they should be the same
+                sanitized_tools.append(actual_tool)
+            elif 'name' in tool and isinstance(tool['name'], dict):
+                # Another corruption pattern where name contains the full config
+                logger.warning(f"Detected corrupted tool configuration (name=dict), fixing: {tool}")
+                actual_tool = tool['name']
+                sanitized_tools.append(actual_tool)
+            elif 'type' in tool and isinstance(tool['type'], str):
+                # Valid tool configuration
+                sanitized_tools.append(tool)
+            else:
+                # Skip invalid/corrupted tools that can't be fixed
+                logger.warning(f"Skipping invalid tool configuration: {tool}")
+        else:
+            logger.warning(f"Skipping non-dict tool: {tool}")
+            # Skip non-dict tools
 
-def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None) -> list:
     prompts = []
     tools = []
+    unhandled_tools = []  # Track tools not handled by main processing
 
-    for tool in tools_list:
-        if tool['type'] == 'datasource':
-            tools.extend(DatasourcesToolkit.get_toolkit(
-                alita_client,
-                datasource_ids=[int(tool['settings']['datasource_id'])],
-                selected_tools=tool['settings']['selected_tools'],
-                toolkit_name=tool.get('toolkit_name', '') or tool.get('name', '')
-            ).get_tools())
-        elif tool['type'] == 'application' and tool.get('agent_type', '') != 'pipeline' :
-            tools.extend(ApplicationToolkit.get_toolkit(
-                alita_client,
-                application_id=int(tool['settings']['application_id']),
-                application_version_id=int(tool['settings']['application_version_id']),
-                selected_tools=[]
-            ).get_tools())
-        elif tool['type'] == 'application' and tool.get('agent_type', '') == 'pipeline':
-            # static get_toolkit returns a list of CompiledStateGraph stubs
-            tools.extend(SubgraphToolkit.get_toolkit(
-                alita_client,
-                application_id=int(tool['settings']['application_id']),
-                application_version_id=int(tool['settings']['application_version_id']),
-                app_api_key=alita_client.auth_token,
-                selected_tools=[],
-                llm=llm
-            ))
-        elif tool['type'] == 'memory':
-            tools += MemoryToolkit.get_toolkit(
-                namespace=tool['settings'].get('namespace', str(tool['id'])),
-                pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
-                store=memory_store,
-            ).get_tools()
-        # TODO: update configuration of internal tools
-        elif tool['type'] == 'internal_tool':
-            if tool['name'] == 'pyodide':
-                tools += SandboxToolkit.get_toolkit(
-                    stateful=False,
-                    allow_net=True,
+    for tool in sanitized_tools:
+        # Flag to track if this tool was processed by the main loop
+        # Used to prevent double processing by fallback systems
+        tool_handled = False
+        try:
+            if tool['type'] == 'datasource':
+                tool_handled = True
+                tools.extend(DatasourcesToolkit.get_toolkit(
+                    alita_client,
+                    datasource_ids=[int(tool['settings']['datasource_id'])],
+                    selected_tools=tool['settings']['selected_tools'],
+                    toolkit_name=tool.get('toolkit_name', '') or tool.get('name', '')
+                ).get_tools())
+            elif tool['type'] == 'application':
+                tool_handled = True
+                # Check if this is a pipeline to enable PrinterNode filtering
+                is_pipeline_subgraph = tool.get('agent_type', '') == 'pipeline'
+
+                tools.extend(ApplicationToolkit.get_toolkit(
+                    alita_client,
+                    application_id=int(tool['settings']['application_id']),
+                    application_version_id=int(tool['settings']['application_version_id']),
+                    selected_tools=[],
+                    ignored_mcp_servers=ignored_mcp_servers,
+                    is_subgraph=is_pipeline_subgraph,  # Pass is_subgraph for pipelines
+                    mcp_tokens=mcp_tokens
+                ).get_tools())
+                # TODO: deprecate next release (1/15/2026)
+                # if is_pipeline_subgraph:
+                #     # static get_toolkit returns a list of CompiledStateGraph stubs
+                #     # Pass is_subgraph=True to enable PrinterNode filtering
+                #     logger.info(f"Processing pipeline as subgraph, will filter PrinterNodes")
+                #     tools.extend(SubgraphToolkit.get_toolkit(
+                #         alita_client,
+                #         application_id=int(tool['settings']['application_id']),
+                #         application_version_id=int(tool['settings']['application_version_id']),
+                #         app_api_key=alita_client.auth_token,
+                #         selected_tools=[],
+                #         llm=llm,
+                #         is_subgraph=True  # Enable PrinterNode filtering for pipelines used as subgraphs
+                #     ))
+            elif tool['type'] == 'memory':
+                tool_handled = True
+                tools += MemoryToolkit.get_toolkit(
+                    namespace=tool['settings'].get('namespace', str(tool['id'])),
+                    pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
+                    store=memory_store,
                 ).get_tools()
-        elif tool['type'] == 'artifact':
-            tools.extend(ArtifactToolkit.get_toolkit(
-                client=alita_client,
-                bucket=tool['settings']['bucket'],
-                toolkit_name=tool.get('toolkit_name', ''),
-                selected_tools=tool['settings'].get('selected_tools', []),
-                llm=llm,
-                # indexer settings
-                pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
-                embedding_model=tool['settings'].get('embedding_model'),
-                collection_name=f"{tool.get('toolkit_name')}",
-            ).get_tools())
-        elif tool['type'] == 'vectorstore':
-            tools.extend(VectorStoreToolkit.get_toolkit(
-                llm=llm,
-                toolkit_name=tool.get('toolkit_name', ''),
-                **tool['settings']).get_tools())
-    
+            # TODO: update configuration of internal tools
+            elif tool['type'] == 'internal_tool':
+                tool_handled = True
+                if tool['name'] == 'pyodide':
+                    tools += SandboxToolkit.get_toolkit(
+                        stateful=False,
+                        allow_net=True,
+                        alita_client=alita_client,
+                    ).get_tools()
+                elif tool['name'] == 'image_generation':
+                    if alita_client and alita_client.model_image_generation:
+                        tools += ImageGenerationToolkit.get_toolkit(
+                            client=alita_client,
+                        ).get_tools()
+                    else:
+                        logger.warning("Image generation internal tool requested "
+                                       "but no image generation model configured")
+                elif tool['name'] == 'planner':
+                    tools += PlanningToolkit.get_toolkit(
+                        pgvector_configuration=tool.get('settings', {}).get('pgvector_configuration'),
+                        conversation_id=conversation_id,
+                    ).get_tools()
+                elif tool['name'] == 'data_analysis':
+                    # Data Analysis internal tool - uses conversation attachment bucket
+                    settings = tool.get('settings', {})
+                    bucket_name = settings.get('bucket_name')
+                    if bucket_name:
+                        tools += DataAnalysisToolkit.get_toolkit(
+                            alita_client=alita_client,
+                            llm=llm,
+                            bucket_name=bucket_name,
+                            toolkit_name="Data Analyst",
+                        ).get_tools()
+                    else:
+                        logger.warning("Data Analysis internal tool requested "
+                                       "but no bucket_name provided in settings")
+            elif tool['type'] == 'artifact':
+                tool_handled = True
+                toolkit_tools = ArtifactToolkit.get_toolkit(
+                    client=alita_client,
+                    bucket=tool['settings']['bucket'],
+                    toolkit_name=tool.get('toolkit_name', ''),
+                    selected_tools=tool['settings'].get('selected_tools', []),
+                    llm=llm,
+                    # indexer settings
+                    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['settings'].get('id', tool.get('id', ''))),
+                ).get_tools()
+                # Inject toolkit_id for artifact tools as well
+                # Pass settings as the tool config since that's where the id field is
+                _inject_toolkit_id(tool['settings'], toolkit_tools)
+                tools.extend(toolkit_tools)
+
+            elif tool['type'] == 'vectorstore':
+                tool_handled = True
+                tools.extend(VectorStoreToolkit.get_toolkit(
+                    llm=llm,
+                    toolkit_name=tool.get('toolkit_name', ''),
+                    **tool['settings']).get_tools())
+            elif tool['type'] == 'planning':
+                tool_handled = True
+                # Planning toolkit for multi-step task tracking
+                settings = tool.get('settings', {})
+                
+                # Check if local mode is enabled (uses filesystem storage, ignores pgvector)
+                use_local = settings.get('local', False)
+                
+                if use_local:
+                    # Local mode - use filesystem storage
+                    logger.info("Planning toolkit using local filesystem storage (local=true)")
+                    pgvector_config = {}
+                else:
+                    # Check if explicit connection_string is provided in pgvector_configuration
+                    explicit_pgvector_config = settings.get('pgvector_configuration', {})
+                    explicit_connstr = explicit_pgvector_config.get('connection_string') if explicit_pgvector_config else None
+                    
+                    if explicit_connstr:
+                        # Use explicitly provided connection string (overrides project secrets)
+                        logger.info("Using explicit connection_string for planning toolkit")
+                        pgvector_config = explicit_pgvector_config
+                    else:
+                        # 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=settings.get('selected_tools', []),
+                    pgvector_configuration=pgvector_config,
+                    conversation_id=conversation_id or settings.get('conversation_id'),
+                ).get_tools())
+            elif tool['type'] == 'mcp':
+                tool_handled = True
+                # remote mcp tool initialization with token injection
+                settings = dict(tool['settings'])
+                url = settings.get('url')
+                
+                # Check if this MCP server should be ignored (user chose to continue without auth)
+                if ignored_mcp_servers and url:
+                    canonical_url = canonical_resource(url)
+                    if canonical_url in ignored_mcp_servers or url in ignored_mcp_servers:
+                        logger.info(f"[MCP Auth] Skipping ignored MCP server: {url}")
+                        continue
+                
+                headers = settings.get('headers')
+                token_data = None
+                session_id = None
+                if mcp_tokens and url:
+                    canonical_url = canonical_resource(url)
+                    logger.info(f"[MCP Auth] Looking for token for URL: {url}")
+                    logger.info(f"[MCP Auth] Canonical URL: {canonical_url}")
+                    logger.info(f"[MCP Auth] Available tokens: {list(mcp_tokens.keys())}")
+                    token_data = mcp_tokens.get(canonical_url)
+                    if token_data:
+                        logger.info(f"[MCP Auth] Found token data for {canonical_url}")
+                        # Handle both old format (string) and new format (dict with access_token and session_id)
+                        if isinstance(token_data, dict):
+                            access_token = token_data.get('access_token')
+                            session_id = token_data.get('session_id')
+                            logger.info(f"[MCP Auth] Token data: access_token={'present' if access_token else 'missing'}, session_id={session_id or 'none'}")
+                        else:
+                            # Backward compatibility: treat as plain token string
+                            access_token = token_data
+                            logger.info(f"[MCP Auth] Using legacy token format (string)")
+                    else:
+                        access_token = None
+                        logger.warning(f"[MCP Auth] No token found for {canonical_url}")
+                else:
+                    access_token = None
+                    
+                if access_token:
+                    merged_headers = dict(headers) if headers else {}
+                    merged_headers.setdefault('Authorization', f'Bearer {access_token}')
+                    settings['headers'] = merged_headers
+                    logger.info(f"[MCP Auth] Added Authorization header for {url}")
+                    
+                # Pass session_id to MCP toolkit if available
+                if session_id:
+                    settings['session_id'] = session_id
+                    logger.info(f"[MCP Auth] Passing session_id to toolkit: {session_id}")
+                tools.extend(McpToolkit.get_toolkit(
+                    toolkit_name=tool.get('toolkit_name', ''),
+                    client=alita_client,
+                    **settings).get_tools())
+            elif tool['type'] == 'skill_router':
+                tool_handled = True
+                # Skills Registry Router Toolkit
+                logger.info(f"Processing skill_router toolkit: {tool}")
+                try:
+                    settings = tool.get('settings', {})
+                    toolkit_name = tool.get('toolkit_name', '')
+                    selected_tools = settings.get('selected_tools', [])
+
+                    toolkit_tools = SkillRouterToolkit.get_toolkit(
+                        client=alita_client,
+                        llm=llm,
+                        toolkit_name=toolkit_name,
+                        selected_tools=selected_tools,
+                        **settings
+                    ).get_tools()
+
+                    tools.extend(toolkit_tools)
+                    logger.info(f"✅ Successfully added {len(toolkit_tools)} tools from SkillRouterToolkit")
+                except Exception as e:
+                    logger.error(f"❌ Failed to initialize SkillRouterToolkit: {e}")
+                    raise
+            elif tool['type'] == 'mcp_config' or tool['type'].startswith('mcp_'):
+                tool_handled = True
+                # MCP Config toolkit - pre-configured MCP servers (stdio or http)
+                # Handle both explicit 'mcp_config' type and dynamic names like 'mcp_playwright'
+                logger.info(f"Processing mcp_config toolkit: {tool}")
+                try:
+                    settings = tool.get('settings', {})
+
+                    # Server name can come from settings or be extracted from type name
+                    server_name = settings.get('server_name')
+                    if not server_name and tool['type'].startswith('mcp_') and tool['type'] != 'mcp_config':
+                        # Extract server name from type (e.g., 'mcp_playwright' -> 'playwright')
+                        server_name = tool['type'][4:]  # Remove 'mcp_' prefix
+
+                    if not server_name:
+                        logger.error(f"❌ No server_name found for mcp_config toolkit: {tool}")
+                        continue
+
+                    toolkit_name = tool.get('toolkit_name', '') or server_name
+                    selected_tools = settings.get('selected_tools', [])
+                    excluded_tools = settings.get('excluded_tools', [])
+
+                    # Get server config (may be in settings or from global config)
+                    server_config = settings.get('server_config')
+
+                    toolkit_tools = McpConfigToolkit.get_toolkit(
+                        server_name=server_name,
+                        server_config=server_config,
+                        user_config=settings,
+                        selected_tools=selected_tools if selected_tools else None,
+                        excluded_tools=excluded_tools if excluded_tools else None,
+                        toolkit_name=toolkit_name,
+                        client=alita_client,
+                    ).get_tools()
+
+                    tools.extend(toolkit_tools)
+                    logger.info(f"✅ Successfully added {len(toolkit_tools)} tools from McpConfigToolkit ({server_name})")
+                except Exception as e:
+                    logger.error(f"❌ Failed to initialize McpConfigToolkit: {e}")
+                    if not debug_mode:
+                        raise
+        except McpAuthorizationRequired:
+            # Re-raise auth required exceptions directly
+            raise
+        except Exception as e:
+            logger.error(f"Error initializing toolkit for tool '{tool.get('name', 'unknown')}': {e}", exc_info=True)
+            if debug_mode:
+                logger.info("Skipping tool initialization error due to debug mode.")
+                continue
+            else:
+                raise ToolException(f"Error initializing toolkit for tool '{tool.get('name', 'unknown')}': {e}")
+
+        # Track unhandled tools (make a copy to avoid reference issues)
+        if not tool_handled:
+            # Ensure we only add valid tool configurations to unhandled_tools
+            if isinstance(tool, dict) and 'type' in tool and isinstance(tool['type'], str):
+                unhandled_tools.append(dict(tool))
+
     if len(prompts) > 0:
         tools += PromptToolkit.get_toolkit(alita_client, prompts).get_tools()
-    
-    # Add community tools
-    tools += community_tools(tools_list, alita_client, llm)
-    # Add alita tools
-    tools += alita_tools(tools_list, alita_client, llm, memory_store)
-    # Add MCP tools
-    tools += _mcp_tools(tools_list, alita_client)
+
+    # Add community tools (only for unhandled tools)
+    tools += community_tools(unhandled_tools, alita_client, llm)
+    # Add alita tools (only for unhandled tools)
+    tools += alita_tools(unhandled_tools, alita_client, llm, memory_store)
+    # Add MCP tools registered via alita-mcp CLI (static registry)
+    # Note: Tools with type='mcp' are already handled in main loop above
+    tools += _mcp_tools(unhandled_tools, alita_client)
     
     # Sanitize tool names to meet OpenAI's function naming requirements
     # tools = _sanitize_tool_names(tools)
@@ -154,6 +425,10 @@ def _sanitize_tool_names(tools: list) -> list:
 
 
 def _mcp_tools(tools_list, alita):
+    """
+    Handle MCP tools registered via alita-mcp CLI (static registry).
+    Skips tools with type='mcp' as those are handled by dynamic discovery.
+    """
     try:
         all_available_toolkits = alita.get_mcp_toolkits()
         toolkit_lookup = {tk["name"]: tk for tk in all_available_toolkits}
@@ -161,6 +436,11 @@ def _mcp_tools(tools_list, alita):
         #
         for selected_toolkit in tools_list:
             server_toolkit_name = selected_toolkit['type']
+            
+            # Skip tools with type='mcp' - they're handled by dynamic discovery
+            if server_toolkit_name == 'mcp':
+                continue
+            
             toolkit_conf = toolkit_lookup.get(server_toolkit_name)
             #
             if not toolkit_conf:
@@ -186,11 +466,18 @@ def _mcp_tools(tools_list, alita):
 
 def _init_single_mcp_tool(server_toolkit_name, toolkit_name, available_tool, alita, toolkit_settings):
     try:
-
-        tool_name = f'{toolkit_name}{TOOLKIT_SPLITTER}{available_tool["name"]}'
+        # Use clean tool name without prefix
+        tool_name = available_tool["name"]
+        # Add toolkit context to description (max 1000 chars)
+        toolkit_context = f" [Toolkit: {clean_string(toolkit_name)}]" if toolkit_name else ''
+        base_description = f"MCP for a tool '{tool_name}': {available_tool.get('description', '')}"
+        description = base_description
+        if toolkit_context and len(base_description + toolkit_context) <= 1000:
+            description = base_description + toolkit_context
+        
         return McpServerTool(
             name=tool_name,
-            description=f"MCP for a tool '{tool_name}': {available_tool.get("description", "")}",
+            description=description,
             args_schema=McpServerTool.create_pydantic_model_from_schema(
                 available_tool.get("inputSchema", {})
             ),

alita_sdk/runtime/toolkits/vectorstore.py

@@ -1,7 +1,7 @@
 from logging import getLogger
 from typing import Any, List, Literal, Optional
 
-from alita_sdk.tools.utils import clean_string, TOOLKIT_SPLITTER
+from alita_sdk.tools.utils import clean_string
 from pydantic import BaseModel, create_model, Field, ConfigDict
 from langchain_core.tools import BaseToolkit, BaseTool
 from alita_sdk.tools.base.tool import BaseAction
@@ -31,7 +31,8 @@ class VectorStoreToolkit(BaseToolkit):
                     toolkit_name: Optional[str] = None,
                     selected_tools: list[str] = []):
         logger.info("Selected tools: %s", selected_tools)
-        prefix = clean_string(toolkit_name) + TOOLKIT_SPLITTER if toolkit_name else ''
+        # Use clean toolkit name for context (max 1000 chars in description)
+        toolkit_context = f" [Toolkit: {clean_string(toolkit_name)}]" if toolkit_name else ''
         if selected_tools is None:
             selected_tools = []
         tools = []
@@ -46,11 +47,16 @@ class VectorStoreToolkit(BaseToolkit):
             # if selected_tools:
             #     if tool["name"] not in selected_tools:
             #         continue
+            # Add toolkit context to description with character limit
+            description = tool["description"]
+            if toolkit_context and len(description + toolkit_context) <= 1000:
+                description = description + toolkit_context
             tools.append(BaseAction(
                 api_wrapper=vectorstore_wrapper,
-                name=f'{prefix}{tool["name"]}',
-                description=tool["description"],
-                args_schema=tool["args_schema"]
+                name=tool["name"],
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={"toolkit_name": toolkit_name, "toolkit_type": "vectorstore"} if toolkit_name else {}
             ))
         return cls(tools=tools)
 

alita_sdk/runtime/tools/__init__.py

@@ -5,7 +5,12 @@ This module provides various tools that can be used within LangGraph agents.
 
 from .sandbox import PyodideSandboxTool, StatefulPyodideSandboxTool, create_sandbox_tool
 from .echo import EchoTool
-from .image_generation import ImageGenerationTool, create_image_generation_tool
+from .image_generation import (
+    ImageGenerationTool,
+    create_image_generation_tool,
+    ImageGenerationToolkit
+)
+from .skill_router import SkillRouterWrapper
 
 __all__ = [
     "PyodideSandboxTool",
@@ -13,5 +18,7 @@ __all__ = [
     "create_sandbox_tool",
     "EchoTool",
     "ImageGenerationTool",
-    "create_image_generation_tool"
-]
+    "ImageGenerationToolkit",
+    "create_image_generation_tool",
+    "SkillRouterWrapper"
+]

alita_sdk/runtime/tools/application.py

@@ -4,7 +4,7 @@ from ..utils.utils import clean_string
 from langchain_core.tools import BaseTool, ToolException
 from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
 from typing import Any, Type, Optional
-from pydantic import create_model, field_validator, BaseModel
+from pydantic import create_model, model_validator, BaseModel
 from pydantic.fields import FieldInfo
 from ..langchain.mixedAgentRenderes import convert_message_to_json
 from logging import getLogger
@@ -41,8 +41,7 @@ def formulate_query(kwargs):
         if key not in ("task", "chat_history"):
             result[key] = value
     return result
-    
-    
+
 
 class Application(BaseTool):
     name: str
@@ -50,11 +49,26 @@ class Application(BaseTool):
     application: Any
     args_schema: Type[BaseModel] = applicationToolSchema
     return_type: str = "str"
+    client: Any
+    args_runnable: dict = {}
+    metadata: dict = {}
+    is_subgraph: Optional[bool] = False
     
-    @field_validator('name', mode='before')
+    @model_validator(mode='before')
     @classmethod
-    def remove_spaces(cls, v):
-        return clean_string(v)
+    def preserve_original_name(cls, data: Any) -> Any:
+        """Preserve the original name in metadata before cleaning."""
+        if isinstance(data, dict):
+            original_name = data.get('name')
+            if original_name:
+                # Initialize metadata if not present
+                if data.get('metadata') is None:
+                    data['metadata'] = {}
+                # Store original name before cleaning
+                data['metadata']['original_name'] = original_name
+                # Clean the name
+                data['name'] = clean_string(original_name)
+        return data
 
     def invoke(self, input: Any, config: Optional[dict] = None, **kwargs: Any) -> Any:
         """Override default invoke to preserve all fields, not just args_schema"""
@@ -66,7 +80,14 @@ class Application(BaseTool):
         return self._run(*config, **all_kwargs)
 
     def _run(self, *args, **kwargs):
+        if self.client and self.args_runnable:
+            # Recreate new LanggraphAgentRunnable in order to reflect the current input_mapping (it can be dynamic for pipelines).
+            # Actually, for pipelines agent toolkits LanggraphAgentRunnable is created (for LLMNode) before pipeline's schema parsing.
+            application_variables = {k: {"name": k, "value": v} for k, v in kwargs.items()}
+            self.application = self.client.application(**self.args_runnable, application_variables=application_variables)
         response = self.application.invoke(formulate_query(kwargs))
+        if self.is_subgraph:
+            return response
         if self.return_type == "str":
             return response["output"]
         else: