65 lines
2.6 KiB
Python
65 lines
2.6 KiB
Python
|
import os
|
||
|
|
from typing import Dict, Any, 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)}
|