rp/pr/tools/agents.py

import os
from typing import Any, Dict, List

from pr.agents.agent_manager import AgentManager
from pr.core.api import call_api


def create_agent(role_name: str, agent_id: str = None) -> Dict[str, Any]:
    """Create a new agent with the specified role."""
    try:
        # Get db_path from environment or default
        db_path = os.environ.get("ASSISTANT_DB_PATH", "~/.assistant_db.sqlite")
        db_path = os.path.expanduser(db_path)

        manager = AgentManager(db_path, call_api)
        agent_id = manager.create_agent(role_name, agent_id)
        return {"status": "success", "agent_id": agent_id, "role": role_name}
    except Exception as e:
        return {"status": "error", "error": str(e)}


def list_agents() -> Dict[str, Any]:
    """List all active agents."""
    try:
        db_path = os.path.expanduser("~/.assistant_db.sqlite")
        manager = AgentManager(db_path, call_api)
        agents = []
        for agent_id, agent in manager.active_agents.items():
            agents.append(
                {
                    "agent_id": agent_id,
                    "role": agent.role.name,
                    "task_count": agent.task_count,
                    "message_count": len(agent.message_history),
                }
            )
        return {"status": "success", "agents": agents}
    except Exception as e:
        return {"status": "error", "error": str(e)}


def execute_agent_task(agent_id: str, task: str, context: Dict[str, Any] = None) -> Dict[str, Any]:
    """Execute a task with the specified agent."""
    try:
        db_path = os.path.expanduser("~/.assistant_db.sqlite")
        manager = AgentManager(db_path, call_api)
        result = manager.execute_agent_task(agent_id, task, context)
        return result
    except Exception as e:
        return {"status": "error", "error": str(e)}


def remove_agent(agent_id: str) -> Dict[str, Any]:
    """Remove an agent."""
    try:
        db_path = os.path.expanduser("~/.assistant_db.sqlite")
        manager = AgentManager(db_path, call_api)
        success = manager.remove_agent(agent_id)
        return {"status": "success" if success else "not_found", "agent_id": agent_id}
    except Exception as e:
        return {"status": "error", "error": str(e)}


def collaborate_agents(orchestrator_id: str, task: str, agent_roles: List[str]) -> Dict[str, Any]:
    """Collaborate multiple agents on a task."""
    try:
        db_path = os.path.expanduser("~/.assistant_db.sqlite")
        manager = AgentManager(db_path, call_api)
        result = manager.collaborate_agents(orchestrator_id, task, agent_roles)
        return result
    except Exception as e:
        return {"status": "error", "error": str(e)}