Merge pull request #157 from kashif/textenv

[textarena] fix how consecutive messages from the same sender and category are handled

Author: init27

src/envs/textarena_env/server/environment.py

@@ -187,6 +187,25 @@ def _legal_players(self) -> List[int]:
 
     def _convert_messages(self, messages: Iterable[Any]) -> List[TextArenaMessage]:
         converted: List[TextArenaMessage] = []
+        buffered_sender: int | None = None
+        buffered_category: str | None = None
+        buffered_content: List[str] = []
+
+        def flush_buffer() -> None:
+            nonlocal buffered_content, buffered_sender, buffered_category
+            if not buffered_content:
+                return
+            converted.append(
+                TextArenaMessage(
+                    sender_id=buffered_sender if buffered_sender is not None else -1,
+                    content="".join(buffered_content),
+                    category=buffered_category or "MESSAGE",
+                )
+            )
+            buffered_content = []
+            buffered_category = None
+            buffered_sender = None
+
         for entry in messages:
             if isinstance(entry, tuple) and len(entry) == 3:
                 sender, content, category = entry
@@ -197,13 +216,18 @@ def _convert_messages(self, messages: Iterable[Any]) -> List[TextArenaMessage]:
                 sender, content, category = -1, str(entry), "MESSAGE"
 
             category_name = getattr(category, "name", str(category))
-            converted.append(
-                TextArenaMessage(
-                    sender_id=int(sender) if isinstance(sender, (int, float)) else -1,
-                    content=str(content),
-                    category=category_name,
-                )
-            )
+            sender_id = int(sender) if isinstance(sender, (int, float)) else -1
+            text = str(content)
+
+            if buffered_content and buffered_category == category_name and buffered_sender == sender_id:
+                buffered_content.append(text)
+            else:
+                flush_buffer()
+                buffered_sender = sender_id
+                buffered_category = category_name
+                buffered_content = [text]
+
+        flush_buffer()
 
         return converted
 
@@ -249,4 +273,3 @@ def _compute_reward_signals(
             for key, value in result.items():
                 aggregated[key] = float(value)
         return aggregated
-

tests/envs/test_textarena_environment.py

@@ -0,0 +1,23 @@
+from envs.textarena_env.server.environment import TextArenaEnvironment
+from envs.textarena_env.models import TextArenaMessage
+
+
+def test_convert_messages_coalesces_consecutive_characters():
+    env = object.__new__(TextArenaEnvironment)
+
+    raw_messages = [
+        (0, "[", "PROMPT"),
+        (0, "GAME", "PROMPT"),
+        (0, "]", "PROMPT"),
+        (1, "A", "MESSAGE"),
+        (1, "B", "MESSAGE"),
+        (2, "!", "MESSAGE"),
+    ]
+
+    converted = env._convert_messages(raw_messages)
+
+    assert converted == [
+        TextArenaMessage(sender_id=0, content="[GAME]", category="PROMPT"),
+        TextArenaMessage(sender_id=1, content="AB", category="MESSAGE"),
+        TextArenaMessage(sender_id=2, content="!", category="MESSAGE"),
+    ]