Add support for single key state, eg just list of messages

langgraph/graph/state.py

@@ -3,7 +3,7 @@ from functools import partial
 from inspect import signature
 from typing import Any, Optional, Type
 
-from langchain_core.runnables import RunnableConfig, RunnableLambda
+from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
@@ -32,6 +32,17 @@ class StateGraph(Graph):
             raise ValueError("Cannot use channel names as node names")
 
         state_keys = list(self.channels)
+        state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
+        update_state = (
+            _update_state_dict
+            if isinstance(state_keys_read, list)
+            else _update_state_root
+        )
+        coerce_state = (
+            partial(_coerce_state, self.schema)
+            if isinstance(state_keys_read, list)
+            else RunnablePassthrough()
+        )
 
         outgoing_edges = defaultdict(list)
         for start, end in self.edges:
@@ -40,9 +51,9 @@ class StateGraph(Graph):
         nodes = {
             key: (
                 Channel.subscribe_to(f"{key}:inbox")
-                | partial(_coerce_state, self.schema)  # coerce/validate using schema
+                | coerce_state  # coerce/validate using schema
                 | node
-                | _update_state
+                | update_state
                 | Channel.write_to(key)
             )
             for key, node in self.nodes.items()
@@ -54,7 +65,7 @@ class StateGraph(Graph):
             if outgoing or key in self.branches:
                 nodes[edges_key] = Channel.subscribe_to(
                     key, tags=["langsmith:hidden"]
-                ) | ChannelRead(state_keys)
+                ) | ChannelRead(state_keys_read)
             if outgoing:
                 nodes[edges_key] |= Channel.write_to(*[dest for dest in outgoing])
             if key in self.branches:
@@ -65,12 +76,12 @@ class StateGraph(Graph):
 
         nodes[START] = (
             Channel.subscribe_to(f"{START}:inbox", tags=["langsmith:hidden"])
-            | _update_state
+            | update_state
             | Channel.write_to(START)
         )
         nodes[f"{START}:edges"] = (
             Channel.subscribe_to(START, tags=["langsmith:hidden"])
-            | ChannelRead(state_keys)
+            | ChannelRead(state_keys_read)
             | Channel.write_to(f"{self.entry_point}:inbox")
         )
 
@@ -88,26 +99,37 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
     return schema(**input)
 
 
-def _update_state(input: dict[str, Any], config: RunnableConfig) -> dict[str, Any]:
+def _update_state_dict(input: dict[str, Any], config: RunnableConfig) -> dict[str, Any]:
     if input is not None:
         ChannelWrite.do_write(config, **input)
     return input
 
 
+def _update_state_root(input: Any, config: RunnableConfig) -> dict[str, Any]:
+    if input is not None:
+        ChannelWrite.do_write(config, __root__=input)
+    return input
+
+
 def _get_channels(schema: Type[dict]) -> dict[str, BaseChannel]:
     if not hasattr(schema, "__annotations__"):
-        raise ValueError("Schema must be a class with type annotations")
+        return {
+            "__root__": _get_channel(schema),
+        }
 
     channels: dict[str, BaseChannel] = {}
     for name, typ in schema.__annotations__.items():
-        if channel := _is_field_binop(typ):
-            channels[name] = channel
-        else:
-            channels[name] = LastValue(typ)
+        channels[name] = _get_channel(typ)
 
     return channels
 
 
+def _get_channel(annotation: Any) -> Optional[BaseChannel]:
+    if channel := _is_field_binop(annotation):
+        return channel
+    return LastValue(annotation)
+
+
 def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__

langgraph/prebuilt/chat_agent_executor.py

@@ -1,6 +1,6 @@
 import json
 import operator
-from typing import Annotated, Sequence, TypedDict
+from typing import Annotated
 
 from langchain.tools.render import format_tool_to_openai_function
 from langchain_core.agents import AgentAction
@@ -23,8 +23,7 @@ def create_function_calling_executor(model, tools):
     )
 
     # Define the function that determines whether to continue or not
-    def should_continue(state):
-        messages = state["messages"]
+    def should_continue(messages):
         last_message = messages[-1]
         # If there is no function call, then we finish
         if "function_call" not in last_message.additional_kwargs:
@@ -34,21 +33,18 @@ def create_function_calling_executor(model, tools):
             return "continue"
 
     # Define the function that calls the model
-    def call_model(state):
-        messages = state["messages"]
+    def call_model(messages):
         response = model.invoke(messages)
         # We return a list, because this will get added to the existing list
-        return {"messages": [response]}
+        return [response]
 
-    async def acall_model(state):
-        messages = state["messages"]
+    async def acall_model(messages):
         response = await model.ainvoke(messages)
         # We return a list, because this will get added to the existing list
-        return {"messages": [response]}
+        return [response]
 
     # Define the function to execute tools
-    def _get_action(state):
-        messages = state["messages"]
+    def _get_action(messages):
         # Based on the continue condition
         # we know the last message involves a function call
         last_message = messages[-1]
@@ -68,7 +64,7 @@ def create_function_calling_executor(model, tools):
         # We use the response to create a FunctionMessage
         function_message = FunctionMessage(content=str(response), name=action.tool)
         # We return a list, because this will get added to the existing list
-        return {"messages": [function_message]}
+        return [function_message]
 
     async def acall_tool(state):
         action = _get_action(state)
@@ -77,17 +73,13 @@ def create_function_calling_executor(model, tools):
         # We use the response to create a FunctionMessage
         function_message = FunctionMessage(content=str(response), name=action.tool)
         # We return a list, because this will get added to the existing list
-        return {"messages": [function_message]}
+        return [function_message]
 
-    # We create the AgentState that we will pass around
+    # Define a new graph with state
     # This simply involves a list of messages
     # We want steps to return messages to append to the list
     # So we annotate the messages attribute with operator.add
-    class AgentState(TypedDict):
-        messages: Annotated[Sequence[BaseMessage], operator.add]
-
-    # Define a new graph
-    workflow = StateGraph(AgentState)
+    workflow = StateGraph(Annotated[list[BaseMessage], operator.add])
 
     # Define the two nodes we will cycle between
     workflow.add_node("agent", RunnableLambda(call_model, acall_model))