From 0386fe5f6a7c8742fd504453795990c82d2f2e1e Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 3 Sep 2025 13:43:47 -0400 Subject: [PATCH] swarm --- libs/langgraph/langgraph/agent/__init__.py | 2 + .../langgraph/agent/middleware/swarm.py | 71 ++++++++++++++++--- libs/langgraph/langgraph/agent/types.py | 6 +- .../langgraph/prebuilt/chat_agent_executor.py | 4 +- 4 files changed, 70 insertions(+), 13 deletions(-) diff --git a/libs/langgraph/langgraph/agent/__init__.py b/libs/langgraph/langgraph/agent/__init__.py index 9959cabab..32de10de4 100644 --- a/libs/langgraph/langgraph/agent/__init__.py +++ b/libs/langgraph/langgraph/agent/__init__.py @@ -144,6 +144,8 @@ def create_agent( first_node = ( f"{middleware_w_before[0].__class__.__name__}.before_model" if middleware_w_before + else f"{middleware_w_modify_model_request[0].__class__.__name__}.modify_model_request" + if middleware_w_modify_model_request else "model_request" ) last_node = ( diff --git a/libs/langgraph/langgraph/agent/middleware/swarm.py b/libs/langgraph/langgraph/agent/middleware/swarm.py index 8af362cb0..ffa966d2c 100644 --- a/libs/langgraph/langgraph/agent/middleware/swarm.py +++ b/libs/langgraph/langgraph/agent/middleware/swarm.py @@ -1,13 +1,66 @@ -from langgraph.agent.types import AgentMiddleware, AgentState, ModelRequest -from typing import Dict, Any, List, Optional, Union -from langgraph.types import interrupt +from langgraph.agent.types import AgentMiddleware, ModelRequest +from langchain_core.tools import BaseTool, tool +from langgraph.agent import create_agent +from langchain_core.messages import AIMessage, ToolMessage +from dataclasses import dataclass +from typing import cast + +@dataclass +class SwarmAgent: + name: str + system_prompt: str + tools: list[BaseTool] -class SwarmMiddleWare(AgentMiddleware): +class SwarmMiddleware(AgentMiddleware): + """Swarm middleware. + + TODOs: + * Support create_agent for handoffs + * Support handoff customization + * do we want to include handoff messages / enable togglging + * default active agent + * handoff tool naming / descriptions + """ + + class State(AgentMiddleware.State): + active_agent: str | None = None - def __init__(self, model_configs: dict[str, dict]): - super().__init__() + @staticmethod + def _create_handoff_tools(agents: list[SwarmAgent]) -> list[BaseTool]: - def modify_model_request( - self, request: ModelRequest, state: AgentState - ) -> ModelRequest: \ No newline at end of file + def handoff_tool(agent_name: str) -> str: + return f"Handing off to {agent_name}" + + return [ + tool(f"handoff_to_{agent.name}", description=f"Handoff tool to trigger a handoff to {agent.name}")(handoff_tool) + for agent in agents + ] + + def __init__(self, agents: list[SwarmAgent]): + self.agents: dict[str, SwarmAgent] = {agent.name: agent for agent in agents} + self.handoff_tools = self._create_handoff_tools(agents) + + def modify_model_request(self, request: ModelRequest, state: State) -> ModelRequest: + if (active_agent := getattr(state, "active_agent", None)) is not None: + agent = self.agents[active_agent] + request.system_prompt = agent.system_prompt + request.tools = agent.tools + + request.tools.extend(self.handoff_tools) + return request + + def after_model(self, state) -> State | None: + # TODO: handle parallel handoffs + + ai_msg: AIMessage = cast(AIMessage, state.messages[-1]) + if ai_msg.tool_calls: + for call in ai_msg.tool_calls: + if call["name"].startswith("handoff_to_"): + active_agent = call["args"]["agent_name"] + return { + "messages": [ToolMessage(name=call["name"], content=f"Successfully transferred to {active_agent}")], + "active_agent": active_agent, + "jump_to": "model", + } + return None \ No newline at end of file diff --git a/libs/langgraph/langgraph/agent/types.py b/libs/langgraph/langgraph/agent/types.py index cda32d6fd..2dc139692 100644 --- a/libs/langgraph/langgraph/agent/types.py +++ b/libs/langgraph/langgraph/agent/types.py @@ -21,16 +21,16 @@ JumpTo = Literal["tools", "model", "__end__"] class ModelRequest: model: BaseChatModel system_prompt: str - messages: Sequence[AnyMessage] # excluding system prompt + messages: list[AnyMessage] # excluding system prompt tool_choice: Any - tools: Sequence[BaseTool] + tools: list[BaseTool] response_format: ResponseFormat | None @dataclass class AgentState: messages: Annotated[list[AnyMessage], add_messages] - model_request: Annotated[ModelRequest | None, EphemeralValue] + model_request: Annotated[ModelRequest | None, EphemeralValue] = None jump_to: Annotated[JumpTo | None, EphemeralValue] = None response: dict | None = None diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index d205062b5..dccba8315 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -38,7 +38,6 @@ from typing_extensions import Annotated, NotRequired, TypedDict from langgraph._internal._runnable import RunnableCallable, RunnableLike from langgraph._internal._typing import MISSING -from langgraph.agent import create_agent from langgraph.agent.types import AgentMiddleware from langgraph.errors import ErrorCode, create_error_message from langgraph.graph import END, StateGraph @@ -472,6 +471,9 @@ def create_react_agent( assert pre_model_hook is None assert post_model_hook is None assert state_schema is None + + from langgraph.agent import create_agent + return create_agent( model=model, tools=tools,