WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-22 15:42:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Make Command accept generic arg for destinations

Browse Source
This commit is contained in:
Nuno Campos
2024-11-12 16:18:23 -08:00
parent 00964b18f6
commit 16bfa80b58
6 changed files with 168 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/state.py
+4 -7
View File
@@ -14,7 +14,6 @@ from typing import (
Optional,
Sequence,
Type,
TypeVar,
Union,
cast,
get_args,
@@ -50,15 +49,13 @@ from langgraph.managed.base import (
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import All, Checkpointer, Command, RetryPolicy
from langgraph.types import All, Checkpointer, Command, N, RetryPolicy
from langgraph.utils.fields import get_field_default
from langgraph.utils.pydantic import create_model
from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
logger = logging.getLogger(__name__)
N = TypeVar("N")
def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None:
if isinstance(schema, type):
@@ -81,7 +78,7 @@ def _get_node_name(node: RunnableLike) -> str:
raise TypeError(f"Unsupported node type: {type(node)}")
class GraphCommand(Command, Generic[N]):
class GraphCommand(Generic[N], Command[N]):
"""One or more commands to update a StateGraph's state and go to, or send messages to nodes."""
__slots__ = ("goto",)
@@ -90,9 +87,9 @@ class GraphCommand(Command, Generic[N]):
self,
*,
update: Optional[dict[str, Any]] = None,
goto: Union[str, Sequence[str]] = (),
send: Union[Send, Sequence[Send]] = (),
resume: Optional[Union[Any, dict[str, Any]]] = None,
goto: Union[str, Sequence[str]] = (),
) -> None:
super().__init__(update=update, send=send, resume=resume)
self.goto = goto
@@ -390,7 +387,7 @@ class StateGraph(Graph):
input = input_hint
if (
(rtn := hints.get("return"))
and get_origin(rtn) is GraphCommand
and get_origin(rtn) in (Command, GraphCommand)
and (rargs := get_args(rtn))
and get_origin(rargs[0]) is Literal
and (vals := get_args(rargs[0]))
libs/langgraph/langgraph/types.py
+7 -1
View File
@@ -4,11 +4,14 @@ from typing import (
TYPE_CHECKING,
Any,
Callable,
Generic,
Hashable,
Literal,
NamedTuple,
Optional,
Sequence,
Type,
TypeVar,
Union,
cast,
)
@@ -221,7 +224,10 @@ class Send:
)
class Command:
N = TypeVar("N", bound=Hashable)
class Command(Generic[N]):
"""One or more commands to update the graph's state and send messages to nodes."""
__slots__ = ("update", "send", "resume")
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+75
View File
@@ -5108,6 +5108,81 @@
'''
# ---
# name: test_send_react_interrupt_control[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_pipe]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_pool]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[sqlite]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_simple_multi_edge
'''
graph TD;
libs/langgraph/tests/__snapshots__/test_pregel_async.ambr
+75
View File
@@ -1302,6 +1302,81 @@
+---------+
'''
# ---
# name: test_send_react_interrupt_control[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_aio_pipe]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_aio_pool]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[sqlite_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[duckdb_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
libs/langgraph/tests/test_pregel.py
+2 -1
View File
@@ -2691,7 +2691,7 @@ def test_send_react_interrupt(
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_send_react_interrupt_control(
request: pytest.FixtureRequest, checkpointer_name: str
request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion
) -> None:
from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage
@@ -2721,6 +2721,7 @@ def test_send_react_interrupt_control(
builder.add_node(foo)
builder.add_edge(START, "agent")
graph = builder.compile()
assert graph.get_graph().draw_mermaid() == snapshot
assert graph.invoke({"messages": [HumanMessage("hello")]}) == {
"messages": [
libs/langgraph/tests/test_pregel_async.py
+5 -2
View File
@@ -2973,7 +2973,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_react_interrupt_control(checkpointer_name: str) -> None:
async def test_send_react_interrupt_control(
checkpointer_name: str, snapshot: SnapshotAssertion
) -> None:
from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage
ai_message = AIMessage(
@@ -2982,7 +2984,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None:
tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())],
)
async def agent(state) -> GraphCommand[Literal["foo"]]:
async def agent(state) -> Command[Literal["foo"]]:
return GraphCommand(
update={"messages": ai_message},
send=[Send(call["name"], call) for call in ai_message.tool_calls],
@@ -3000,6 +3002,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None:
builder.add_node(foo)
builder.add_edge(START, "agent")
graph = builder.compile()
assert graph.get_graph().draw_mermaid() == snapshot
assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == {
"messages": [