Adding tests to check message serialization/deserialization

Avoiding python type yaml tag for message type

Signed-off-by: Emmanuel Hugonnet <[email protected]>

docs/communication.md

@@ -4,13 +4,13 @@ The communication between agents happen on STOMP protocol. The message exchanged
 
 ## STOMP Queue
 
-Per convention each agent listen on 2 queues, normally sharing the name with the Agent using them (not mandatory, its a configuration of the transport):
+Per convention each agent listens on 2 queues, normally sharing the name with the Agent using them (not mandatory, its a configuration of the transport):
 
 ```plain-text
 /queue/request/AgentName
 /queue/response/AgentName
 ```
-The first one is used to send request to the Agent and the second one to send answer back.
+The first one is used to send requests to the Agent and the second one to send answers back.
 
 ## WiseAgentMessage Schema
 
@@ -92,7 +92,7 @@ schema:
 ### `_message_type`
 - **Type**: `string`
 - **Description**: 
-  (Optional) The type of the message, represented as an integer.
+  (Optional) The type of the message, represented as an string.
 - **Required**: false
 
   **Enum**: 
@@ -113,7 +113,7 @@ schema:
 ### `_sender`
 - **Type**: `string`
 - **Description**: 
-  (Optional) he agent that is sending this message.
+  (Optional) The agent that is sending this message.
 - **Required**: false
 
 ### `_tool_id`

src/wiseagents/wise_agent_messaging.py

@@ -1,24 +1,31 @@
+import logging
 from abc import *
-from enum import Enum, auto
+from enum import StrEnum
 from typing import Callable, Optional
 
+import yaml
 from yaml import YAMLObject
-from wiseagents.yaml import WiseAgentsLoader
-
-class WiseAgentMessageType(Enum):
-    ACK = auto()
-    ALERT = auto()
-    CANNOT_ANSWER = auto()
-    QUERY = auto()
-    RESPONSE = auto()
-    ACTION_REQUEST = auto()
-    HUMAN = auto()
-
+from wiseagents.yaml import WiseAgentsLoader 
+from yaml.resolver import BaseResolver
+
+
+class WiseAgentMessageType(StrEnum):
+    ACK = "ACK"
+    ALERT = "ALERT"
+    CANNOT_ANSWER = "CANNOT_ANSWER"
+    QUERY = "QUERY"
+    RESPONSE = "RESPONSE"
+    ACTION_REQUEST = "ACTION_REQUEST"
+    HUMAN = "HUMAN"
+
 class WiseAgentEvent:
     """
     TODO
     """
 
+def wiseAgentMessageType_representer(dumper, data):
+    return dumper.represent_scalar(BaseResolver.DEFAULT_SCALAR_TAG, str(data.value))
+
 
 class WiseAgentMessage(YAMLObject):
     ''' A message that can be sent between agents. '''
@@ -47,6 +54,24 @@ def __init__(self, message: str, sender: Optional[str] = None, message_type: Opt
             self._context_name = context_name
         else:
             self._context_name = 'default'
+        self.__class__.yaml_dumper.add_representer(WiseAgentMessageType, wiseAgentMessageType_representer)
+
+    def __setstate__(self, state):
+        self._message = state["_message"]
+        self._sender =  state["_sender"]
+        if state["_message_type"] is not None and state["_message_type"] != "":
+            logging.debug(f"__setstate__ state[_message_type]: {state['_message_type']}")
+            self._message_type = WiseAgentMessageType(state["_message_type"])
+        else:
+            self._message_type = None
+        self._chat_id =  state["_chat_id"]
+        self._tool_id =  state["_tool_id"]
+        self._route_response_to =  state["_route_response_to"]
+        if  state["_context_name"] is not None:
+            self._context_name = state["_context_name"]
+        else:
+            self._context_name = 'default'
+
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(message={self.message}, sender={self.sender}, message_type={self.message_type}, id={self.chat_id}, tool_id={self.tool_id}, context_name={self.context_name}, route_response_to={self.route_response_to}, route_response_to={self.route_response_to})"

test_message.yaml → tests/wiseagents/test_message.yaml

@@ -2,8 +2,7 @@
 _chat_id: '12345'
 _context_name: Weather
 _message: Hello
-_message_type: !!python/object/apply:wiseagents.wise_agent_messaging.WiseAgentMessageType
-- 1
+_message_type: ACK
 _route_response_to: Agent1
 _sender: Agent1
 _tool_id: WeatherAgent

tests/wiseagents/test_yaml_deserializer.py

@@ -4,7 +4,7 @@
 import pytest
 import yaml
 
-from wiseagents import WiseAgent, WiseAgentRegistry
+from wiseagents import WiseAgent, WiseAgentRegistry, WiseAgentMessage, WiseAgentMessageType
 from wiseagents.yaml import WiseAgentsLoader
 
 
@@ -100,3 +100,24 @@ def test_assistant_desiralizer():
             deserialized_agent = yaml.load(stream, Loader=yaml.Loader)
         except yaml.YAMLError as exc:
             print(exc)
+def test_deserialize_message():
+    message = [];
+    with open(pathlib.Path().resolve() / "tests/wiseagents/test_message.yaml") as stream:
+        try:
+            message = yaml.load(stream, Loader=yaml.Loader)
+            logging.info(str(message))
+        except yaml.YAMLError as exc:
+            print(exc)
+    assert isinstance(message, WiseAgentMessage)
+    assert message.message == "Hello"
+    assert message.sender == "Agent1"
+    msgType = type(message.message_type)
+    logging.info(str(msgType))
+    assert isinstance(message.message_type, WiseAgentMessageType)
+    assert message.message_type == WiseAgentMessageType.ACK
+    assert message.chat_id == "12345"
+    assert message.tool_id == "WeatherAgent"
+    assert message.context_name =="Weather"
+    assert message.route_response_to =="Agent1"
+
+

tests/wiseagents/test_yaml_serializtion.py

@@ -1,4 +1,7 @@
 import logging
+import os
+import pathlib
+import unittest
 
 import pytest
 import yaml
@@ -149,6 +152,16 @@ def test_serialize_message():
                                tool_id="WeatherAgent", 
                                context_name="Weather", 
                                route_response_to="Agent1")
-    with open("test_message.yaml", "w") as stream:
+    with open(pathlib.Path().resolve() / "tests/wiseagents/test_serialized_message.yaml", "w") as stream:
         yaml.dump(message, stream)
+    unittest.TestCase().assertListEqual(
+        list(open(pathlib.Path().resolve() / "tests/wiseagents/test_serialized_message.yaml", "r")),
+        list(open(pathlib.Path().resolve() / "tests/wiseagents/test_message.yaml", "r")))
+    with open(pathlib.Path().resolve() / "tests/wiseagents/test_serialized_message.yaml") as stream:
+        try:
+            deserialized_message = yaml.load(stream, Loader=yaml.Loader)
+        except yaml.YAMLError as exc:
+            print(exc)  
+    os.remove(pathlib.Path().resolve() / "tests/wiseagents/test_serialized_message.yaml")
+