alita-sdk 0.3.428.post2__py3-none-any.whl → 0.3.430__py3-none-any.whl

alita_sdk/runtime/tools/mcp_inspect_tool.py

@@ -0,0 +1,248 @@
+"""
+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:
+            # Run the async inspection
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # If we're in an async context, we need to create a new event loop
+                import concurrent.futures
+                with concurrent.futures.ThreadPoolExecutor() as executor:
+                    future = executor.submit(self._run_async_inspection, resource_type)
+                    return future.result(timeout=self.timeout)
+            else:
+                return loop.run_until_complete(self._run_async_inspection(resource_type))
+        except Exception as e:
+            logger.error(f"Error inspecting MCP server '{self.server_name}': {e}")
+            return f"Error inspecting MCP server: {e}"
+
+    def _run_async_inspection(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)
+
+    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", **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()}")
+
+            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", **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()}")
+
+            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", **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()}")
+
+            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,10 +15,60 @@ 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
 
+    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
     def create_pydantic_model_from_schema(schema: dict, model_name: str = "ArgsSchema"):

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,17 @@ def get_tools(tools_list, alita, llm, store: Optional[BaseStore] = None, *args,
                 logger.error(f"Error getting ADO repos tools: {e}")
             continue
 
+        # Handle MCP toolkit (located in runtime/toolkits)
+        if tool_type == 'mcp':
+            try:
+                from alita_sdk.runtime.toolkits.mcp import get_tools as mcp_get_tools
+                tools.extend(mcp_get_tools(tool, alita, llm))
+                logger.debug(f"Successfully loaded MCP tools")
+            except Exception as e:
+                logger.error(f"Error getting MCP tools: {e}")
+                raise ToolException(f"Error getting MCP tools: {e}")
+            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/figma/__init__.py

@@ -1,16 +1,10 @@
-from typing import Dict, List, Literal, Optional
+from typing import List, Literal, Optional
 
 from langchain_core.tools import BaseTool, BaseToolkit
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
 from ..base.tool import BaseAction
-from .api_wrapper import (
-    FigmaApiWrapper,
-    GLOBAL_LIMIT,
-    DEFAULT_FIGMA_IMAGES_PROMPT,
-    DEFAULT_FIGMA_SUMMARY_PROMPT,
-    DEFAULT_NUMBER_OF_THREADS,
-)
+from .api_wrapper import FigmaApiWrapper, GLOBAL_LIMIT
 from ..elitea_base import filter_missconfigured_index_tools
 from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
 from ...configurations.figma import FigmaConfiguration
@@ -34,14 +28,7 @@ def get_tools(tool):
             collection_name=str(tool['toolkit_name']),
             doctype='doc',
             embedding_model=tool['settings'].get('embedding_model'),
-            vectorstore_type="PGVector",
-            # figma summary/image prompt settings (toolkit-level)
-            # TODO disabled until new requirements
-            # apply_images_prompt=tool["settings"].get("apply_images_prompt"),
-            # images_prompt=tool["settings"].get("images_prompt"),
-            # apply_summary_prompt=tool["settings"].get("apply_summary_prompt"),
-            # summary_prompt=tool["settings"].get("summary_prompt"),
-            # number_of_threads=tool["settings"].get("number_of_threads"),
+            vectorstore_type="PGVector"
         )
         .get_tools()
     )
@@ -60,39 +47,6 @@ class FigmaToolkit(BaseToolkit):
         FigmaToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         return create_model(
             name,
-            # TODO disabled until new requirements
-            # apply_images_prompt=(Optional[bool], Field(
-            #     description="Enable advanced image processing instructions for Figma image nodes.",
-            #     default=True,
-            # )),
-            # images_prompt=(Optional[Dict[str, str]], Field(
-            #     description=(
-            #         "Instruction for how to analyze image-based nodes "
-            #         "(screenshots, diagrams, etc.) during Figma file retrieving. "
-            #         "Must contain a single 'prompt' key with the text."
-            #     ),
-            #     default=DEFAULT_FIGMA_IMAGES_PROMPT,
-            # )),
-            # apply_summary_prompt=(Optional[bool], Field(
-            #     description="Enable LLM-based summarization over loaded Figma data.",
-            #     default=True,
-            # )),
-            # summary_prompt=(Optional[Dict[str, str]], Field(
-            #     description=(
-            #         "Instruction for the LLM on how to summarize loaded Figma data. "
-            #         "Must contain a single 'prompt' key with the text."
-            #     ),
-            #     default=DEFAULT_FIGMA_SUMMARY_PROMPT,
-            # )),
-            number_of_threads=(Optional[int], Field(
-                description=(
-                    "Number of worker threads to use when downloading and processing Figma images. "
-                    f"Valid values are from 1 to 5. Default is {DEFAULT_NUMBER_OF_THREADS}."
-                ),
-                default=DEFAULT_NUMBER_OF_THREADS,
-                ge=1,
-                le=5,
-            )),
             global_limit=(Optional[int], Field(description="Global limit", default=GLOBAL_LIMIT)),
             global_regexp=(Optional[str], Field(description="Global regex pattern", default=None)),
             selected_tools=(