Add Pydantic Field descriptions to improve LLM JSON schema parsing

mesa_llm/memory/episodic_memory.py

@@ -2,7 +2,7 @@
 from collections import deque
 from typing import TYPE_CHECKING
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from mesa_llm.memory.memory import Memory, MemoryEntry, _format_message_entry
 
@@ -11,7 +11,9 @@
 
 
 class EventGrade(BaseModel):
-    grade: int
+    grade: int = Field(
+        description="Integer score representing the importance of the event"
+    )
 
 
 def normalize_dict_values(scores: dict, min_target: float, max_target: float) -> dict:

mesa_llm/reasoning/react.py

@@ -1,7 +1,7 @@
 import json
 from typing import TYPE_CHECKING
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from mesa_llm.reasoning.reasoning import Observation, Plan, Reasoning
 
@@ -10,8 +10,10 @@
 
 
 class ReActOutput(BaseModel):
-    reasoning: str
-    action: str
+    reasoning: str = Field(
+        description="Step-by-step reasoning about the situation based on memory and observation"
+    )
+    action: str = Field(description="The specific action to take without using tools")
 
 
 class ReActReasoning(Reasoning):

tests/test_memory/test_episodic_memory.py

@@ -4,7 +4,11 @@
 
 import pytest
 
-from mesa_llm.memory.episodic_memory import EpisodicMemory, normalize_dict_values
+from mesa_llm.memory.episodic_memory import (
+    EpisodicMemory,
+    EventGrade,
+    normalize_dict_values,
+)
 from mesa_llm.memory.memory import MemoryEntry
 
 
@@ -38,6 +42,14 @@ def test_normalize_dict_floats_logic_when_empty():
 class TestEpisodicMemory:
     """Core functionality test"""
 
+    def test_event_grade_schema_includes_field_description(self):
+        """Structured output schema should keep the grading instructions."""
+        schema = EventGrade.model_json_schema()
+
+        assert schema["properties"]["grade"]["description"] == (
+            "Integer score representing the importance of the event"
+        )
+
     def test_memory_init(self, episodic_mock_agent):
         """Test EpisodicMemory class initialization with defaults and custom values"""
         memory = EpisodicMemory(

tests/test_reasoning/test_react.py

@@ -22,6 +22,17 @@ def test_react_output_creation(self):
         assert output.reasoning == "I need to move to a better position"
         assert output.action == "move_north"
 
+    def test_react_output_schema_includes_field_descriptions(self):
+        """Structured output schema should keep the field guidance text."""
+        schema = ReActOutput.model_json_schema()
+
+        assert schema["properties"]["reasoning"]["description"] == (
+            "Step-by-step reasoning about the situation based on memory and observation"
+        )
+        assert schema["properties"]["action"]["description"] == (
+            "The specific action to take without using tools"
+        )
+
 
 class TestReActReasoning:
     """Test the ReActReasoning class."""