Skip to content

Commit 51dcfc4

Browse files
committed
fix: store unique_ids instead of Agent objects in send_message memory content
1 parent 4ea91ed commit 51dcfc4

File tree

2 files changed

+96
-4
lines changed

2 files changed

+96
-4
lines changed

mesa_llm/llm_agent.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -267,8 +267,8 @@ async def asend_message(self, message: str, recipients: list[Agent]) -> str:
267267
type="message",
268268
content={
269269
"message": message,
270-
"sender": self,
271-
"recipients": recipients,
270+
"sender": self.unique_id,
271+
"recipients": [r.unique_id for r in recipients],
272272
},
273273
)
274274

@@ -283,8 +283,8 @@ def send_message(self, message: str, recipients: list[Agent]) -> str:
283283
type="message",
284284
content={
285285
"message": message,
286-
"sender": self,
287-
"recipients": recipients,
286+
"sender": self.unique_id,
287+
"recipients": [r.unique_id for r in recipients],
288288
},
289289
)
290290

tests/test_llm_agent.py

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# tests/test_llm_agent.py
22

3+
import json
34
import re
45

56
import pytest
@@ -612,3 +613,94 @@ def __init__(self):
612613
obs = agent.generate_obs()
613614

614615
assert len(obs.local_state) == 0
616+
617+
618+
# ---------------------------------------------------------------------------
619+
# send_message / asend_message – store unique_ids, not Agent objects (#156)
620+
# ---------------------------------------------------------------------------
621+
622+
623+
def _make_send_message_model(monkeypatch):
624+
"""Shared setup: two-agent MultiGrid model with ShortTermMemory."""
625+
monkeypatch.setenv("GEMINI_API_KEY", "dummy")
626+
627+
class DummyModel(Model):
628+
def __init__(self):
629+
super().__init__(seed=45)
630+
self.grid = MultiGrid(3, 3, torus=False)
631+
632+
def add_agent(self, pos):
633+
agents = LLMAgent.create_agents(
634+
self,
635+
n=1,
636+
reasoning=lambda agent: None,
637+
system_prompt="Test",
638+
vision=-1,
639+
internal_state=[],
640+
)
641+
agent = agents.to_list()[0]
642+
self.grid.place_agent(agent, pos)
643+
return agent
644+
645+
model = DummyModel()
646+
647+
sender = model.add_agent((0, 0))
648+
sender.memory = ShortTermMemory(agent=sender, n=5, display=True)
649+
sender.unique_id = 10
650+
651+
recipient = model.add_agent((1, 1))
652+
recipient.memory = ShortTermMemory(agent=recipient, n=5, display=True)
653+
recipient.unique_id = 20
654+
655+
return sender, recipient
656+
657+
658+
def test_send_message_stores_serializable_ids(monkeypatch):
659+
"""send_message stores sender/recipients as unique_ids, not Agent objects."""
660+
sender, recipient = _make_send_message_model(monkeypatch)
661+
662+
captured = {}
663+
664+
def capture_content(type, content):
665+
captured.update(content)
666+
667+
monkeypatch.setattr(recipient.memory, "add_to_memory", capture_content)
668+
monkeypatch.setattr(sender.memory, "add_to_memory", lambda *a, **kw: None)
669+
670+
sender.send_message("hello", recipients=[recipient])
671+
672+
assert captured["sender"] == 10
673+
assert captured["recipients"] == [20]
674+
assert captured["message"] == "hello"
675+
676+
# Must not raise TypeError when serializing
677+
data = json.loads(json.dumps(captured))
678+
assert data["sender"] == 10
679+
assert data["recipients"] == [20]
680+
681+
682+
@pytest.mark.asyncio
683+
async def test_asend_message_stores_serializable_ids(monkeypatch):
684+
"""asend_message stores sender/recipients as unique_ids, not Agent objects."""
685+
sender, recipient = _make_send_message_model(monkeypatch)
686+
687+
captured = {}
688+
689+
async def capture_content(type, content):
690+
captured.update(content)
691+
692+
async def noop(*a, **kw):
693+
pass
694+
695+
monkeypatch.setattr(recipient.memory, "aadd_to_memory", capture_content)
696+
monkeypatch.setattr(sender.memory, "aadd_to_memory", noop)
697+
698+
await sender.asend_message("hello", recipients=[recipient])
699+
700+
assert captured["sender"] == 10
701+
assert captured["recipients"] == [20]
702+
assert captured["message"] == "hello"
703+
704+
data = json.loads(json.dumps(captured))
705+
assert data["sender"] == 10
706+
assert data["recipients"] == [20]

0 commit comments

Comments
 (0)