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 from ..memory.knowledge_store import KnowledgeStore @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.knowledge_store = KnowledgeStore(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) knowledge_matches = self.knowledge_store.search_entries(task, top_k=3) agent.task_count += 1 messages = agent.get_messages_for_api() if knowledge_matches: knowledge_content = "Knowledge base matches based on your query:\\n" for i, entry in enumerate(knowledge_matches, 1): shortened_content = entry.content[:2000] knowledge_content += f"{i}. {shortened_content}\\n\\n" messages.insert(-1, {'role': 'user', 'content': knowledge_content}) 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]