improve: add comprehensive type hints to LLMAgent and ModuleLLM#170
Open
abhinavk0220 wants to merge 11 commits intomesa:mainfrom
Open
improve: add comprehensive type hints to LLMAgent and ModuleLLM#170abhinavk0220 wants to merge 11 commits intomesa:mainfrom
abhinavk0220 wants to merge 11 commits intomesa:mainfrom
Conversation
…rely, breaking the full test suite.This commit migrates all affected files to the new APIs.Changes:- Replace MultiGrid/SingleGrid imports with OrthogonalMooreGrid from mesa.discrete_space- Replace old ContinuousSpace(x_max, y_max) with new dimensions-based API- Fix agent placement to use cell-based API (agent.cell = grid._cells[pos])- Fix neighbor lookup using get_neighborhood + coordinate set matching- Fix model.steps -> model.step (renamed in Mesa 4.x)- Remove seed= kwarg from Model.__init__ calls (mesa_signals compat fix)- Fix record_model.py self.steps -> self.stepAll 203 tests now passing.
…ools.py to use Mesa 4.x cell-based API- Replace place_agent() with agent.cell = grid._cells[pos]- Fix ContinuousSpace boundary/torus tests for new API behavior- Remove SingleGrid|MultiGrid references from inbuilt_tools.py- All 262 tests passing
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
…return type annotations to all methods in LLMAgent and ModuleLLM- Add explicit type annotations to all instance attributes in __init__- Improve _build_observation() return type: tuple[dict[str, Any], dict[str, Any]]- Improve apply_plan/aapply_plan return type: list[dict[str, Any]]- Add Any import to both files for complete typing coverage- Expand docstrings with Args/Returns sections where missing- All 262 tests passing
Contributor
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughThis PR migrates the codebase to Mesa 4.x by replacing grid types (MultiGrid/SingleGrid → OrthogonalMooreGrid/OrthogonalVonNeumannGrid), updating ContinuousSpace imports, refactoring agent positioning logic, adding comprehensive type annotations across LLMAgent and ModuleLLM, and expanding test coverage for movement and space handling across multiple environments. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Member
|
@CodeRabbit review |
Contributor
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
mesa_llm/llm_agent.py (3)
105-113:⚠️ Potential issue | 🟠 MajorPreserve all tool-call results in memory.
The dict comprehension collapses every tool result into the same
name/responsekeys, so only the last tool call survives in memory when a plan executes more than one tool. Store a list under a stable key instead of flattening the response list.🛠️ Suggested fix
await self.memory.aadd_to_memory( type="action", content={ - k: v - for tool_call in tool_call_resp - for k, v in tool_call.items() - if k not in ["tool_call_id", "role"] + "tool_calls": [ + { + k: v + for k, v in tool_call.items() + if k not in {"tool_call_id", "role"} + } + for tool_call in tool_call_resp + ] }, ) ... self.memory.add_to_memory( type="action", content={ - k: v - for tool_call in tool_call_resp - for k, v in tool_call.items() - if k not in ["tool_call_id", "role"] + "tool_calls": [ + { + k: v + for k, v in tool_call.items() + if k not in {"tool_call_id", "role"} + } + for tool_call in tool_call_resp + ] }, )Also applies to: 133-141
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@mesa_llm/llm_agent.py` around lines 105 - 113, The current memory.aadd_to_memory call flattens multiple tool_call_resp dicts into a single dict, losing earlier tool results; instead, change the payload to store the entire list under a stable key (e.g., "tool_call_results" or "tool_calls") so all entries are preserved (reference: memory.aadd_to_memory and variable tool_call_resp), and make the same change to the similar block handling tool_call_resp at the later section (lines 133-141) so both places append the full list rather than collapsing into one dict.
344-349:⚠️ Potential issue | 🔴 CriticalBalance pre/post-step hooks exactly once, even on errors.
astep()opens an async pre/post pair and then callsself.step(), but overriddenstep()methods are wrapped below with their own pre/post hooks. For subclasses that only implementstep, that double-opens the lifecycle and leaves memory bookkeeping unbalanced. The wrappers also skip the post hook when user code raises.🛠️ Suggested fix
async def astep(self) -> None: - await self.apre_step() - if hasattr(self, "step") and self.__class__.step != LLMAgent.step: self.step() + return - await self.apost_step() + await self.apre_step() + try: + return + finally: + await self.apost_step() ... def wrapped(self, *args: Any, **kwargs: Any) -> Any: - LLMAgent.pre_step(self, *args, **kwargs) - result = user_step(self, *args, **kwargs) - LLMAgent.post_step(self, *args, **kwargs) - return result + LLMAgent.pre_step(self) + try: + return user_step(self, *args, **kwargs) + finally: + LLMAgent.post_step(self) ... async def awrapped(self, *args: Any, **kwargs: Any) -> Any: await self.apre_step() - result = await user_astep(self, *args, **kwargs) - await self.apost_step() - return result + try: + return await user_astep(self, *args, **kwargs) + finally: + await self.apost_step()Also applies to: 360-378
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@mesa_llm/llm_agent.py` around lines 344 - 349, The pre/post-step lifecycle is being opened twice for subclasses that override step and may skip post on exceptions; fix by calling apre_step() once, executing self.step() inside a try/finally, and always calling apost_step() in the finally so the pair is balanced even on errors; update the same pattern for the other wrapper region (the methods around lines 360-378) and ensure any wrapper implementations of step do not themselves call apre_step/apost_step so the lifecycle is only managed in the outer caller.
229-249:⚠️ Potential issue | 🟠 MajorReplace
self.model.stepwithself.model.stepsin two locations.
Model.step()is the simulation method;Model.stepsis the integer counter tracking executed steps. Usingself.model.stepassigns a bound method object tostep: int, violating the Observation type contract and causing runtime errors.Affected locations
- Line 239:
step: int = self.model.step→step: int = self.model.steps- Line 253:
step: int = self.model.step→step: int = self.model.steps🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@mesa_llm/llm_agent.py` around lines 229 - 249, The code assigns a bound method object to the Observation.step field by using self.model.step instead of the integer counter; update both occurrences to use the integer counter property self.model.steps (e.g., in agenerate_obs and the corresponding generate_obs method) so step: int = self.model.steps and the returned Observation(step=step, ...) uses an int, not a bound method.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@hell`:
- Around line 1-275: Remove the stray merge/diff artifact file named "hell" that
contains the checked-in patch text for tests/test_tools/test_inbuilt_tools.py;
delete this file from the PR (or revert the commit that introduced it) so only
the real tests/test_tools/test_inbuilt_tools.py remains in the changeset and no
duplicated diff text is checked in.
In `@mesa_llm/recording/record_model.py`:
- Line 104: The recorder is logging the step method object because it uses
self.step; update the two calls to self.recorder.record_model_event that pass
the step counter (the "step_start" and corresponding "step_end"/other event
call) to use self.steps instead of self.step so the numeric step count from the
Mesa model (self.steps) is recorded; locate the calls to
self.recorder.record_model_event("step_start", {"step": self.step}) and the
analogous call later and replace self.step with self.steps.
In `@tests/test_llm_agent.py`:
- Around line 96-98: The ShortTermMemory instance is being constructed with
agent=agent instead of the neighbor, causing both memories to claim the same
owner; update the neighbor.memory assignment so the ShortTermMemory is created
with agent=neighbor (or otherwise attach the memory to the neighbor by passing
neighbor into ShortTermMemory) so neighbor.memory.agent correctly references
neighbor rather than agent.
In `@tests/test_tools/test_inbuilt_tools.py`:
- Around line 333-362: Change the tests so they actually move past the upper
bound: in test_move_one_step_boundary_on_continuousspace and
test_move_one_step_torus_wrap_on_continuousspace set the agent start position to
y=10.0 (e.g., agent.pos = (2.0, 10.0)) so move_one_step("North") attempts to go
to y=11.0; then assert the non-torus case leaves the agent at (2.0, 10.0) and
the torus case wraps to (2.0, 0.0), and update the expected result strings to
match those outcomes; look for the functions/test names
test_move_one_step_boundary_on_continuousspace,
test_move_one_step_torus_wrap_on_continuousspace, the move_one_step helper,
DummyAgent and ContinuousSpace to locate and change the assertions.
---
Outside diff comments:
In `@mesa_llm/llm_agent.py`:
- Around line 105-113: The current memory.aadd_to_memory call flattens multiple
tool_call_resp dicts into a single dict, losing earlier tool results; instead,
change the payload to store the entire list under a stable key (e.g.,
"tool_call_results" or "tool_calls") so all entries are preserved (reference:
memory.aadd_to_memory and variable tool_call_resp), and make the same change to
the similar block handling tool_call_resp at the later section (lines 133-141)
so both places append the full list rather than collapsing into one dict.
- Around line 344-349: The pre/post-step lifecycle is being opened twice for
subclasses that override step and may skip post on exceptions; fix by calling
apre_step() once, executing self.step() inside a try/finally, and always calling
apost_step() in the finally so the pair is balanced even on errors; update the
same pattern for the other wrapper region (the methods around lines 360-378) and
ensure any wrapper implementations of step do not themselves call
apre_step/apost_step so the lifecycle is only managed in the outer caller.
- Around line 229-249: The code assigns a bound method object to the
Observation.step field by using self.model.step instead of the integer counter;
update both occurrences to use the integer counter property self.model.steps
(e.g., in agenerate_obs and the corresponding generate_obs method) so step: int
= self.model.steps and the returned Observation(step=step, ...) uses an int, not
a bound method.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b4f3f488-5303-46f9-bf19-201359e324b9
📒 Files selected for processing (10)
hellmesa_llm/llm_agent.pymesa_llm/module_llm.pymesa_llm/recording/record_model.pymesa_llm/tools/inbuilt_tools.pytests/conftest.pytests/test_llm_agent.pytests/test_parallel_stepping.pytests/test_reasoning/test_cot.pytests/test_tools/test_inbuilt_tools.py
Comment on lines
+1
to
+275
| [1mdiff --cc tests/test_tools/test_inbuilt_tools.py[m | ||
| [1mindex fce291c,7a9c367..0000000[m | ||
| [1m--- a/tests/test_tools/test_inbuilt_tools.py[m | ||
| [1m+++ b/tests/test_tools/test_inbuilt_tools.py[m | ||
| [36m@@@ -3,9 -3,9 +3,10 @@@[m [mfrom __future__ import annotation[m | ||
| from types import SimpleNamespace[m | ||
| [m | ||
| import pytest[m | ||
| [31m -from mesa.discrete_space import OrthogonalMooreGrid, OrthogonalVonNeumannGrid[m | ||
| [31m -from mesa.space import ContinuousSpace, MultiGrid, SingleGrid[m | ||
| [32m +from mesa.discrete_space import OrthogonalMooreGrid[m | ||
| [32m +from mesa.experimental.continuous_space import ContinuousSpace[m | ||
| [32m +[m | ||
| [32m+ [m | ||
| from mesa_llm.tools.inbuilt_tools import ([m | ||
| move_one_step,[m | ||
| speak_to,[m | ||
| [36m@@@ -155,3 -329,256 +334,257 @@@[m [mdef test_move_one_step_on_continuousspa[m | ||
| [m | ||
| assert agent.pos == (2.0, 3.0)[m | ||
| assert result == "agent 6 moved to (2.0, 3.0)."[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_boundary_on_continuousspace():[m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = None[m | ||
| [32m+ model.space = ContinuousSpace(x_max=10.0, y_max=10.0, torus=False)[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=30, model=model)[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ model.space.place_agent(agent, (2.0, 9.0))[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.pos == (2.0, 9.0)[m | ||
| [32m+ assert "boundary" in result.lower()[m | ||
| [32m+ assert "North" in result[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_torus_wrap_on_continuousspace():[m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = None[m | ||
| [32m+ model.space = ContinuousSpace(x_max=10.0, y_max=10.0, torus=True)[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=31, model=model)[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ model.space.place_agent(agent, (2.0, 9.0))[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.pos == (2.0, 0.0)[m | ||
| [32m+ assert result == "agent 31 moved to (2.0, 0.0)."[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_boundary_singlegrid_north():[m | ||
| [32m+ """Agent at top edge of SingleGrid trying to go North gets a clear message."""[m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = SingleGrid(width=5, height=5, torus=False)[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=20, model=model)[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ model.grid.place_agent(agent, (2, 4)) # y=4 is the top edge[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ # agent should not have moved[m | ||
| [32m+ assert agent.pos == (2, 4)[m | ||
| [32m+ assert "boundary" in result.lower()[m | ||
| [32m+ assert "North" in result[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_torus_wrap_singlegrid_north():[m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = SingleGrid(width=5, height=5, torus=True)[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=23, model=model)[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ model.grid.place_agent(agent, (2, 4))[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.pos == (2, 0)[m | ||
| [32m+ assert result == "agent 23 moved to (2, 0)."[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_boundary_multigrid_west():[m | ||
| [32m+ """Agent at left edge of MultiGrid trying to go West gets a clear message."""[m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = MultiGrid(width=5, height=5, torus=False)[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=21, model=model)[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ model.grid.place_agent(agent, (0, 2)) # x=0 is the left edge[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "West")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.pos == (0, 2)[m | ||
| [32m+ assert "boundary" in result.lower()[m | ||
| [32m+ assert "West" in result[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_torus_wrap_multigrid_west():[m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = MultiGrid(width=5, height=5, torus=True)[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=24, model=model)[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ model.grid.place_agent(agent, (0, 2))[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "West")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.pos == (4, 2)[m | ||
| [32m+ assert result == "agent 24 moved to (4, 2)."[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_singlegrid_occupied_target():[m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = SingleGrid(width=5, height=5, torus=False)[m | ||
| [32m+ [m | ||
| [32m+ moving_agent = DummyAgent(unique_id=25, model=model)[m | ||
| [32m+ blocking_agent = DummyAgent(unique_id=26, model=model)[m | ||
| [32m+ model.agents.extend([moving_agent, blocking_agent])[m | ||
| [32m+ model.grid.place_agent(moving_agent, (2, 2))[m | ||
| [32m+ model.grid.place_agent(blocking_agent, (2, 3))[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(moving_agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ assert moving_agent.pos == (2, 2)[m | ||
| [32m+ assert blocking_agent.pos == (2, 3)[m | ||
| [32m+ assert "occupied" in result.lower()[m | ||
| [32m+ assert "North" in result[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_boundary_orthogonal_grid():[m | ||
| [32m+ """Agent at edge of OrthogonalMooreGrid with no cell in that direction gets a clear message."""[m | ||
| [32m+ [m | ||
| [32m+ class _DummyOrthogonalGrid(OrthogonalMooreGrid):[m | ||
| [32m+ pass[m | ||
| [32m+ [m | ||
| [32m+ orth_grid = object.__new__(_DummyOrthogonalGrid)[m | ||
| [32m+ orth_grid.torus = False[m | ||
| [32m+ orth_grid.dimensions = (5, 5)[m | ||
| [32m+ start = (0, 1)[m | ||
| [32m+ start_cell = SimpleNamespace([m | ||
| [32m+ coordinate=start, agents=[], connections={}, is_full=False[m | ||
| [32m+ )[m | ||
| [32m+ orth_grid._cells = {start: start_cell}[m | ||
| [32m+ [m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = orth_grid[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=22, model=model)[m | ||
| [32m+ agent.cell = start_cell[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ # cell should be unchanged[m | ||
| [32m+ assert agent.cell is start_cell[m | ||
| [32m+ assert "boundary" in result.lower()[m | ||
| [32m+ assert "North" in result[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_boundary_orthogonal_torus_missing_wrapped_cell():[m | ||
| [32m+ class _DummyOrthogonalGrid(OrthogonalMooreGrid):[m | ||
| [32m+ pass[m | ||
| [32m+ [m | ||
| [32m+ orth_grid = object.__new__(_DummyOrthogonalGrid)[m | ||
| [32m+ orth_grid.torus = True[m | ||
| [32m+ orth_grid.dimensions = (3, 3)[m | ||
| [32m+ start = (0, 0)[m | ||
| [32m+ start_cell = SimpleNamespace(coordinate=start, agents=[], is_full=False)[m | ||
| [32m+ # Wrapped target for North would be (2, 0), but it is intentionally absent.[m | ||
| [32m+ orth_grid._cells = {start: start_cell}[m | ||
| [32m+ [m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = orth_grid[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=38, model=model)[m | ||
| [32m+ agent.cell = start_cell[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.cell is start_cell[m | ||
| [32m+ assert "boundary" in result.lower()[m | ||
| [32m+ assert "North" in result[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_full_target_orthogonal_grid():[m | ||
| [32m+ class _DummyOrthogonalGrid(OrthogonalMooreGrid):[m | ||
| [32m+ pass[m | ||
| [32m+ [m | ||
| [32m+ orth_grid = object.__new__(_DummyOrthogonalGrid)[m | ||
| [32m+ orth_grid.torus = False[m | ||
| [32m+ orth_grid.dimensions = (5, 5)[m | ||
| [32m+ start = (1, 1)[m | ||
| [32m+ end = (0, 1)[m | ||
| [32m+ start_cell = SimpleNamespace([m | ||
| [32m+ coordinate=start, agents=[], connections={}, is_full=False[m | ||
| [32m+ )[m | ||
| [32m+ full_target_cell = SimpleNamespace([m | ||
| [32m+ coordinate=end,[m | ||
| [32m+ agents=[SimpleNamespace(unique_id=99)],[m | ||
| [32m+ connections={},[m | ||
| [32m+ is_full=True,[m | ||
| [32m+ )[m | ||
| [32m+ start_cell.connections[(-1, 0)] = full_target_cell[m | ||
| [32m+ orth_grid._cells = {start: start_cell, end: full_target_cell}[m | ||
| [32m+ [m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = orth_grid[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=27, model=model)[m | ||
| [32m+ agent.cell = start_cell[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "North")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.cell is start_cell[m | ||
| [32m+ assert "full" in result.lower()[m | ||
| [32m+ assert "North" in result[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_diagonal_on_orthogonal_vonneumann_grid():[m | ||
| [32m+ class _DummyOrthogonalVonNeumannGrid(OrthogonalVonNeumannGrid):[m | ||
| [32m+ pass[m | ||
| [32m+ [m | ||
| [32m+ orth_grid = object.__new__(_DummyOrthogonalVonNeumannGrid)[m | ||
| [32m+ orth_grid.torus = False[m | ||
| [32m+ orth_grid.dimensions = (5, 5)[m | ||
| [32m+ start = (2, 2)[m | ||
| [32m+ end = (1, 3) # NorthEast[m | ||
| [32m+ start_cell = SimpleNamespace(coordinate=start, agents=[], is_full=False)[m | ||
| [32m+ end_cell = SimpleNamespace(coordinate=end, agents=[], is_full=False)[m | ||
| [32m+ orth_grid._cells = {start: start_cell, end: end_cell}[m | ||
| [32m+ [m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = orth_grid[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=28, model=model)[m | ||
| [32m+ agent.cell = start_cell[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "NorthEast")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.cell is end_cell[m | ||
| [32m+ assert result == "agent 28 moved to (1, 3)."[m | ||
| [32m+ [m | ||
| [32m+ [m | ||
| [32m+ def test_move_one_step_torus_wrap_orthogonal_grid():[m | ||
| [32m+ class _DummyOrthogonalGrid(OrthogonalMooreGrid):[m | ||
| [32m+ pass[m | ||
| [32m+ [m | ||
| [32m+ orth_grid = object.__new__(_DummyOrthogonalGrid)[m | ||
| [32m+ orth_grid.torus = True[m | ||
| [32m+ orth_grid.dimensions = (3, 3)[m | ||
| [32m+ start = (0, 0)[m | ||
| [32m+ end = (2, 2) # NorthWest wraps on torus[m | ||
| [32m+ start_cell = SimpleNamespace(coordinate=start, agents=[], is_full=False)[m | ||
| [32m+ wrapped_cell = SimpleNamespace(coordinate=end, agents=[], is_full=False)[m | ||
| [32m+ orth_grid._cells = {start: start_cell, end: wrapped_cell}[m | ||
| [32m+ [m | ||
| [32m+ model = DummyModel()[m | ||
| [32m+ model.grid = orth_grid[m | ||
| [32m+ [m | ||
| [32m+ agent = DummyAgent(unique_id=29, model=model)[m | ||
| [32m+ agent.cell = start_cell[m | ||
| [32m+ model.agents.append(agent)[m | ||
| [32m+ [m | ||
| [32m+ result = move_one_step(agent, "NorthWest")[m | ||
| [32m+ [m | ||
| [32m+ assert agent.cell is wrapped_cell[m | ||
| [32m+ assert result == "agent 29 moved to (2, 2)."[m | ||
| [32m++[m |
Contributor
There was a problem hiding this comment.
Remove the stray patch artifact.
This file is checked-in diff text for tests/test_tools/test_inbuilt_tools.py, not source. It duplicates content already under tests/ and looks like an unresolved patch/merge artifact, so it should be dropped from the PR.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@hell` around lines 1 - 275, Remove the stray merge/diff artifact file named
"hell" that contains the checked-in patch text for
tests/test_tools/test_inbuilt_tools.py; delete this file from the PR (or revert
the commit that introduced it) so only the real
tests/test_tools/test_inbuilt_tools.py remains in the changeset and no
duplicated diff text is checked in.
Comment on lines
96
to
98
| neighbor = model.add_agent((1, 2)) | ||
| neighbor.memory = ShortTermMemory( | ||
| agent=agent, | ||
| n=5, | ||
| display=True, | ||
| ) | ||
| neighbor.memory = ShortTermMemory(agent=agent, n=5, display=True) | ||
| neighbor.unique_id = 2 |
Contributor
There was a problem hiding this comment.
Attach the neighbor memory to neighbor.
Line 97 builds neighbor.memory with agent=agent, so both ShortTermMemory instances believe they belong to the same owner. That keeps this fixture internally inconsistent and can hide bugs once memory ownership is used during observation or formatting.
Suggested fix
- neighbor.memory = ShortTermMemory(agent=agent, n=5, display=True)
+ neighbor.memory = ShortTermMemory(agent=neighbor, n=5, display=True)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| neighbor = model.add_agent((1, 2)) | |
| neighbor.memory = ShortTermMemory( | |
| agent=agent, | |
| n=5, | |
| display=True, | |
| ) | |
| neighbor.memory = ShortTermMemory(agent=agent, n=5, display=True) | |
| neighbor.unique_id = 2 | |
| neighbor = model.add_agent((1, 2)) | |
| neighbor.memory = ShortTermMemory(agent=neighbor, n=5, display=True) | |
| neighbor.unique_id = 2 |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_llm_agent.py` around lines 96 - 98, The ShortTermMemory instance
is being constructed with agent=agent instead of the neighbor, causing both
memories to claim the same owner; update the neighbor.memory assignment so the
ShortTermMemory is created with agent=neighbor (or otherwise attach the memory
to the neighbor by passing neighbor into ShortTermMemory) so
neighbor.memory.agent correctly references neighbor rather than agent.
Comment on lines
333
to
+362
| def test_move_one_step_boundary_on_continuousspace(): | ||
| """In Mesa 4.x ContinuousSpace has no boundary checking — agent moves freely.""" | ||
| model = DummyModel() | ||
| model.grid = None | ||
| model.space = ContinuousSpace(x_max=10.0, y_max=10.0, torus=False) | ||
| model.space = ContinuousSpace(dimensions=[[0, 10.0], [0, 10.0]], torus=False) | ||
|
|
||
| agent = DummyAgent(unique_id=30, model=model) | ||
| model.agents.append(agent) | ||
| model.space.place_agent(agent, (2.0, 9.0)) | ||
| agent.pos = (2.0, 9.0) | ||
|
|
||
| result = move_one_step(agent, "North") | ||
|
|
||
| assert agent.pos == (2.0, 9.0) | ||
| assert "boundary" in result.lower() | ||
| assert "North" in result | ||
| assert agent.pos == (2.0, 10.0) | ||
| assert result == "agent 30 moved to (2.0, 10.0)." | ||
|
|
||
|
|
||
| def test_move_one_step_torus_wrap_on_continuousspace(): | ||
| """In Mesa 4.x ContinuousSpace torus wrapping is not built-in — agent moves freely.""" | ||
| model = DummyModel() | ||
| model.grid = None | ||
| model.space = ContinuousSpace(x_max=10.0, y_max=10.0, torus=True) | ||
| model.space = ContinuousSpace(dimensions=[[0, 10.0], [0, 10.0]], torus=True) | ||
|
|
||
| agent = DummyAgent(unique_id=31, model=model) | ||
| model.agents.append(agent) | ||
| model.space.place_agent(agent, (2.0, 9.0)) | ||
| agent.pos = (2.0, 9.0) | ||
|
|
||
| result = move_one_step(agent, "North") | ||
|
|
||
| assert agent.pos == (2.0, 0.0) | ||
| assert result == "agent 31 moved to (2.0, 0.0)." | ||
| assert agent.pos == (2.0, 10.0) | ||
| assert result == "agent 31 moved to (2.0, 10.0)." |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's find the Mesa ContinuousSpace implementation
find . -type f -name "*.py" | xargs grep -l "class ContinuousSpace" | head -5Repository: mesa/mesa-llm
Length of output: 1865
🏁 Script executed:
# Let's also check the move_one_step implementation
fd -e py "inbuilt_tools"Repository: mesa/mesa-llm
Length of output: 1865
🏁 Script executed:
# Check if mesa is available or if there's local implementation
find . -type f -name "*.py" -path "*/continuous*" | head -10Repository: mesa/mesa-llm
Length of output: 1865
🌐 Web query:
Mesa ContinuousSpace in_bounds upper boundary inclusive exclusive
💡 Result:
In Mesa’s experimental ContinuousSpace (experimental.continuous_space.continuous_space.ContinuousSpace), in_bounds(point) treats the upper boundary as inclusive (it checks point <= dimensions[:, 1]). (mesa.readthedocs.io)
In Mesa’s legacy/maintenance mesa.space.ContinuousSpace, there isn’t an in_bounds() method; instead out_of_bounds(pos) treats the upper boundary as exclusive (x >= x_max or y >= y_max is out of bounds). (mesa.readthedocs.io)
Citations:
- 1: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space.html
- 2: https://mesa.readthedocs.io/stable/_modules/mesa/space.html
🌐 Web query:
Mesa experimental continuous_space ContinuousSpace bounds checking torus
💡 Result:
In Mesa’s experimental ContinuousSpace, bounds checking (and torus wrapping) happens when you set an agent’s position:
ContinuousSpaceAgent.position = valuefirst callsspace.in_bounds(value). (mesa.readthedocs.io)- If
valueis out of bounds:- if
space.torusisTrue, it wraps the point withspace.torus_correct(value)(per-dimension modulo into[min, min+size)). (mesa.readthedocs.io) - otherwise it raises
ValueError. (mesa.readthedocs.io)
- if
Key details / gotchas:
in_bounds(point)is implemented as(point >= min) & (point <= max)(note the inclusive upper bound<= max). (mesa.readthedocs.io)torus_correct(point)maps any coordinate via:min + mod(point - min, size). (mesa.readthedocs.io)- So
point == maxwould mathematically wrap tomin(becausemod(size, size) == 0) but it will not be wrapped in practice, sincein_boundstreatsmaxas already valid and skipstorus_correct. (mesa.readthedocs.io) - This means you can end up with both
minandmaxpositions existing simultaneously; in torus distance calculations, those edges are treated as adjacent (effectively the same location). (mesa.readthedocs.io)
- So
Citations:
- 1: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space_agents.html?utm_source=openai
- 2: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space_agents.html?utm_source=openai
- 3: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space_agents.html?utm_source=openai
- 4: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space.html
- 5: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space.html
- 6: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space_agents.html?utm_source=openai
- 7: https://mesa.readthedocs.io/stable/_modules/experimental/continuous_space/continuous_space.html
Test cases do not trigger boundary or torus behavior.
Both tests start at y=9.0 and move north to y=10.0, which is within the valid bounds [0, 10.0]. Mesa's experimental ContinuousSpace.in_bounds() treats the upper endpoint as inclusive, so these assertions pass without exercising the behavior their names describe. To test boundary rejection, move past the upper limit (e.g., 10.0 to 11.0). To test torus wrapping, move beyond the bounds to trigger the wrap logic.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_tools/test_inbuilt_tools.py` around lines 333 - 362, Change the
tests so they actually move past the upper bound: in
test_move_one_step_boundary_on_continuousspace and
test_move_one_step_torus_wrap_on_continuousspace set the agent start position to
y=10.0 (e.g., agent.pos = (2.0, 10.0)) so move_one_step("North") attempts to go
to y=11.0; then assert the non-torus case leaves the agent at (2.0, 10.0) and
the torus case wraps to (2.0, 0.0), and update the expected result strings to
match those outcomes; look for the functions/test names
test_move_one_step_boundary_on_continuousspace,
test_move_one_step_torus_wrap_on_continuousspace, the move_one_step helper,
DummyAgent and ContinuousSpace to locate and change the assertions.
|
@abhinavk0220, I was looking at some of the changes. Just to confirm if this is intended, here the generate function will be sending a And as of curiosity, are the change written to align with Mesa 4? |
…PI- Replace MultiGrid/SingleGrid with OrthogonalMooreGrid- Replace ContinuousSpace(x_max=, y_max=) with dimensions=[[]] API- Replace Model(seed=42) with Model()- Replace grid.place_agent() with cell.add_agent() + agent.cell assignment- Fix continuous space: set agent.pos directly instead of space.place_agent()- All 287 tests passing
for more information, see https://pre-commit.ci
Author
@IlamaranMagesh Yes, that was intentional! litellm returns ModelResponse | CustomStreamWrapper depending on the streaming mode. Using Any avoids importing litellm internals as a hard dependency. I'm happy to tighten it to ModelResponse | CustomStreamWrapper if you prefer! And yes, to answer your second question, the recent commits were specifically written to align the codebase with the new Mesa 4.x API. |
added 2 commits
March 8, 2026 00:26
…p (agent=neighbor, not agent=agent)- Fix step counter in record_model.py: use self._time instead of self.step (Mesa 4.x uses _time as the internal step counter, model.step is the method)
…hub.com/abhinavk0220/mesa-llm into improve/type-hints-llm-agent-module-llm
|
I understand, you are right. Tight coupling doesn't make sense now. @colinfrisch what's your say? |
Author
Thanks for understanding! If the preference is explicit typing, I can update it to: from litellm import ModelResponse, CustomStreamWrapper def generate(...) -> ModelResponse | CustomStreamWrapper: Just say the word and I'll push the update! 😊 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds comprehensive type hints and improved docstrings to the two core
files in mesa-llm:
LLMAgentandModuleLLM.Changes
mesa_llm/llm_agent.py-> None,-> str,-> Observation, etc.)__init___build_observation()return type:tuple[dict[str, Any], dict[str, Any]]apply_plan/aapply_planreturn type:list[dict[str, Any]]mesa_llm/module_llm.py_build_messages()return type:list[dict[str, str]]generate/ageneratereturn type:Any(litellm response object)__init__Testing
All 262 tests passing.
Related
Closes #[issue if any]
Part of GSoC 2026 contributions — #153
Summary by CodeRabbit
Bug Fixes
New Features
Tests