⚙️ Session 3 Advanced: Enterprise Agent Patterns¶

async def start_all_servers(self) -> Dict[str, bool]:
    """The startup orchestration - Bringing your digital workforce online."""
    results = {}

    for name, config in self.server_configs.items():
        result = await self._start_single_server(name, config)
        results[name] = result

    return results

async def _start_single_server(self, name: str, config: MCPServerConfig) -> bool:
    """Individual server startup with comprehensive validation."""
    try:
        logger.info(f"Starting MCP server: {name}")
        adapter = MCPAdapter(
            command=config.command,
            args=config.args,
            timeout=config.timeout
        )

        # The critical handshake - Test connection and discover tools
        await adapter.start()
        tools = await adapter.list_tools()

        # Success - Store adapter and update status
        self.adapters[name] = adapter
        self.health_status[name] = True

        logger.info(f"Server '{name}' started with {len(tools)} tools")
        return True

    except Exception as e:
        logger.error(f"Failed to start server '{name}': {e}")
        self.health_status[name] = False
        return False

async def get_adapter(self, server_name: str) -> Optional[MCPAdapter]:
    """Smart adapter access with automatic health checking."""
    if server_name not in self.adapters:
        logger.warning(f"Server '{server_name}' not found")
        return None

    if not self.health_status.get(server_name, False):
        logger.warning(f"Server '{server_name}' is unhealthy")
        # The resilience factor - Attempt automatic restart
        await self._restart_server(server_name)

    return self.adapters.get(server_name)

async def _restart_server(self, name: str):
    """Automatic recovery - Attempt to restart a failed server."""
    config = self.server_configs.get(name)
    if config:
        await self._start_single_server(name, config)

class MultiToolMCPAgent:
    """Advanced ReAct agent using multiple MCP servers - Your coordination engine."""

    def __init__(self, mcp_manager: MCPServerManager):
        self.mcp_manager = mcp_manager
        self.llm = None
        self.agent_executor = None
        self.memory = None
        self.available_tools = {}

    async def initialize(self) -> bool:
        """The comprehensive initialization - Building your multi-tool intelligence."""
        try:
            # Intelligence foundation - Initialize LLM
            self.llm = ChatOpenAI(
                model=Config.LLM.model,
                temperature=Config.LLM.temperature,
                api_key=Config.OPENAI_API_KEY
            )

            # Memory system - Initialize conversation memory
            self.memory = ConversationBufferWindowMemory(
                k=10,  # Remember last 10 exchanges
                memory_key="chat_history",
                return_messages=True
            )

            # Tool ecosystem - Collect tools from all available servers
            langchain_tools = await self._collect_all_tools()

            if not langchain_tools:
                logger.error("No tools available from MCP servers")
                return False

            # Agent creation - Create enhanced agent executor
            self.agent_executor = self._create_enhanced_agent(langchain_tools)

            tool_count = len(langchain_tools)
            server_count = len(self.available_tools)
            logger.info(f"Multi-tool agent: {tool_count} tools from {server_count} servers")

            return True

        except Exception as e:
            logger.error(f"Failed to initialize multi-tool agent: {e}")
            return False

async def _collect_all_tools(self) -> List[Tool]:
    """The tool aggregator - Collect tools from all available MCP servers."""
    all_tools = []

    # Discover all available servers
    for server_name in self.mcp_manager.server_configs.keys():
        try:
            adapter = await self.mcp_manager.get_adapter(server_name)
            if not adapter:
                logger.warning(f"Skipping unavailable server: {server_name}")
                continue

            # Get tools from this server
            mcp_tools = await adapter.list_tools()
            server_tools = self._create_langchain_tools(mcp_tools, adapter, server_name)

            # Track tools per server for management
            self.available_tools[server_name] = {
                "tool_count": len(server_tools),
                "tool_names": [tool.name for tool in server_tools],
                "adapter": adapter
            }

            all_tools.extend(server_tools)
            logger.info(f"Collected {len(server_tools)} tools from {server_name}")

        except Exception as e:
            logger.error(f"Failed to collect tools from {server_name}: {e}")
            continue

    return all_tools

def _create_langchain_tools(self, mcp_tools, adapter, server_name):
    """The enhanced translation layer - Converting MCP tools with comprehensive error handling."""
    langchain_tools = []

    for mcp_tool in mcp_tools:
        # Enhanced wrapper with server context
        async def tool_wrapper(tool_input: str, tool_name=mcp_tool.name, srv_name=server_name):
            """Enhanced wrapper with retry logic and error context."""
            max_retries = 3

            for attempt in range(max_retries):
                try:
                    # Health check before tool call
                    if not self.mcp_manager.health_status.get(srv_name, False):
                        await self.mcp_manager._restart_server(srv_name)

                    result = await adapter.call_tool(tool_name, {"input": tool_input})
                    return f"[{srv_name}] {str(result)}"

                except Exception as e:
                    if attempt < max_retries - 1:
                        logger.warning(f"Tool call failed (attempt {attempt + 1}), retrying: {e}")
                        await asyncio.sleep(1)  # Brief delay before retry
                        continue
                    else:
                        error_msg = f"Tool {tool_name} from {srv_name} failed after {max_retries} attempts: {str(e)}"
                        logger.error(error_msg)
                        return error_msg

        # Enhanced LangChain tool with server context
        langchain_tool = Tool(
            name=f"{server_name}_{mcp_tool.name}",
            description=f"[{server_name}] {mcp_tool.description or f'Tool from {server_name}'}",
            func=lambda x, tn=mcp_tool.name, sn=server_name: asyncio.create_task(tool_wrapper(x, tn, sn))
        )
        langchain_tools.append(langchain_tool)

    return langchain_tools

def _create_enhanced_agent(self, langchain_tools):
    """The intelligence amplifier - Create enhanced ReAct agent with sophisticated prompting."""
    react_prompt = PromptTemplate.from_template("""
You are an intelligent AI assistant with access to multiple specialized tools across different servers.
You can use weather information, file system operations, database queries, and other tools as available.

STRATEGIC APPROACH:
1. Analyze the user's request carefully to understand their true intent
2. Identify which tools might be helpful from your available toolkit
3. Use tools in logical sequence, building on previous results
4. If a tool from one server fails, try equivalent tools from other servers
5. Provide comprehensive, helpful responses that synthesize information from multiple sources

AVAILABLE TOOL SERVERS:
{tool_servers}

TOOL COORDINATION GUIDELINES:
- Tools are prefixed with server names (e.g., weather_get_current_weather)
- If a tool fails, the error message will indicate the server and issue
- You can use tools from multiple servers to cross-validate information
- Always explain your reasoning when switching between different tool servers

Available tools: {tools}

Use this format:
Question: {input}
Thought: Let me analyze what information I need and which tools can help
Action: [tool_name_with_server_prefix]
Action Input: [input_for_tool]
Observation: [result_from_tool]
... (repeat Thought/Action/Observation as needed)
Thought: I have gathered sufficient information from multiple sources
Final Answer: [comprehensive_response_synthesizing_all_information]

Conversation history: {chat_history}
Question: {input}
{agent_scratchpad}
""")

    # Generate server summary for prompt
    tool_servers = []
    for server_name, info in self.available_tools.items():
        tool_servers.append(f"- {server_name}: {info['tool_count']} tools ({', '.join(info['tool_names'])})")

    # Create enhanced agent
    agent = create_react_agent(
        llm=self.llm,
        tools=langchain_tools,
        prompt=react_prompt.partial(tool_servers="\n".join(tool_servers))
    )

    return AgentExecutor(
        agent=agent,
        tools=langchain_tools,
        memory=self.memory,
        verbose=Config.AGENT_CONFIG["verbose"],
        max_iterations=Config.AGENT_CONFIG["max_iterations"],
        handle_parsing_errors=True,
        return_intermediate_steps=True,
        max_execution_time=Config.AGENT_CONFIG.get("timeout", 300)
    )

async def run(self, query: str, context: Dict[str, Any] = None) -> Dict[str, Any]:
    """Enhanced agent execution with comprehensive monitoring and error handling."""
    if not self.agent_executor:
        return {
            "success": False,
            "error": "Agent not initialized. Call initialize() first.",
            "query": query
        }

    start_time = asyncio.get_event_loop().time()

    try:
        # Prepare enhanced input with context
        enhanced_input = {"input": query}
        if context:
            enhanced_input.update(context)

        # Execute with timeout and monitoring
        result = await asyncio.wait_for(
            self.agent_executor.ainvoke(enhanced_input),
            timeout=Config.AGENT_CONFIG.get("timeout", 300)
        )

        execution_time = asyncio.get_event_loop().time() - start_time

        return {
            "success": True,
            "query": query,
            "output": result["output"],
            "execution_time": execution_time,
            "steps": len(result.get("intermediate_steps", [])),
            "tools_used": self._extract_tools_used(result.get("intermediate_steps", []))
        }

    except asyncio.TimeoutError:
        execution_time = asyncio.get_event_loop().time() - start_time
        return {
            "success": False,
            "query": query,
            "error": f"Agent execution timed out after {execution_time:.2f} seconds",
            "execution_time": execution_time
        }

    except Exception as e:
        execution_time = asyncio.get_event_loop().time() - start_time
        logger.error(f"Agent execution failed: {e}")
        return {
            "success": False,
            "query": query,
            "error": f"Agent execution error: {str(e)}",
            "execution_time": execution_time
        }

def _extract_tools_used(self, intermediate_steps) -> List[str]:
    """Extract list of tools used during execution for monitoring."""
    tools_used = []
    for step in intermediate_steps:
        if hasattr(step, 'tool') and step.tool:
            tools_used.append(step.tool)
    return list(set(tools_used))  # Remove duplicates