Add Agriculture decision Model

examples/agriculture_model/__init__.py

@@ -0,0 +1 @@
+import examples.agriculture_model.tools

examples/agriculture_model/agent.py

@@ -0,0 +1,132 @@
+from enum import Enum
+
+from mesa_llm.llm_agent import LLMAgent
+from mesa_llm.memory.st_lt_memory import STLTMemory
+from mesa_llm.tools.tool_manager import ToolManager
+
+FARMER_TOOL_MANAGER = ToolManager()
+
+
+class CropState(Enum):
+    IDLE = "IDLE"
+    PLANTED = "PLANTED"
+    GROWING = "GROWING"
+    READY = "READY"
+
+
+class FarmerAgent(LLMAgent):
+    def __init__(
+        self,
+        model,
+        reasoning,
+        llm_model,
+        system_prompt,
+        vision,
+        internal_state,
+        step_prompt,
+    ):
+        super().__init__(
+            model=model,
+            reasoning=reasoning,
+            llm_model=llm_model,
+            system_prompt=system_prompt,
+            vision=vision,
+            internal_state=internal_state,
+            step_prompt=step_prompt,
+        )
+
+        self.tool_manager = FARMER_TOOL_MANAGER
+        self.land_size = self.random.randint(1, 5)
+        self.wealth = self.random.randint(1000, 50000)
+        self._base_internal_state = list(self.internal_state)
+
+        self.crop_type = self.random.choice(["wheat", "rice", "maize"])
+        self.crop_state = CropState.IDLE
+
+        self.plant_date = None
+        self.harvest_date = None
+        self.fertilizer = 0.0
+
+        self.yield_output = 0.0
+        self.profit = 0.0
+
+        self.memory = STLTMemory(
+            agent=self,
+            llm_model=llm_model,
+            display=True,
+        )
+
+    def observe_environment(self):
+        self.internal_state = [
+            *self._base_internal_state,
+            f"My land size is {self.land_size}",
+            f"On a scale of 1000 - 50000, my wealth is {self.wealth}",
+            f"My crop type is {self.crop_type}",
+            f"My crop state is {self.crop_state.value}",
+            f"Rainfall condition is {self.model.rainfall}",
+        ]
+
+    def update_crop_state(self):
+        if (
+            self.crop_state in {CropState.PLANTED, CropState.GROWING}
+            and self.harvest_date is not None
+            and self.model.current_day >= self.harvest_date
+        ):
+            self.crop_state = CropState.READY
+
+    def decide(self):
+        observation = self.generate_obs()
+
+        prompt = f"""
+        You are a farmer.
+        wealth: {self.wealth}
+        Crop type: {self.crop_type}
+        Crop state: {self.crop_state}
+        land size: {self.land_size}
+
+        Decide actions:
+        - plant_crop(days_to_harvest)
+        - apply_fertilizer(level)
+        - harvest_crop()
+        - speak_to()
+
+        Use tools smartly.
+        """
+
+        plan = self.reasoning.plan(
+            prompt=prompt,
+            obs=observation,
+            selected_tools=[
+                "plant_crop",
+                "apply_fertilizer",
+                "harvest_crop",
+                "speak_to",
+            ],
+        )
+
+        self.apply_plan(plan)
+
+    def compute_yield(self):
+        base_yield = 1000 * self.land_size
+
+        if self.model.rainfall == "LOW":
+            rain_factor = 0.7
+        elif self.model.rainfall == "HIGH":
+            rain_factor = 1.3
+        else:
+            rain_factor = 1.0
+
+        fert_factor = 1 + 0.3 * self.fertilizer
+        noise = self.random.uniform(0.9, 1.1)
+
+        self.profit = 0
+
+        self.yield_output = base_yield * rain_factor * fert_factor * noise
+        price = self.model.market_price[self.crop_type]
+        self.profit += self.yield_output * price
+
+    def step(self):
+        self.observe_environment()
+        self.update_crop_state()
+        self.decide()
+        self.update_crop_state()

examples/agriculture_model/app.py

@@ -0,0 +1,126 @@
+import logging
+import warnings
+
+from dotenv import load_dotenv
+from mesa.visualization import (
+    Slider,
+    SolaraViz,
+    make_plot_component,
+    make_space_component,
+)
+
+from examples.agriculture_model.agent import CropState, FarmerAgent
+from examples.agriculture_model.model import FarmerModel
+from mesa_llm.parallel_stepping import enable_automatic_parallel_stepping
+from mesa_llm.reasoning.react import ReActReasoning
+
+warnings.filterwarnings(
+    "ignore",
+    category=UserWarning,
+    module="pydantic.main",
+    message=r".*Pydantic serializer warnings.*",
+)
+logging.getLogger("pydantic").setLevel(logging.ERROR)
+
+enable_automatic_parallel_stepping()
+load_dotenv()
+
+
+CROP_COLORS = {
+    CropState.IDLE: "#D7C7A3",
+    CropState.PLANTED: "#A3D977",
+    CropState.GROWING: "#6FBF73",
+    CropState.READY: "#FFBD66",
+}
+
+
+model_params = {
+    "initial_farmers": Slider("Number of Farmers", 10, 1, 50, 1),
+    "rainfall": Slider("Rainfall Level", 0.5, 0.0, 1.0, 0.1),
+    "seed": {
+        "type": "InputText",
+        "value": 42,
+        "label": "Random Seed",
+    },
+    "width": 20,
+    "height": 20,
+    "reasoning": ReActReasoning,
+    "llm_model": "ollama/llama3.1:latest",
+    "vision": 3,
+    "parallel_stepping": False,
+}
+
+model = FarmerModel(
+    initial_farmers=10,
+    rainfall=0.5,
+    width=model_params["width"],
+    height=model_params["height"],
+    reasoning=model_params["reasoning"],
+    llm_model=model_params["llm_model"],
+    vision=model_params["vision"],
+    seed=model_params["seed"]["value"],
+    parallel_stepping=model_params["parallel_stepping"],
+)
+
+
+def farmer_portrayal(agent):
+    if agent is None:
+        return None
+
+    portrayal = {"size": 50}
+
+    if isinstance(agent, FarmerAgent):
+        portrayal["color"] = CROP_COLORS.get(agent.crop_state, "#888888")
+
+    return portrayal
+
+
+def post_process(ax):
+    ax.set_aspect("equal")
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.get_figure().set_size_inches(8, 8)
+
+
+space_component = make_space_component(
+    farmer_portrayal,
+    post_process=post_process,
+    draw_grid=False,
+)
+
+crop_chart = make_plot_component(
+    {
+        "Rice": "#4CAF50",
+        "Wheat": "#FFC107",
+        "Maize": "#FF5722",
+    }
+)
+
+profit_chart = make_plot_component(
+    {
+        "Total Profit": "#2196F3",
+        "Average Profit": "#9C27B0",
+    }
+)
+
+state_chart = make_plot_component(
+    {
+        "Planted": "#A3D977",
+        "Growing": "#6FBF73",
+        "Ready": "#FFBD66",
+        "Idle": "#D7C7A3",
+    }
+)
+
+
+page = SolaraViz(
+    model,
+    components=[
+        space_component,
+        crop_chart,
+        profit_chart,
+        state_chart,
+    ],
+    model_params=model_params,
+    name="Agriculture Decision Model (Mesa-LLM)",
+)

examples/agriculture_model/model.py

@@ -0,0 +1,135 @@
+from datetime import datetime, timedelta
+
+from mesa import Model
+from mesa.datacollection import DataCollector
+from mesa.space import MultiGrid
+
+from examples.agriculture_model.agent import CropState, FarmerAgent
+from mesa_llm.reasoning.reasoning import Reasoning
+from mesa_llm.recording.record_model import record_model
+
+
+@record_model(output_dir="recordings")
+class FarmerModel(Model):
+    def __init__(
+        self,
+        initial_farmers: int,
+        rainfall: int,
+        width: int,
+        height: int,
+        reasoning: type[Reasoning],
+        llm_model: str,
+        vision: int,
+        parallel_stepping=True,
+        seed=None,
+    ):
+        normalized_seed = None if seed in (None, "") else int(seed)
+        super().__init__(seed=normalized_seed)
+        self.width = width
+        self.height = height
+        self.parallel_stepping = parallel_stepping
+        self.grid = MultiGrid(width, height, torus=False)
+
+        self.start_date = datetime(2024, 6, 1)
+        self.current_day = self.start_date
+
+        self.rainfall = self._normalize_rainfall(rainfall=rainfall)
+        self.market_price = {
+            "wheat": 2.0,
+            "rice": 3.0,
+            "maize": 1.5,
+        }
+
+        for _ in range(initial_farmers):
+            agent = FarmerAgent(
+                model=self,
+                reasoning=reasoning,
+                llm_model=llm_model,
+                system_prompt="You are a smart farmer.",
+                step_prompt="Decide actions.",
+                vision=vision,
+                internal_state=["hardworking"],
+            )
+            x = self.random.randrange(self.grid.width)
+            y = self.random.randrange(self.grid.height)
+            self.grid.place_agent(agent, (x, y))
+
+        self.datacollector = DataCollector(
+            model_reporters={
+                "Rice": lambda m: sum(
+                    1
+                    for a in m.agents
+                    if isinstance(a, FarmerAgent) and a.crop_type == "rice"
+                ),
+                "Wheat": lambda m: sum(
+                    1
+                    for a in m.agents
+                    if isinstance(a, FarmerAgent) and a.crop_type == "wheat"
+                ),
+                "Maize": lambda m: sum(
+                    1
+                    for a in m.agents
+                    if isinstance(a, FarmerAgent) and a.crop_type == "maize"
+                ),
+                "Total Profit": lambda m: sum(
+                    a.profit for a in m.agents if isinstance(a, FarmerAgent)
+                ),
+                "Average Profit": lambda m: (
+                    sum(a.profit for a in m.agents if isinstance(a, FarmerAgent))
+                    / max(
+                        1,
+                        sum(1 for a in m.agents if isinstance(a, FarmerAgent)),
+                    )
+                ),
+                "Planted": lambda m: sum(
+                    1
+                    for a in m.agents
+                    if isinstance(a, FarmerAgent) and a.crop_state == CropState.PLANTED
+                ),
+                "Growing": lambda m: sum(
+                    1
+                    for a in m.agents
+                    if isinstance(a, FarmerAgent) and a.crop_state == CropState.GROWING
+                ),
+                "Ready": lambda m: sum(
+                    1
+                    for a in m.agents
+                    if isinstance(a, FarmerAgent) and a.crop_state == CropState.READY
+                ),
+                "Idle": lambda m: sum(
+                    1
+                    for a in m.agents
+                    if isinstance(a, FarmerAgent) and a.crop_state == CropState.IDLE
+                ),
+            }
+        )
+
+    def _normalize_rainfall(self, rainfall: float | str | None) -> str:
+        if isinstance(rainfall, str):
+            rainfall_upper = rainfall.upper()
+            if rainfall_upper in {"LOW", "NORMAL", "HIGH"}:
+                return rainfall_upper
+            raise ValueError(f"Unsupported rainfall value: {rainfall}")
+
+        if rainfall is None:
+            return "NORMAL"
+        if rainfall < 0.33:
+            return "LOW"
+        if rainfall > 0.66:
+            return "HIGH"
+        return "NORMAL"
+
+    def step(self):
+        self.current_day += timedelta(days=1)
+        self.agents.shuffle_do("step")
+        self.datacollector.collect(self)
+
+        if (self.current_day - self.start_date).days >= 120:
+            self.running = False
+
+
+if __name__ == "__main__":
+    from examples.agriculture_model.app import model
+
+    for _ in range(5):
+        model.step()