alita-sdk 0.3.447__tar.gz → 0.3.449__tar.gz

{alita_sdk-0.3.447 → alita_sdk-0.3.449}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.447
+Version: 0.3.449
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedj27@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0
@@ -24,6 +24,7 @@ Requires-Dist: fastapi==0.115.9
 Requires-Dist: httpcore==1.0.7
 Requires-Dist: urllib3>=2
 Requires-Dist: certifi==2024.8.30
+Requires-Dist: aiohttp>=3.9.0
 Provides-Extra: runtime
 Requires-Dist: streamlit>=1.28.0; extra == "runtime"
 Requires-Dist: langchain_core<0.4.0,>=0.3.76; extra == "runtime"

{alita_sdk-0.3.447 → alita_sdk-0.3.449}/alita_sdk/runtime/toolkits/mcp.py

@@ -7,6 +7,7 @@ Following MCP specification: https://modelcontextprotocol.io/specification/2025-
 import logging
 import re
 import requests
+import asyncio
 from typing import List, Optional, Any, Dict, Literal, ClassVar, Union
 
 from langchain_core.tools import BaseToolkit, BaseTool
@@ -17,6 +18,7 @@ from ..tools.mcp_remote_tool import McpRemoteTool
 from ..tools.mcp_inspect_tool import McpInspectTool
 from ...tools.utils import TOOLKIT_SPLITTER, clean_string
 from ..models.mcp_models import McpConnectionConfig
+from ..utils.mcp_sse_client import McpSseClient
 from ..utils.mcp_oauth import (
     McpAuthorizationRequired,
     canonical_resource,
@@ -402,12 +404,9 @@ class McpToolkit(BaseToolkit):
                 ]
 
             # Create BaseTool instances from discovered metadata
-            # Use session_id from discovery (which may have been created if server requires it)
-            # If discovery returned a session_id, it means a new session was created
-            # Log it so UI can capture and store it
+            # Use session_id from frontend (passed via connection_config)
             if session_id:
-                logger.info(f"[MCP Session] Using session from discovery: {session_id}")
-                logger.info(f"[MCP Session] Session created - server: {canonical_resource(connection_config.url)}, session_id: {session_id}")
+                logger.info(f"[MCP Session] Using session_id from frontend: {session_id}")
             
             for tool_metadata in tool_metadata_list:
                 server_tool = cls._create_tool_from_dict(
@@ -462,67 +461,76 @@ class McpToolkit(BaseToolkit):
         timeout: int
     ) -> List[Dict[str, Any]]:
         """
-        Synchronously discover tools and prompts from MCP server using HTTP requests.
+        Discover tools and prompts from MCP server using SSE client.
         Returns list of tool/prompt dictionaries with name, description, and inputSchema.
         Prompts are converted to tools that can be invoked.
         """
-        all_tools = []
-        
-        # Check if we have a session_id from UI (passed via connection_config)
-        # If not, we'll try discovery without session first, and create one if needed
         session_id = connection_config.session_id
         
-        # Discover regular tools - if server requires session and we don't have one, this will fail
-        # and we'll catch it below to initialize a session
+        if not session_id:
+            logger.warning(f"[MCP Session] No session_id provided for '{toolkit_name}' - server may require it")
+            logger.warning(f"[MCP Session] Frontend should generate a UUID and include it with mcp_tokens")
+        
+        # Run async discovery in sync context
         try:
-            tools_data = cls._discover_mcp_endpoint(
-                endpoint="tools/list",
-                toolkit_name=toolkit_name,
-                connection_config=connection_config,
-                timeout=timeout,
-                session_id=session_id
-            )
-        except Exception as discovery_error:
-            # Check if error is due to missing session
-            if "Missing sessionId" in str(discovery_error) and not session_id:
-                logger.info(f"[MCP Session] Server requires session for discovery, initializing...")
-                # Try to initialize session if we have OAuth headers
-                session_id = cls._initialize_mcp_session(
+            all_tools = asyncio.run(
+                cls._discover_tools_async(
                     toolkit_name=toolkit_name,
                     connection_config=connection_config,
-                    timeout=timeout,
-                    extra_headers=connection_config.headers  # Use headers from connection_config
+                    timeout=timeout
                 )
-                
-                if session_id:
-                    logger.info(f"[MCP Session] Created session for discovery: {session_id}")
-                    # Retry discovery with session
-                    tools_data = cls._discover_mcp_endpoint(
-                        endpoint="tools/list",
-                        toolkit_name=toolkit_name,
-                        connection_config=connection_config,
-                        timeout=timeout,
-                        session_id=session_id
-                    )
-                else:
-                    logger.error(f"[MCP Session] Failed to create session for discovery")
-                    raise discovery_error
-            else:
-                raise
-        all_tools.extend(tools_data)
-        logger.info(f"Discovered {len(tools_data)} tools from MCP toolkit '{toolkit_name}'")
+            )
+            return all_tools, session_id
+        except Exception as e:
+            logger.error(f"[MCP SSE] Discovery failed for '{toolkit_name}': {e}")
+            raise
+    
+    @classmethod
+    async def _discover_tools_async(
+        cls,
+        toolkit_name: str,
+        connection_config: McpConnectionConfig,
+        timeout: int
+    ) -> List[Dict[str, Any]]:
+        """
+        Async implementation of tool discovery using SSE client.
+        """
+        all_tools = []
+        session_id = connection_config.session_id
+        
+        if not session_id:
+            logger.error(f"[MCP SSE] session_id is required for SSE servers")
+            raise ValueError("session_id is required. Frontend must generate UUID.")
+        
+        logger.info(f"[MCP SSE] Discovering from {connection_config.url} with session {session_id}")
         
-        # Discover prompts and convert them to tools
+        # Prepare headers
+        headers = {}
+        if connection_config.headers:
+            headers.update(connection_config.headers)
+        
+        # Create SSE client
+        client = McpSseClient(
+            url=connection_config.url,
+            session_id=session_id,
+            headers=headers,
+            timeout=timeout
+        )
+        
+        # Initialize MCP session
+        await client.initialize()
+        logger.info(f"[MCP SSE] Session initialized for '{toolkit_name}'")
+        
+        # Discover tools
+        tools = await client.list_tools()
+        all_tools.extend(tools)
+        logger.info(f"[MCP SSE] Discovered {len(tools)} tools from '{toolkit_name}'")
+        
+        # Discover prompts
         try:
-            prompts_data = cls._discover_mcp_endpoint(
-                endpoint="prompts/list",
-                toolkit_name=toolkit_name,
-                connection_config=connection_config,
-                timeout=timeout,
-                session_id=session_id
-            )
+            prompts = await client.list_prompts()
             # Convert prompts to tool format
-            for prompt in prompts_data:
+            for prompt in prompts:
                 prompt_tool = {
                     "name": f"prompt_{prompt.get('name', 'unnamed')}",
                     "description": prompt.get('description', f"Execute prompt: {prompt.get('name')}"),
@@ -547,252 +555,12 @@ class McpToolkit(BaseToolkit):
                     "_mcp_prompt_name": prompt.get('name')
                 }
                 all_tools.append(prompt_tool)
-            logger.info(f"Discovered {len(prompts_data)} prompts from MCP toolkit '{toolkit_name}'")
+            logger.info(f"[MCP SSE] Discovered {len(prompts)} prompts from '{toolkit_name}'")
         except Exception as e:
-            logger.warning(f"Failed to discover prompts from MCP toolkit '{toolkit_name}': {e}")
+            logger.warning(f"[MCP SSE] Failed to discover prompts: {e}")
         
-        logger.info(f"Total discovered {len(all_tools)} tools+prompts from MCP toolkit '{toolkit_name}'")
-        return all_tools, session_id
-
-    @classmethod
-    def _initialize_mcp_session(
-        cls,
-        toolkit_name: str,
-        connection_config: McpConnectionConfig,
-        timeout: int,
-        extra_headers: Optional[Dict[str, str]] = None
-    ) -> Optional[str]:
-        """
-        Initialize an MCP session for stateful SSE servers.
-        Returns sessionId if successful, None if server doesn't require sessions.
-        
-        Note: This session should be stored by the caller (e.g., UI sessionStorage)
-        and passed back for subsequent requests. Sessions should NOT be cached
-        on the backend for security reasons (multi-tenant environment).
-        
-        Args:
-            toolkit_name: Name of the toolkit
-            connection_config: MCP connection configuration
-            timeout: Request timeout in seconds
-            extra_headers: Additional headers (e.g., Authorization) to include
-        """
-        import time
-
-        mcp_request = {
-            "jsonrpc": "2.0",
-            "id": f"initialize_{int(time.time())}",
-            "method": "initialize",
-            "params": {
-                "protocolVersion": "2024-11-05",
-                "capabilities": {
-                    "roots": {"listChanged": True},
-                    "sampling": {}
-                },
-                "clientInfo": {
-                    "name": "Alita MCP Client",
-                    "version": "1.0.0"
-                }
-            }
-        }
-
-        headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json, text/event-stream"
-        }
-        if connection_config.headers:
-            headers.update(connection_config.headers)
-        if extra_headers:
-            headers.update(extra_headers)
-            logger.debug(f"[MCP Session] Added extra headers: {list(extra_headers.keys())}")
-
-        try:
-            logger.info(f"[MCP Session] Attempting to initialize session for {connection_config.url}")
-            logger.debug(f"[MCP Session] Initialize request: {mcp_request}")
-            response = requests.post(
-                connection_config.url,
-                json=mcp_request,
-                headers=headers,
-                timeout=timeout
-            )
-
-            logger.info(f"[MCP Session] Initialize response status: {response.status_code}")
-            if response.status_code == 200:
-                # Parse the response to extract sessionId
-                content_type = response.headers.get('Content-Type', '')
-                logger.debug(f"[MCP Session] Response Content-Type: {content_type}")
-                
-                if 'text/event-stream' in content_type:
-                    data = cls._parse_sse_response(response.text)
-                elif 'application/json' in content_type:
-                    data = response.json()
-                else:
-                    logger.warning(f"[MCP Session] Unexpected Content-Type during initialize: {content_type}")
-                    logger.debug(f"[MCP Session] Response text: {response.text[:500]}")
-                    return None
-
-                logger.debug(f"[MCP Session] Parsed response: {data}")
-                
-                # Extract sessionId from response
-                result = data.get("result", {})
-                session_id = result.get("sessionId")
-                if session_id:
-                    logger.info(f"[MCP Session] ✓ Session initialized for '{toolkit_name}': {session_id}")
-                    logger.info(f"[MCP Session] This sessionId should be stored by the caller and sent back for subsequent requests")
-                    return session_id
-                else:
-                    logger.info(f"[MCP Session] No sessionId in initialize response for '{toolkit_name}' - server may not require sessions")
-                    logger.debug(f"[MCP Session] Full result: {result}")
-                    return None
-            else:
-                logger.warning(f"[MCP Session] Initialize returned {response.status_code} for '{toolkit_name}' - server may not support sessions")
-                logger.debug(f"[MCP Session] Response: {response.text[:500]}")
-                return None
-
-        except Exception as e:
-            logger.warning(f"[MCP Session] Failed to initialize MCP session for '{toolkit_name}': {e} - proceeding without session")
-            import traceback
-            logger.debug(f"[MCP Session] Traceback: {traceback.format_exc()}")
-            return None
-
-    @classmethod
-    def _discover_mcp_endpoint(
-        cls,
-        endpoint: str,
-        toolkit_name: str,
-        connection_config: McpConnectionConfig,
-        timeout: int,
-        session_id: Optional[str] = None
-    ) -> List[Dict[str, Any]]:
-        """
-        Discover items from a specific MCP endpoint (tools/list or prompts/list).
-        Returns list of dictionaries.
-        """
-        import time
-
-        # MCP protocol request
-        mcp_request = {
-            "jsonrpc": "2.0",
-            "id": f"discover_{endpoint.replace('/', '_')}_{int(time.time())}",
-            "method": endpoint,
-            "params": {}
-        }
-
-        headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json, text/event-stream"
-        }
-        if connection_config.headers:
-            headers.update(connection_config.headers)
-
-        # Add sessionId to URL if provided (for stateful SSE servers)
-        url = connection_config.url
-        if session_id:
-            separator = '&' if '?' in url else '?'
-            url = f"{url}{separator}sessionId={session_id}"
-            logger.info(f"[MCP Session] Using session {session_id} for {endpoint} request")
-        else:
-            logger.debug(f"[MCP Session] No session ID available for {endpoint} request - server may not require sessions")
-
-        try:
-            logger.debug(f"Sending MCP {endpoint} request to {url}")
-            response = requests.post(
-                url,
-                json=mcp_request,
-                headers=headers,
-                timeout=timeout
-            )
-
-            auth_header = response.headers.get('WWW-Authenticate', '')
-            if response.status_code == 401:
-                resource_metadata_url = extract_resource_metadata_url(auth_header, connection_config.url)
-                metadata = fetch_resource_metadata(resource_metadata_url, timeout=timeout) if resource_metadata_url else None
-                
-                # If we couldn't get metadata from the resource_metadata endpoint,
-                # infer authorization servers from the WWW-Authenticate header and server URL
-                if not metadata or not metadata.get('authorization_servers'):
-                    inferred_servers = infer_authorization_servers_from_realm(auth_header, connection_config.url)
-                    if inferred_servers:
-                        if not metadata:
-                            metadata = {}
-                        metadata['authorization_servers'] = inferred_servers
-                        logger.info(f"Inferred authorization servers for {connection_config.url}: {inferred_servers}")
-                        
-                        # Fetch OAuth authorization server metadata from the inferred server
-                        # This avoids CORS issues in the frontend
-                        from alita_sdk.runtime.utils.mcp_oauth import fetch_oauth_authorization_server_metadata
-                        auth_server_metadata = fetch_oauth_authorization_server_metadata(inferred_servers[0], timeout=timeout)
-                        if auth_server_metadata:
-                            metadata['oauth_authorization_server'] = auth_server_metadata
-                            logger.info(f"Fetched OAuth metadata for {inferred_servers[0]}")
-                
-                raise McpAuthorizationRequired(
-                    message=f"MCP server {connection_config.url} requires OAuth authorization",
-                    server_url=canonical_resource(connection_config.url),
-                    resource_metadata_url=resource_metadata_url,
-                    www_authenticate=auth_header,
-                    resource_metadata=metadata,
-                    status=response.status_code,
-                    tool_name=toolkit_name,
-                )
-
-            if response.status_code != 200:
-                logger.error(f"MCP server returned non-200 status: {response.status_code}")
-                raise Exception(f"HTTP {response.status_code}: {response.text}")
-
-            # Check content type and parse accordingly
-            content_type = response.headers.get('Content-Type', '')
-            
-            if 'text/event-stream' in content_type:
-                # Parse SSE (Server-Sent Events) format
-                data = cls._parse_sse_response(response.text)
-            elif 'application/json' in content_type:
-                # Parse regular JSON
-                try:
-                    data = response.json()
-                except ValueError as json_err:
-                    raise Exception(f"Invalid JSON response: {json_err}. Response text: {response.text[:200]}")
-            else:
-                raise Exception(f"Unexpected Content-Type: {content_type}. Response: {response.text[:200]}")
-
-            if "error" in data:
-                raise Exception(f"MCP Error: {data['error']}")
-
-            # Parse MCP response - different endpoints return different keys
-            result = data.get("result", {})
-            if endpoint == "tools/list":
-                return result.get("tools", [])
-            elif endpoint == "prompts/list":
-                return result.get("prompts", [])
-            else:
-                return result.get("items", [])
-
-        except Exception as e:
-            logger.error(f"Failed to discover from {endpoint} on MCP toolkit '{toolkit_name}': {e}")
-            raise
-
-    @staticmethod
-    def _parse_sse_response(sse_text: str) -> Dict[str, Any]:
-        """
-        Parse Server-Sent Events (SSE) format response.
-        SSE format: event: message\ndata: {json}\n\n
-        """
-        import json
-        
-        lines = sse_text.strip().split('\n')
-        data_line = None
-        
-        for line in lines:
-            if line.startswith('data:'):
-                data_line = line[5:].strip()  # Remove 'data:' prefix
-                break
-        
-        if not data_line:
-            raise Exception(f"No data found in SSE response: {sse_text[:200]}")
-        
-        try:
-            return json.loads(data_line)
-        except json.JSONDecodeError as e:
-            raise Exception(f"Failed to parse SSE data as JSON: {e}. Data: {data_line[:200]}")
+        logger.info(f"[MCP SSE] Total discovered {len(all_tools)} items from '{toolkit_name}'")
+        return all_tools
 
     @classmethod
     def _create_tool_from_dict(

alita_sdk-0.3.449/alita_sdk/runtime/tools/mcp_remote_tool.py

@@ -0,0 +1,166 @@
+"""
+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_sse_client import McpSseClient
+
+logger = logging.getLogger(__name__)
+
+
+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()
+    
+    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 __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."""
+        from ...tools.utils import TOOLKIT_SPLITTER
+        
+        # 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.rsplit(TOOLKIT_SPLITTER, 1)[-1] if TOOLKIT_SPLITTER in self.name else self.name
+            logger.warning(f"original_tool_name not set for '{self.name}', using extracted: {tool_name_for_server}")
+        
+        logger.info(f"[MCP SSE] 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 SSE client
+            client = McpSseClient(
+                url=self.server_url,
+                session_id=self.session_id,
+                headers=headers,
+                timeout=self.tool_timeout_sec
+            )
+            
+            # Execute tool call via SSE
+            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 SSE] 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 {}