|
import time
|
|
import json
|
|
import uuid
|
|
from typing import Dict, List, Any, Optional, Callable
|
|
from dataclasses import dataclass, field
|
|
from .agent_roles import AgentRole, get_agent_role
|
|
from .agent_communication import AgentMessage, AgentCommunicationBus, MessageType
|
|
|
|
@dataclass
|
|
class AgentInstance:
|
|
agent_id: str
|
|
role: AgentRole
|
|
message_history: List[Dict[str, Any]] = field(default_factory=list)
|
|
context: Dict[str, Any] = field(default_factory=dict)
|
|
created_at: float = field(default_factory=time.time)
|
|
task_count: int = 0
|
|
|
|
def add_message(self, role: str, content: str):
|
|
self.message_history.append({
|
|
'role': role,
|
|
'content': content,
|
|
'timestamp': time.time()
|
|
})
|
|
|
|
def get_system_message(self) -> Dict[str, str]:
|
|
return {'role': 'system', 'content': self.role.system_prompt}
|
|
|
|
def get_messages_for_api(self) -> List[Dict[str, str]]:
|
|
return [self.get_system_message()] + [
|
|
{'role': msg['role'], 'content': msg['content']}
|
|
for msg in self.message_history
|
|
]
|
|
|
|
class AgentManager:
|
|
def __init__(self, db_path: str, api_caller: Callable):
|
|
self.db_path = db_path
|
|
self.api_caller = api_caller
|
|
self.communication_bus = AgentCommunicationBus(db_path)
|
|
self.active_agents: Dict[str, AgentInstance] = {}
|
|
self.session_id = str(uuid.uuid4())[:16]
|
|
|
|
def create_agent(self, role_name: str, agent_id: Optional[str] = None) -> str:
|
|
if agent_id is None:
|
|
agent_id = f"{role_name}_{str(uuid.uuid4())[:8]}"
|
|
|
|
role = get_agent_role(role_name)
|
|
agent = AgentInstance(
|
|
agent_id=agent_id,
|
|
role=role
|
|
)
|
|
|
|
self.active_agents[agent_id] = agent
|
|
return agent_id
|
|
|
|
def get_agent(self, agent_id: str) -> Optional[AgentInstance]:
|
|
return self.active_agents.get(agent_id)
|
|
|
|
def remove_agent(self, agent_id: str) -> bool:
|
|
if agent_id in self.active_agents:
|
|
del self.active_agents[agent_id]
|
|
return True
|
|
return False
|
|
|
|
def execute_agent_task(self, agent_id: str, task: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
|
agent = self.get_agent(agent_id)
|
|
if not agent:
|
|
return {'error': f'Agent {agent_id} not found'}
|
|
|
|
if context:
|
|
agent.context.update(context)
|
|
|
|
agent.add_message('user', task)
|
|
agent.task_count += 1
|
|
|
|
messages = agent.get_messages_for_api()
|
|
|
|
try:
|
|
response = self.api_caller(
|
|
messages=messages,
|
|
temperature=agent.role.temperature,
|
|
max_tokens=agent.role.max_tokens
|
|
)
|
|
|
|
if response and 'choices' in response:
|
|
assistant_message = response['choices'][0]['message']['content']
|
|
agent.add_message('assistant', assistant_message)
|
|
|
|
return {
|
|
'success': True,
|
|
'agent_id': agent_id,
|
|
'response': assistant_message,
|
|
'role': agent.role.name,
|
|
'task_count': agent.task_count
|
|
}
|
|
else:
|
|
return {'error': 'Invalid API response', 'agent_id': agent_id}
|
|
|
|
except Exception as e:
|
|
return {'error': str(e), 'agent_id': agent_id}
|
|
|
|
def send_agent_message(self, from_agent_id: str, to_agent_id: str,
|
|
content: str, message_type: MessageType = MessageType.REQUEST,
|
|
metadata: Optional[Dict[str, Any]] = None):
|
|
message = AgentMessage(
|
|
from_agent=from_agent_id,
|
|
to_agent=to_agent_id,
|
|
message_type=message_type,
|
|
content=content,
|
|
metadata=metadata or {},
|
|
timestamp=time.time(),
|
|
message_id=str(uuid.uuid4())[:16]
|
|
)
|
|
|
|
self.communication_bus.send_message(message, self.session_id)
|
|
return message.message_id
|
|
|
|
def get_agent_messages(self, agent_id: str, unread_only: bool = True) -> List[AgentMessage]:
|
|
return self.communication_bus.get_messages(agent_id, unread_only)
|
|
|
|
def collaborate_agents(self, orchestrator_id: str, task: str, agent_roles: List[str]):
|
|
orchestrator = self.get_agent(orchestrator_id)
|
|
if not orchestrator:
|
|
orchestrator_id = self.create_agent('orchestrator')
|
|
orchestrator = self.get_agent(orchestrator_id)
|
|
|
|
worker_agents = []
|
|
for role in agent_roles:
|
|
agent_id = self.create_agent(role)
|
|
worker_agents.append({
|
|
'agent_id': agent_id,
|
|
'role': role
|
|
})
|
|
|
|
orchestration_prompt = f'''Task: {task}
|
|
|
|
Available specialized agents:
|
|
{chr(10).join([f"- {a['agent_id']} ({a['role']})" for a in worker_agents])}
|
|
|
|
Break down the task and delegate subtasks to appropriate agents. Coordinate their work and integrate results.'''
|
|
|
|
orchestrator_result = self.execute_agent_task(orchestrator_id, orchestration_prompt)
|
|
|
|
results = {
|
|
'orchestrator': orchestrator_result,
|
|
'agents': []
|
|
}
|
|
|
|
for agent_info in worker_agents:
|
|
agent_id = agent_info['agent_id']
|
|
messages = self.get_agent_messages(agent_id)
|
|
|
|
for msg in messages:
|
|
subtask = msg.content
|
|
result = self.execute_agent_task(agent_id, subtask)
|
|
results['agents'].append(result)
|
|
|
|
self.send_agent_message(
|
|
from_agent_id=agent_id,
|
|
to_agent_id=orchestrator_id,
|
|
content=result.get('response', ''),
|
|
message_type=MessageType.RESPONSE
|
|
)
|
|
self.communication_bus.mark_as_read(msg.message_id)
|
|
|
|
return results
|
|
|
|
def get_session_summary(self) -> str:
|
|
summary = {
|
|
'session_id': self.session_id,
|
|
'active_agents': len(self.active_agents),
|
|
'agents': [
|
|
{
|
|
'agent_id': agent_id,
|
|
'role': agent.role.name,
|
|
'task_count': agent.task_count,
|
|
'message_count': len(agent.message_history)
|
|
}
|
|
for agent_id, agent in self.active_agents.items()
|
|
]
|
|
}
|
|
return json.dumps(summary)
|
|
|
|
def clear_session(self):
|
|
self.active_agents.clear()
|
|
self.communication_bus.clear_messages(session_id=self.session_id)
|
|
self.session_id = str(uuid.uuid4())[:16] |