alita-sdk 0.3.379__py3-none-any.whl → 0.3.462__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
@@ -11,11 +12,14 @@ from .datasource import DatasourcesToolkit
 from .prompt import PromptToolkit
 from .subgraph import SubgraphToolkit
 from .vectorstore import VectorStoreToolkit
+from .mcp import McpToolkit
 from ..tools.mcp_server_tool import McpServerTool
 from ..tools.sandbox import SandboxToolkit
+from ..tools.image_generation import ImageGenerationToolkit
 # Import community tools
 from ...community import get_toolkits as community_toolkits, get_tools as community_tools
 from ...tools.memory import MemoryToolkit
+from ..utils.mcp_oauth import canonical_resource, McpAuthorizationRequired
 from ...tools.utils import TOOLKIT_SPLITTER
 
 logger = logging.getLogger(__name__)
@@ -26,71 +30,138 @@ def get_toolkits():
         ArtifactToolkit.toolkit_config_schema(),
         MemoryToolkit.toolkit_config_schema(),
         VectorStoreToolkit.toolkit_config_schema(),
-        SandboxToolkit.toolkit_config_schema()
+        SandboxToolkit.toolkit_config_schema(),
+        ImageGenerationToolkit.toolkit_config_schema(),
+        McpToolkit.toolkit_config_schema()
     ]
 
     return core_toolkits + community_toolkits() + alita_toolkits()
 
 
-def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None) -> list:
+def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None, debug_mode: Optional[bool] = False, mcp_tokens: Optional[dict] = None) -> list:
     prompts = []
     tools = []
 
     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,
+        try:
+            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':
+                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())
+                # backward compatibility for pipeline application type as subgraph node
+                if 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()
-        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':
+                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['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')}",
+                    collection_schema = str(tool['id'])
+                ).get_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'] == 'mcp':
+                # remote mcp tool initialization with token injection
+                settings = dict(tool['settings'])
+                url = settings.get('url')
+                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())
+        except Exception as e:
+            if isinstance(e, McpAuthorizationRequired):
+                raise
+            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}")
     
     if len(prompts) > 0:
         tools += PromptToolkit.get_toolkit(alita_client, prompts).get_tools()
@@ -99,7 +170,8 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
     tools += community_tools(tools_list, alita_client, llm)
     # Add alita tools
     tools += alita_tools(tools_list, alita_client, llm, memory_store)
-    # Add MCP tools
+    # 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(tools_list, alita_client)
     
     # Sanitize tool names to meet OpenAI's function naming requirements
@@ -154,6 +226,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 +237,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:

alita_sdk/runtime/tools/__init__.py

@@ -5,7 +5,11 @@ 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
+)
 
 __all__ = [
     "PyodideSandboxTool",
@@ -13,5 +17,6 @@ __all__ = [
     "create_sandbox_tool",
     "EchoTool",
     "ImageGenerationTool",
+    "ImageGenerationToolkit",
     "create_image_generation_tool"
-]
+]

alita_sdk/runtime/tools/application.py

@@ -50,6 +50,8 @@ class Application(BaseTool):
     application: Any
     args_schema: Type[BaseModel] = applicationToolSchema
     return_type: str = "str"
+    client: Any
+    args_runnable: dict = {}
     
     @field_validator('name', mode='before')
     @classmethod
@@ -66,6 +68,11 @@ 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.return_type == "str":
             return response["output"]

alita_sdk/runtime/tools/function.py

@@ -7,7 +7,7 @@ from langchain_core.callbacks import dispatch_custom_event
 from langchain_core.messages import ToolCall
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool, ToolException
-from typing import Any, Optional, Union, Annotated
+from typing import Any, Optional, Union
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from pydantic import ValidationError
 
@@ -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,29 +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"alita_state = {state_copy}"
-        # add classes related to sandbox client
-        # read the content of alita_sdk/runtime/cliens/sandbox_client.py
-        try:
-            import os
-            from pathlib import Path
-
-            # Get the directory of the current file and construct the path to sandbox_client.py
-            current_dir = Path(__file__).parent
-            sandbox_client_path = current_dir.parent / 'clients' / 'sandbox_client.py'
-
-            with open(sandbox_client_path, 'r') as f:
-                sandbox_client_code = f.read()
-            pyodide_predata += f"\n{sandbox_client_code}\n"
-            pyodide_predata += (f"alita_client = SandboxClient(base_url='{self.alita_client.base_url}',"
-                                f"project_id={self.alita_client.project_id},"
-                                f"auth_token='{self.alita_client.auth_token}')")
-        except FileNotFoundError:
-            logger.error(f"sandbox_client.py not found at {sandbox_client_path}. Ensure the file exists.")
+        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:
@@ -132,14 +128,22 @@ class FunctionTool(BaseTool):
             if not self.output_variables:
                 return {"messages": [{"role": "assistant", "content": dumps(tool_result)}]}
             else:
-                if self.output_variables[0] == "messages":
-                    return {
+                if "messages" in self.output_variables:
+                    messages_dict = {
                         "messages": [{
                             "role": "assistant",
-                            "content": dumps(tool_result) if not isinstance(tool_result, ToolException) else str(
-                                tool_result)
+                            "content": dumps(tool_result)
+                            if not isinstance(tool_result, ToolException) and not isinstance(tool_result, str)
+                            else str(tool_result)
                         }]
                     }
+                    for var in self.output_variables:
+                        if var != "messages":
+                            if isinstance(tool_result, dict) and var in tool_result:
+                                messages_dict[var] = tool_result[var]
+                            else:
+                                messages_dict[var] = tool_result
+                    return messages_dict
                 else:
                     return { self.output_variables[0]: tool_result }
         except ValidationError:

alita_sdk/runtime/tools/graph.py

@@ -47,8 +47,8 @@ def formulate_query(kwargs):
 
 
 class GraphTool(BaseTool):
-    name: str
-    description: str
+    name: str = 'GraphTool'
+    description: str = 'Graph tool for tools'
     graph: CompiledStateGraph
     args_schema: Type[BaseModel] = graphToolSchema
     return_type: str = "str"
@@ -65,10 +65,16 @@ class GraphTool(BaseTool):
         all_kwargs = {**kwargs, **extras, **schema_values}
         if config is None:
             config = {}
-        return self._run(*config, **all_kwargs)
+        # Pass the config to the _run empty or the one passed from the parent executor.
+        return self._run(config, **all_kwargs)
 
     def _run(self, *args, **kwargs):
-        response = self.graph.invoke(formulate_query(kwargs))
+        config = None
+        # From invoke method we are passing only 1 arg so it is safe to do this condition and config assignment.
+        # Default to None is safe because it will be checked also on the langchain side.
+        if args:
+            config = args[0]
+        response = self.graph.invoke(formulate_query(kwargs), config=config)
         if self.return_type == "str":
             return response["output"]
         else:

alita_sdk/runtime/tools/image_generation.py

@@ -2,16 +2,59 @@
 Image generation tool for Alita SDK.
 """
 import logging
-from typing import Optional, Type, Any
-from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field
+from typing import Optional, Type, Any, List, Literal
+from langchain_core.tools import BaseTool, BaseToolkit
+from pydantic import BaseModel, Field, create_model, ConfigDict
 
 logger = logging.getLogger(__name__)
 
+name = "image_generation"
+
+
+def get_tools(tools_list: list, alita_client=None, llm=None,
+              memory_store=None):
+    """
+    Get image generation tools for the provided tool configurations.
+
+    Args:
+        tools_list: List of tool configurations
+        alita_client: Alita client instance (required for image generation)
+        llm: LLM client instance (unused for image generation)
+        memory_store: Optional memory store instance (unused)
+
+    Returns:
+        List of image generation tools
+    """
+    all_tools = []
+
+    for tool in tools_list:
+        if (tool.get('type') == 'image_generation' or
+                tool.get('toolkit_name') == 'image_generation'):
+            try:
+                if not alita_client:
+                    logger.error("Alita client is required for image "
+                                 "generation tools")
+                    continue
+
+                toolkit_instance = ImageGenerationToolkit.get_toolkit(
+                    client=alita_client,
+                    toolkit_name=tool.get('toolkit_name', '')
+                )
+                all_tools.extend(toolkit_instance.get_tools())
+            except Exception as e:
+                logger.error(f"Error in image generation toolkit "
+                             f"get_tools: {e}")
+                logger.error(f"Tool config: {tool}")
+                raise
+
+    return all_tools
+
 
 class ImageGenerationInput(BaseModel):
     """Input schema for image generation tool."""
-    prompt: str = Field(description="Text prompt describing the image to generate")
+    prompt: str = Field(
+        description="Text prompt describing the image to generate"
+    )
     n: int = Field(
         default=1, description="Number of images to generate (1-10)",
         ge=1, le=10
@@ -22,7 +65,7 @@ class ImageGenerationInput(BaseModel):
     )
     quality: str = Field(
         default="auto",
-        description="Quality of the generated image ('low', 'medium', 'high', 'auto')"
+        description="Quality of the generated image ('low', 'medium', 'high')"
     )
     style: Optional[str] = Field(
         default=None, description="Style of the generated image (optional)"
@@ -69,7 +112,8 @@ class ImageGenerationTool(BaseTool):
                 else:
                     content_chunks.append({
                         "type": "text",
-                        "text": f"Generated {len(images)} images for prompt: '{prompt}'"
+                        "text": f"Generated {len(images)} images for "
+                                f"prompt: '{prompt}'"
                     })
                 
                 # Add image content for each generated image
@@ -85,7 +129,8 @@ class ImageGenerationTool(BaseTool):
                         content_chunks.append({
                             "type": "image_url",
                             "image_url": {
-                                "url": f"data:image/png;base64,{image_data['b64_json']}"
+                                "url": f"data:image/png;base64,"
+                                       f"{image_data['b64_json']}"
                             }
                         })
                 
@@ -94,7 +139,8 @@ class ImageGenerationTool(BaseTool):
             # Fallback to text response if no images in result
             return [{
                 "type": "text",
-                "text": f"Image generation completed but no images returned: {result}"
+                "text": f"Image generation completed but no images "
+                        f"returned: {result}"
             }]
             
         except Exception as e:
@@ -114,3 +160,53 @@ class ImageGenerationTool(BaseTool):
 def create_image_generation_tool(client):
     """Create an image generation tool with the provided Alita client."""
     return ImageGenerationTool(client=client)
+
+
+class ImageGenerationToolkit(BaseToolkit):
+    """Toolkit for image generation tools."""
+    tools: List[BaseTool] = []
+
+    @staticmethod
+    def toolkit_config_schema() -> BaseModel:
+        """Get the configuration schema for the image generation toolkit."""
+        # Create sample tool to get schema
+        sample_tool = ImageGenerationTool(client=None)
+        selected_tools = {sample_tool.name: sample_tool.args_schema.schema()}
+
+        return create_model(
+            'image_generation',
+            selected_tools=(
+                List[Literal[tuple(selected_tools)]],
+                Field(
+                    default=[],
+                    json_schema_extra={'args_schemas': selected_tools}
+                )
+            ),
+            __config__=ConfigDict(json_schema_extra={
+                'metadata': {
+                    "label": "Image Generation",
+                    "icon_url": "image_generation.svg",
+                    "hidden": True,
+                    "categories": ["internal_tool"],
+                    "extra_categories": ["image generation"],
+                }
+            })
+        )
+
+    @classmethod
+    def get_toolkit(cls, client=None, **kwargs):
+        """
+        Get toolkit with image generation tools.
+
+        Args:
+            client: Alita client instance (required)
+            **kwargs: Additional arguments
+        """
+        if not client:
+            raise ValueError("Alita client is required for image generation")
+            
+        tools = [ImageGenerationTool(client=client)]
+        return cls(tools=tools)
+
+    def get_tools(self):
+        return self.tools