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

alita_sdk/runtime/tools/loop.py

@@ -102,7 +102,9 @@ Input Data:
         logger.debug(f"LoopNode input: {predict_input}")
         completion = self.client.invoke(predict_input, config=config)
         logger.debug(f"LoopNode pure output: {completion}")
-        loop_data = _old_extract_json(completion.content.strip())
+        from ..langchain.utils import extract_text_from_completion
+        content_text = extract_text_from_completion(completion)
+        loop_data = _old_extract_json(content_text.strip())
         logger.debug(f"LoopNode output: {loop_data}")
         if self.return_type == "str":
             accumulated_response = ''

alita_sdk/runtime/tools/loop_output.py

@@ -93,7 +93,9 @@ Answer must be JSON only extractable by JSON.LOADS."""
         else:
             input_[-1].content += self.unstructured_output
             completion = self.client.invoke(input_, config=config)
-            result = _extract_json(completion.content.strip())
+            from ..langchain.utils import extract_text_from_completion
+            content_text = extract_text_from_completion(completion)
+            result = _extract_json(content_text.strip())
         try:
             tool_result: dict | List[dict] = self.tool.invoke(result, config=config, kwargs=kwargs)
             dispatch_custom_event(

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_remote_tool.py

@@ -0,0 +1,181 @@
+"""
+MCP Remote Tool for direct HTTP/SSE invocation.
+This tool is used for remote MCP servers accessed via HTTP/SSE.
+"""
+
+import asyncio
+import json
+import logging
+import time
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Optional
+
+from .mcp_server_tool import McpServerTool
+from pydantic import Field
+from ..utils.mcp_oauth import (
+    McpAuthorizationRequired,
+    canonical_resource,
+    extract_resource_metadata_url,
+    fetch_resource_metadata_async,
+    infer_authorization_servers_from_realm,
+)
+from ..utils.mcp_client import McpClient
+
+logger = logging.getLogger(__name__)
+
+# Global registry to store MCP tool session metadata by tool name
+# This is used to pass session info to callbacks since LangChain's serialization doesn't include all fields
+MCP_TOOL_SESSION_REGISTRY: Dict[str, Dict[str, Any]] = {}
+
+
+class McpRemoteTool(McpServerTool):
+    """
+    Tool for invoking remote MCP server tools via HTTP/SSE.
+    Extends McpServerTool and overrides _run to use direct HTTP calls instead of client.mcp_tool_call.
+    """
+    
+    # Remote MCP connection details
+    server_url: str = Field(..., description="URL of the remote MCP server")
+    server_headers: Optional[Dict[str, str]] = Field(default=None, description="HTTP headers for authentication")
+    original_tool_name: Optional[str] = Field(default=None, description="Original tool name from MCP server (before optimization)")
+    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
+    session_id: Optional[str] = Field(default=None, description="MCP session ID for stateful SSE servers")
+    
+    def model_post_init(self, __context: Any) -> None:
+        """Update metadata with session info after model initialization."""
+        super().model_post_init(__context)
+        self._update_metadata_with_session()
+        self._register_session_metadata()
+    
+    def _update_metadata_with_session(self):
+        """Update the metadata dict with current session information."""
+        if self.session_id:
+            if self.metadata is None:
+                self.metadata = {}
+            self.metadata.update({
+                'mcp_session_id': self.session_id,
+                'mcp_server_url': canonical_resource(self.server_url)
+            })
+    
+    def _register_session_metadata(self):
+        """Register session metadata in global registry for callback access."""
+        if self.session_id and self.server_url:
+            MCP_TOOL_SESSION_REGISTRY[self.name] = {
+                'mcp_session_id': self.session_id,
+                'mcp_server_url': canonical_resource(self.server_url)
+            }
+            logger.debug(f"[MCP] Registered session metadata for tool '{self.name}': session={self.session_id}")
+    
+    def __getstate__(self):
+        """Custom serialization for pickle compatibility."""
+        state = super().__getstate__()
+        # Ensure headers are serializable
+        if 'server_headers' in state and state['server_headers'] is not None:
+            state['server_headers'] = dict(state['server_headers'])
+        return state
+
+    def _run(self, *args, **kwargs):
+        """
+        Execute the MCP tool via direct HTTP/SSE call to the remote server.
+        Overrides the parent method to avoid using client.mcp_tool_call.
+        """
+        try:
+            # Always create a new event loop for sync context
+            with ThreadPoolExecutor() as executor:
+                future = executor.submit(self._run_in_new_loop, kwargs)
+                return future.result(timeout=self.tool_timeout_sec)
+        except McpAuthorizationRequired:
+            # Bubble up so LangChain can surface a tool error with useful metadata
+            raise
+        except Exception as e:
+            logger.error(f"Error executing remote MCP tool '{self.name}': {e}")
+            return f"Error executing tool: {e}"
+
+    def _run_in_new_loop(self, kwargs: Dict[str, Any]) -> str:
+        """Run the async tool invocation in a new event loop."""
+        return asyncio.run(self._execute_remote_tool(kwargs))
+
+    async def _execute_remote_tool(self, kwargs: Dict[str, Any]) -> str:
+        """Execute the actual remote MCP tool call using SSE client."""
+        
+        # Check for session_id requirement
+        if not self.session_id:
+            logger.error(f"[MCP Session] Missing session_id for tool '{self.name}'")
+            raise Exception("sessionId required. Frontend must generate UUID and send with mcp_tokens.")
+        
+        # Use the original tool name from discovery for MCP server invocation
+        tool_name_for_server = self.original_tool_name
+        if not tool_name_for_server:
+            tool_name_for_server = self.name
+            logger.warning(f"original_tool_name not set for '{self.name}', using: {tool_name_for_server}")
+        
+        logger.info(f"[MCP] Executing tool '{tool_name_for_server}' with session {self.session_id}")
+        
+        try:
+            # Prepare headers
+            headers = {}
+            if self.server_headers:
+                headers.update(self.server_headers)
+            
+            # Create unified MCP client (auto-detects transport)
+            client = McpClient(
+                url=self.server_url,
+                session_id=self.session_id,
+                headers=headers,
+                timeout=self.tool_timeout_sec
+            )
+            
+            # Execute tool call (client auto-detects SSE vs Streamable HTTP)
+            async with client:
+                await client.initialize()
+                result = await client.call_tool(tool_name_for_server, kwargs)
+            
+            # Format the result
+            if isinstance(result, dict):
+                # Check for content array (common in MCP responses)
+                if "content" in result:
+                    content_items = result["content"]
+                    if isinstance(content_items, list):
+                        # Extract text from content items
+                        text_parts = []
+                        for item in content_items:
+                            if isinstance(item, dict):
+                                if item.get("type") == "text" and "text" in item:
+                                    text_parts.append(item["text"])
+                                elif "text" in item:
+                                    text_parts.append(item["text"])
+                                else:
+                                    text_parts.append(json.dumps(item))
+                            else:
+                                text_parts.append(str(item))
+                        return "\n".join(text_parts)
+                
+                # Return formatted JSON if no content field
+                return json.dumps(result, indent=2)
+            
+            # Return as string for other types
+            return str(result)
+            
+        except Exception as e:
+            logger.error(f"[MCP] Tool execution failed: {e}", exc_info=True)
+            raise
+
+    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")
+    
+    def get_session_metadata(self) -> dict:
+        """Return session metadata to be included in tool responses."""
+        if self.session_id:
+            return {
+                'mcp_session_id': self.session_id,
+                'mcp_server_url': canonical_resource(self.server_url)
+            }
+        return {}

alita_sdk/runtime/tools/mcp_server_tool.py

@@ -1,11 +1,12 @@
 import uuid
 from logging import getLogger
-from typing import Any, Type, Literal, Optional, Union, List
+from typing import Any, Type, Literal, Optional, Union, List, Annotated
 
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field, create_model, EmailStr, constr
+from pydantic import BaseModel, Field, create_model, ConfigDict, StringConstraints
 
-from ...tools.utils import TOOLKIT_SPLITTER
+# EmailStr moved to pydantic_extra_types in pydantic v2, use str for simplicity
+EmailStr = str
 
 logger = getLogger(__name__)
 
@@ -19,6 +20,7 @@ class McpServerTool(BaseTool):
     server: str
     tool_timeout_sec: int = 60
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @staticmethod
     def create_pydantic_model_from_schema(schema: dict, model_name: str = "ArgsSchema"):
@@ -60,7 +62,7 @@ class McpServerTool(BaseTool):
                 if field.get("format") == "email":
                     return EmailStr
                 if "pattern" in field:
-                    return constr(regex=field["pattern"])
+                    return Annotated[str, StringConstraints(pattern=field["pattern"])]
                 return str
             if t == "integer":
                 return int
@@ -90,12 +92,13 @@ class McpServerTool(BaseTool):
         return create_model(model_name, **fields)
 
     def _run(self, *args, **kwargs):
+        # Use the tool name directly (no prefix extraction needed)
         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,
+                "name": self.name,
                 "arguments": kwargs
             }
         }

alita_sdk/runtime/tools/planning/__init__.py

@@ -0,0 +1,36 @@
+"""
+Planning tools for runtime agents.
+
+Provides plan management for multi-step task execution with progress tracking.
+Supports two storage backends:
+1. PostgreSQL - when connection_string is provided (production/indexer_worker)
+2. Filesystem - when no connection string (local CLI usage)
+"""
+
+from .wrapper import (
+    PlanningWrapper,
+    PlanStep,
+    PlanState,
+    FilesystemStorage,
+    PostgresStorage,
+)
+from .models import (
+    AgentPlan, 
+    PlanStatus, 
+    ensure_plan_tables,
+    delete_plan_by_conversation_id,
+    cleanup_on_graceful_completion
+)
+
+__all__ = [
+    "PlanningWrapper",
+    "PlanStep",
+    "PlanState",
+    "FilesystemStorage",
+    "PostgresStorage",
+    "AgentPlan",
+    "PlanStatus",
+    "ensure_plan_tables",
+    "delete_plan_by_conversation_id",
+    "cleanup_on_graceful_completion",
+]