rp/agent_communication.py at 9a5bf46a54616d7d21091745c5e29f2b07bfae83

 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