|
import asyncio
|
|
import os
|
|
from typing import Any, Dict, List
|
|
|
|
from pr.agents.agent_manager import AgentManager
|
|
from pr.core.api import call_api
|
|
from pr.config import DEFAULT_MODEL, DEFAULT_API_URL
|
|
from pr.tools.base import get_tools_definition
|
|
|
|
|
|
def _create_api_wrapper():
|
|
"""Create a wrapper function for call_api that matches AgentManager expectations."""
|
|
model = os.environ.get("AI_MODEL", DEFAULT_MODEL)
|
|
api_url = os.environ.get("API_URL", DEFAULT_API_URL)
|
|
api_key = os.environ.get("OPENROUTER_API_KEY", "")
|
|
use_tools = int(os.environ.get("USE_TOOLS", "0"))
|
|
tools_definition = get_tools_definition() if use_tools else []
|
|
|
|
def api_wrapper(messages, temperature=None, max_tokens=None, **kwargs):
|
|
try:
|
|
# Try to get the current event loop
|
|
loop = asyncio.get_running_loop()
|
|
# If we're in an event loop, use run_coroutine_threadsafe
|
|
import concurrent.futures
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
call_api(
|
|
messages=messages,
|
|
model=model,
|
|
api_url=api_url,
|
|
api_key=api_key,
|
|
use_tools=use_tools,
|
|
tools_definition=tools_definition,
|
|
verbose=False,
|
|
),
|
|
loop,
|
|
)
|
|
return future.result()
|
|
except RuntimeError:
|
|
# No event loop running, use asyncio.run
|
|
return asyncio.run(
|
|
call_api(
|
|
messages=messages,
|
|
model=model,
|
|
api_url=api_url,
|
|
api_key=api_key,
|
|
use_tools=use_tools,
|
|
tools_definition=tools_definition,
|
|
verbose=False,
|
|
)
|
|
)
|
|
|
|
return api_wrapper
|
|
|
|
|
|
def create_agent(role_name: str, agent_id: str = None) -> Dict[str, Any]:
|
|
"""Create a new agent with the specified role."""
|
|
try:
|
|
db_path = os.environ.get("ASSISTANT_DB_PATH", "~/.assistant_db.sqlite")
|
|
db_path = os.path.expanduser(db_path)
|
|
|
|
api_wrapper = _create_api_wrapper()
|
|
manager = AgentManager(db_path, api_wrapper)
|
|
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")
|
|
api_wrapper = _create_api_wrapper()
|
|
manager = AgentManager(db_path, api_wrapper)
|
|
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")
|
|
api_wrapper = _create_api_wrapper()
|
|
manager = AgentManager(db_path, api_wrapper)
|
|
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")
|
|
api_wrapper = _create_api_wrapper()
|
|
manager = AgentManager(db_path, api_wrapper)
|
|
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")
|
|
api_wrapper = _create_api_wrapper()
|
|
manager = AgentManager(db_path, api_wrapper)
|
|
result = manager.collaborate_agents(orchestrator_id, task, agent_roles)
|
|
return result
|
|
except Exception as e:
|
|
return {"status": "error", "error": str(e)}
|