rp/pr/agents/agent_manager.py

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]