feat: add LLM Schelling Segregation example using mesa-llm

llm/llm_schelling/README.md

@@ -0,0 +1,71 @@
+# LLM Schelling Segregation
+
+An LLM-powered implementation of Schelling's (1971) classic segregation model,
+built with [Mesa](https://github.com/projectmesa/mesa) and
+[Mesa-LLM](https://github.com/projectmesa/mesa-llm).
+
+## Overview
+
+The Schelling segregation model is one of the most influential agent-based
+models ever published. It demonstrates that even mild individual preferences
+for same-group neighbors produce strong global segregation — a striking example
+of emergent behavior from simple rules.
+
+**Classical model:** An agent moves if fewer than a fixed threshold (e.g. 30%)
+of its neighbors share its group.
+
+**This model:** Agents reason in natural language about their neighborhood
+composition and decide whether they feel comfortable staying or want to move.
+The LLM can weigh contextual factors, producing richer dynamics than a fixed
+threshold allows.
+
+## The Model
+
+Agents of two groups (A and B) are placed on a grid. Each step:
+1. Each agent observes its Moore neighborhood (up to 8 neighbors)
+2. It describes the neighborhood composition in natural language to the LLM
+3. The LLM decides: `happy` (stay) or `unhappy` (move)
+4. Unhappy agents relocate to a random empty cell
+
+The simulation tracks happiness levels and a segregation index over time.
+
+### Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `width` | Grid width | 10 |
+| `height` | Grid height | 10 |
+| `density` | Fraction of cells occupied | 0.8 |
+| `minority_fraction` | Fraction of agents in Group B | 0.4 |
+| `llm_model` | LLM model string | `gemini/gemini-2.0-flash` |
+
+## Running the Model
+
+Set your API key:
+```bash
+export GEMINI_API_KEY=your_key_here
+```
+
+Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+Run the visualization:
+```bash
+solara run app.py
+```
+
+## Comparison with Classical Schelling
+
+| Feature | Classical Schelling | LLM Schelling |
+|---------|--------------------|--------------------|
+| Decision rule | Fixed threshold (e.g. 30%) | LLM natural language reasoning |
+| Agent memory | None | Short-term memory of interactions |
+| Flexibility | Rigid | Emergent from reasoning |
+| Interpretability | Mathematical | Natural language explanations |
+
+## Reference
+
+Schelling, T.C. (1971). Dynamic models of segregation.
+*Journal of Mathematical Sociology*, 1(2), 143–186.

llm/llm_schelling/app.py

@@ -0,0 +1,90 @@
+import matplotlib.pyplot as plt
+import solara
+from llm_schelling.model import LLMSchellingModel
+from mesa.visualization import SolaraViz, make_plot_component
+from mesa.visualization.utils import update_counter
+
+GROUP_COLORS = {0: "#2196F3", 1: "#FF5722"}  # Blue, Orange
+
+model_params = {
+    "width": {
+        "type": "SliderInt",
+        "value": 10,
+        "label": "Grid width",
+        "min": 5,
+        "max": 20,
+        "step": 1,
+    },
+    "height": {
+        "type": "SliderInt",
+        "value": 10,
+        "label": "Grid height",
+        "min": 5,
+        "max": 20,
+        "step": 1,
+    },
+    "density": {
+        "type": "SliderFloat",
+        "value": 0.8,
+        "label": "Population density",
+        "min": 0.1,
+        "max": 1.0,
+        "step": 0.05,
+    },
+    "minority_fraction": {
+        "type": "SliderFloat",
+        "value": 0.4,
+        "label": "Minority fraction",
+        "min": 0.1,
+        "max": 0.5,
+        "step": 0.05,
+    },
+}
+
+
+def SchellingGridPlot(model):
+    """Visualize the grid showing agent groups and happiness."""
+    update_counter.get()
+
+    width = model.grid.dimensions[0]
+    height = model.grid.dimensions[1]
+
+    fig, ax = plt.subplots(figsize=(6, 6))
+    ax.set_xlim(0, width)
+    ax.set_ylim(0, height)
+    ax.set_aspect("equal")
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.set_title("Schelling Segregation (LLM)\nBlue=Group A, Orange=Group B, X=Unhappy")
+
+    for agent in model.agents:
+        x, y = agent.pos
+        color = GROUP_COLORS[agent.group]
+        marker = "o" if agent.is_happy else "x"
+        ax.plot(
+            x + 0.5,
+            y + 0.5,
+            marker=marker,
+            color=color,
+            markersize=8,
+            markeredgewidth=2,
+        )
+
+    return solara.FigureMatplotlib(fig)
+
+
+HappinessPlot = make_plot_component({"happy": "#4CAF50", "unhappy": "#F44336"})
+SegregationPlot = make_plot_component("segregation_index")
+
+model = LLMSchellingModel()
+
+page = SolaraViz(
+    model,
+    components=[
+        SchellingGridPlot,
+        HappinessPlot,
+        SegregationPlot,
+    ],
+    model_params=model_params,
+    name="LLM Schelling Segregation",
+)

llm/llm_schelling/llm_schelling/agent.py

@@ -0,0 +1,100 @@
+import logging
+
+from mesa_llm.llm_agent import LLMAgent
+from mesa_llm.reasoning.reasoning import Reasoning
+
+logger = logging.getLogger(__name__)
+
+
+class SchellingAgent(LLMAgent):
+    """
+    An LLM-powered agent in the Schelling Segregation model.
+
+    Unlike the classical Schelling model where agents move if fewer than
+    a fixed fraction of neighbors share their group, this agent reasons
+    about its neighborhood using an LLM and decides whether it feels
+    comfortable staying or wants to move.
+
+    Attributes:
+        group (int): The agent's group identity (0 or 1).
+        is_happy (bool): Whether the agent is satisfied with its location.
+    """
+
+    def __init__(self, model, reasoning: type[Reasoning], group: int):
+        group_label = "Group A" if group == 0 else "Group B"
+        other_label = "Group B" if group == 0 else "Group A"
+
+        system_prompt = f"""You are an agent in a social simulation. You belong to {group_label}.
+You are deciding whether you feel comfortable in your current neighborhood.
+Look at your neighbors: if too many belong to {other_label} and too few to {group_label},
+you may feel uncomfortable and want to move.
+Respond with ONLY one word: 'happy' if you want to stay, or 'unhappy' if you want to move."""
+
+        super().__init__(
+            model=model,
+            reasoning=reasoning,
+            system_prompt=system_prompt,
+            vision=1,
+            internal_state=["group", "is_happy"],
+        )
+        self.group = group
+        self.is_happy = True
+
+    def step(self):
+        """Decide whether to move based on LLM reasoning about neighborhood."""
+        obs = self.generate_obs()
+
+        # Count neighbors by group
+        neighbors = list(self.cell.neighborhood.agents) if self.cell else []
+        same = sum(
+            1 for n in neighbors if hasattr(n, "group") and n.group == self.group
+        )
+        total = len(neighbors)
+        different = total - same
+
+        if total == 0:
+            self.is_happy = True
+            self.internal_state = [f"group:{self.group}", "is_happy:True"]
+            return
+
+        group_label = "Group A" if self.group == 0 else "Group B"
+        other_label = "Group B" if self.group == 0 else "Group A"
+
+        step_prompt = f"""You belong to {group_label}.
+Your current neighborhood has:
+- {same} neighbors from {group_label} (your group)
+- {different} neighbors from {other_label} (the other group)
+- {total} neighbors total
+
+Do you feel comfortable here, or do you want to move to a different location?
+Respond with ONLY one word: 'happy' or 'unhappy'."""
+
+        plan = self.reasoning.plan(obs, step_prompt=step_prompt)
+
+        # Parse LLM response
+        response_text = ""
+        try:
+            if hasattr(plan, "llm_plan") and plan.llm_plan:
+                for block in plan.llm_plan:
+                    if hasattr(block, "text"):
+                        response_text += block.text.lower()
+        except Exception as e:
+            logger.warning("Failed to parse LLM response: %s", e)
+
+        self.is_happy = "unhappy" not in response_text
+        self.internal_state = [
+            f"group:{self.group}",
+            f"is_happy:{self.is_happy}",
+        ]
+
+        # If unhappy, move to a random empty cell
+        if not self.is_happy:
+            empty_cells = [
+                cell
+                for cell in self.model.grid.all_cells
+                if len(list(cell.agents)) == 0
+            ]
+            if empty_cells:
+                new_cell = self.model.random.choice(empty_cells)
+                self.cell = new_cell
+                self.pos = new_cell.coordinate

llm/llm_schelling/llm_schelling/model.py

@@ -0,0 +1,94 @@
+import mesa
+from mesa.discrete_space import OrthogonalMooreGrid
+from mesa_llm.reasoning.cot import CoTReasoning
+
+from .agent import SchellingAgent
+
+
+class LLMSchellingModel(mesa.Model):
+    """
+    An LLM-powered Schelling Segregation model.
+
+    The classical Schelling (1971) segregation model shows that even mild
+    individual preferences for same-group neighbors lead to strong global
+    segregation. In the classical model, agents move if fewer than a fixed
+    threshold fraction of their neighbors share their group.
+
+    This model replaces the threshold rule with LLM reasoning: agents
+    describe their neighborhood in natural language and decide whether they
+    feel comfortable staying or want to move. This produces richer dynamics
+    where the decision depends on framing, context, and reasoning — not
+    just a number.
+
+    Reference:
+        Schelling, T.C. (1971). Dynamic models of segregation.
+        Journal of Mathematical Sociology, 1(2), 143-186.
+
+    Args:
+        width (int): Grid width.
+        height (int): Grid height.
+        density (float): Fraction of cells that are occupied.
+        minority_fraction (float): Fraction of agents in group 1 (minority).
+        llm_model (str): LLM model string in 'provider/model' format.
+        rng: Random number generator.
+    """
+
+    def __init__(
+        self,
+        width: int = 10,
+        height: int = 10,
+        density: float = 0.8,
+        minority_fraction: float = 0.4,
+        llm_model: str = "gemini/gemini-2.0-flash",
+        rng=None,
+    ):
+        super().__init__(rng=rng)
+
+        self.grid = OrthogonalMooreGrid((width, height), torus=True, random=self.random)
+
+        self.datacollector = mesa.DataCollector(
+            model_reporters={
+                "happy": lambda m: sum(1 for a in m.agents if a.is_happy),
+                "unhappy": lambda m: sum(1 for a in m.agents if not a.is_happy),
+                "segregation_index": lambda m: self._segregation_index(m),
+            }
+        )
+
+        # Place agents on grid
+        for cell in self.grid.all_cells:
+            if self.random.random() < density:
+                group = 1 if self.random.random() < minority_fraction else 0
+                agent = SchellingAgent(
+                    model=self,
+                    reasoning=CoTReasoning,
+                    group=group,
+                )
+                agent.cell = cell
+                agent.pos = cell.coordinate
+
+        self.running = True
+        self.datacollector.collect(self)
+
+    def step(self):
+        """Advance the model by one step."""
+        self.agents.shuffle_do("step")
+        self.datacollector.collect(self)
+
+        # Stop if everyone is happy
+        if all(a.is_happy for a in self.agents):
+            self.running = False
+
+    @staticmethod
+    def _segregation_index(model):
+        """
+        Measure segregation as the average fraction of same-group neighbors.
+        Higher values indicate more segregation.
+        """
+        scores = []
+        for agent in model.agents:
+            neighbors = list(agent.cell.neighborhood.agents)
+            if not neighbors:
+                continue
+            same = sum(1 for n in neighbors if n.group == agent.group)
+            scores.append(same / len(neighbors))
+        return sum(scores) / len(scores) if scores else 0.0

llm/llm_schelling/requirements.txt

@@ -0,0 +1,2 @@
+mesa[viz]>=3.0
+mesa-llm>=0.1.0