From d07ce8ad12ec4d970d646a00fa3fd30f601dc0d4 Mon Sep 17 00:00:00 2001 From: Michael Siu Date: Wed, 2 Apr 2025 04:33:28 -0700 Subject: [PATCH 1/2] add llm graph-based orchestrator prototype --- mesa/experimental/llm_layer/__init__.py | 0 mesa/experimental/llm_layer/agents.py | 27 +++++++++++++++++ mesa/experimental/llm_layer/models.py | 24 +++++++++++++++ mesa/experimental/llm_layer/orchestrator.py | 33 +++++++++++++++++++++ mesa/experimental/llm_layer/run.py | 20 +++++++++++++ 5 files changed, 104 insertions(+) create mode 100644 mesa/experimental/llm_layer/__init__.py create mode 100644 mesa/experimental/llm_layer/agents.py create mode 100644 mesa/experimental/llm_layer/models.py create mode 100644 mesa/experimental/llm_layer/orchestrator.py create mode 100644 mesa/experimental/llm_layer/run.py diff --git a/mesa/experimental/llm_layer/__init__.py b/mesa/experimental/llm_layer/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/mesa/experimental/llm_layer/agents.py b/mesa/experimental/llm_layer/agents.py new file mode 100644 index 00000000000..a054c91205e --- /dev/null +++ b/mesa/experimental/llm_layer/agents.py @@ -0,0 +1,27 @@ +from mesa import Agent + +class BasicAgent(Agent): + def __init__(self, model): + # Pass the parameters to the parent class. + super().__init__(model) + + def step(self): + if self.random.random() > 0.5: + print(f"[Basic {self.unique_id}] Collecting nearby resource.") + else: + print(f"[Basic {self.unique_id}] Waiting...") + +class CognitiveAgent(Agent): + def __init__(self, model, orchestrator): + super().__init__(model) + self.orchestrator = orchestrator + self.memory = [] + + def step(self): + context = { + "goal": "collect", + "memory": self.memory + } + action = self.orchestrator.execute_graph("plan", self, context) + self.memory.append(action) + print(f"[Cognitive {self.unique_id}] {action}") \ No newline at end of file diff --git a/mesa/experimental/llm_layer/models.py b/mesa/experimental/llm_layer/models.py new file mode 100644 index 00000000000..d95a5d02cad --- /dev/null +++ b/mesa/experimental/llm_layer/models.py @@ -0,0 +1,24 @@ +from mesa import Model +from agents import BasicAgent, CognitiveAgent + +class HybridModel(Model): + def __init__(self, num_agents, orchestrator): + super().__init__() + self.num_agents = num_agents + self.orchestrator = orchestrator + + # Create 3 basic agents and 2 cognitive agents + BasicAgent.create_agents(model=self, n=3) + CognitiveAgent.create_agents(model=self, n=2, orchestrator=orchestrator) + + def step(self): + self.agents.shuffle_do("step") + + + + +def llm_collect(agent, state): + return f"[LLM] Reasoned to collect based on memory length {len(state['memory'])}" + +def wait_tool(agent, state): + return "[RULE] Wait due to uncertainty or cooldown" \ No newline at end of file diff --git a/mesa/experimental/llm_layer/orchestrator.py b/mesa/experimental/llm_layer/orchestrator.py new file mode 100644 index 00000000000..fbdfbc903ae --- /dev/null +++ b/mesa/experimental/llm_layer/orchestrator.py @@ -0,0 +1,33 @@ +# LangGraph-inspired orchestrator (future-proof structure) +class Orchestrator: + def __init__(self): + self.nodes = {} + self.edges = {} + + def add_node(self, name, func): + self.nodes[name] = func + + def add_edge(self, from_node, to_node): + self.edges.setdefault(from_node, []).append(to_node) + + def add_conditional_edges(self, from_node, condition_fn): + self.edges[from_node] = [(condition_fn, target) for target in self.nodes if target != from_node] + + def execute_graph(self, start_node, agent, state): + current_node = start_node + while current_node: + result = self.nodes[current_node](agent, state) + state["last_output"] = result + next_node = self._resolve_next_node(current_node, state) + if not next_node: + break + current_node = next_node + return result + + def _resolve_next_node(self, current_node, state): + if current_node not in self.edges: + return None + for condition_fn, target in self.edges[current_node]: + if condition_fn(state): + return target + return None \ No newline at end of file diff --git a/mesa/experimental/llm_layer/run.py b/mesa/experimental/llm_layer/run.py new file mode 100644 index 00000000000..3faae0016df --- /dev/null +++ b/mesa/experimental/llm_layer/run.py @@ -0,0 +1,20 @@ +from models import HybridModel, llm_collect, wait_tool +from orchestrator import Orchestrator + + +# Set up the orchestrator and graph +orchestrator = Orchestrator() +orchestrator.add_node("plan", llm_collect) +orchestrator.add_node("wait", wait_tool) + +# Add conditional edge: alternate between collecting and waiting +orchestrator.add_conditional_edges( + "plan", + lambda state: "wait" if len(state["memory"]) % 2 == 0 else None +) + +# Run the model +model = HybridModel(num_agents=5, orchestrator=orchestrator) +for step in range(3): + print(f"\n--- Step {step + 1} ---") + model.step() \ No newline at end of file From 799a55da877a47477f87aa0bcfa5b1131a70211a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 2 Apr 2025 12:05:06 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mesa/experimental/llm_layer/agents.py | 11 +++++------ mesa/experimental/llm_layer/models.py | 9 +++++---- mesa/experimental/llm_layer/orchestrator.py | 6 ++++-- mesa/experimental/llm_layer/run.py | 6 ++---- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/mesa/experimental/llm_layer/agents.py b/mesa/experimental/llm_layer/agents.py index a054c91205e..2be6c9635bb 100644 --- a/mesa/experimental/llm_layer/agents.py +++ b/mesa/experimental/llm_layer/agents.py @@ -1,16 +1,18 @@ from mesa import Agent + class BasicAgent(Agent): def __init__(self, model): # Pass the parameters to the parent class. super().__init__(model) - + def step(self): if self.random.random() > 0.5: print(f"[Basic {self.unique_id}] Collecting nearby resource.") else: print(f"[Basic {self.unique_id}] Waiting...") + class CognitiveAgent(Agent): def __init__(self, model, orchestrator): super().__init__(model) @@ -18,10 +20,7 @@ def __init__(self, model, orchestrator): self.memory = [] def step(self): - context = { - "goal": "collect", - "memory": self.memory - } + context = {"goal": "collect", "memory": self.memory} action = self.orchestrator.execute_graph("plan", self, context) self.memory.append(action) - print(f"[Cognitive {self.unique_id}] {action}") \ No newline at end of file + print(f"[Cognitive {self.unique_id}] {action}") diff --git a/mesa/experimental/llm_layer/models.py b/mesa/experimental/llm_layer/models.py index d95a5d02cad..f42664c16c2 100644 --- a/mesa/experimental/llm_layer/models.py +++ b/mesa/experimental/llm_layer/models.py @@ -1,6 +1,8 @@ -from mesa import Model from agents import BasicAgent, CognitiveAgent +from mesa import Model + + class HybridModel(Model): def __init__(self, num_agents, orchestrator): super().__init__() @@ -15,10 +17,9 @@ def step(self): self.agents.shuffle_do("step") - - def llm_collect(agent, state): return f"[LLM] Reasoned to collect based on memory length {len(state['memory'])}" + def wait_tool(agent, state): - return "[RULE] Wait due to uncertainty or cooldown" \ No newline at end of file + return "[RULE] Wait due to uncertainty or cooldown" diff --git a/mesa/experimental/llm_layer/orchestrator.py b/mesa/experimental/llm_layer/orchestrator.py index fbdfbc903ae..1038f16bc29 100644 --- a/mesa/experimental/llm_layer/orchestrator.py +++ b/mesa/experimental/llm_layer/orchestrator.py @@ -11,7 +11,9 @@ def add_edge(self, from_node, to_node): self.edges.setdefault(from_node, []).append(to_node) def add_conditional_edges(self, from_node, condition_fn): - self.edges[from_node] = [(condition_fn, target) for target in self.nodes if target != from_node] + self.edges[from_node] = [ + (condition_fn, target) for target in self.nodes if target != from_node + ] def execute_graph(self, start_node, agent, state): current_node = start_node @@ -30,4 +32,4 @@ def _resolve_next_node(self, current_node, state): for condition_fn, target in self.edges[current_node]: if condition_fn(state): return target - return None \ No newline at end of file + return None diff --git a/mesa/experimental/llm_layer/run.py b/mesa/experimental/llm_layer/run.py index 3faae0016df..06671986e94 100644 --- a/mesa/experimental/llm_layer/run.py +++ b/mesa/experimental/llm_layer/run.py @@ -1,7 +1,6 @@ from models import HybridModel, llm_collect, wait_tool from orchestrator import Orchestrator - # Set up the orchestrator and graph orchestrator = Orchestrator() orchestrator.add_node("plan", llm_collect) @@ -9,12 +8,11 @@ # Add conditional edge: alternate between collecting and waiting orchestrator.add_conditional_edges( - "plan", - lambda state: "wait" if len(state["memory"]) % 2 == 0 else None + "plan", lambda state: "wait" if len(state["memory"]) % 2 == 0 else None ) # Run the model model = HybridModel(num_agents=5, orchestrator=orchestrator) for step in range(3): print(f"\n--- Step {step + 1} ---") - model.step() \ No newline at end of file + model.step()