feat: add LLM Schelling Segregation example using mesa-llm

llm/llm_schelling/README.md

@@ -0,0 +1,73 @@
+# LLM Schelling Segregation
+
+## Summary
+
+An LLM-powered reimplementation of Schelling's (1971) classic segregation model where agents
+**reason in natural language** about their neighborhood before deciding to stay or move —
+instead of following a fixed satisfaction threshold.
+
+## The Model
+
+Agents of two groups (A and B) are placed on a 5×5 grid. Each step:
+
+1. Each agent observes its Moore neighborhood (up to 8 neighbors on a torus grid)
+2. It describes the neighborhood composition to the LLM
+3. The LLM decides: `happy` (stay) or `unhappy` (move to a random empty cell)
+
+The simulation tracks happiness levels and a segregation index over time. The model stops
+when all agents are happy.
+
+## What makes this different from classical Schelling
+
+Classical Schelling uses a fixed threshold rule — an agent moves if fewer than X% of its
+neighbors share its group. The outcome is mathematically determined by that threshold.
+
+Here, agents **reason** at each step:
+
+> "I belong to Group A. My neighborhood has 3 Group A and 4 Group B neighbors.
+> The mix is manageable — I feel comfortable here. I'll stay."
+
+This produces **qualitatively different dynamics**. The LLM agents weigh context,
+not just ratios — and the result shows it.
+
+## Visualization
+
+**Initial state (Step 0) — random placement:**
+
+![Initial state — blue (Group A) and orange (Group B) agents distributed randomly](schelling_initial.png)
+
+Random mix across the 5×5 grid. No clustering, no preference expressed yet. Happiness and
+segregation charts are empty.
+
+**After 2 LLM reasoning steps — stable integration:**
+
+![Step 2 — all agents happy, zero segregation, model stopped](schelling_dashboard.png)
+
+| Step | Happy | Unhappy | Segregation Index | What happened |
+|------|-------|---------|-------------------|---------------|
+| 0 | — | — | — | Random initial placement |
+| 1 | ~16 | ~0 | 0.00 | LLM agents assess neighbors, most decide to stay |
+| 2 | 19 | 0 | 0.00 | All agents happy — simulation stops |
+
+**Why this matters:** The classical Schelling model *always* produces segregation from even
+mild preferences. The LLM version produced **zero segregation** — agents reasoned their way
+to comfort in a diverse neighborhood. No hardcoded tolerance parameter, no rule relaxation.
+The LLM considered context (neighborhood quality, stability, social mix) and decided the
+mixed state was acceptable. This is something a fixed-threshold agent cannot do.
+
+## How to Run
+
+```bash
+cp .env.example .env  # fill in your API key
+pip install -r requirements.txt
+solara run app.py
+```
+
+## Supported LLM Providers
+
+Gemini, OpenAI, Anthropic, Ollama (local) — configured via `.env`.
+
+## 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,134 @@
+import os
+import sys
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+load_dotenv(Path(__file__).parent / ".env")
+os.environ["PYTHONIOENCODING"] = "utf-8"
+if sys.stdout.encoding != "utf-8":
+    sys.stdout.reconfigure(encoding="utf-8")
+if sys.stderr.encoding != "utf-8":
+    sys.stderr.reconfigure(encoding="utf-8")
+
+import matplotlib.pyplot as plt  # noqa: E402
+import solara  # noqa: E402
+from llm_schelling.model import LLMSchellingModel  # noqa: E402
+from mesa.visualization import SolaraViz  # noqa: E402
+from mesa.visualization.utils import update_counter  # noqa: E402
+
+GROUP_COLORS = {0: "#2196F3", 1: "#FF5722"}  # Blue, Orange
+
+model_params = {
+    "width": {
+        "type": "SliderInt",
+        "value": 5,
+        "label": "Grid width",
+        "min": 5,
+        "max": 20,
+        "step": 1,
+    },
+    "height": {
+        "type": "SliderInt",
+        "value": 5,
+        "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)
+
+
+def SchellingStatsPlot(model):
+    """Combined happiness + segregation index chart."""
+    update_counter.get()
+
+    df = model.datacollector.get_model_vars_dataframe()
+
+    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 5), facecolor="#1a1a2e")
+    fig.patch.set_facecolor("#1a1a2e")
+
+    for ax in (ax1, ax2):
+        ax.set_facecolor("#16213e")
+        ax.tick_params(colors="white")
+        ax.xaxis.label.set_color("white")
+        ax.yaxis.label.set_color("white")
+        ax.title.set_color("white")
+        for spine in ax.spines.values():
+            spine.set_edgecolor("#444")
+
+    if not df.empty:
+        ax1.plot(df.index, df["happy"], color="#4CAF50", label="Happy", linewidth=2)
+        ax1.plot(df.index, df["unhappy"], color="#F44336", label="Unhappy", linewidth=2)
+    ax1.set_title("Agent Happiness", fontsize=11)
+    ax1.set_ylabel("Count", color="white")
+    ax1.legend(facecolor="#16213e", labelcolor="white", fontsize=8)
+
+    if not df.empty and "segregation_index" in df.columns:
+        ax2.plot(df.index, df["segregation_index"], color="#2196F3", linewidth=2)
+    ax2.set_title("Segregation Index", fontsize=11)
+    ax2.set_ylabel("Index", color="white")
+    ax2.set_xlabel("Step", color="white")
+
+    fig.tight_layout(pad=1.5)
+    return solara.FigureMatplotlib(fig)
+
+
+model = LLMSchellingModel()
+
+page = SolaraViz(
+    model,
+    components=[
+        SchellingGridPlot,
+        SchellingStatsPlot,
+    ],
+    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