update langgraph & pass state explicitly in handoffs (#20)

Author: vbarda

README.md

@@ -116,9 +116,6 @@ By default, the agents in the swarm are assumed to use handoff tools created wit
 * add tool call arguments for the LLM to populate, for example a task description for the next agent
 * change what data is passed to the next agent as part of the handoff: by default `create_handoff_tool` passes **full** message history (all of the messages generated in the swarm up to this point), as well as the contents of `Command.update` to the next agent
 
-> [!IMPORTANT]
-> If you want to change what messages are passed to the next agent, you **must** use a different state schema key for `messages` in your agent implementation (e.g., `alice_messages`). By default, all agent (subgraph) state updates are applied to the swarm (parent) graph state during the handoff. Since all of the agents by default are assumed to communicate over a single `messages` key, this means that the agent's messages are **automatically combined** into the parent graph's `messages`, unless an agent uses a different key for `messages`. See more on this in the [customizing agent implementation](#customizing-agent-implementation) section.
-
 Here is an example of what a custom handoff tool might look like:
 
 ```python
@@ -170,7 +167,7 @@ def create_custom_handoff_tool(*, agent_name: str, tool_name: str, tool_descript
 
 ### Customizing agent implementation
 
-By default, individual agents are expected to communicate over a single `messages` key that is shared by all agents and the overall multi-agent swarm graph. This means that **all** of the messages from **all** of the agents will be combined into a single, shared list of messages. This might not be desirable if you don't want to expose an agent's internal history of messages. To change this, you can customize the agent by taking the following steps:
+By default, individual agents are expected to communicate over a single `messages` key that is shared by all agents and the overall multi-agent swarm graph. This means that messages from **all** of the agents will be combined into a single, shared list of messages. This might not be desirable if you don't want to expose an agent's internal history of messages. To change this, you can customize the agent by taking the following steps:
 
 1.  use custom [state schema](https://langchain-ai.github.io/langgraph/concepts/low_level#schema) with a different key for messages, for example `alice_messages`
 1.  write a wrapper that converts the parent graph state to the child agent state and back (see this [how-to](https://langchain-ai.github.io/langgraph/how-tos/subgraph-transform-state/) guide)

langgraph_swarm/handoff.py

@@ -3,7 +3,7 @@
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import BaseTool, InjectedToolCallId, tool
 from langgraph.graph.state import CompiledStateGraph
-from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt import InjectedState, ToolNode
 from langgraph.types import Command
 from typing_extensions import Annotated
 
@@ -34,6 +34,7 @@ def create_handoff_tool(*, agent_name: str, description: str | None = None) -> B
 
     @tool(name, description=description)
     def handoff_to_agent(
+        state: Annotated[dict, InjectedState],
         tool_call_id: Annotated[str, InjectedToolCallId],
     ):
         tool_message = ToolMessage(
@@ -44,7 +45,7 @@ def handoff_to_agent(
         return Command(
             goto=agent_name,
             graph=Command.PARENT,
-            update={"messages": [tool_message], "active_agent": agent_name},
+            update={"messages": state["messages"] + [tool_message], "active_agent": agent_name},
         )
 
     handoff_to_agent.metadata = {METADATA_KEY_HANDOFF_DESTINATION: agent_name}

pyproject.toml

@@ -4,15 +4,15 @@ build-backend = "pdm.backend"
 
 [project]
 name = "langgraph-swarm"
-version = "0.0.4"
+version = "0.0.5"
 description = "An implementation of a multi-agent swarm using LangGraph"
 authors = [
     {name = "Vadym Barda", email = "[email protected] "}
 ]
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-    "langgraph>=0.3.1,<0.4.0",
+    "langgraph>=0.3.5,<0.4.0",
     "langchain-core>=0.3.40,<0.4.0"
 ]