WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-05 01:07:48 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #50 from langchain-ai/nc/20jan/root-state

Browse Source 
Add support for single key state, eg just list of messages
This commit is contained in:
Nuno Campos
2024-01-21 09:07:49 -08:00
committed by GitHub
parent cc764f8423 60482b7adb
commit bd6490443f
5 changed files with 719 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/__init__.py
+2 -1
View File
@@ -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"]
langgraph/graph/message.py
+24
View File
@@ -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])
langgraph/graph/state.py
+34 -12
View File
@@ -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__
tests/test_pregel.py
+327 -4
View File
@@ -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"),
]
},
]
tests/test_pregel_async.py
+332 -5
View File
@@ -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"),
]
},
]