alita-sdk 0.3.423__py3-none-any.whl → 0.3.435__py3-none-any.whl

alita_sdk/runtime/toolkits/tools.py

@@ -12,6 +12,7 @@ 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
@@ -29,7 +30,8 @@ def get_toolkits():
         MemoryToolkit.toolkit_config_schema(),
         VectorStoreToolkit.toolkit_config_schema(),
         SandboxToolkit.toolkit_config_schema(),
-        ImageGenerationToolkit.toolkit_config_schema()
+        ImageGenerationToolkit.toolkit_config_schema(),
+        McpToolkit.toolkit_config_schema()
     ]
 
     return core_toolkits + community_toolkits() + alita_toolkits()
@@ -106,6 +108,11 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                     llm=llm,
                     toolkit_name=tool.get('toolkit_name', ''),
                     **tool['settings']).get_tools())
+            elif tool['type'] == 'mcp':
+                tools.extend(McpToolkit.get_toolkit(
+                    toolkit_name=tool.get('toolkit_name', ''),
+                    client=alita_client,
+                    **tool['settings']).get_tools())
         except Exception as e:
             logger.error(f"Error initializing toolkit for tool '{tool.get('name', 'unknown')}': {e}", exc_info=True)
             if debug_mode:
@@ -121,7 +128,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
@@ -176,6 +184,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}
@@ -183,6 +195,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/mcp_inspect_tool.py

@@ -0,0 +1,284 @@
+"""
+MCP Server Inspection Tool.
+Allows inspecting available tools, prompts, and resources on an MCP server.
+"""
+
+import asyncio
+import json
+import logging
+import time
+from typing import Any, Type, Dict, List, Optional
+
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field, ConfigDict
+import aiohttp
+
+logger = logging.getLogger(__name__)
+
+
+class McpInspectInput(BaseModel):
+    """Input schema for MCP server inspection tool."""
+
+    resource_type: str = Field(
+        default="all",
+        description="What to inspect: 'tools', 'prompts', 'resources', or 'all'"
+    )
+
+
+class McpInspectTool(BaseTool):
+    """Tool for inspecting available tools, prompts, and resources on an MCP server."""
+
+    name: str = "mcp_inspect"
+    description: str = "List available tools, prompts, and resources from the MCP server"
+    args_schema: Type[BaseModel] = McpInspectInput
+    return_type: str = "str"
+
+    # MCP server connection details
+    server_name: str = Field(..., description="Name of the MCP server")
+    server_url: str = Field(..., description="URL of the MCP server")
+    server_headers: Optional[Dict[str, str]] = Field(default=None, description="HTTP headers for authentication")
+    timeout: int = Field(default=30, description="Request timeout in seconds")
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def __getstate__(self):
+        """Custom serialization for pickle compatibility."""
+        state = self.__dict__.copy()
+        # Convert headers dict to regular dict to avoid any reference issues
+        if 'server_headers' in state and state['server_headers'] is not None:
+            state['server_headers'] = dict(state['server_headers'])
+        return state
+
+    def __setstate__(self, state):
+        """Custom deserialization for pickle compatibility."""
+        # Initialize Pydantic internal attributes if needed
+        if '__pydantic_fields_set__' not in state:
+            state['__pydantic_fields_set__'] = set(state.keys())
+        if '__pydantic_extra__' not in state:
+            state['__pydantic_extra__'] = None
+        if '__pydantic_private__' not in state:
+            state['__pydantic_private__'] = None
+
+        # Update object state
+        self.__dict__.update(state)
+
+    def _run(self, resource_type: str = "all") -> str:
+        """Inspect the MCP server for available resources."""
+        try:
+            # Always create a new event loop for sync context
+            # This avoids issues with existing event loops in threads
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                future = executor.submit(self._run_in_new_loop, resource_type)
+                return future.result(timeout=self.timeout)
+        except Exception as e:
+            logger.error(f"Error inspecting MCP server '{self.server_name}': {e}")
+            return f"Error inspecting MCP server: {e}"
+
+    def _run_in_new_loop(self, resource_type: str) -> str:
+        """Run the async inspection in a new event loop."""
+        return asyncio.run(self._inspect_server(resource_type))
+
+    async def _inspect_server(self, resource_type: str) -> str:
+        """Perform the actual MCP server inspection."""
+        results = {}
+
+        # Determine what to inspect
+        inspect_tools = resource_type in ["all", "tools"]
+        inspect_prompts = resource_type in ["all", "prompts"]
+        inspect_resources = resource_type in ["all", "resources"]
+
+        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout)) as session:
+
+            # List tools
+            if inspect_tools:
+                try:
+                    tools = await self._list_tools(session)
+                    results["tools"] = tools
+                except Exception as e:
+                    logger.warning(f"Failed to list tools from {self.server_name}: {e}")
+                    results["tools"] = {"error": str(e)}
+
+            # List prompts
+            if inspect_prompts:
+                try:
+                    prompts = await self._list_prompts(session)
+                    results["prompts"] = prompts
+                except Exception as e:
+                    logger.warning(f"Failed to list prompts from {self.server_name}: {e}")
+                    results["prompts"] = {"error": str(e)}
+
+            # List resources
+            if inspect_resources:
+                try:
+                    resources = await self._list_resources(session)
+                    results["resources"] = resources
+                except Exception as e:
+                    logger.warning(f"Failed to list resources from {self.server_name}: {e}")
+                    results["resources"] = {"error": str(e)}
+
+        return self._format_results(results, resource_type)
+
+    def _parse_sse(self, text: str) -> Dict[str, Any]:
+        """Parse Server-Sent Events (SSE) format response."""
+        for line in text.split('\n'):
+            line = line.strip()
+            if line.startswith('data:'):
+                json_str = line[5:].strip()
+                return json.loads(json_str)
+        raise ValueError("No data found in SSE response")
+
+    async def _list_tools(self, session: aiohttp.ClientSession) -> Dict[str, Any]:
+        """List available tools from the MCP server."""
+        request = {
+            "jsonrpc": "2.0",
+            "id": f"list_tools_{int(time.time())}",
+            "method": "tools/list",
+            "params": {}
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+            **self.server_headers
+        }
+
+        async with session.post(self.server_url, json=request, headers=headers) as response:
+            if response.status != 200:
+                raise Exception(f"HTTP {response.status}: {await response.text()}")
+
+            # Handle both JSON and SSE responses
+            content_type = response.headers.get('Content-Type', '')
+            if 'text/event-stream' in content_type:
+                # Parse SSE format
+                text = await response.text()
+                data = self._parse_sse(text)
+            else:
+                data = await response.json()
+
+            if "error" in data:
+                raise Exception(f"MCP Error: {data['error']}")
+
+            return data.get("result", {})
+
+    async def _list_prompts(self, session: aiohttp.ClientSession) -> Dict[str, Any]:
+        """List available prompts from the MCP server."""
+        request = {
+            "jsonrpc": "2.0",
+            "id": f"list_prompts_{int(time.time())}",
+            "method": "prompts/list",
+            "params": {}
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+            **self.server_headers
+        }
+
+        async with session.post(self.server_url, json=request, headers=headers) as response:
+            if response.status != 200:
+                raise Exception(f"HTTP {response.status}: {await response.text()}")
+
+            # Handle both JSON and SSE responses
+            content_type = response.headers.get('Content-Type', '')
+            if 'text/event-stream' in content_type:
+                text = await response.text()
+                data = self._parse_sse(text)
+            else:
+                data = await response.json()
+
+            if "error" in data:
+                raise Exception(f"MCP Error: {data['error']}")
+
+            return data.get("result", {})
+
+    async def _list_resources(self, session: aiohttp.ClientSession) -> Dict[str, Any]:
+        """List available resources from the MCP server."""
+        request = {
+            "jsonrpc": "2.0",
+            "id": f"list_resources_{int(time.time())}",
+            "method": "resources/list",
+            "params": {}
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+            **self.server_headers
+        }
+
+        async with session.post(self.server_url, json=request, headers=headers) as response:
+            if response.status != 200:
+                raise Exception(f"HTTP {response.status}: {await response.text()}")
+
+            # Handle both JSON and SSE responses
+            content_type = response.headers.get('Content-Type', '')
+            if 'text/event-stream' in content_type:
+                text = await response.text()
+                data = self._parse_sse(text)
+            else:
+                data = await response.json()
+
+            if "error" in data:
+                raise Exception(f"MCP Error: {data['error']}")
+
+            return data.get("result", {})
+
+    def _format_results(self, results: Dict[str, Any], resource_type: str) -> str:
+        """Format the inspection results for display."""
+        output_lines = [f"=== MCP Server Inspection: {self.server_name} ==="]
+        output_lines.append(f"Server URL: {self.server_url}")
+        output_lines.append("")
+
+        # Format tools
+        if "tools" in results:
+            if "error" in results["tools"]:
+                output_lines.append(f"❌ TOOLS: Error - {results['tools']['error']}")
+            else:
+                tools = results["tools"].get("tools", [])
+                output_lines.append(f"🔧 TOOLS ({len(tools)} available):")
+                if tools:
+                    for tool in tools:
+                        name = tool.get("name", "Unknown")
+                        desc = tool.get("description", "No description")
+                        output_lines.append(f"  • {name}: {desc}")
+                else:
+                    output_lines.append("  (No tools available)")
+                output_lines.append("")
+
+        # Format prompts
+        if "prompts" in results:
+            if "error" in results["prompts"]:
+                output_lines.append(f"❌ PROMPTS: Error - {results['prompts']['error']}")
+            else:
+                prompts = results["prompts"].get("prompts", [])
+                output_lines.append(f"💬 PROMPTS ({len(prompts)} available):")
+                if prompts:
+                    for prompt in prompts:
+                        name = prompt.get("name", "Unknown")
+                        desc = prompt.get("description", "No description")
+                        output_lines.append(f"  • {name}: {desc}")
+                else:
+                    output_lines.append("  (No prompts available)")
+                output_lines.append("")
+
+        # Format resources
+        if "resources" in results:
+            if "error" in results["resources"]:
+                output_lines.append(f"❌ RESOURCES: Error - {results['resources']['error']}")
+            else:
+                resources = results["resources"].get("resources", [])
+                output_lines.append(f"📁 RESOURCES ({len(resources)} available):")
+                if resources:
+                    for resource in resources:
+                        uri = resource.get("uri", "Unknown")
+                        name = resource.get("name", uri)
+                        desc = resource.get("description", "No description")
+                        output_lines.append(f"  • {name}: {desc}")
+                        output_lines.append(f"    URI: {uri}")
+                else:
+                    output_lines.append("  (No resources available)")
+                output_lines.append("")
+
+        return "\n".join(output_lines)

alita_sdk/runtime/tools/mcp_server_tool.py

@@ -3,7 +3,7 @@ from logging import getLogger
 from typing import Any, Type, Literal, Optional, Union, List
 
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field, create_model, EmailStr, constr
+from pydantic import BaseModel, Field, create_model, EmailStr, constr, ConfigDict
 
 from ...tools.utils import TOOLKIT_SPLITTER
 
@@ -15,9 +15,61 @@ class McpServerTool(BaseTool):
     description: str
     args_schema: Optional[Type[BaseModel]] = None
     return_type: str = "str"
-    client: Any
+    client: Any = Field(default=None, exclude=True)  # Exclude from serialization
     server: str
     tool_timeout_sec: int = 60
+    is_prompt: bool = False  # Flag to indicate if this is a prompt tool
+    prompt_name: Optional[str] = None  # Original prompt name if this is a prompt
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def __getstate__(self):
+        """Custom serialization to exclude non-serializable objects."""
+        state = self.__dict__.copy()
+        # Remove the client since it contains threading objects that can't be pickled
+        state['client'] = None
+        # Store args_schema as a schema dict instead of the dynamic class
+        if hasattr(self, 'args_schema') and self.args_schema is not None:
+            # Convert the Pydantic model back to schema dict for pickling
+            try:
+                state['_args_schema_dict'] = self.args_schema.model_json_schema()
+                state['args_schema'] = None
+            except Exception as e:
+                logger.warning(f"Failed to serialize args_schema: {e}")
+                # If conversion fails, just remove it
+                state['args_schema'] = None
+                state['_args_schema_dict'] = {}
+        return state
+
+    def __setstate__(self, state):
+        """Custom deserialization to handle missing objects."""
+        # Restore the args_schema from the stored schema dict
+        args_schema_dict = state.pop('_args_schema_dict', {})
+
+        # Initialize required Pydantic internal attributes
+        if '__pydantic_fields_set__' not in state:
+            state['__pydantic_fields_set__'] = set(state.keys())
+        if '__pydantic_extra__' not in state:
+            state['__pydantic_extra__'] = None
+        if '__pydantic_private__' not in state:
+            state['__pydantic_private__'] = None
+
+        # Directly update the object's __dict__ to bypass Pydantic validation
+        self.__dict__.update(state)
+
+        # Recreate the args_schema from the stored dict if available
+        if args_schema_dict:
+            try:
+                recreated_schema = self.create_pydantic_model_from_schema(args_schema_dict)
+                self.__dict__['args_schema'] = recreated_schema
+            except Exception as e:
+                logger.warning(f"Failed to recreate args_schema: {e}")
+                self.__dict__['args_schema'] = None
+        else:
+            self.__dict__['args_schema'] = None
+
+        # Note: client will be None after unpickling
+        # The toolkit should reinitialize the client when needed
 
 
     @staticmethod
@@ -90,14 +142,31 @@ class McpServerTool(BaseTool):
         return create_model(model_name, **fields)
 
     def _run(self, *args, **kwargs):
-        call_data = {
-            "server": self.server,
-            "tool_timeout_sec": self.tool_timeout_sec,
-            "tool_call_id": str(uuid.uuid4()),
-            "params": {
-                "name": self.name.rsplit(TOOLKIT_SPLITTER)[1] if TOOLKIT_SPLITTER in self.name else self.name,
-                "arguments": kwargs
+        # Extract the actual tool/prompt name (remove toolkit prefix)
+        actual_name = self.name.rsplit(TOOLKIT_SPLITTER)[1] if TOOLKIT_SPLITTER in self.name else self.name
+        
+        if self.is_prompt:
+            # For prompts, use prompts/get endpoint
+            call_data = {
+                "server": self.server,
+                "tool_timeout_sec": self.tool_timeout_sec,
+                "tool_call_id": str(uuid.uuid4()),
+                "method": "prompts/get",
+                "params": {
+                    "name": self.prompt_name or actual_name.replace("prompt_", ""),
+                    "arguments": kwargs.get("arguments", kwargs)
+                }
+            }
+        else:
+            # For regular tools, use tools/call endpoint
+            call_data = {
+                "server": self.server,
+                "tool_timeout_sec": self.tool_timeout_sec,
+                "tool_call_id": str(uuid.uuid4()),
+                "params": {
+                    "name": actual_name,
+                    "arguments": kwargs
+                }
             }
-        }
         
         return self.client.mcp_tool_call(call_data)

alita_sdk/runtime/utils/streamlit.py

@@ -868,10 +868,24 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                 label = f"{'🔒 ' if is_secret else ''}{'*' if is_required else ''}{field_name.replace('_', ' ').title()}"
                                 
                                 if field_type == 'string':
-                                    if is_secret:
+                                    # Check if this is an enum field
+                                    if field_schema.get('enum'):
+                                        # Dropdown for enum values
+                                        options = field_schema['enum']
+                                        default_index = 0
+                                        if default_value and str(default_value) in options:
+                                            default_index = options.index(str(default_value))
+                                        toolkit_config_values[field_name] = st.selectbox(
+                                            label,
+                                            options=options,
+                                            index=default_index,
+                                            help=field_description,
+                                            key=f"config_{field_name}_{selected_toolkit_idx}"
+                                        )
+                                    elif is_secret:
                                         toolkit_config_values[field_name] = st.text_input(
                                             label,
-                                            value=str(default_value) if default_value else '', 
+                                            value=str(default_value) if default_value else '',
                                             help=field_description,
                                             type="password",
                                             key=f"config_{field_name}_{selected_toolkit_idx}"
@@ -879,7 +893,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                     else:
                                         toolkit_config_values[field_name] = st.text_input(
                                             label,
-                                            value=str(default_value) if default_value else '', 
+                                            value=str(default_value) if default_value else '',
                                             help=field_description,
                                             key=f"config_{field_name}_{selected_toolkit_idx}"
                                         )
@@ -971,6 +985,23 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                                 key=f"config_{field_name}_{selected_toolkit_idx}"
                                             )
                                         toolkit_config_values[field_name] = [line.strip() for line in array_input.split('\n') if line.strip()]
+                                elif field_type == 'object':
+                                    # Handle object/dict types (like headers)
+                                    obj_input = st.text_area(
+                                        f"{label} (JSON object)",
+                                        value=json.dumps(default_value) if isinstance(default_value, dict) else str(default_value) if default_value else '',
+                                        help=f"{field_description} - Enter as JSON object, e.g. {{\"Authorization\": \"Bearer token\"}}",
+                                        placeholder='{"key": "value"}',
+                                        key=f"config_{field_name}_{selected_toolkit_idx}"
+                                    )
+                                    try:
+                                        if obj_input.strip():
+                                            toolkit_config_values[field_name] = json.loads(obj_input)
+                                        else:
+                                            toolkit_config_values[field_name] = None
+                                    except json.JSONDecodeError as e:
+                                        st.error(f"Invalid JSON format for {field_name}: {e}")
+                                        toolkit_config_values[field_name] = None
                         else:
                             st.info("This toolkit doesn't require additional configuration.")
                         

alita_sdk/runtime/utils/toolkit_utils.py

@@ -46,11 +46,14 @@ def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any],
         # Log the configuration being used
         logger.info(f"Instantiating toolkit {toolkit_name} with LLM client")
         logger.debug(f"Toolkit {toolkit_name} configuration: {toolkit_config}")
-        
+
+        # Use toolkit type from config, or fall back to lowercase toolkit name
+        toolkit_type = toolkit_config.get('type', toolkit_name.lower())
+
         # Create a tool configuration dict with required fields
         tool_config = {
             'id': toolkit_config.get('id', random.randint(1, 1000000)),
-            'type': toolkit_config.get('type', toolkit_name.lower()),
+            'type': toolkit_config.get('type', toolkit_type),
             'settings': settings,
             'toolkit_name': toolkit_name
         }

alita_sdk/tools/__init__.py

@@ -131,6 +131,11 @@ def get_tools(tools_list, alita, llm, store: Optional[BaseStore] = None, *args,
                 logger.error(f"Error getting ADO repos tools: {e}")
             continue
 
+        # Skip MCP toolkit - it's handled by runtime/toolkits/tools.py to avoid duplicate loading
+        if tool_type == 'mcp':
+            logger.debug(f"Skipping MCP toolkit '{tool.get('toolkit_name')}' - handled by runtime toolkit system")
+            continue
+
         # Handle standard tools
         if tool_type in AVAILABLE_TOOLS and 'get_tools' in AVAILABLE_TOOLS[tool_type]:
             try:

alita_sdk/tools/chunkers/sematic/proposal_chunker.py

@@ -6,7 +6,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from langchain.text_splitter import TokenTextSplitter
 
 from typing import Optional, List
-from langchain_core.pydantic_v1 import BaseModel
+from pydantic import BaseModel
 from ..utils import tiktoken_length
 
 logger = getLogger(__name__)

alita_sdk/tools/gitlab/api_wrapper.py

@@ -115,6 +115,11 @@ class GitLabAPIWrapper(CodeIndexerToolkit):
         """Remove trailing slash from URL if present."""
         return url.rstrip('/') if url else url
 
+    @model_validator(mode='before')
+    @classmethod
+    def validate_toolkit_before(cls, values: Dict) -> Dict:
+        return super().validate_toolkit(values)
+
     @model_validator(mode='after')
     def validate_toolkit(self):
         try: