langcode_agent/main.py

#!/usr/bin/env python3
"""
MaxClaw Agent - A configurable agent harness using DeepAgents
Supports multiple LLM providers via LangChain's init_chat_model
Enhanced with custom tools for Proxmox and Ollama management
"""
import os
import sys
from typing import Optional, List
from deepagents import create_deep_agent
from langchain.chat_models import init_chat_model
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.tools import Tool

# Import custom tools
try:
    from tools import get_proxmox_tools, get_ollama_tools
    TOOLS_AVAILABLE = True
except ImportError:
    TOOLS_AVAILABLE = False
    print("Warning: Custom tools not available. Install required dependencies.")


def load_model_from_config():
    """
    Load the LLM model based on configuration.
    Configuration can come from environment variables or a config file.
    For simplicity, we'll use environment variables.
    
    Expected environment variables:
    - OPENAI_API_KEY: For OpenAI models
    - OLLAMA_BASE_URL: For Ollama (e.g., https://ollama.solarpi.net)
    - MODEL_PROVIDER: One of "openai", "ollama", "openrouter", etc.
    - MODEL_NAME: The specific model name (e.g., "gpt-4o", "gemma2:9b")
    """
    provider = os.getenv("MODEL_PROVIDER", "openai").lower()
    model_name = os.getenv("MODEL_NAME", "gpt-4o")
    
    # Configure based on provider
    if provider == "openai":
        # Uses OPENAI_API_KEY from environment
        return init_chat_model(f"openai:{model_name}")
    elif provider == "ollama":
        # For Ollama, we need to specify the base URL
        base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
        return init_chat_model(f"ollama:{model_name}", base_url=base_url)
    elif provider == "openrouter":
        # For OpenRouter, we need OPENROUTER_API_KEY and optionally base URL
        # Note: LangChain's init_chat_model for openrouter might require specific setup
        # As of now, we assume it's supported via openai-compatible API
        base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
        api_key = os.getenv("OPENROUTER_API_KEY")
        if not api_key:
            raise ValueError("OPENROUTER_API_KEY environment variable is required for OpenRouter")
        # Using openai-compatible API
        return init_chat_model(
            f"openai:{model_name}", 
            base_url=base_url, 
            api_key=api_key
        )
    else:
        # Fallback to openai for unknown providers (or add more as needed)
        return init_chat_model(f"openai:{model_name}")


def get_custom_tools() -> List[Tool]:
    """
    Get custom tools for the MaxClaw agent
    
    Returns:
        List of custom tools
    """
    custom_tools = []
    
    if not TOOLS_AVAILABLE:
        return custom_tools
    
    # Add Proxmox tools if configured
    proxmox_host = os.getenv("PROXMOX_HOST")
    if proxmox_host:
        proxmox_user = os.getenv("PROXMOX_USER", "root")
        proxmox_password = os.getenv("PROXMOX_PASSWORD")
        proxmox_token = os.getenv("PROXMOX_TOKEN")
        
        try:
            proxmox_tools = get_proxmox_tools(
                host=proxmox_host,
                user=proxmox_user,
                password=proxmox_password,
                api_token=proxmox_token
            )
            custom_tools.extend(proxmox_tools)
            print(f"✅ Added {len(proxmox_tools)} Proxmox tools")
        except Exception as e:
            print(f"⚠️  Could not load Proxmox tools: {e}")
    
    # Add Ollama tools if configured
    ollama_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
    try:
        ollama_tools = get_ollama_tools(base_url=ollama_url)
        custom_tools.extend(ollama_tools)
        print(f"✅ Added {len(ollama_tools)} Ollama tools")
    except Exception as e:
        print(f"⚠️  Could not load Ollama tools: {e}")
    
    return custom_tools


def create_maxclaw_agent(system_prompt: Optional[str] = None):
    """
    Create and return a DeepAgent instance with configurable model.
    """
    if system_prompt is None:
        system_prompt = """You are MaxClaw, a helpful AI agent that can plan, use tools, 
        manage context, and delegate to subagents. You are designed to be configurable 
        to work with various LLM providers.
        
        You have access to custom tools for managing Proxmox virtual machines and Ollama models.
        Use these tools when asked about system administration, virtualization, or AI model management."""
    
    model = load_model_from_config()
    
    # Get custom tools
    custom_tools = get_custom_tools()
    
    # Create the agent with custom tools
    agent = create_deep_agent(
        model=model,
        tools=custom_tools,  # Add custom tools to the agent
        system_prompt=system_prompt
    )
    
    return agent


def main():
    """
    Main entry point for the MaxClaw agent.
    """
    print("🚀 Initializing MaxClaw Agent...")
    
    try:
        agent = create_maxclaw_agent()
        print("✅ Agent initialized successfully!")
        
        # Simple test interaction
        print("\n💬 You can now interact with the agent. Type 'exit' to quit.")
        print("💡 Try asking about Proxmox VMs or Ollama models!")
        
        while True:
            user_input = input("\n👤 You: ").strip()
            if user_input.lower() in ['exit', 'quit']:
                print("👋 Goodbye!")
                break
            
            if not user_input:
                continue
                
            # Invoke the agent
            try:
                result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
                # Extract the AI's response
                ai_message = result["messages"][-1]
                print(f"\n🤖 MaxClaw: {ai_message.content}")
            except Exception as e:
                print(f"\n❌ Error: {str(e)}")
                
    except Exception as e:
        print(f"❌ Failed to initialize agent: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()