py/tools/agent_tools.py

"""
Agent tools and information management

This module contains:
- Agent metadata and configuration
- Agent singleton registry
- Functions to get available agents
- Agent switching tool creation
- Tool conversion utilities for OpenAI function format
"""

import inspect
import logging
from typing import Dict, List, Optional, Callable, Any, get_type_hints
from langchain.agents import Tool
from langchain.tools import StructuredTool
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


# Agent information dictionary - single source of truth for all agents
AGENT_INFO = {
    "document_reader": {
        "display_name": "Document Reader",
        "description": "Handles documentation queries with RAG capabilities",
        "tools": {
            "read": ["search_version_documentation", "search_all_versions", "list_available_versions"],
            "write": []  # No write operations
        }
    },
    "profile_settings": {
        "display_name": "Profile Settings",
        "description": "Manages user preferences and settings",
        "tools": {
            "read": ["view_profile", "list_preferences", "view_notification_settings"],
            "write": ["update_preference", "update_notifications", "reset_profile"]
        }
    }
}


# Global agent singleton registry
_AGENT_INSTANCES: Dict[str, Any] = {}
_LLM: Optional[ChatOpenAI] = None
_API_GATEWAY = None
_CURRENT_CONTEXT: Dict[str, Any] = {}  # Store context for switching tools


def initialize_agents(llm: ChatOpenAI, api_gateway=None) -> None:
    """
    Initialize all agent singletons
    
    Args:
        llm: The language model to use for all agents
        api_gateway: Optional API gateway for certain agents
    """
    global _AGENT_INSTANCES, _LLM, _API_GATEWAY
    
    # Avoid circular imports by importing here
    from agents.document_reader import DocumentReaderAgent
    from agents.profile_settings import ProfileSettingsAgent
    
    _LLM = llm
    _API_GATEWAY = api_gateway
    
    logger.info("Initializing agent singletons...")
    
    # Create agent instances (only implemented ones)
    _AGENT_INSTANCES = {
        "document_reader": DocumentReaderAgent(llm),
        "profile_settings": ProfileSettingsAgent(llm),
    }
    
    logger.info(f"Initialized {len(_AGENT_INSTANCES)} agents")


def get_agent(agent_id: str) -> Optional[Any]:
    """
    Get agent instance by ID
    
    Args:
        agent_id: The agent identifier
        
    Returns:
        Agent instance or None if not found
    """
    return _AGENT_INSTANCES.get(agent_id)


def run_agent(agent_id: str, message: str, context: Optional[Dict] = None) -> Dict[str, Any]:
    """
    Run an agent with a message
    
    Args:
        agent_id: The agent identifier
        message: The message to process
        context: Optional context dictionary
        
    Returns:
        Agent response dictionary
    """
    agent = get_agent(agent_id)
    if not agent:
        raise ValueError(f"Unknown agent: {agent_id}")
    
    # Store context globally for switching tools to access
    global _CURRENT_CONTEXT
    _CURRENT_CONTEXT = {
        "message": message,
        "context": context or {}
    }
    
    return agent.run(message, context)


def get_available_agents(exclude_agent_id: Optional[str] = None) -> Dict[str, Dict]:
    """
    Get all available agents, optionally excluding one
    
    Args:
        exclude_agent_id: Agent ID to exclude from results
        
    Returns:
        Dictionary of agent_id -> agent_info
    """
    return {
        agent_id: info 
        for agent_id, info in AGENT_INFO.items() 
        if agent_id != exclude_agent_id
    }


def create_agent_switching_tool(target_agent_id: str, target_agent_name: str) -> Callable:
    """
    Create a tool function for switching to a specific agent
    
    Args:
        target_agent_id: ID of the agent to switch to
        target_agent_name: Display name of the agent
        
    Returns:
        Callable function that performs the agent switch
    """
    def switch_to_agent(reason: str = "") -> str:
        """Switch conversation to another agent. 
        
        Args:
            reason: Brief explanation of why switching (e.g., 'User wants to update profile settings')
        """
        # Get current context
        global _CURRENT_CONTEXT
        message = _CURRENT_CONTEXT.get("message", "")
        context = _CURRENT_CONTEXT.get("context", {})
        
        logger.info(f"Switching from current agent to {target_agent_id}")
        
        # Log the transfer
        transfer_msg = f"Transferring to {target_agent_name}"
        if reason:
            transfer_msg += f": {reason}"
        
        try:
            # Run the target agent with the current message
            result = run_agent(target_agent_id, message, context)
            
            # Return a special marker with the target agent's response
            # This tells BaseAgent to update agent_id and use this response
            return f"__SWITCH_AGENT__|{target_agent_id}|{transfer_msg}\n\n{result.get('output', '')}"
        except Exception as e:
            logger.error(f"Error during agent switch: {e}")
            return f"__SWITCH_AGENT__|{target_agent_id}|{transfer_msg}\n\n[Error: Failed to run {target_agent_name} - {str(e)}]"
    
    # Set function metadata for LangChain
    switch_to_agent.__name__ = f"switch_to_{target_agent_id}"
    switch_to_agent.__doc__ = f"Transfer conversation to {target_agent_name} agent. Always provide a reason parameter explaining why you're switching (e.g., 'User wants to update profile settings')."
    
    return switch_to_agent


def create_switching_tools_for_agent(current_agent_id: str) -> List[StructuredTool]:
    """
    Create all agent switching tools for a given agent
    
    Args:
        current_agent_id: ID of the current agent
        
    Returns:
        List of StructuredTool objects for switching to other agents
    """
    tools = []
    other_agents = get_available_agents(exclude_agent_id=current_agent_id)
    
    for agent_id, agent_info in other_agents.items():
        # Create a dynamic Pydantic model for this specific agent
        class SwitchToAgentInput(BaseModel):
            reason: str = Field(
                default="", 
                description=f"Brief explanation of why switching to {agent_info['display_name']} (e.g., 'User wants to update profile settings')"
            )
        
        # Set the class name dynamically for better debugging
        SwitchToAgentInput.__name__ = f"SwitchTo{agent_id.title().replace('_', '')}Input"
        
        switch_func = create_agent_switching_tool(
            agent_id, 
            agent_info['display_name']
        )
        
        tool = StructuredTool.from_function(
            func=switch_func,
            name=f"switch_to_{agent_id}",
            description=f"Transfer to {agent_info['display_name']} - {agent_info['description']}",
            args_schema=SwitchToAgentInput
        )
        tools.append(tool)
    
    return tools


def get_agent_tools(agent_id: str) -> Dict[str, List[str]]:
    """
    Get the tools (read and write) for a specific agent
    
    Args:
        agent_id: ID of the agent
        
    Returns:
        Dictionary with 'read' and 'write' tool lists
    """
    if agent_id not in AGENT_INFO:
        raise ValueError(f"Unknown agent ID: {agent_id}")
    
    return AGENT_INFO[agent_id]["tools"]


def is_write_tool(agent_id: str, tool_name: str) -> bool:
    """
    Check if a tool is a write operation (requires confirmation)
    
    Args:
        agent_id: ID of the agent
        tool_name: Name of the tool
        
    Returns:
        True if the tool is a write operation
    """
    agent_tools = get_agent_tools(agent_id)
    return tool_name in agent_tools.get("write", [])


def convert_tool_to_openai_function(tool: Any) -> Dict[str, Any]:
    """
    Convert a LangChain Tool or StructuredTool to OpenAI function format.
    
    Handles both Tool and StructuredTool instances appropriately.
    """
    # Check if this is a StructuredTool with args_schema
    if hasattr(tool, 'args_schema') and tool.args_schema:
        # Use the Pydantic schema directly
        schema = tool.args_schema.schema()
        
        # Extract properties and required fields from Pydantic schema
        properties = {}
        required = []
        
        for field_name, field_info in schema.get('properties', {}).items():
            properties[field_name] = field_info
            
        # Get required fields from schema
        required = schema.get('required', [])
        
        return {
            "name": tool.name,
            "description": tool.description,
            "parameters": {
                "type": "object",
                "properties": properties,
                "required": required
            }
        }
    
    # Fall back to signature introspection for regular Tool
    func = tool.func
    sig = inspect.signature(func)
    type_hints = get_type_hints(func)
    
    # Build properties from function parameters
    properties = {}
    required = []
    
    for param_name, param in sig.parameters.items():
        # Skip 'self' parameter if it exists
        if param_name == 'self':
            continue
            
        # Determine the type
        param_type = type_hints.get(param_name, type(param.default) if param.default != param.empty else str)
        
        # Convert Python types to JSON Schema types
        if param_type == str:
            json_type = "string"
        elif param_type == int:
            json_type = "integer"
        elif param_type == float:
            json_type = "number"
        elif param_type == bool:
            json_type = "boolean"
        elif param_type == list or str(param_type).startswith('typing.List'):
            json_type = "array"
        elif param_type == dict or str(param_type).startswith('typing.Dict'):
            json_type = "object"
        else:
            json_type = "string"  # Default fallback
        
        # Build property definition
        prop_def = {"type": json_type}
        
        # Add description from docstring if available
        if func.__doc__:
            # Simple extraction - could be enhanced with docstring parsing
            prop_def["description"] = f"Parameter: {param_name}"
        
        # Add default value if present
        if param.default != param.empty:
            prop_def["default"] = param.default
        else:
            # If no default, it's required
            required.append(param_name)
        
        properties[param_name] = prop_def
    
    # Return the OpenAI function schema
    return {
        "name": tool.name,
        "description": tool.description,
        "parameters": {
            "type": "object",
            "properties": properties,
            "required": required
        }
    }