traia-iatp 0.1.1__py3-none-any.whl

traia_iatp/server/templates/agent_executor.py.j2

@@ -0,0 +1,264 @@
+"""
+{{ agent_name }} - A2A Server Agent Executor Implementation
+
+This module implements the agent executor for {{ agent_name }}.
+Supports both synchronous responses and SSE streaming.
+"""
+
+import asyncio
+import json
+import logging
+import os
+from typing import AsyncGenerator, Optional, Dict, Any
+from datetime import datetime
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events.event_queue import EventQueue
+from a2a.types import Message, TextPart
+from a2a.utils import new_agent_text_message
+from crewai import Task
+from traia_iatp.mcp import MCPServerConfig, MCPAgentBuilder, run_with_mcp_tools, MCPServerInfo
+
+logger = logging.getLogger(__name__)
+
+# Create a thread pool for CPU-bound CrewAI operations
+executor = ThreadPoolExecutor(max_workers=10)
+
+# Get MCP server API key if required
+{% if requires_api_key %}
+# Check for API keys required by the MCP server
+MCP_API_KEY = None
+{% for api_key_name in api_keys %}
+if not MCP_API_KEY and os.getenv("{{ api_key_name }}"):
+    MCP_API_KEY = os.getenv("{{ api_key_name }}")
+    logger.info(f"Using API key from {{ api_key_name }} environment variable")
+{% endfor %}
+
+if not MCP_API_KEY:
+    logger.warning("No API key found for MCP server authentication.")
+    logger.warning("The MCP server requires one of these environment variables to be set:")
+    {% for api_key_name in api_keys %}
+    logger.warning("  - {{ api_key_name }}")
+    {% endfor %}
+else:
+    logger.info("MCP server API key loaded successfully")
+{% else %}
+MCP_API_KEY = None
+{% endif %}
+
+
+class CustomEvent:
+    """Custom event class for SSE streaming."""
+    def __init__(self, event_type: str, data: Dict[str, Any]):
+        self.type = event_type
+        self.data = data
+
+
+class {{ class_name }}AgentExecutor(AgentExecutor):
+    """Agent executor for {{ agent_name }}."""
+    
+    def __init__(self, mcp_config: MCPServerConfig, supports_streaming: bool = False):
+        self.mcp_config = mcp_config
+        self.supports_streaming = supports_streaming
+        self.mcp_server_info = MCPServerInfo(
+            id="",  # Not needed for direct usage
+            name=mcp_config.name,
+            url=mcp_config.url,
+            description=mcp_config.description,
+            server_type=mcp_config.server_type,
+            capabilities=mcp_config.capabilities,
+            metadata=mcp_config.metadata,
+            tags=mcp_config.metadata.get("tags", [])
+        )
+    
+    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Process a request using the {{ agent_name }} capabilities."""
+        try:
+            # Get the user's request from context
+            request_text = context.get_user_input()
+            if not request_text:
+                # Send empty response with task ID if available
+                msg = new_agent_text_message("No user message provided")
+                if hasattr(context, 'task_id') and context.task_id:
+                    msg.taskId = context.task_id
+                await event_queue.enqueue_event(msg)
+                return
+            
+            # Check if client requested streaming
+            stream_requested = False
+            if hasattr(context, 'configuration') and context.configuration:
+                output_mode = context.configuration.get('output_mode', '')
+                stream_requested = output_mode == 'text/event-stream'
+            
+            # Execute the request
+            if stream_requested and self.supports_streaming:
+                await self._execute_streaming(context, event_queue, request_text)
+            else:
+                await self._execute_standard(context, event_queue, request_text)
+                
+        except Exception as e:
+            logger.error(f"Error processing request: {e}")
+            msg = new_agent_text_message(f"Error processing request: {str(e)}")
+            if hasattr(context, 'task_id') and context.task_id:
+                msg.taskId = context.task_id
+            await event_queue.enqueue_event(msg)
+    
+    async def _execute_standard(self, context: RequestContext, event_queue: EventQueue, request_text: str) -> None:
+        """Execute standard (non-streaming) request."""
+        # Get additional context if provided
+        task_context = {}
+        if hasattr(context, 'metadata'):
+            task_context = context.metadata or {}
+        
+        # Create an agent for this request
+        agent = MCPAgentBuilder.create_agent(
+            role=f"{{ agent_name }} Specialist",
+            goal=f"Process the request using {self.mcp_config.name} capabilities",
+            backstory=f"You are an expert at using {self.mcp_config.name}. {self.mcp_config.description}"
+        )
+        
+        # Create a task
+        task = Task(
+            description=request_text,
+            expected_output="The processed result based on the request",
+            agent=agent
+        )
+        
+        # Create a wrapper function to handle the arguments properly
+        def run_crew_task():
+            # Build kwargs for run_with_mcp_tools
+            kwargs = {
+                "tasks": [task],
+                "mcp_server": self.mcp_server_info,
+                "inputs": task_context,
+                "skip_health_check": True
+            }
+            
+            # Only add api_key if MCP server requires authentication
+            if self.mcp_config.metadata.get("requires_api_key", False) and MCP_API_KEY:
+                kwargs["api_key"] = MCP_API_KEY
+                logger.debug("Including API key for authenticated MCP connection")
+            
+            return run_with_mcp_tools(**kwargs)
+        
+        # Run CrewAI in thread pool to avoid blocking the event loop
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(executor, run_crew_task)
+        
+        # Send the result as agent message with task ID if available
+        msg = new_agent_text_message(str(result))
+        if hasattr(context, 'task_id') and context.task_id:
+            msg.taskId = context.task_id
+        await event_queue.enqueue_event(msg)
+    
+    async def _execute_streaming(self, context: RequestContext, event_queue: EventQueue, request_text: str) -> None:
+        """Execute streaming request using SSE."""
+        try:
+            # Send initial event to indicate streaming has started
+            await event_queue.enqueue_event(
+                CustomEvent("stream_start", {"message": "Starting streaming response"})
+            )
+            
+            # Stream chunks as they become available
+            chunk_count = 0
+            async for chunk in self._stream_mcp_response(request_text, context):
+                # Send each chunk as a separate SSE event
+                await event_queue.enqueue_event(
+                    CustomEvent("stream_chunk", {
+                        "chunk_id": chunk_count,
+                        "content": chunk
+                    })
+                )
+                
+                # Also emit as message event for standard SSE subscribers
+                await event_queue.enqueue_event(
+                    CustomEvent("message", {
+                        "role": "agent",
+                        "content": chunk,
+                        "chunk_id": chunk_count
+                    })
+                )
+                
+                chunk_count += 1
+                
+                # Yield control to allow other tasks to run
+                await asyncio.sleep(0)
+            
+            # Send completion event
+            await event_queue.enqueue_event(
+                CustomEvent("stream_complete", {
+                    "total_chunks": chunk_count,
+                    "message": "Streaming completed successfully"
+                })
+            )
+            
+        except Exception as e:
+            logger.error(f"Error in streaming execution: {e}")
+            await event_queue.enqueue_event(
+                CustomEvent("stream_error", {
+                    "error": str(e),
+                    "message": "Streaming encountered an error"
+                })
+            )
+    
+    async def _stream_mcp_response(self, request_text: str, context: RequestContext) -> AsyncGenerator[str, None]:
+        """
+        Stream responses from MCP server.
+        This is a placeholder that should be implemented based on specific MCP server capabilities.
+        """
+        # For now, simulate streaming by breaking response into chunks
+        # In real implementation, this would connect to MCP server's streaming endpoint
+        
+        # Get full response first (in real implementation, this would be streamed)
+        task_context = {}
+        if hasattr(context, 'metadata'):
+            task_context = context.metadata or {}
+        
+        agent = MCPAgentBuilder.create_agent(
+            role=f"{{ agent_name }} Streaming Specialist",
+            goal=f"Process the streaming request using {self.mcp_config.name} capabilities",
+            backstory=f"You are an expert at using {self.mcp_config.name} for streaming data. {self.mcp_config.description}"
+        )
+        
+        task = Task(
+            description=request_text,
+            expected_output="The processed streaming result",
+            agent=agent
+        )
+        
+        # For demonstration, get the full result and stream it in chunks
+        def run_streaming_task():
+            # Build kwargs for run_with_mcp_tools
+            kwargs = {
+                "tasks": [task],
+                "mcp_server": self.mcp_server_info,
+                "inputs": task_context,
+                "skip_health_check": True
+            }
+            
+            # Only add api_key if MCP server requires authentication
+            if self.mcp_config.metadata.get("requires_api_key", False) and MCP_API_KEY:
+                kwargs["api_key"] = MCP_API_KEY
+                logger.debug("Including API key for authenticated MCP streaming connection")
+            
+            return run_with_mcp_tools(**kwargs)
+        
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(executor, run_streaming_task)
+        
+        # Simulate streaming by chunking the response
+        result_str = str(result)
+        chunk_size = 100  # Characters per chunk
+        
+        for i in range(0, len(result_str), chunk_size):
+            chunk = result_str[i:i + chunk_size]
+            yield chunk
+            await asyncio.sleep(0.1)  # Simulate network delay
+    
+    async def cancel(self, task_id: str) -> None:
+        """Cancel a running task."""
+        logger.info(f"Cancelling task: {task_id}")
+        # Implementation depends on MCP server capabilities
+        # For now, just log the cancellation request
+        pass 

traia_iatp/server/templates/docker-compose.yml.j2

@@ -0,0 +1,23 @@
+version: '3.8'
+
+services:
+  {{ agent_id }}:
+  build: .
+  container_name: {{ agent_id }}-local
+    ports:
+      - "${PORT:-8000}:${PORT:-8000}"
+    environment:
+      - PORT=${PORT:-8000}
+      - HOST=0.0.0.0
+      - LLM_MODEL=${LLM_MODEL:-openai/gpt-4.1}
+      {% for env_var in environment_variables %}
+      - {{ env_var.name }}={{ env_var.value }}
+      {% endfor %}
+    volumes:
+      - ./.env:/app/.env:ro
+    {% if mcp_server_url %}
+    # MCP Server configuration
+    extra_hosts:
+      - "host.docker.internal:host-gateway"  # For accessing local MCP servers
+    {% endif %}
+    restart: unless-stopped 

traia_iatp/server/templates/env.example.j2

@@ -0,0 +1,67 @@
+# {{ agent_name }} Environment Variables
+# Copy this file to .env and fill in your values
+
+# Server Configuration
+PORT=8000
+HOST=0.0.0.0
+
+# Language Model Configuration
+LLM_MODEL=openai/gpt-4.1
+# You can also use other models:
+# LLM_MODEL=vertex_ai/gemini-2.0-flash
+# LLM_MODEL=anthropic/claude-4-sonnetf
+
+# API Keys (add as needed based on your LLM choice)
+# OPENAI_API_KEY=your-openai-key-here
+# ANTHROPIC_API_KEY=your-anthropic-key-here
+# GOOGLE_API_KEY=your-google-key-here
+# VERTEX_PROJECT_ID=your-gcp-project-id
+# VERTEX_LOCATION=us-central1
+
+# MCP Server Configuration
+MCP_SERVER_URL={{ mcp_server_url }}
+{% if requires_api_key and api_keys %}
+# MCP Server API Keys (REQUIRED - The MCP server requires authentication)
+# Set one of the following environment variables:
+{% for api_key_name in api_keys %}
+{{ api_key_name }}=your-api-key-here
+{% endfor %}
+{% endif %}
+{% if mcp_server_metadata %}
+# Additional MCP configuration
+{% for key, value in mcp_server_metadata.items() %}
+{% if key.upper() == key %}
+{{ key }}={{ value }}
+{% endif %}
+{% endfor %}
+{% endif %}
+
+# Logging
+LOG_LEVEL=INFO
+
+# MongoDB Configuration (if using registry)
+# MONGODB_USER=your-mongodb-user
+# MONGODB_PASSWORD=your-mongodb-password
+# MONGODB_HOST=your-mongodb-host
+# MONGODB_PORT=27017
+# MONGODB_DATABASE=iatp
+
+{% if auth_required %}
+# Authentication Configuration
+{% for scheme in auth_schemes %}
+# Configure {{ scheme }} authentication
+{% if scheme == "bearer" %}
+# AUTH_TOKEN=your-bearer-token
+{% elif scheme == "api_key" %}
+# API_KEY=your-api-key
+{% elif scheme == "oauth2" %}
+# OAUTH2_CLIENT_ID=your-client-id
+# OAUTH2_CLIENT_SECRET=your-client-secret
+{% endif %}
+{% endfor %}
+{% endif %}
+
+# Additional custom environment variables
+{% for env_var in environment_variables %}
+# {{ env_var.name }}={{ env_var.value }}
+{% endfor %} 

traia_iatp/server/templates/gitignore.j2

@@ -0,0 +1,78 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+venv/
+.venv/
+*.egg-info/
+dist/
+build/
+.pytest_cache/
+.coverage
+htmlcov/
+.mypy_cache/
+.ruff_cache/
+
+# Virtual Environments
+bin/
+include/
+lib/
+lib64/
+pyvenv.cfg
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+.project
+.pydevproject
+
+# Environment Variables
+.env
+.env.*
+!.env.example
+
+# Logs
+*.log
+logs/
+
+# OS
+.DS_Store
+Thumbs.db
+desktop.ini
+
+# UV Package Manager
+.uv/
+uv.lock
+
+# Docker
+*.pid
+docker-compose.override.yml
+
+# {{ agent_name }} specific
+{% if additional_ignores %}
+# Custom ignores
+{% for ignore in additional_ignores %}
+{{ ignore }}
+{% endfor %}
+{% endif %}
+
+# Temporary files
+*.tmp
+*.temp
+*.bak
+.cache/
+
+# Database
+*.db
+*.sqlite
+*.sqlite3
+
+# Generated files
+generated/
+output/ 

traia_iatp/server/templates/grpc_server.py.j2

@@ -0,0 +1,218 @@
+"""
+{{ agency_name }} - gRPC Server Implementation
+
+This module provides a gRPC server for {{ agency_name }}.
+This is an optional component that can run alongside the HTTP/2 server for high-performance scenarios.
+"""
+
+import asyncio
+import logging
+import os
+import json
+from concurrent import futures
+from typing import AsyncIterator, Optional
+
+import grpc
+from grpc import aio
+from traia_iatp.mcp import MCPServerConfig
+from .agent_executor import {{ class_name }}AgentExecutor
+
+# Note: These would be generated from .proto files
+# from . import a2a_pb2, a2a_pb2_grpc
+
+logger = logging.getLogger(__name__)
+
+
+class {{ class_name }}GrpcService:  # (a2a_pb2_grpc.A2AServiceServicer):
+    """gRPC service implementation for {{ agency_name }}."""
+    
+    def __init__(self, mcp_config: MCPServerConfig):
+        self.mcp_config = mcp_config
+        self.executor = {{ class_name }}AgentExecutor(
+            mcp_config,
+            supports_streaming=True  # gRPC always supports streaming
+        )
+    
+    async def SendMessage(self, request, context):
+        """Handle unary RPC: single request, single response."""
+        try:
+            # Extract message content
+            message_content = request.message.content
+            
+            # Process with executor
+            result = await self.executor.process_unary_request(message_content)
+            
+            # Create response
+            # response = a2a_pb2.SendMessageResponse()
+            # response.result.content = result
+            # return response
+            
+            # Placeholder return
+            return {"result": {"content": result}}
+            
+        except Exception as e:
+            logger.error(f"Error in SendMessage: {e}")
+            await context.abort(
+                grpc.StatusCode.INTERNAL,
+                f"Internal error: {str(e)}"
+            )
+    
+    async def StreamMessage(self, request, context):
+        """Handle server streaming RPC: single request, stream of responses."""
+        try:
+            # Extract message content
+            message_content = request.message.content
+            
+            # Stream responses
+            async for chunk in self.executor.process_streaming_request(message_content):
+                # response = a2a_pb2.StreamMessageResponse()
+                # response.chunk.content = chunk
+                # yield response
+                
+                # Placeholder yield
+                yield {"chunk": {"content": chunk}}
+                
+        except Exception as e:
+            logger.error(f"Error in StreamMessage: {e}")
+            await context.abort(
+                grpc.StatusCode.INTERNAL,
+                f"Streaming error: {str(e)}"
+            )
+    
+    async def BidirectionalStream(self, request_iterator, context):
+        """Handle bidirectional streaming RPC."""
+        try:
+            async for request in request_iterator:
+                # Process each incoming request
+                message_content = request.message.content
+                
+                # Stream responses for this request
+                async for chunk in self.executor.process_streaming_request(message_content):
+                    # response = a2a_pb2.StreamMessageResponse()
+                    # response.chunk.content = chunk
+                    # yield response
+                    
+                    # Placeholder yield
+                    yield {"chunk": {"content": chunk}}
+                    
+        except Exception as e:
+            logger.error(f"Error in BidirectionalStream: {e}")
+            await context.abort(
+                grpc.StatusCode.INTERNAL,
+                f"Bidirectional streaming error: {str(e)}"
+            )
+    
+    async def GetCapabilities(self, request, context):
+        """Return agent capabilities."""
+        try:
+            # capabilities = a2a_pb2.CapabilitiesResponse()
+            # capabilities.name = "{{ agency_id }}"
+            # capabilities.description = "{{ agency_description }}"
+            # capabilities.version = "{{ agency_version }}"
+            # capabilities.supports_streaming = True
+            # capabilities.supports_bidirectional = True
+            # return capabilities
+            
+            # Placeholder return
+            return {
+                "name": "{{ agency_id }}",
+                "description": "{{ agency_description }}",
+                "version": "{{ agency_version }}",
+                "supports_streaming": True,
+                "supports_bidirectional": True,
+                "mcp_capabilities": self.mcp_config.capabilities
+            }
+            
+        except Exception as e:
+            logger.error(f"Error in GetCapabilities: {e}")
+            await context.abort(
+                grpc.StatusCode.INTERNAL,
+                f"Error getting capabilities: {str(e)}"
+            )
+
+
+async def serve_grpc():
+    """Start the gRPC server."""
+    # Load configuration
+    config_path = "agency_config.json"
+    if os.path.exists(config_path):
+        with open(config_path, "r") as f:
+            config_data = json.load(f)
+        mcp_data = config_data.get("mcp_server", {})
+    else:
+        mcp_data = {
+            "name": "{{ mcp_server_name }}",
+            "url": "{{ mcp_server_url }}",
+            "description": "{{ mcp_server_description }}",
+            "server_type": "{{ mcp_server_type }}",
+            "capabilities": {{ mcp_server_capabilities | tojson }},
+            "metadata": {{ mcp_server_metadata | tojson }}
+        }
+    
+    # Create MCP config
+    mcp_config = MCPServerConfig(
+        name=mcp_data["name"],
+        url=mcp_data["url"],
+        description=mcp_data["description"],
+        server_type=mcp_data.get("server_type", "streamable-http"),
+        capabilities=mcp_data.get("capabilities", []),
+        metadata=mcp_data.get("metadata", {})
+    )
+    
+    # Create gRPC server
+    server = aio.server(
+        futures.ThreadPoolExecutor(max_workers=10),
+        options=[
+            ('grpc.max_send_message_length', 4 * 1024 * 1024),  # 4MB
+            ('grpc.max_receive_message_length', 4 * 1024 * 1024),  # 4MB
+            ('grpc.keepalive_time_ms', 10000),
+            ('grpc.keepalive_timeout_ms', 5000),
+            ('grpc.http2.max_concurrent_streams', 100),
+            ('grpc.http2.max_frame_size', 16384),
+        ]
+    )
+    
+    # Add service
+    service = {{ class_name }}GrpcService(mcp_config)
+    # a2a_pb2_grpc.add_A2AServiceServicer_to_server(service, server)
+    
+    # Enable reflection for debugging
+    # from grpc_reflection.v1alpha import reflection
+    # SERVICE_NAMES = (
+    #     a2a_pb2.DESCRIPTOR.services_by_name['A2AService'].full_name,
+    #     reflection.SERVICE_NAME,
+    # )
+    # reflection.enable_server_reflection(SERVICE_NAMES, server)
+    
+    # Get port from environment
+    port = int(os.environ.get("GRPC_PORT", 50051))
+    
+    # Add insecure port (TLS can be added later)
+    server.add_insecure_port(f'[::]:{port}')
+    
+    logger.info(f"Starting {{ agency_name }} gRPC Server")
+    logger.info(f"Listening on port {port}")
+    logger.info(f"MCP Server: {{ mcp_server_name }}")
+    
+    # Start server
+    await server.start()
+    
+    try:
+        await server.wait_for_termination()
+    except KeyboardInterrupt:
+        logger.info("Shutting down gRPC server...")
+        await server.stop(5)
+
+
+def main():
+    """Main entry point for gRPC server."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    )
+    
+    asyncio.run(serve_grpc())
+
+
+if __name__ == "__main__":
+    main()