diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 1a3f5530e..a11288d91 100644 --- a/langgraph/graph/state.py +++ b/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__ diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index 253473416..df6588dc1 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/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)) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 1dc57988e..30842bd44 100644 --- a/tests/test_pregel.py +++ b/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"), + ] + }, + ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 1b162cbb8..40d3739e1 100644 --- a/tests/test_pregel_async.py +++ b/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"), + ] + }, + ]