diff --git a/langgraph/graph/__init__.py b/langgraph/graph/__init__.py index 1173035b5..8fac44cfa 100644 --- a/langgraph/graph/__init__.py +++ b/langgraph/graph/__init__.py @@ -1,4 +1,5 @@ from langgraph.graph.graph import END, Graph +from langgraph.graph.message import MessageGraph from langgraph.graph.state import StateGraph -__all__ = ["END", "Graph", "StateGraph"] +__all__ = ["END", "Graph", "StateGraph", "MessageGraph"] diff --git a/langgraph/graph/message.py b/langgraph/graph/message.py new file mode 100644 index 000000000..d41a2e747 --- /dev/null +++ b/langgraph/graph/message.py @@ -0,0 +1,24 @@ +from typing import Annotated, Union + +from langchain_core.messages import AnyMessage + +from langgraph.graph.state import StateGraph + +Messages = Union[list[AnyMessage], AnyMessage] + + +def add_messages(left: Messages, right: Messages) -> Messages: + if not isinstance(left, list): + left = [left] + if not isinstance(right, list): + right = [right] + return left + right + + +class MessageGraph(StateGraph): + """A StateGraph where every node + - receives a list of messages as input + - returns one or more messages as output.""" + + def __init__(self) -> None: + super().__init__(Annotated[list[AnyMessage], add_messages]) 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/tests/test_pregel.py b/tests/test_pregel.py index 1dc57988e..ac198490b 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 @@ -16,7 +17,10 @@ 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 +from langgraph.graph.message import MessageGraph from langgraph.graph.state import StateGraph +from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor +from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -788,8 +792,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 +896,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 +975,324 @@ 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( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) == { + "messages": [ + 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({"messages": [HumanMessage(content="what is weather in sf")]}) + ] == [ + { + "agent": { + "messages": [ + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"query"', + } + }, + ) + ] + } + }, + { + "action": { + "messages": [ + FunctionMessage(content="result for query", name="search_api") + ] + } + }, + { + "agent": { + "messages": [ + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + ) + ] + } + }, + { + "action": { + "messages": [ + FunctionMessage(content="result for another", name="search_api") + ] + } + }, + {"agent": {"messages": [AIMessage(content="answer")]}}, + { + "__end__": { + "messages": [ + 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"), + ] + } + }, + ] + + +def test_message_graph() -> None: + from langchain.chat_models.fake import FakeMessagesListChatModel + from langchain_community.tools import tool + from langchain_core.agents import AgentAction + 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] + + model = 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"), + ] + ) + + tool_executor = ToolExecutor(tools) + + # Define the function that determines whether to continue or not + 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: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + def call_tool(messages): + # Based on the continue condition + # we know the last message involves a function call + last_message = messages[-1] + # We construct an AgentAction from the function_call + action = AgentAction( + tool=last_message.additional_kwargs["function_call"]["name"], + tool_input=json.loads( + last_message.additional_kwargs["function_call"]["arguments"] + ), + log="", + ) + # We call the tool_executor and get back a response + response = tool_executor.invoke(action) + # We use the response to create a FunctionMessage + return FunctionMessage(content=str(response), name=action.tool) + + # Define a new graph + workflow = MessageGraph() + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("action", call_tool) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges( + # First, we define the start node. We use `agent`. + # This means these are the edges taken after the `agent` node is called. + "agent", + # Next, we pass in the function that will determine which node is called next. + should_continue, + # Finally we pass in a mapping. + # The keys are strings, and the values are other nodes. + # END is a special node marking that the graph should finish. + # What will happen is we will call `should_continue`, and then the output of that + # will be matched against the keys in this mapping. + # Based on which one it matches, that node will then be called. + { + # If `tools`, then we call the tool node. + "continue": "action", + # Otherwise we finish. + "end": END, + }, + ) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("action", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + 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..ce50a2271 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,9 @@ 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.graph.message import MessageGraph +from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor +from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -834,8 +839,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 +943,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 +1022,329 @@ 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( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) == { + "messages": [ + 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( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "agent": { + "messages": [ + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"query"', + } + }, + ) + ] + } + }, + { + "action": { + "messages": [ + FunctionMessage(content="result for query", name="search_api") + ] + } + }, + { + "agent": { + "messages": [ + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + ) + ] + } + }, + { + "action": { + "messages": [ + FunctionMessage(content="result for another", name="search_api") + ] + } + }, + {"agent": {"messages": [AIMessage(content="answer")]}}, + { + "__end__": { + "messages": [ + 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"), + ] + } + }, + ] + + +async def test_message_graph() -> None: + from langchain.chat_models.fake import FakeMessagesListChatModel + from langchain_community.tools import tool + from langchain_core.agents import AgentAction + 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] + + model = 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"), + ] + ) + + tool_executor = ToolExecutor(tools) + + # Define the function that determines whether to continue or not + 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: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + async def call_tool(messages): + # Based on the continue condition + # we know the last message involves a function call + last_message = messages[-1] + # We construct an AgentAction from the function_call + action = AgentAction( + tool=last_message.additional_kwargs["function_call"]["name"], + tool_input=json.loads( + last_message.additional_kwargs["function_call"]["arguments"] + ), + log="", + ) + # We call the tool_executor and get back a response + response = await tool_executor.ainvoke(action) + # We use the response to create a FunctionMessage + return FunctionMessage(content=str(response), name=action.tool) + + # Define a new graph + workflow = MessageGraph() + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("action", call_tool) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges( + # First, we define the start node. We use `agent`. + # This means these are the edges taken after the `agent` node is called. + "agent", + # Next, we pass in the function that will determine which node is called next. + should_continue, + # Finally we pass in a mapping. + # The keys are strings, and the values are other nodes. + # END is a special node marking that the graph should finish. + # What will happen is we will call `should_continue`, and then the output of that + # will be matched against the keys in this mapping. + # Based on which one it matches, that node will then be called. + { + # If `tools`, then we call the tool node. + "continue": "action", + # Otherwise we finish. + "end": END, + }, + ) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("action", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + 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"), + ] + }, + ]