traia-iatp 0.1.1__py3-none-any.whl

traia_iatp/mcp/traia_mcp_adapter.py

@@ -0,0 +1,381 @@
+#!/usr/bin/env python
+"""
+Traia MCP Adapter
+
+A custom MCP adapter that extends CrewAI's MCPServerAdapter to support
+passing headers (like Authorization) to streamable-http MCP servers.
+
+This adapter transparently handles both authenticated and non-authenticated
+MCP connections by injecting headers when provided.
+"""
+
+import logging
+from typing import Dict, Any, List, Optional, Union
+import httpx
+from functools import wraps
+
+from crewai_tools import MCPServerAdapter
+from crewai.tools import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+class TraiaMCPAdapter(MCPServerAdapter):
+    """
+    MCP adapter that supports custom headers for streamable-http transport.
+    
+    This adapter extracts headers from server_params and injects them into
+    all HTTP requests made to the MCP server. It works transparently for
+    both authenticated and non-authenticated connections.
+    
+    Example:
+        ```python
+        # With authentication headers
+        server_params = {
+            "url": "https://mcp.example.com/mcp",
+            "transport": "streamable-http",
+            "headers": {"Authorization": "Bearer YOUR_API_KEY"}
+        }
+        
+        # Without headers (standard connection)
+        server_params = {
+            "url": "https://mcp.example.com/mcp",
+            "transport": "streamable-http"
+        }
+        
+        with TraiaMCPAdapter(server_params) as tools:
+            # Use tools with or without authentication
+            agent = Agent(tools=tools)
+        ```
+    """
+    
+    def __init__(
+        self,
+        server_params: Union[Dict[str, Any], Any],
+        *tool_names: str
+    ):
+        """
+        Initialize the adapter with optional header support.
+        
+        Args:
+            server_params: Server configuration. For streamable-http, can include:
+                          - url: The MCP server URL
+                          - transport: "streamable-http"
+                          - headers: Optional dict of headers to include in requests
+            *tool_names: Optional tool names to filter
+        """
+        # Handle dict params
+        if isinstance(server_params, dict):
+            # Extract headers if present (don't modify original)
+            server_params_copy = server_params.copy()
+            self._auth_headers = server_params_copy.pop("headers", {})
+            
+            if self._auth_headers:
+                logger.info(f"TraiaMCPAdapter: Headers configured: {list(self._auth_headers.keys())}")
+                # Apply monkey patch before parent initialization
+                self._apply_httpx_patch()
+            
+            # Pass clean params to parent
+            super().__init__(server_params_copy, *tool_names)
+        else:
+            # Non-dict params, no headers possible
+            self._auth_headers = {}
+            super().__init__(server_params, *tool_names)
+    
+    def _apply_httpx_patch(self):
+        """Monkey-patch httpx.AsyncClient to inject headers."""
+        original_init = httpx.AsyncClient.__init__
+        auth_headers = self._auth_headers
+        
+        @wraps(original_init)
+        def patched_init(client_self, *args, **kwargs):
+            # Get existing headers
+            existing_headers = kwargs.get('headers', {})
+            if isinstance(existing_headers, dict):
+                # Merge our headers
+                merged_headers = {**auth_headers, **existing_headers}
+                kwargs['headers'] = merged_headers
+                logger.debug(f"Injected headers into httpx.AsyncClient: {list(auth_headers.keys())}")
+            
+            # Also ensure headers are preserved on redirects
+            if 'follow_redirects' not in kwargs:
+                kwargs['follow_redirects'] = True
+            
+            # Call original init
+            original_init(client_self, *args, **kwargs)
+        
+        # Apply the patch
+        httpx.AsyncClient.__init__ = patched_init
+        self._original_httpx_init = original_init
+        logger.debug("Applied httpx.AsyncClient monkey patch")
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager and restore original httpx if patched."""
+        # Restore original httpx.AsyncClient.__init__ if we patched it
+        if hasattr(self, '_original_httpx_init'):
+            httpx.AsyncClient.__init__ = self._original_httpx_init
+            logger.debug("Restored original httpx.AsyncClient")
+        
+        # Call parent exit
+        return super().__exit__(exc_type, exc_val, exc_tb)
+    
+    def __enter__(self) -> List[BaseTool]:
+        """Enter context manager."""
+        if self._auth_headers:
+            logger.debug(f"TraiaMCPAdapter: Using authenticated connection")
+        
+        return super().__enter__()
+
+
+def create_mcp_adapter(
+    url: str,
+    headers: Optional[Dict[str, str]] = None,
+    transport: str = "streamable-http",
+    tool_names: Optional[List[str]] = None
+) -> TraiaMCPAdapter:
+    """
+    Create a Traia MCP adapter with optional headers.
+    
+    Args:
+        url: MCP server URL
+        headers: Optional dictionary of headers to include in requests
+        transport: Transport type (default: streamable-http)
+        tool_names: Optional list of tool names to filter
+    
+    Returns:
+        TraiaMCPAdapter configured with or without headers
+    
+    Example:
+        ```python
+        # With headers
+        adapter = create_mcp_adapter(
+            url="https://news-mcp.example.com/mcp",
+            headers={"Authorization": "Bearer YOUR_API_KEY"}
+        )
+        
+        # Without headers
+        adapter = create_mcp_adapter(
+            url="https://news-mcp.example.com/mcp"
+        )
+        
+        with adapter as tools:
+            agent = Agent(tools=tools)
+        ```
+    """
+    server_params = {
+        "url": url,
+        "transport": transport
+    }
+    
+    if headers:
+        server_params["headers"] = headers
+    
+    if tool_names:
+        return TraiaMCPAdapter(server_params, *tool_names)
+    else:
+        return TraiaMCPAdapter(server_params)
+
+
+def create_mcp_adapter_with_auth(
+    url: str,
+    api_key: str,
+    auth_header: str = "Authorization",
+    auth_prefix: str = "Bearer",
+    transport: str = "streamable-http",
+    tool_names: Optional[List[str]] = None
+) -> TraiaMCPAdapter:
+    """
+    Create a Traia MCP adapter with authentication.
+    
+    Args:
+        url: MCP server URL
+        api_key: API key for authentication
+        auth_header: Header name for auth (default: Authorization)
+        auth_prefix: Auth scheme prefix (default: Bearer)
+        transport: Transport type (default: streamable-http)
+        tool_names: Optional list of tool names to filter
+    
+    Returns:
+        TraiaMCPAdapter configured with auth headers
+    
+    Example:
+        ```python
+        adapter = create_mcp_adapter_with_auth(
+            url="https://news-mcp.example.com/mcp",
+            api_key="your-api-key"
+        )
+        
+        with adapter as tools:
+            # Tools will include Authorization header in all requests
+            agent = Agent(tools=tools)
+        ```
+    """
+    headers = {}
+    if auth_prefix:
+        headers[auth_header] = f"{auth_prefix} {api_key}"
+    else:
+        headers[auth_header] = api_key
+    
+    return create_mcp_adapter(url, headers, transport, tool_names)
+
+
+# Backwards compatibility aliases
+HeaderMCPAdapter = TraiaMCPAdapter
+create_mcp_adapter_with_headers = create_mcp_adapter
+
+
+# Usage examples
+USAGE_GUIDE = """
+TraiaMCPAdapter Usage Guide
+==========================
+
+The TraiaMCPAdapter seamlessly handles both authenticated and non-authenticated
+MCP connections. It extracts headers from server_params and injects them into
+all HTTP requests when using streamable-http transport.
+
+Basic Usage
+-----------
+```python
+from iatp.mcp import TraiaMCPAdapter
+
+# Standard connection (no authentication)
+server_params = {
+    "url": "http://localhost:8000/mcp",
+    "transport": "streamable-http"
+}
+
+# Authenticated connection
+server_params = {
+    "url": "http://localhost:8000/mcp",
+    "transport": "streamable-http",
+    "headers": {
+        "Authorization": "Bearer YOUR_API_KEY"
+    }
+}
+
+# Use the same way regardless of authentication
+with TraiaMCPAdapter(server_params) as tools:
+    agent = Agent(
+        role="My Agent",
+        tools=tools
+    )
+```
+
+Using Helper Functions
+----------------------
+```python
+from iatp.mcp import create_mcp_adapter_with_auth
+
+# Create authenticated adapter
+adapter = create_mcp_adapter_with_auth(
+    url="http://localhost:8000/mcp",
+    api_key="your-api-key"
+)
+
+with adapter as tools:
+    # Use authenticated tools
+    pass
+```
+
+Multiple Headers
+----------------
+```python
+adapter = create_mcp_adapter(
+    url="http://localhost:8000/mcp",
+    headers={
+        "Authorization": "Bearer YOUR_API_KEY",
+        "X-API-Version": "v1",
+        "X-Client-ID": "my-client"
+    }
+)
+```
+
+With Tool Filtering
+-------------------
+```python
+adapter = create_mcp_adapter_with_auth(
+    url="http://localhost:8000/mcp",
+    api_key="your-api-key",
+    tool_names=["search_news", "get_api_info"]
+)
+```
+
+Server-Side Authentication
+--------------------------
+For MCP servers that require authentication, implement FastMCP middleware:
+
+```python
+from fastmcp import FastMCP, Context
+from fastmcp.server.middleware import Middleware, MiddlewareContext
+from fastmcp.server.dependencies import get_http_request, get_context
+from starlette.requests import Request
+
+class AuthMiddleware(Middleware):
+    async def on_request(self, context: MiddlewareContext, call_next):
+        try:
+            # Access the raw HTTP request
+            request: Request = get_http_request()
+            
+            # Extract bearer token from Authorization header
+            auth = request.headers.get("Authorization", "")
+            token = auth[7:].strip() if auth.lower().startswith("bearer ") else None
+            
+            if not token:
+                # Check X-API-KEY header as alternative
+                token = request.headers.get("X-API-KEY", "")
+            
+            if token:
+                # Store the API key in the context state
+                if hasattr(context, 'state'):
+                    context.state.api_key = token
+                else:
+                    # Try to store it in the request state as fallback
+                    request.state.api_key = token
+        except Exception as e:
+            logger.debug(f"Could not extract API key from request: {e}")
+        
+        return await call_next(context)
+
+mcp = FastMCP("My Server", middleware=[AuthMiddleware()])
+
+def get_session_api_key(context: Context) -> Optional[str]:
+    '''Get the API key for the current session.'''
+    try:
+        # Try to get the API key from the context state
+        if hasattr(context, 'state') and hasattr(context.state, 'api_key'):
+            return context.state.api_key
+        
+        # Fallback: try to get it from the current HTTP request
+        try:
+            request: Request = get_http_request()
+            if hasattr(request.state, 'api_key'):
+                return request.state.api_key
+        except Exception:
+            pass
+        
+        # If we're in a tool context, try to get the context using the dependency
+        try:
+            ctx = get_context()
+            if hasattr(ctx, 'state') and hasattr(ctx.state, 'api_key'):
+                return ctx.state.api_key
+        except Exception:
+            pass
+    except Exception as e:
+        logger.debug(f"Could not retrieve API key from context: {e}")
+    
+    return None
+```
+
+Note: Headers are only applicable for streamable-http transport.
+For stdio or SSE transports, authentication must be handled differently.
+"""
+
+
+__all__ = [
+    'TraiaMCPAdapter',
+    'HeaderMCPAdapter',  # Alias for backward compatibility
+    'create_mcp_adapter',
+    'create_mcp_adapter_with_auth',
+    'create_mcp_adapter_with_headers',  # Alias
+    'USAGE_GUIDE'
+] 

traia_iatp/preview_diagrams.html

@@ -0,0 +1,181 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>IATP Mermaid Diagrams Preview</title>
+    <script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
+    <style>
+        body {
+            font-family: Arial, sans-serif;
+            max-width: 1200px;
+            margin: 0 auto;
+            padding: 20px;
+            background-color: #f5f5f5;
+        }
+        h1, h2 {
+            color: #333;
+        }
+        .diagram-container {
+            background-color: white;
+            border: 1px solid #ddd;
+            border-radius: 8px;
+            padding: 20px;
+            margin-bottom: 30px;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }
+        .mermaid {
+            text-align: center;
+        }
+    </style>
+</head>
+<body>
+    <h1>IATP (Inter-Agent Transfer Protocol) Diagrams</h1>
+    
+    <div class="diagram-container">
+        <h2>Standard Request-Response Flow</h2>
+        <div class="mermaid">
+sequenceDiagram
+    participant CrewAI as CrewAI Agent
+    participant Client as A2A Client
+    participant Server as A2A Server
+    participant MCP as MCP Server
+    
+    CrewAI->>Client: Execute request
+    Client->>Client: Ensure HTTP/2 connection
+    Client->>Server: Send message (JSON-RPC)
+    Server->>Server: Process request
+    Server->>MCP: Execute MCP tools
+    MCP-->>Server: Return results
+    Server-->>Client: Send response/task
+    
+    alt Task-based response
+        loop Poll for completion
+            Client->>Server: Get task status
+            Server-->>Client: Task state
+        end
+        Client->>Server: Get final result
+        Server-->>Client: Task messages/artifacts
+    else Direct response
+        Note over Client: Extract message content
+    end
+    
+    Client-->>CrewAI: Return formatted result
+        </div>
+    </div>
+
+    <div class="diagram-container">
+        <h2>SSE Streaming Flow</h2>
+        <div class="mermaid">
+sequenceDiagram
+    participant CrewAI as CrewAI Agent
+    participant Client as A2A Client
+    participant Server as A2A Server (HTTP/2)
+    participant MCP as MCP Server
+    
+    CrewAI->>Client: Stream request
+    Client->>Client: Check streaming support
+    Client->>Server: POST /a2a/stream (SSE)
+    Server->>Server: Initialize SSE stream
+    Server->>Client: stream_start event
+    
+    loop Streaming chunks
+        Server->>MCP: Get data chunk
+        MCP-->>Server: Data
+        Server->>Client: stream_chunk event
+        Client->>CrewAI: Yield chunk
+    end
+    
+    Server->>Client: stream_complete event
+    Client->>Client: Close stream
+    Client-->>CrewAI: Complete
+        </div>
+    </div>
+
+    <div class="diagram-container">
+        <h2>Client Decision Flow</h2>
+        <div class="mermaid">
+graph TD
+    subgraph "CrewAI Agent Request"
+        A[Agent calls tool._run<br/>with request]
+    end
+    
+    subgraph "A2ATool Decision Logic"
+        B{Check context}
+        C[_execute_with_retry]
+        D{streaming=True<br/>AND<br/>supports_streaming?}
+    end
+    
+    subgraph "Standard Mode"
+        E[_execute_standard]
+        F[Send JSON-RPC message]
+        G{Response type?}
+        H[Direct Message<br/>Extract text parts]
+        I[Task Response<br/>Poll with _wait_for_task_completion]
+        J[_get_task_status<br/>Loop until complete]
+    end
+    
+    subgraph "Streaming Mode"
+        K[_execute_streaming]
+        L[Open SSE connection<br/>via _stream_sse]
+        M[Collect all chunks]
+        N[Return combined result]
+    end
+    
+    subgraph "Async Iterator Mode"
+        O[stream method called]
+        P[Force streaming=True]
+        Q[_stream_sse context manager]
+        R[Yield chunks individually]
+    end
+    
+    A --> B
+    B --> C
+    C --> D
+    D -->|No| E
+    D -->|Yes| K
+    
+    E --> F
+    F --> G
+    G -->|Message| H
+    G -->|Task| I
+    I --> J
+    J -->|Loop| I
+    
+    K --> L
+    L --> M
+    M --> N
+    
+    O --> P
+    P --> Q
+    Q --> R
+    
+    H --> S[Return result]
+    J --> S
+    N --> S
+    R --> T[Async iteration]
+    
+    style A fill:#e1f5fe
+    style S fill:#c8e6c9
+    style T fill:#c8e6c9
+    style D fill:#fff3cd
+    style G fill:#fff3cd
+        </div>
+    </div>
+
+    <script>
+        mermaid.initialize({ 
+            startOnLoad: true,
+            theme: 'default',
+            themeVariables: {
+                primaryColor: '#f5f5f5',
+                primaryTextColor: '#333',
+                primaryBorderColor: '#7C0000',
+                lineColor: '#5D5D5D',
+                secondaryColor: '#006100',
+                tertiaryColor: '#fff'
+            }
+        });
+    </script>
+</body>
+</html> 

traia_iatp/registry/__init__.py

@@ -0,0 +1,26 @@
+"""IATP registry module for managing utility agents and MCP servers."""
+
+from .mongodb_registry import UtilityAgentRegistry, MCPServerRegistry
+from .iatp_search_api import (
+    find_utility_agent, 
+    list_utility_agents, 
+    search_utility_agents,
+    find_mcp_server,
+    list_mcp_servers,
+    search_mcp_servers,
+    get_mcp_server
+)
+from .embeddings import get_embedding_service
+
+__all__ = [
+    "UtilityAgentRegistry",
+    "MCPServerRegistry",
+    "find_utility_agent",
+    "list_utility_agents", 
+    "search_utility_agents",
+    "find_mcp_server",
+    "list_mcp_servers",
+    "search_mcp_servers", 
+    "get_mcp_server",
+    "get_embedding_service",
+]