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))

tests/test_pregel.py

@@ -1,9 +1,10 @@
+import json
 import operator
 import time
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from typing import Annotated, Generator, Optional, TypedDict, Union
+from typing import Annotated, Generator, Optional, Self, TypedDict, Union
 
 import pytest
 from langchain_core.runnables import RunnablePassthrough
@@ -17,6 +18,7 @@ from langgraph.channels.topic import Topic
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, Graph
 from langgraph.graph.state import StateGraph
+from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor
 from langgraph.pregel import Channel, GraphRecursionError, Pregel
 from langgraph.pregel.reserved import ReservedChannels
 
@@ -788,8 +790,6 @@ def test_conditional_graph() -> None:
 
 
 def test_conditional_graph_state() -> None:
-    from copy import deepcopy
-
     from langchain.llms.fake import FakeStreamingListLLM
     from langchain_community.tools import tool
     from langchain_core.agents import AgentAction, AgentFinish
@@ -894,7 +894,7 @@ def test_conditional_graph_state() -> None:
         ),
     }
 
-    assert [deepcopy(c) for c in app.stream({"input": "what is weather in sf"})] == [
+    assert [*app.stream({"input": "what is weather in sf"})] == [
         {
             "agent": {
                 "agent_outcome": AgentAction(
@@ -973,3 +973,118 @@ def test_conditional_graph_state() -> None:
             }
         },
     ]
+
+
+def test_prebuilt_chat() -> None:
+    from langchain.chat_models.fake import FakeMessagesListChatModel
+    from langchain_community.tools import tool
+    from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage
+
+    class FakeFuntionChatModel(FakeMessagesListChatModel):
+        def bind_functions(self, functions: list) -> Self:
+            return self
+
+    @tool()
+    def search_api(query: str) -> str:
+        """Searches the API for the query."""
+        return f"result for {query}"
+
+    tools = [search_api]
+
+    app = create_function_calling_executor(
+        FakeFuntionChatModel(
+            responses=[
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": json.dumps("query"),
+                        }
+                    },
+                ),
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": json.dumps("another"),
+                        }
+                    },
+                ),
+                AIMessage(content="answer"),
+            ]
+        ),
+        tools,
+    )
+
+    assert app.invoke([HumanMessage(content="what is weather in sf")]) == [
+        HumanMessage(content="what is weather in sf"),
+        AIMessage(
+            content="",
+            additional_kwargs={
+                "function_call": {"name": "search_api", "arguments": '"query"'}
+            },
+        ),
+        FunctionMessage(content="result for query", name="search_api"),
+        AIMessage(
+            content="",
+            additional_kwargs={
+                "function_call": {"name": "search_api", "arguments": '"another"'}
+            },
+        ),
+        FunctionMessage(content="result for another", name="search_api"),
+        AIMessage(content="answer"),
+    ]
+
+    assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [
+        {
+            "agent": [
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {"name": "search_api", "arguments": '"query"'}
+                    },
+                )
+            ]
+        },
+        {"action": [FunctionMessage(content="result for query", name="search_api")]},
+        {
+            "agent": [
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": '"another"',
+                        }
+                    },
+                )
+            ]
+        },
+        {"action": [FunctionMessage(content="result for another", name="search_api")]},
+        {"agent": [AIMessage(content="answer")]},
+        {
+            "__end__": [
+                HumanMessage(content="what is weather in sf"),
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {"name": "search_api", "arguments": '"query"'}
+                    },
+                ),
+                FunctionMessage(content="result for query", name="search_api"),
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": '"another"',
+                        }
+                    },
+                ),
+                FunctionMessage(content="result for another", name="search_api"),
+                AIMessage(content="answer"),
+            ]
+        },
+    ]

tests/test_pregel_async.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import operator
 from contextlib import asynccontextmanager, contextmanager
 from typing import (
@@ -8,6 +9,7 @@ from typing import (
     AsyncIterator,
     Generator,
     Optional,
+    Self,
     TypedDict,
     Union,
 )
@@ -23,6 +25,7 @@ from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, Graph, StateGraph
+from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor
 from langgraph.pregel import Channel, GraphRecursionError, Pregel
 from langgraph.pregel.reserved import ReservedChannels
 
@@ -834,8 +837,6 @@ async def test_conditional_graph() -> None:
 
 
 async def test_conditional_graph_state() -> None:
-    from copy import deepcopy
-
     from langchain.llms.fake import FakeStreamingListLLM
     from langchain_community.tools import tool
     from langchain_core.agents import AgentAction, AgentFinish
@@ -940,9 +941,7 @@ async def test_conditional_graph_state() -> None:
         ),
     }
 
-    assert [
-        deepcopy(c) async for c in app.astream({"input": "what is weather in sf"})
-    ] == [
+    assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
         {
             "agent": {
                 "agent_outcome": AgentAction(
@@ -1021,3 +1020,120 @@ async def test_conditional_graph_state() -> None:
             }
         },
     ]
+
+
+async def test_prebuilt_chat() -> None:
+    from langchain.chat_models.fake import FakeMessagesListChatModel
+    from langchain_community.tools import tool
+    from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage
+
+    class FakeFuntionChatModel(FakeMessagesListChatModel):
+        def bind_functions(self, functions: list) -> Self:
+            return self
+
+    @tool()
+    def search_api(query: str) -> str:
+        """Searches the API for the query."""
+        return f"result for {query}"
+
+    tools = [search_api]
+
+    app = create_function_calling_executor(
+        FakeFuntionChatModel(
+            responses=[
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": json.dumps("query"),
+                        }
+                    },
+                ),
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": json.dumps("another"),
+                        }
+                    },
+                ),
+                AIMessage(content="answer"),
+            ]
+        ),
+        tools,
+    )
+
+    assert await app.ainvoke([HumanMessage(content="what is weather in sf")]) == [
+        HumanMessage(content="what is weather in sf"),
+        AIMessage(
+            content="",
+            additional_kwargs={
+                "function_call": {"name": "search_api", "arguments": '"query"'}
+            },
+        ),
+        FunctionMessage(content="result for query", name="search_api"),
+        AIMessage(
+            content="",
+            additional_kwargs={
+                "function_call": {"name": "search_api", "arguments": '"another"'}
+            },
+        ),
+        FunctionMessage(content="result for another", name="search_api"),
+        AIMessage(content="answer"),
+    ]
+
+    assert [
+        c async for c in app.astream([HumanMessage(content="what is weather in sf")])
+    ] == [
+        {
+            "agent": [
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {"name": "search_api", "arguments": '"query"'}
+                    },
+                )
+            ]
+        },
+        {"action": [FunctionMessage(content="result for query", name="search_api")]},
+        {
+            "agent": [
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": '"another"',
+                        }
+                    },
+                )
+            ]
+        },
+        {"action": [FunctionMessage(content="result for another", name="search_api")]},
+        {"agent": [AIMessage(content="answer")]},
+        {
+            "__end__": [
+                HumanMessage(content="what is weather in sf"),
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {"name": "search_api", "arguments": '"query"'}
+                    },
+                ),
+                FunctionMessage(content="result for query", name="search_api"),
+                AIMessage(
+                    content="",
+                    additional_kwargs={
+                        "function_call": {
+                            "name": "search_api",
+                            "arguments": '"another"',
+                        }
+                    },
+                ),
+                FunctionMessage(content="result for another", name="search_api"),
+                AIMessage(content="answer"),
+            ]
+        },
+    ]