from __future__ import annotations
from typing import Any, Callable, Literal, Optional, Sequence, Type, Union, cast, get_args, AsyncIterator, Iterator
from uuid import UUID, uuid5
try:
from langchain.agents import AgentState, AgentStateWithStructuredResponse
except ImportError: # LangChain < 1.1.4
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage
from typing_extensions import Annotated, NotRequired, TypedDict
class AgentState(TypedDict):
"""Minimal supervisor state compatible with LangGraph 1.0."""
messages: Annotated[Sequence[AnyMessage], add_messages]
class AgentStateWithStructuredResponse(AgentState):
structured_response: Any
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AnyMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import BaseTool
from langgraph._internal._config import patch_configurable
from langgraph._internal._runnable import Runnable, RunnableCallable, RunnableLike
from langchain.messages import SystemMessage
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langgraph.typing import InputT, NodeInputT
from langgraph.types import StreamMode
from langchain.agents.factory import create_agent
from langchain.agents.structured_output import ResponseFormat
from langchain.agents.middleware import AgentMiddleware
from langchain.agents.middleware.types import StateT_co, ContextT, ResponseT
from langgraph.pregel import Pregel
from langgraph.pregel.remote import RemoteGraph
from typing_extensions import Annotated, TypedDict, Unpack
from aidputils.agents.toolkit.agent_name import AgentNameMode, with_agent_name
from aidputils.agents.toolkit.handoff import (
METADATA_KEY_HANDOFF_DESTINATION,
_normalize_agent_name,
create_handoff_back_messages,
create_handoff_tool,
)
OutputMode = Literal["full_history", "last_message"]
"""Mode for adding agent outputs to the message history in the multi-agent workflow
- `full_history`: add the entire agent message history
- `last_message`: add only the last message
"""
MODELS_NO_PARALLEL_TOOL_CALLS = {"o3-mini", "o3", "o4-mini"}
def _make_call_agent(
agent: Pregel[Any],
output_mode: OutputMode,
add_handoff_back_messages: bool,
supervisor_name: str,
) -> Callable[[dict], dict] | Runnable[NodeInputT, Any]:
if output_mode not in get_args(OutputMode):
raise ValueError(
f"Invalid agent output mode: {output_mode}. Needs to be one of {get_args(OutputMode)}"
)
def _process_output(output: dict | None) -> dict:
if not output:
return {"messages": []}
messages = output.get("messages") or output.get("agent", {}).get("messages")
if not messages:
return output
if output_mode == "full_history":
pass
elif output_mode == "last_message":
if isinstance(messages[-1], ToolMessage):
messages = messages[-2:]
else:
messages = messages[-1:]
else:
raise ValueError(
f"Invalid agent output mode: {output_mode}. "
f"Needs to be one of {OutputMode.__args__}"
)
if add_handoff_back_messages:
messages.extend(create_handoff_back_messages(agent.name, supervisor_name))
updated_output = dict(output)
updated_output["messages"] = messages
return updated_output
# langgraph protocol.py agent i nvoke signature
# def astream(
# self,
# input: InputT | Command | None,
# config: RunnableConfig | None = None,
# *,
# context: ContextT | None = None,
# stream_mode: StreamMode | list[StreamMode] | None = None,
# interrupt_before: All | Sequence[str] | None = None,
# interrupt_after: All | Sequence[str] | None = None,
# subgraphs: bool = False,
# ) -> AsyncIterator[dict[str, Any] | Any]:
def call_agent(input: InputT, config: RunnableConfig,
context: ContextT = None, stream_mode: StreamMode = "updates") -> dict:
thread_id = (context or {}).get("thread_id") if isinstance(context, dict) else None
if thread_id and isinstance(agent, RemoteGraph):
config = patch_configurable(
config,
{"thread_id": str(uuid5(UUID(str(thread_id)), agent.name))},
)
stream_out = agent.stream(input=input, config=config, context=context, stream_mode=stream_mode)
output = read_dict_from_stream(stream_out)
return _process_output(output or {})
async def acall_agent(input: InputT, config: RunnableConfig,
context: ContextT = None, stream_mode: StreamMode = "updates") -> dict:
thread_id = (context or {}).get("thread_id") if isinstance(context, dict) else None
if thread_id and isinstance(agent, RemoteGraph):
config = patch_configurable(
config,
{"thread_id": str(uuid5(UUID(str(thread_id)), agent.name))},
)
stream_out = agent.astream(input=input, config=config, context=context, stream_mode=stream_mode)
output = await read_async_dict_from_stream(stream_out)
return _process_output(output or {})
return RunnableCallable(name=agent.name, func=call_agent, afunc=acall_agent)
[docs]
async def read_async_dict_from_stream(
stream: AsyncIterator[dict[str, Any] | Any]
) -> dict[str, Any]:
"""
Reads all chunks from an async stream and aggregates them into a single dictionary.
This assumes that each chunk can be safely merged into the final dictionary
using dict.update().
"""
full_data: dict[str, Any] = {}
async for chunk in stream:
if isinstance(chunk, dict):
full_data.update(chunk)
else:
# Handle cases where the 'Any' part of the type hint might appear.
# This logic depends heavily on what 'Any' represents in your specific stream.
# For demonstration, we raise an error if it's not a dict chunk:
raise TypeError(f"Expected a dictionary chunk, but received type: {type(chunk)}")
return full_data
[docs]
def read_dict_from_stream(
stream: Iterator[dict[str, Any] | Any]
) -> dict[str, Any]:
"""
Reads all chunks from an async stream and aggregates them into a single dictionary.
This assumes that each chunk can be safely merged into the final dictionary
using dict.update().
"""
full_data: dict[str, Any] = {}
for chunk in stream:
if isinstance(chunk, dict):
full_data.update(chunk)
else:
# Handle cases where the 'Any' part of the type hint might appear.
# This logic depends heavily on what 'Any' represents in your specific stream.
# For demonstration, we raise an error if it's not a dict chunk:
raise TypeError(f"Expected a dictionary chunk, but received type: {type(chunk)}")
return full_data
def _get_handoff_destinations(tools: Sequence[BaseTool | Callable]) -> list[str]:
"""Extract handoff destinations from provided tools.
Args:
tools: List of tools to inspect.
Returns:
List of agent names that are handoff destinations.
"""
return [
tool.metadata[METADATA_KEY_HANDOFF_DESTINATION]
for tool in tools
if isinstance(tool, BaseTool)
and tool.metadata is not None
and METADATA_KEY_HANDOFF_DESTINATION in tool.metadata
]
def _prepare_tool_node(
tools: list[BaseTool | Callable] | ToolNode | None,
handoff_tool_prefix: Optional[str],
add_handoff_messages: bool,
agent_names: set[str],
) -> ToolNode:
"""Prepare the tool node used by the supervisor.
The supervisor can work with either a pre-built ``ToolNode`` or a plain list of tool
definitions. This helper normalizes those inputs and guarantees that handoff tools are
present for every managed agent.
This function takes in a list of tools, a handoff tool prefix, a flag to add handoff messages,
and a set of agent names. It returns a ToolNode that can be used in the supervisor agent.
Args:
tools: A list of tools or a ToolNode instance. Can be None.
handoff_tool_prefix: An optional prefix for the handoff tools.
add_handoff_messages: A boolean flag to indicate whether to add handoff messages.
agent_names: A set of agent names.
Returns:
A ToolNode instance. The ToolNode instance contains a Sequence list of
tools that can be called.
Raises:
ValueError: If custom handoff tools are provided but not for all subagents.
"""
if isinstance(tools, ToolNode):
input_tool_node = tools
tool_classes = list(tools.tools_by_name.values())
elif tools:
input_tool_node = ToolNode(tools)
# get the tool functions wrapped in a tool class from the ToolNode
tool_classes = list(input_tool_node.tools_by_name.values())
else:
input_tool_node = None
tool_classes = []
handoff_destinations = _get_handoff_destinations(tool_classes)
if handoff_destinations:
if missing_handoff_destinations := set(agent_names) - set(handoff_destinations):
raise ValueError(
"When providing custom handoff tools, you must provide them for all subagents. "
f"Missing handoff tools for agents '{missing_handoff_destinations}'."
)
# Handoff tools should be already provided here
tool_node = cast(ToolNode, input_tool_node)
else:
handoff_tools = [
create_handoff_tool(
agent_name=agent_name,
name=(
None
if handoff_tool_prefix is None
else f"{handoff_tool_prefix}{_normalize_agent_name(agent_name)}"
),
add_handoff_messages=add_handoff_messages,
)
for agent_name in agent_names
]
all_tools = tool_classes + list(handoff_tools)
# re-wrap the combined tools in a ToolNode
# if the original input was a ToolNode, apply the same params
if input_tool_node is not None:
tool_node = ToolNode(
all_tools,
name=str(input_tool_node.name),
tags=list(input_tool_node.tags) if input_tool_node.tags else None,
handle_tool_errors=input_tool_node._handle_tool_errors,
messages_key=input_tool_node._messages_key,
)
else:
tool_node = ToolNode(all_tools)
return tool_node
class _OuterState(TypedDict):
"""The state of the supervisor workflow."""
messages: Annotated[Sequence[AnyMessage], add_messages]
[docs]
def create_supervisor(
agents: list[Pregel],
*,
model: BaseChatModel,
tools: Sequence[BaseTool | Callable | dict[str, Any]] | None = None,
system_prompt: str | SystemMessage | None = None,
middleware: Sequence[AgentMiddleware[StateT_co, ContextT]] = (),
response_format: ResponseFormat[ResponseT] | type[ResponseT] | None = None,
state_schema: Type[AgentState] | None = None,
context_schema: Type[Any] | None = None,
supervisor_name: str = "supervisor",
# cache: BaseCache | None = None,
output_mode: OutputMode = "last_message", # "last_message",
add_handoff_messages: bool = True,
handoff_tool_prefix: Optional[str] = None,
include_agent_name: AgentNameMode | None = None,
add_handoff_back_messages: Optional[bool] = None,
) -> StateGraph:
"""Create a multi-agent supervisor workflow.
The generated graph places a supervisor agent in front of a set of worker agents.
The supervisor decides when to answer directly, when to call tools, and when to hand
work off to one of the managed agents.
Args:
agents: List of agents to manage.
An agent can be a LangGraph [`CompiledStateGraph`](https://reference.langchain.com/python/langgraph/graphs/#langgraph.graph.state.CompiledStateGraph),
a functional API workflow, or any other [Pregel](https://reference.langchain.com/python/langgraph/pregel/#langgraph.pregel.Pregel)
object.
model: Language model to use for the supervisor
tools: Tools to use for the supervisor
system_prompt: Optional prompt to use for the supervisor. Can be one of:
- `str`: This is converted to a `SystemMessage` and added to the beginning of the list of messages in `state["messages"]`.
- `SystemMessage`: this is added to the beginning of the list of messages in `state["messages"]`.
- `Callable`: This function should take in full graph state and the output is then passed to the language model.
- `Runnable`: This runnable should take in full graph state and the output is then passed to the language model.
response_format: An optional schema for the final supervisor output.
If provided, output will be formatted to match the given schema and returned in the `'structured_response'` state key.
If not provided, `structured_response` will not be present in the output state.
Can be passed in as:
- An OpenAI function/tool schema,
- A JSON Schema,
- A TypedDict class,
- A Pydantic class.
- A tuple `(prompt, schema)`, where schema is one of the above.
The prompt will be used together with the model that is being used to generate the structured response.
!!! Important
`response_format` requires the model to support `.with_structured_output`
!!! Note
`response_format` requires `structured_response` key in your state schema.
You can use the prebuilt `langgraph.prebuilt.chat_agent_executor.AgentStateWithStructuredResponse`.
middleware: A sequence of middleware instances to apply to the agent.
Middleware can intercept and modify agent behavior at various stages.
!!! tip ""
See the [Middleware](https://docs.langchain.com/oss/python/langchain/middleware)
docs for more information.
parallel_tool_calls: Whether to allow the supervisor LLM to call tools in parallel (only OpenAI and Anthropic).
Use this to control whether the supervisor can hand off to multiple agents at once.
If `True`, will enable parallel tool calls.
If `False`, will disable parallel tool calls.
!!! Important
This is currently supported only by OpenAI and Anthropic models.
To control parallel tool calling for other providers, add explicit instructions for tool use to the system prompt.
state_schema: State schema to use for the supervisor graph.
context_schema: Specifies the schema for the context object that will be passed to the workflow.
output_mode: Mode for adding managed agents' outputs to the message history in the multi-agent workflow.
Can be one of:
- `full_history`: Add the entire agent message history
- `last_message`: Add only the last message
add_handoff_messages: Whether to add a pair of `(AIMessage, ToolMessage)` to the message history
when a handoff occurs.
handoff_tool_prefix: Optional prefix for the handoff tools (e.g., `'delegate_to_'` or `'transfer_to_'`)
If provided, the handoff tools will be named `handoff_tool_prefix_agent_name`.
If not provided, the handoff tools will be named `transfer_to_agent_name`.
add_handoff_back_messages: Whether to add a pair of `(AIMessage, ToolMessage)` to the message history
when returning control to the supervisor to indicate that a handoff has occurred.
supervisor_name: Name of the supervisor node.
include_agent_name: Use to specify how to expose the agent name to the underlying supervisor LLM.
- `None`: Relies on the LLM provider using the name attribute on the AI message. Currently, only OpenAI supports this.
- `'inline'`: Add the agent name directly into the content field of the AI message using XML-style tags.
Example: `"How can I help you"` -> `"<name>agent_name</name><content>How can I help you?</content>"`
Example:
```python
from langchain_openai import ChatOpenAI
from langgraph_supervisor import create_supervisor
from langgraph.prebuilt import create_react_agent
# Create specialized agents
def add(a: float, b: float) -> float:
'''Add two numbers.'''
return a + b
def web_search(query: str) -> str:
'''Search the web for information.'''
return 'Here are the headcounts for each of the FAANG companies in 2024...'
math_agent = create_react_agent(
model="openai:gpt-4o",
tools=[add],
name="math_expert",
)
research_agent = create_react_agent(
model="openai:gpt-4o",
tools=[web_search],
name="research_expert",
)
# Create supervisor workflow
workflow = create_supervisor(
[research_agent, math_agent],
model=ChatOpenAI(model="gpt-4o"),
)
# Compile and run
app = workflow.compile()
result = app.invoke({
"messages": [
{
"role": "user",
"content": "what's the combined headcount of the FAANG companies in 2024?"
}
]
})
```
"""
# if (config_schema := deprecated_kwargs.get("config_schema", None)) is not None:
# warn(
# "`config_schema` is deprecated. Please use `context_schema` instead.",
# DeprecationWarning,
# stacklevel=2,
# )
# context_schema = config_schema
#
if add_handoff_back_messages is None:
add_handoff_back_messages = add_handoff_messages
supervisor_schema = state_schema or (
AgentStateWithStructuredResponse if response_format is not None else AgentState # type: ignore[deprecated]
)
workflow_schema = state_schema or _OuterState
agent_names = set()
for agent in agents:
if agent.name is None or agent.name == "LangGraph":
raise ValueError(
"Please specify a name when you create your agent, either via `create_react_agent(..., name=agent_name)` "
"or via `graph.compile(name=name)`."
)
if agent.name in agent_names:
raise ValueError(
f"Agent with name '{agent.name}' already exists. Agent names must be unique."
)
agent_names.add(agent.name)
tool_node = _prepare_tool_node(
tools,
handoff_tool_prefix,
add_handoff_messages,
agent_names,
)
all_tools = list(tool_node.tools_by_name.values())
if include_agent_name:
model = with_agent_name(model, include_agent_name)
supervisor_agent = create_agent( # type: ignore[deprecated]
name=supervisor_name,
model=model,
tools=all_tools,
system_prompt=system_prompt,
state_schema=supervisor_schema, # type: ignore[invalid-argument-type]
middleware=middleware,
response_format=response_format,
)
builder = StateGraph(workflow_schema, context_schema=context_schema)
builder.add_node(supervisor_agent, destinations=tuple(agent_names) + (END,))
builder.add_edge(START, supervisor_agent.name)
for agent in agents:
builder.add_node(
node=agent.name,
action=_make_call_agent(
agent,
output_mode,
add_handoff_back_messages=add_handoff_back_messages,
supervisor_name=supervisor_name,
),
)
builder.add_edge(agent.name, supervisor_agent.name)
return builder