more swarm progress

This commit is contained in:
Sydney Runkle
2025-09-03 15:14:33 -04:00
parent 0386fe5f6a
commit f6d0382d66
2 changed files with 56 additions and 29 deletions
+28 -15
View File
@@ -66,15 +66,6 @@ def create_agent(
if m.__class__.after_model is not AgentMiddleware.after_model
]
default_model_request = ModelRequest(
model=model,
tools=list(tool_node.tools_by_name.values()),
system_prompt=system_prompt,
response_format=response_format,
messages=[],
tool_choice=None,
)
# create graph, add nodes
graph = StateGraph(
AgentState,
@@ -85,7 +76,14 @@ def create_agent(
def model_request(state: AgentState) -> AgentState:
request = state.model_request or default_model_request
request = state.model_request or ModelRequest(
model=model,
tools=list(tool_node.tools_by_name.values()),
system_prompt=system_prompt,
response_format=response_format,
messages=state.messages,
tool_choice=None,
)
# prepare messages
if request.system_prompt:
@@ -102,10 +100,8 @@ def create_agent(
)
return {"messages": output["raw"], "response": output["parsed"]}
else:
model_ = request.model
output = model_.invoke(
messages, tools=request.tools, tool_choice=request.tool_choice
)
model_ = request.model.bind_tools(request.tools, tool_choice=request.tool_choice)
output = model_.invoke(messages)
if state.response is not None:
return {"messages": output, "response": None}
else:
@@ -125,6 +121,16 @@ def create_agent(
def modify_model_request_node(state: AgentState) -> dict[str, ModelRequest]:
# TODO assert request.tools in tools, or pass them to tool node
default_model_request = ModelRequest(
model=model,
tools=list(tool_node.tools_by_name.values()),
system_prompt=system_prompt,
response_format=response_format,
messages=state.messages,
tool_choice=None,
)
return {"model_request": m.modify_model_request(state.model_request or default_model_request, state)}
graph.add_node(
@@ -162,7 +168,7 @@ def create_agent(
[first_node, END],
)
graph.add_conditional_edges(
last_node, _make_model_to_tools_edge(first_node), ["tools", END]
last_node, _make_model_to_tools_edge(first_node), [first_node, "tools", END]
)
# add before model edges
@@ -216,6 +222,13 @@ def create_agent(
f"{m2.__class__.__name__}.after_model",
first_node,
)
# _add_middleware_edge(
# graph,
# middleware_w_after[-1].after_model,
# f"{middleware_w_after[-1].__class__.__name__}.after_model",
# "model_request",
# first_node,
# )
return graph
@@ -1,10 +1,11 @@
from langgraph.agent.types import AgentMiddleware, ModelRequest
from langgraph.agent.types import AgentMiddleware, ModelRequest, AgentJump
from langchain_core.tools import BaseTool, tool
from langgraph.agent import create_agent
from langchain_core.messages import AIMessage, ToolMessage
from langchain_core.messages import AIMessage, ToolMessage, HumanMessage
from dataclasses import dataclass
from typing import cast
@dataclass
class SwarmAgent:
name: str
@@ -14,28 +15,35 @@ class SwarmAgent:
class SwarmMiddleware(AgentMiddleware):
"""Swarm middleware.
TODOs:
* Support create_agent for handoffs
* Support handoff customization
* 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
@staticmethod
def _create_handoff_tools(agents: list[SwarmAgent]) -> list[BaseTool]:
handoff_tools: list[BaseTool] = []
def handoff_tool(agent_name: str) -> str:
return f"Handing off to {agent_name}"
for agent in agents:
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 handoff_tool() -> str:
return f"Handing off to {agent.name}"
handoff_tools.append(
tool(
f"handoff_to_{agent.name}",
description=f"Handoff tool to trigger a handoff to {agent.name}",
)(handoff_tool)
)
return handoff_tools
def __init__(self, agents: list[SwarmAgent]):
self.agents: dict[str, SwarmAgent] = {agent.name: agent for agent in agents}
@@ -51,15 +59,21 @@ class SwarmMiddleware(AgentMiddleware):
return request
def after_model(self, state) -> State | None:
# TODO: handle parallel handoffs
# TODO: handle parallel handoffs, we don't do this currently
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"]
active_agent = call["name"].replace("handoff_to_", "")
return {
"messages": [ToolMessage(name=call["name"], content=f"Successfully transferred to {active_agent}")],
"messages": [
ToolMessage(
name=call["name"],
content=f"Successfully transferred to {active_agent}",
tool_call_id=call["id"],
)
],
"active_agent": active_agent,
"jump_to": "model",
}