Make Command accept generic arg for destinations

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
@@ -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;
@@ -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'}}}%%
+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": [
+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": [