WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-10 19:57:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

For cond edges support specifying list of possible destinations as a list or typing annotation

Browse Source
This commit is contained in:
Nuno Campos
2024-04-29 10:38:06 -07:00
parent 6519d72837
commit a83fd19b29
3 changed files with 85 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+13 -1
View File
@@ -5,11 +5,15 @@ from typing import (
Awaitable,
Callable,
Dict,
Literal,
NamedTuple,
Optional,
Sequence,
Union,
cast,
get_args,
get_origin,
get_type_hints,
)
from langchain_core.runnables import Runnable
@@ -142,7 +146,7 @@ class Graph:
Callable[..., Awaitable[Union[str, list[str]]]],
Runnable[Any, Union[str, list[str]]],
],
path_map: Optional[dict[str, str]] = None,
path_map: Optional[Union[dict[str, str], list[str]]] = None,
then: Optional[str] = None,
) -> None:
"""Add a conditional edge from the starting node to any number of destination nodes.
@@ -166,6 +170,14 @@ class Graph:
"Adding an edge to a graph that has already been compiled. This will "
"not be reflected in the compiled graph."
)
# coerce path_map to a dictionary
if isinstance(path_map, dict):
pass
elif isinstance(path_map, list):
path_map = {name: name for name in path_map}
elif rtn_type := get_type_hints(path).get("return"):
if get_origin(rtn_type) is Literal:
path_map = {name: name for name in get_args(rtn_type)}
# find a name for the condition
path = coerce_to_runnable(path, name=None, trace=True)
name = path.name or "condition"
tests/__snapshots__/test_pregel.ambr
+66 -66
View File
@@ -1541,76 +1541,76 @@
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_run]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +-----------+
| analyzer_one | | condition |
+--------------+ +-----------+
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +--------------------+
| analyzer_one | | rewrite_query_then |
+--------------+ +--------------------+
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_step]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +-----------+
| analyzer_one | | condition |
+--------------+ +-----------+
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +--------------------+
| analyzer_one | | rewrite_query_then |
+--------------+ +--------------------+
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_message_graph[end_of_run]
tests/test_pregel.py
+6 -4
View File
@@ -4,7 +4,7 @@ import time
import warnings
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from typing import Annotated, Any, Generator, Optional, TypedDict, Union
from typing import Annotated, Any, Generator, Literal, Optional, TypedDict, Union
import pytest
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
@@ -3115,6 +3115,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
invalid_graph.add_conditional_edges(
source="prepare",
path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
path_map=["tool_two_slow", "tool_two_fast"],
)
invalid_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
invalid_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
@@ -3387,6 +3388,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
def rewrite_query_then(data: State) -> Literal["retriever_two"]:
return "retriever_two"
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
@@ -3398,9 +3402,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
workflow.set_entry_point("rewrite_query")
workflow.add_edge("rewrite_query", "analyzer_one")
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_conditional_edges(
"rewrite_query", lambda _: "retriever_two", {"retriever_two": "retriever_two"}
)
workflow.add_conditional_edges("rewrite_query", rewrite_query_then)
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
workflow.set_finish_point("qa")