rp/pr/agents/agent_communication.py

import json
import sqlite3
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional


class MessageType(Enum):
    REQUEST = "request"
    RESPONSE = "response"
    NOTIFICATION = "notification"


@dataclass
class AgentMessage:
    message_id: str
    from_agent: str
    to_agent: str
    message_type: MessageType
    content: str
    metadata: dict
    timestamp: float

    def to_dict(self) -> dict:
        return {
            "message_id": self.message_id,
            "from_agent": self.from_agent,
            "to_agent": self.to_agent,
            "message_type": self.message_type.value,
            "content": self.content,
            "metadata": self.metadata,
            "timestamp": self.timestamp,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "AgentMessage":
        return cls(
            message_id=data["message_id"],
            from_agent=data["from_agent"],
            to_agent=data["to_agent"],
            message_type=MessageType(data["message_type"]),
            content=data["content"],
            metadata=data["metadata"],
            timestamp=data["timestamp"],
        )


class AgentCommunicationBus:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        self._create_tables()

    def _create_tables(self):
        cursor = self.conn.cursor()
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS agent_messages (
                message_id TEXT PRIMARY KEY,
                from_agent TEXT,
                to_agent TEXT,
                message_type TEXT,
                content TEXT,
                metadata TEXT,
                timestamp REAL,
                session_id TEXT,
                read INTEGER DEFAULT 0
            )
        """
        )
        self.conn.commit()

    def send_message(self, message: AgentMessage, session_id: Optional[str] = None):
        cursor = self.conn.cursor()

        cursor.execute(
            """
            INSERT INTO agent_messages
            (message_id, from_agent, to_agent, message_type, content, metadata, timestamp, session_id)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """,
            (
                message.message_id,
                message.from_agent,
                message.to_agent,
                message.message_type.value,
                message.content,
                json.dumps(message.metadata),
                message.timestamp,
                session_id,
            ),
        )

        self.conn.commit()

    def get_messages(self, agent_id: str, unread_only: bool = True) -> List[AgentMessage]:
        cursor = self.conn.cursor()
        if unread_only:
            cursor.execute(
                """
                SELECT message_id, from_agent, to_agent, message_type, content, metadata, timestamp
                FROM agent_messages
                WHERE to_agent = ? AND read = 0
                ORDER BY timestamp ASC
            """,
                (agent_id,),
            )
        else:
            cursor.execute(
                """
                SELECT message_id, from_agent, to_agent, message_type, content, metadata, timestamp
                FROM agent_messages
                WHERE to_agent = ?
                ORDER BY timestamp ASC
            """,
                (agent_id,),
            )

        messages = []
        for row in cursor.fetchall():
            messages.append(
                AgentMessage(
                    message_id=row[0],
                    from_agent=row[1],
                    to_agent=row[2],
                    message_type=MessageType(row[3]),
                    content=row[4],
                    metadata=json.loads(row[5]) if row[5] else {},
                    timestamp=row[6],
                )
            )
        return messages

    def mark_as_read(self, message_id: str):
        cursor = self.conn.cursor()
        cursor.execute("UPDATE agent_messages SET read = 1 WHERE message_id = ?", (message_id,))
        self.conn.commit()

    def clear_messages(self, session_id: Optional[str] = None):
        cursor = self.conn.cursor()
        if session_id:
            cursor.execute("DELETE FROM agent_messages WHERE session_id = ?", (session_id,))
        else:
            cursor.execute("DELETE FROM agent_messages")
        self.conn.commit()

    def close(self):
        self.conn.close()

    def receive_messages(self, agent_id: str) -> List[AgentMessage]:
        return self.get_messages(agent_id, unread_only=True)

    def get_conversation_history(self, agent_a: str, agent_b: str) -> List[AgentMessage]:
        cursor = self.conn.cursor()
        cursor.execute(
            """
            SELECT message_id, from_agent, to_agent, message_type, content, metadata, timestamp
            FROM agent_messages
            WHERE (from_agent = ? AND to_agent = ?) OR (from_agent = ? AND to_agent = ?)
            ORDER BY timestamp ASC
        """,
            (agent_a, agent_b, agent_b, agent_a),
        )

        messages = []
        for row in cursor.fetchall():
            messages.append(
                AgentMessage(
                    message_id=row[0],
                    from_agent=row[1],
                    to_agent=row[2],
                    message_type=MessageType(row[3]),
                    content=row[4],
                    metadata=json.loads(row[5]) if row[5] else {},
                    timestamp=row[6],
                )
            )
        return messages