agent-pard0x/tools_integration_demo.py

#!/usr/bin/env python3
"""
Example usage of SSH and Shell tools in both simple-react-agent and multi-agent approaches.
This file demonstrates how to integrate both local shell commands and remote SSH commands.
"""

import os
import warnings
from langchain.chat_models import init_chat_model
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage

# Suppress warnings
warnings.filterwarnings("ignore", message="The shell tool has no safeguards by default. Use at your own risk.")
warnings.filterwarnings("ignore", message="The SSH tool has no safeguards by default. Use at your own risk.")


def demo_simple_react_agent():
    """Demonstrate using both Shell and SSH tools in simple react agent."""
    from simple_react_agent.custom_tools import ShellTool, SSHTool, print_poem
    
    print("\n" + "="*60)
    print("🤖 Simple React Agent Demo - Shell & SSH Tools")
    print("="*60)
    
    # Initialize the chat model
    llm = init_chat_model("openai:gpt-4o-mini")
    
    # Create shell tool (local commands)
    shell_tool = ShellTool()
    
    # Create SSH tool (uncomment and configure for your server)
    # ssh_tool = SSHTool(
    #     host="your-server.com",
    #     username="your-username",
    #     key_filename="~/.ssh/id_rsa",  # or use password="your-password"
    #     ask_human_input=True  # Ask for confirmation before SSH commands
    # )
    
    # Define tools (SSH tool commented out for demo)
    tools = [shell_tool, print_poem]  # , ssh_tool]
    
    # Create agent with enhanced system prompt
    system_prompt = """You are an expert system administrator with access to both local and remote system tools.

    ## AVAILABLE TOOLS
    1. **Shell Tool (terminal)**: Execute commands on the local machine
    2. **SSH Tool (ssh)**: Execute commands on remote servers (when configured)
    3. **Poem Tool (print_poem)**: Generate motivational poems for debugging sessions

    ## SAFETY PROTOCOLS
    - Always start with read-only diagnostic commands
    - Explain what each command does before executing
    - Ask for confirmation before running potentially harmful commands
    - Use appropriate tools based on whether you need local or remote system access

    ## EXAMPLES
    - Use shell tool for local system diagnostics: `df -h`, `ps aux`, `top`
    - Use SSH tool for remote server management: connecting to production servers
    - Use poem tool when users need motivation during difficult troubleshooting

    Be helpful, safe, and efficient in your system administration assistance.
    """
    
    # Create the agent
    agent = create_react_agent(llm, tools, state_modifier=system_prompt)
    
    # Demo interaction
    print("\n📝 Demo Query: 'Check local disk usage and system load'")
    response = agent.invoke({
        "messages": [
            HumanMessage(content="Check local disk usage and system load")
        ]
    })
    
    print("\n🤖 Agent Response:")
    for message in response["messages"]:
        if hasattr(message, 'content') and message.content:
            print(f"  {message.content}")


def demo_multi_agent_with_tools():
    """Demonstrate integrating SSH and Shell tools with multi-agent supervisor."""
    from multi_agent_supervisor.custom_tools import ShellTool, SSHTool, print_poem
    from multi_agent_supervisor.agents import (
        create_os_detector_worker,
        create_logs_analyzer_worker,
        create_performance_analyzer_worker
    )
    
    print("\n" + "="*60)
    print("🤖 Multi-Agent Supervisor Demo - Enhanced with SSH")
    print("="*60)
    
    # Initialize the chat model
    llm = init_chat_model("openai:gpt-4o-mini")
    
    # Create tools
    shell_tool = ShellTool()
    
    # SSH tool configuration (commented for demo)
    # ssh_tool = SSHTool(
    #     host="production-server.com",
    #     username="admin",
    #     key_filename="~/.ssh/production_key",
    #     ask_human_input=True
    # )
    
    # Enhanced tools list
    enhanced_tools = [shell_tool, print_poem]  # , ssh_tool]
    
    # Create enhanced agents with both local and remote capabilities
    os_detector = create_os_detector_worker(
        llm=llm,
        tools=enhanced_tools,
        system_message_suffix="\nYou can use SSH tool to detect OS on remote servers when needed.",
    )
    
    logs_analyzer = create_logs_analyzer_worker(
        llm=llm,
        tools=enhanced_tools,
        system_message_suffix="\nYou can analyze logs on both local and remote systems using appropriate tools.",
    )
    
    performance_analyzer = create_performance_analyzer_worker(
        llm=llm,
        tools=enhanced_tools,
        system_message_suffix="\nYou can monitor performance on both local and remote systems.",
    )
    
    print("\n✅ Multi-agent system enhanced with SSH capabilities!")
    print("  - Agents can now work with both local and remote systems")
    print("  - SSH tool provides secure remote command execution")
    print("  - Shell tool handles local system operations")


def show_tool_usage_examples():
    """Show examples of how to use the tools directly."""
    print("\n" + "="*60)
    print("📚 Tool Usage Examples")
    print("="*60)
    
    print("\n🔧 Shell Tool Usage:")
    print("""
    from custom_tools import ShellTool
    
    shell_tool = ShellTool(ask_human_input=True)  # Ask for confirmation
    result = shell_tool.run(commands=["df -h", "free -m", "uptime"])
    """)
    
    print("\n🌐 SSH Tool Usage:")
    print("""
    from custom_tools import SSHTool
    
    # Using SSH key authentication
    ssh_tool = SSHTool(
        host="server.example.com",
        username="admin",
        key_filename="~/.ssh/id_rsa",
        ask_human_input=True
    )
    result = ssh_tool.run(commands=["df -h", "systemctl status nginx"])
    
    # Using password authentication
    ssh_tool = SSHTool(
        host="192.168.1.100",
        username="user",
        password="secure_password",
        port=22
    )
    result = ssh_tool.run(commands="ps aux | grep python")
    """)
    
    print("\n🎭 Poem Tool Usage:")
    print("""
    from custom_tools import print_poem
    
    result = print_poem("debugging")  # Theme: debugging
    """)


def main():
    """Main demonstration function."""
    print("🚀 SSH + Shell Tools Integration Demo")
    print("This demo shows how to use both local shell and remote SSH tools")
    
    # Check if OpenAI API key is set
    if not os.getenv("OPENAI_API_KEY"):
        print("\n⚠️  Warning: OPENAI_API_KEY not set. Some demos may not work.")
        print("   Set your API key: export OPENAI_API_KEY='your-key-here'")
    
    # Show usage examples
    show_tool_usage_examples()
    
    # Uncomment to run live demos (requires OpenAI API key)
    # demo_simple_react_agent()
    # demo_multi_agent_with_tools()
    
    print("\n" + "="*60)
    print("✅ Integration Complete!")
    print("="*60)
    print("\nBoth approaches now support:")
    print("  • Local shell command execution (ShellTool)")
    print("  • Remote SSH command execution (SSHTool)")
    print("  • Motivational poems (print_poem)")
    print("\nTo use SSH tools, uncomment and configure the SSH connection parameters.")
    print("Make sure to install paramiko: uv add paramiko")


if __name__ == "__main__":
    main()