daita-agents 0.1.0__py3-none-any.whl

daita/agents/substrate.py

@@ -0,0 +1,895 @@
+"""
+Substrate Agent - The foundational agent for Daita Agents.
+
+This agent provides a blank slate that users can build upon to create
+custom agents for any task, with simplified error handling and retry capabilities.
+All operations are automatically traced without any configuration required.
+"""
+import asyncio
+import logging
+import os
+from datetime import datetime
+from typing import Dict, Any, Optional, List, Union, Callable
+
+from ..config.base import AgentConfig, AgentType
+from ..core.interfaces import LLMProvider
+from ..core.exceptions import (
+    DaitaError, AgentError, LLMError, PluginError,
+    ValidationError, InvalidDataError, NotFoundError
+)
+from .base import BaseAgent
+
+logger = logging.getLogger(__name__)
+
+# Import unified plugin access
+from ..plugins import PluginAccess
+from ..llm.factory import create_llm_provider
+from ..config.settings import settings
+from ..core.tools import AgentTool, ToolRegistry
+
+class SubstrateAgent(BaseAgent):
+    """
+    The foundational agent with unified prompt system and handler-based task execution.
+
+    SubstrateAgent uses a two-tier prompt system for composable AI behavior:
+    - Agent-level prompt (WHO the agent is) - set at creation
+    - Call-level instructions (WHAT to do) - set per process() call
+
+    All operations are automatically traced without any configuration required.
+
+    Args:
+        name: Agent name
+        prompt: Agent identity and default behavior. Can be:
+            - str: Single prompt for all tasks
+            - dict: Task-specific prompts {"analyze": "...", "transform": "..."}
+        focus: Column/field filtering for data operations
+        handlers: Custom task handlers {task_name: handler_function}
+        tools: Plugin instances or AgentTool objects
+        llm_provider: LLM provider name or instance
+        model: Model name to use
+        api_key: API key for LLM provider (auto-detected if not provided)
+
+    Examples:
+        Basic agent:
+            >>> agent = SubstrateAgent(name="processor")
+            >>> result = await agent.process("analyze", data)
+
+        Agent with identity:
+            >>> agent = SubstrateAgent(
+            ...     name="validator",
+            ...     prompt="You are a data quality expert. Always return JSON."
+            ... )
+            >>> result = await agent.process("analyze", data)
+
+        Agent with call-specific instructions:
+            >>> agent = SubstrateAgent(prompt="You validate data quality.")
+            >>> result = await agent.process(
+            ...     "analyze",
+            ...     data,
+            ...     context={"instructions": "Focus on email validation"}
+            ... )
+
+        Task-specific prompts:
+            >>> agent = SubstrateAgent(
+            ...     prompt={
+            ...         "analyze": "You are a data analyst.",
+            ...         "transform": "You are a data transformer."
+            ...     }
+            ... )
+
+        Composition (agent identity + specific instructions):
+            >>> agent = SubstrateAgent(
+            ...     prompt="You are a data quality expert."
+            ... )
+            >>> result = await agent.process(
+            ...     "analyze",
+            ...     batch_data,
+            ...     context={"instructions": "This is legacy data - be careful"}
+            ... )
+            >>> # LLM receives both the agent prompt AND the specific instructions
+
+        Data operations example:
+            >>> validator = SubstrateAgent(
+            ...     name="ETL Validator",
+            ...     prompt="You validate data quality for ETL pipelines."
+            ... )
+            >>> result = await validator.process(
+            ...     "analyze",
+            ...     staging_data,
+            ...     context={"instructions": "Validate against production schema"}
+            ... )
+    """
+    
+    # Class-level defaults for smart constructor
+    _default_llm_provider = "openai"
+    _default_model = "gpt-4"
+
+    @classmethod
+    def configure_defaults(cls, **kwargs):
+        """Set global defaults for all SubstrateAgent instances."""
+        for key, value in kwargs.items():
+            setattr(cls, f'_default_{key}', value)
+
+    def __new__(cls, name=None, **kwargs):
+        """Smart constructor with auto-configuration."""
+        # Auto-configuration from environment and defaults
+        if not kwargs.get('llm_provider'):
+            kwargs['llm_provider'] = getattr(cls, '_default_llm_provider', 'openai')
+        if not kwargs.get('model'):
+            kwargs['model'] = getattr(cls, '_default_model', 'gpt-4')
+        if not kwargs.get('api_key'):
+            provider = kwargs.get('llm_provider', 'openai')
+            # Only try to get API key if provider is a string (not an object)
+            if isinstance(provider, str):
+                kwargs['api_key'] = settings.get_llm_api_key(provider)
+
+        return super().__new__(cls)
+    
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        llm_provider: Optional[Union[str, LLMProvider]] = None,
+        model: Optional[str] = None,
+        api_key: Optional[str] = None,
+        config: Optional[AgentConfig] = None,
+        agent_id: Optional[str] = None,
+        prompt: Optional[Union[str, Dict[str, str]]] = None,
+        focus: Optional[Union[List[str], str, Dict[str, Any]]] = None,
+        handlers: Optional[Dict[str, Callable]] = None,
+        relay: Optional[str] = None,
+        mcp: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
+        display_reasoning: bool = False,
+        **kwargs
+    ):
+        """
+        Initialize the Substrate Agent with smart constructor pattern.
+
+        This constructor auto-creates LLM providers and provides sensible
+        defaults while preserving all functionality.
+
+        Args:
+            name: Agent name (required for direct instantiation)
+            llm_provider: LLM provider name ("openai", "anthropic") or instance
+            model: Model name ("gpt-4", "claude-3-sonnet-20240229")
+            api_key: API key for LLM provider (auto-detected from env if not provided)
+            config: Agent configuration (auto-generated if not provided)
+            agent_id: Unique identifier for the agent
+            prompt: Custom prompt or prompt templates
+            focus: Default focus configuration for data processing
+            handlers: Dictionary of task handlers {task_name: handler_function}
+            relay: Name of relay channel for publishing results
+            mcp: MCP server(s) for tool integration - single dict or list of dicts
+            display_reasoning: Enable minimal decision display in console
+            **kwargs: Additional configuration options (can include 'tools' parameter)
+        """
+        # Auto-create LLM provider if needed
+        if isinstance(llm_provider, str) or llm_provider is None:
+            provider_name = llm_provider or self._default_llm_provider
+            model_name = model or self._default_model
+            api_key_to_use = api_key or settings.get_llm_api_key(provider_name)
+            
+            if api_key_to_use:
+                llm_provider = create_llm_provider(
+                    provider=provider_name,
+                    model=model_name,
+                    api_key=api_key_to_use,
+                    agent_id=agent_id
+                )
+            else:
+                logger.warning(f"No API key found for {provider_name}. LLM functionality will be disabled.")
+                llm_provider = None
+        # Create default config if none provided
+        if config is None:
+            config = AgentConfig(
+                name=name or "Substrate Agent",
+                type=AgentType.SUBSTRATE,
+                **kwargs
+            )
+        
+        # Initialize base agent (which handles automatic tracing)
+        super().__init__(config, llm_provider, agent_id, name)
+        
+        # Store customization options
+        self.prompt = prompt
+        self.default_focus = focus
+        self.relay = relay
+        
+        # Decision display setup
+        self.display_reasoning = display_reasoning
+        self._decision_display = None
+        
+        if display_reasoning:
+            self._setup_decision_display()
+
+        # Tool management (unified system)
+        self.tool_registry = ToolRegistry()
+        self.tool_sources = kwargs.get('tools', [])  # Plugins, AgentTool instances, or callables
+        self._tools_setup = False
+
+        # MCP server integration
+        self.mcp_registry = None
+        self.mcp_tools = []
+        if mcp is not None:
+            # Normalize to list
+            mcp_servers = [mcp] if isinstance(mcp, dict) else mcp
+            self._mcp_server_configs = mcp_servers
+            # MCP setup happens lazily on first use to avoid blocking init
+        else:
+            self._mcp_server_configs = []
+
+        # Task handlers - users can define custom behavior
+        self.handlers: Dict[str, Callable] = handlers or {}
+
+        # Built-in handlers for common operations
+        self._register_builtin_handlers()
+
+        # Plugin access for direct plugin usage
+        self.plugins = PluginAccess()
+
+        logger.debug(f"Substrate Agent {self.name} initialized with {len(self.handlers)} handlers")
+    
+    
+    def _setup_decision_display(self):
+        """Setup minimal decision display for local development."""
+        try:
+            from ..display.console import create_console_decision_display
+            from ..core.decision_tracing import register_agent_decision_stream
+
+            # Create display
+            self._decision_display = create_console_decision_display(
+                agent_name=self.name,
+                agent_id=self.agent_id
+            )
+
+            # Register with decision streaming system
+            register_agent_decision_stream(
+                agent_id=self.agent_id,
+                callback=self._decision_display.handle_event
+            )
+
+            logger.debug(f"Decision display enabled for agent {self.name}")
+
+        except Exception as e:
+            logger.warning(f"Failed to setup decision display: {e}")
+            self.display_reasoning = False
+            self._decision_display = None
+
+    async def _setup_mcp_tools(self):
+        """
+        Setup MCP servers and discover available tools.
+
+        This is called lazily on first agent.process() to avoid blocking
+        agent initialization with MCP server connections.
+        """
+        if self.mcp_registry is not None:
+            # Already setup
+            return
+
+        if not self._mcp_server_configs:
+            # No MCP servers configured
+            return
+
+        try:
+            from ..plugins.mcp import MCPServer, MCPToolRegistry
+
+            logger.info(f"Setting up {len(self._mcp_server_configs)} MCP server(s) for {self.name}")
+
+            # Create registry
+            self.mcp_registry = MCPToolRegistry()
+
+            # Connect to each server and register tools
+            for server_config in self._mcp_server_configs:
+                server = MCPServer(
+                    command=server_config.get("command"),
+                    args=server_config.get("args", []),
+                    env=server_config.get("env", {}),
+                    server_name=server_config.get("name")
+                )
+
+                # Add to registry (automatically connects and discovers tools)
+                await self.mcp_registry.add_server(server)
+
+            # Get all tools from registry
+            self.mcp_tools = self.mcp_registry.get_all_tools()
+
+            logger.info(f"MCP setup complete: {self.mcp_registry.tool_count} tools from {self.mcp_registry.server_count} server(s)")
+
+        except ImportError:
+            logger.error(
+                "MCP SDK not installed. Install with: pip install mcp\n"
+                "See: https://github.com/modelcontextprotocol/python-sdk"
+            )
+            raise
+
+        except Exception as e:
+            logger.error(f"Failed to setup MCP servers: {str(e)}")
+            raise
+
+    async def _setup_tools(self):
+        """
+        Discover and register tools from all sources.
+
+        Called lazily on first process() call to avoid blocking initialization.
+        Sources can be:
+        - Plugin instances with get_tools() method
+        - AgentTool instances directly
+        - MCP server configurations
+        """
+        if self._tools_setup:
+            return  # Already setup
+
+        # 1. Setup MCP tools first
+        if self._mcp_server_configs and self.mcp_registry is None:
+            await self._setup_mcp_tools()
+            # Convert MCP tools to AgentTool format
+            for mcp_tool in self.mcp_tools:
+                agent_tool = AgentTool.from_mcp_tool(mcp_tool, self.mcp_registry)
+                self.tool_registry.register(agent_tool)
+
+        # 2. Register plugin tools
+        for source in self.tool_sources:
+            if isinstance(source, AgentTool):
+                # Direct AgentTool registration
+                self.tool_registry.register(source)
+                logger.debug(f"Registered tool: {source.name}")
+
+            elif hasattr(source, 'get_tools'):
+                # Plugin with get_tools() method
+                plugin_tools = source.get_tools()
+                if plugin_tools:
+                    self.tool_registry.register_many(plugin_tools)
+                    logger.info(
+                        f"Registered {len(plugin_tools)} tools from "
+                        f"{source.__class__.__name__}"
+                    )
+            else:
+                logger.warning(
+                    f"Invalid tool source: {source}. "
+                    f"Expected AgentTool or plugin with get_tools() method."
+                )
+
+        self._tools_setup = True
+        logger.info(
+            f"Agent {self.name} initialized with {self.tool_registry.tool_count} tools"
+        )
+
+    async def _call_llm(
+        self,
+        task: str,
+        data: Any,
+        context: Dict[str, Any],
+        default_system_prompt: Optional[str] = None
+    ) -> str:
+        """
+        Centralized LLM gateway - ALL LLM calls go through here.
+
+        Composes: agent prompt + context instructions + data
+
+        Args:
+            task: Task name (for task-specific prompts)
+            data: Input data to process
+            context: Call-level context with optional 'instructions'
+            default_system_prompt: Fallback if no agent prompt configured
+
+        Returns:
+            LLM response string
+
+        Raises:
+            AgentError: If no LLM configured
+        """
+        if not self.llm:
+            raise AgentError(f"No LLM provider configured for task: {task}")
+
+        # 1. Build system message (agent identity/behavior)
+        system_prompt = None
+        if self.prompt:
+            if isinstance(self.prompt, dict):
+                # Task-specific prompt
+                system_prompt = self.prompt.get(task, self.prompt.get('default'))
+            else:
+                # Single prompt for all tasks
+                system_prompt = str(self.prompt)
+
+        # Use default if no agent prompt
+        if not system_prompt and default_system_prompt:
+            system_prompt = default_system_prompt
+
+        # 2. Build user message (instructions + data)
+        user_parts = []
+
+        # Add call-specific instructions
+        instructions = context.get('instructions')
+        if instructions:
+            user_parts.append(str(instructions))
+
+        # Add data
+        if data is not None:
+            # Smart data formatting
+            if isinstance(data, str):
+                data_str = data
+            elif isinstance(data, (list, dict)):
+                import json
+                try:
+                    data_str = json.dumps(data, indent=2, default=str)
+                except:
+                    data_str = str(data)
+            else:
+                data_str = str(data)
+
+            # Truncate very long data
+            if len(data_str) > 10000:
+                data_str = data_str[:10000] + "\n... (truncated)"
+
+            user_parts.append(f"Data:\n{data_str}")
+
+        user_message = "\n\n".join(user_parts) if user_parts else ""
+
+        # 3. Call LLM with proper structure
+        # Check if provider supports system messages (most do)
+        if hasattr(self.llm, 'generate_with_system'):
+            response = await self.llm.generate_with_system(
+                system=system_prompt,
+                user=user_message
+            )
+        else:
+            # Fallback: concatenate into single prompt
+            if system_prompt:
+                full_prompt = f"{system_prompt}\n\n{user_message}"
+            else:
+                full_prompt = user_message
+            response = await self.llm.generate(full_prompt)
+
+        return response
+
+    def _register_builtin_handlers(self):
+        """Register built-in handlers for common tasks."""
+        self.handlers.update({
+            'analyze': self._handle_analyze,
+            'transform': self._handle_transform,
+            'process_data': self._handle_process_data,
+            'custom': self._handle_custom,
+            'relay_message': self._handle_relay_message,
+            'llm_query': self._handle_llm_query
+        })
+    
+    async def _process_once(
+        self,
+        task: str,
+        data: Any,
+        context: Dict[str, Any],
+        attempt: int,
+        max_attempts: int
+    ) -> Dict[str, Any]:
+        """
+        Execute the task once using the SubstrateAgent's handler system.
+
+        This override provides the SubstrateAgent's specific task routing logic
+        while letting the parent handle retry and tracing automatically.
+        """
+        # Setup all tools (lazy initialization - includes MCP and plugin tools)
+        await self._setup_tools()
+
+        # Apply focus if specified
+        focus_to_use = context.get('focus') or self.default_focus
+        if focus_to_use and data is not None:
+            try:
+                from ..core.focus import apply_focus
+                data = apply_focus(data, focus_to_use)
+            except Exception as e:
+                logger.warning(f"Focus application failed: {str(e)}")
+                # Continue with original data if focus fails
+        
+        # Find the appropriate handler
+        handler = None
+        if task in self.handlers:
+            handler = self.handlers[task]
+        elif self.llm:
+            # Fall back to LLM query if no specific handler found
+            handler = self._handle_llm_query
+        else:
+            # No handler and no LLM available - use default behavior
+            return self._handle_default(data, context, self)
+        
+        # Execute the handler
+        try:
+            if asyncio.iscoroutinefunction(handler):
+                result = await handler(data, context, self)
+            else:
+                result = handler(data, context, self)
+            
+            # Ensure result is a dictionary
+            if not isinstance(result, dict):
+                result = {"handler_result": result}
+            
+            # Add metadata
+            result.update({
+                "task": task,
+                "agent_id": self.agent_id,
+                "agent_name": self.name,
+                "attempt": attempt,
+                "handler_used": handler.__name__ if hasattr(handler, '__name__') else str(handler),
+                "timestamp": datetime.utcnow().isoformat()
+            })
+            
+            # Publish to relay if configured and result is successful
+            if self.relay:
+                await self._publish_to_relay(result, context)
+            
+            return result
+            
+        except Exception as e:
+            logger.error(f"Handler execution failed in {self.name}: {str(e)}")
+            raise
+    
+    # Built-in Handlers
+    
+    async def _handle_llm_query(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
+        """Handle direct LLM queries using unified gateway."""
+        if not self.llm:
+            raise AgentError("No LLM provider configured for LLM query")
+
+        # Use gateway for consistent behavior
+        response = await self._call_llm(
+            task=context.get('task', 'llm_query'),
+            data=data,
+            context=context,
+            default_system_prompt="You are a helpful AI assistant. Process the user's request."
+        )
+
+        return {
+            "agent_response": response,
+            "llm_provider": self.llm.provider_name if hasattr(self.llm, 'provider_name') else 'unknown'
+        }
+
+    async def _llm_with_tools(self, prompt: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Handle LLM query with MCP tool calling support.
+
+        This implements the tool calling flow:
+        1. Send prompt with available tools to LLM
+        2. If LLM requests tool call, execute it via MCP
+        3. Return tool result to LLM for final response
+        4. Return final response to user
+        """
+        # Convert MCP tools to LLM function format
+        tool_definitions = [tool.to_llm_function() for tool in self.mcp_tools]
+
+        logger.debug(f"LLM query with {len(tool_definitions)} MCP tools available")
+
+        # For now, we'll do a simple implementation that appends tool info to the prompt
+        # Full function calling integration requires LLM provider updates
+        # This is MVP: tools are described, LLM can reference them in response
+
+        # Build tool descriptions for prompt
+        tool_descriptions = "\n\nAvailable tools:\n"
+        for tool in self.mcp_tools:
+            tool_descriptions += f"- {tool.name}: {tool.description}\n"
+
+        enhanced_prompt = f"{prompt}{tool_descriptions}"
+
+        response = await self.llm.generate(enhanced_prompt)
+
+        return {
+            "agent_response": response,
+            "prompt_used": enhanced_prompt,
+            "mcp_tools_available": [t.name for t in self.mcp_tools],
+            "llm_provider": self.llm.provider_name if hasattr(self.llm, 'provider_name') else 'unknown'
+        }
+
+    async def call_mcp_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Manually call an MCP tool.
+
+        This allows users to explicitly call MCP tools from their handlers
+        or for testing purposes.
+
+        Args:
+            tool_name: Name of the MCP tool to call
+            arguments: Arguments for the tool
+
+        Returns:
+            Tool execution result
+
+        Example:
+            ```python
+            result = await agent.call_mcp_tool("read_file", {"path": "/data/file.txt"})
+            ```
+        """
+        if not self.mcp_registry:
+            raise RuntimeError("No MCP servers configured. Add mcp parameter to SubstrateAgent.")
+
+        return await self.mcp_registry.call_tool(tool_name, arguments)
+    
+    async def _handle_analyze(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
+        """Handle analysis tasks with unified prompt system."""
+        try:
+            # Basic data type analysis (non-LLM)
+            analysis_result = {
+                "data_type": type(data).__name__,
+                "data_size": len(data) if hasattr(data, '__len__') else None,
+                "keys": list(data.keys()) if isinstance(data, dict) else None,
+                "sample": str(data)[:200] if data else None
+            }
+
+            # LLM analysis using gateway
+            if self.llm:
+                llm_analysis = await self._call_llm(
+                    task="analyze",
+                    data=data,
+                    context=context,
+                    default_system_prompt="You are a data analyst. Analyze the provided data and provide detailed insights."
+                )
+                analysis_result["llm_analysis"] = llm_analysis
+
+            return analysis_result
+
+        except Exception as e:
+            logger.error(f"Analysis failed: {str(e)}")
+            return {
+                "analysis_error": str(e),
+                "data_type": type(data).__name__,
+                "message": "Analysis completed with errors"
+            }
+    
+    async def _handle_transform(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
+        """Handle transformation tasks with LLM support."""
+        try:
+            # Use LLM if available
+            if self.llm:
+                transformed = await self._call_llm(
+                    task="transform",
+                    data=data,
+                    context=context,
+                    default_system_prompt="You are a data transformer. Transform the data according to instructions."
+                )
+
+                return {
+                    "transformed": transformed,
+                    "transformation_type": "llm_guided",
+                    "original_type": type(data).__name__
+                }
+
+            # Fallback: basic transformations if no LLM
+            if isinstance(data, dict):
+                transformed = {k: str(v).upper() if isinstance(v, str) else v for k, v in data.items()}
+                transformation_type = "dict_string_uppercase"
+            elif isinstance(data, list):
+                transformed = [str(item).title() if isinstance(item, str) else item for item in data]
+                transformation_type = "list_string_title"
+            elif isinstance(data, str):
+                transformed = data.title()
+                transformation_type = "string_title"
+            else:
+                transformed = str(data)
+                transformation_type = "to_string"
+
+            return {
+                "transformed": transformed,
+                "transformation_type": transformation_type,
+                "original_type": type(data).__name__
+            }
+
+        except Exception as e:
+            logger.error(f"Transformation failed: {str(e)}")
+            return {
+                "transformation_error": str(e),
+                "original_data": data,
+                "message": "Transformation completed with errors"
+            }
+    
+    async def _handle_process_data(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
+        """Handle generic data processing with optional LLM."""
+        result = {
+            'status': 'success',
+            'message': 'Data processed',
+            'input_type': type(data).__name__,
+            'processed_at': datetime.utcnow().isoformat(),
+            'task': 'process_data'
+        }
+
+        # Add LLM processing if available
+        if self.llm:
+            processed = await self._call_llm(
+                task="process_data",
+                data=data,
+                context=context,
+                default_system_prompt="Process this data as requested."
+            )
+            result['llm_processed'] = processed
+
+        return result
+    
+    async def _handle_custom(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
+        """Handle custom tasks."""
+        return {
+            'status': 'success',
+            'message': 'Custom task executed',
+            'data_received': data is not None,
+            'context_keys': list(context.keys()),
+            'task': 'custom'
+        }
+    
+    async def _handle_relay_message(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
+        """Handle incoming relay messages from other agents."""
+        return {
+            'status': 'success',
+            'message': 'Received relay message',
+            'source_agent': context.get('source_agent', 'unknown'),
+            'data_preview': str(data)[:200] if data else None,
+            'task': 'relay_message'
+        }
+    
+    def _handle_default(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
+        """Default handler when no specific handler is found."""
+        return {
+            'status': 'success',
+            'message': f'Default handler executed for task: {context.get("task", "unknown")}',
+            'data_type': type(data).__name__,
+            'data_received': data is not None,
+            'task': context.get('task', 'unknown')
+        }
+    
+    # User customization methods
+    
+    def add_handler(self, task: str, handler: Callable):
+        """Add a custom task handler."""
+        if not task:
+            raise ValidationError("Task name cannot be empty")
+        if not callable(handler):
+            raise ValidationError("Handler must be callable")
+        
+        self.handlers[task] = handler
+        logger.debug(f"Added handler for task: {task}")
+    
+    def remove_handler(self, task: str):
+        """Remove a task handler."""
+        if task in self.handlers:
+            del self.handlers[task]
+            logger.debug(f"Removed handler for task: {task}")
+    
+    def add_plugin(self, plugin: Any):
+        """
+        Add a plugin to the agent's tool sources.
+
+        The plugin's tools will be registered on next tool setup.
+        """
+        self.tool_sources.append(plugin)
+        logger.debug(f"Added plugin: {plugin.__class__.__name__}")
+
+    def register_tool(self, tool: AgentTool) -> None:
+        """
+        Register a single tool manually.
+
+        Useful for adding custom tools after agent initialization.
+
+        Args:
+            tool: AgentTool instance to register
+
+        Example:
+            ```python
+            agent = SubstrateAgent(name="my_agent")
+
+            custom_tool = AgentTool.from_function(my_custom_function)
+            agent.register_tool(custom_tool)
+            ```
+        """
+        self.tool_registry.register(tool)
+
+    def register_tools(self, tools: List[AgentTool]) -> None:
+        """
+        Register multiple tools manually.
+
+        Args:
+            tools: List of AgentTool instances
+        """
+        self.tool_registry.register_many(tools)
+
+    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Execute a tool by name.
+
+        Provides manual tool execution for testing or custom handlers.
+
+        Args:
+            name: Tool name
+            arguments: Tool arguments dict
+
+        Returns:
+            Tool execution result
+
+        Example:
+            ```python
+            result = await agent.call_tool("query_database", {"sql": "SELECT 1"})
+            ```
+        """
+        await self._setup_tools()
+        return await self.tool_registry.execute(name, arguments)
+
+    @property
+    def available_tools(self) -> List[AgentTool]:
+        """
+        Get list of all available tools.
+
+        Returns:
+            List of AgentTool instances
+        """
+        return self.tool_registry.tools.copy()
+
+    @property
+    def tool_names(self) -> List[str]:
+        """Get list of all tool names"""
+        return self.tool_registry.tool_names
+
+    async def stop(self) -> None:
+        """Stop agent and clean up all resources including MCP connections."""
+        # Clean up MCP connections first
+        if self.mcp_registry:
+            try:
+                await self.mcp_registry.disconnect_all()
+                logger.info(f"Cleaned up MCP connections for agent {self.name}")
+            except Exception as e:
+                logger.warning(f"Error cleaning up MCP connections: {e}")
+
+        # Call parent stop for standard cleanup
+        await super().stop()
+
+    def get_token_usage(self) -> Dict[str, int]:
+        """
+        Get token usage for this agent using automatic tracing.
+
+        Returns comprehensive token statistics from the unified tracing system.
+        """
+        if not self.llm or not hasattr(self.llm, 'get_token_stats'):
+            # Fallback for agents without LLM or tracing
+            return {
+                'total_tokens': 0,
+                'prompt_tokens': 0,
+                'completion_tokens': 0,
+                'requests': 0
+            }
+
+        return self.llm.get_token_stats()
+    
+    async def _publish_to_relay(self, result: Dict[str, Any], context: Dict[str, Any]):
+        """Publish result to relay channel."""
+        try:
+            from ..core.relay import publish
+            
+            await publish(
+                channel=self.relay,
+                agent_response=result,
+                publisher=self.name
+            )
+            logger.debug(f"Published result to relay channel: {self.relay}")
+        except Exception as e:
+            logger.warning(f"Failed to publish to relay channel {self.relay}: {str(e)}")
+            # Don't re-raise - relay failures shouldn't break main processing
+    
+    @property
+    def health(self) -> Dict[str, Any]:
+        """Enhanced health information for SubstrateAgent."""
+        base_health = super().health
+
+        # Add SubstrateAgent-specific health info
+        base_health.update({
+            'handlers': {
+                'count': len(self.handlers),
+                'builtin': len([h for h in self.handlers.keys() if not h.startswith('_')]),
+                'custom': len([h for h in self.handlers.keys() if h.startswith('_')])
+            },
+            'tools': {
+                'count': self.tool_registry.tool_count,
+                'setup': self._tools_setup,
+                'names': self.tool_registry.tool_names if self._tools_setup else []
+            },
+            'relay': {
+                'enabled': self.relay is not None,
+                'channel': self.relay
+            },
+            'llm': {
+                'available': self.llm is not None,
+                'provider': self.llm.provider_name if self.llm and hasattr(self.llm, 'provider_name') else None
+            }
+        })
+
+        return base_health