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

lib: Add support for graphs without edges

Browse Source 
- Return Control(update_state=, trigger=, send=) from your nodes instead
- Annotate nodes with Control[Literal["destination"]] to see your graph connections drawn
This commit is contained in:
Nuno Campos
2024-10-31 09:38:43 -07:00
parent f7a93f99ef
commit 0c5c2e6370
6 changed files with 180 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/constants.py
+2
View File
@@ -20,6 +20,8 @@ START = sys.intern("__start__")
"""The first (maybe virtual) node in graph-style Pregel."""
END = sys.intern("__end__")
"""The last (maybe virtual) node in graph-style Pregel."""
SELF = sys.intern("__self__")
"""The implicit branch that handles each node's Control values."""
# --- Reserved write keys ---
INPUT = sys.intern("__input__")
libs/langgraph/langgraph/graph/graph.py
+13 -1
View File
@@ -27,6 +27,7 @@ from typing_extensions import Self
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.constants import (
EMPTY_SEQ,
END,
NS_END,
NS_SEP,
@@ -47,6 +48,7 @@ logger = logging.getLogger(__name__)
class NodeSpec(NamedTuple):
runnable: Runnable
metadata: Optional[dict[str, Any]] = None
ends: Optional[tuple[str, ...]] = EMPTY_SEQ
class Branch(NamedTuple):
@@ -123,7 +125,7 @@ class Branch(NamedTuple):
result: Any,
config: RunnableConfig,
) -> Union[Runnable, Any]:
if not isinstance(result, list):
if not isinstance(result, (list, tuple)):
result = [result]
if self.ends:
destinations: Sequence[Union[Send, str]] = [
@@ -364,6 +366,9 @@ class Graph:
for node in self.nodes:
if node != start and node != branch.then:
all_sources.add(node)
for name, spec in self.nodes.items():
if spec.ends:
all_sources.add(name)
# validate sources
for source in all_sources:
if source not in self.nodes and source != START:
@@ -387,6 +392,9 @@ class Graph:
for node in self.nodes:
if node != start and node != branch.then:
all_targets.add(node)
for name, spec in self.nodes.items():
if spec.ends:
all_targets.update(spec.ends)
# validate targets
for node in self.nodes:
if node not in all_targets:
@@ -620,5 +628,9 @@ class CompiledGraph(Pregel):
)
if branch.then is not None:
add_edge(end, branch.then)
for key, n in self.builder.nodes.items():
if n.ends:
for end in n.ends:
add_edge(key, end, conditional=True)
return graph
libs/langgraph/langgraph/graph/state.py
+86 -27
View File
@@ -14,6 +14,7 @@ from typing import (
Type,
Union,
cast,
get_args,
get_origin,
get_type_hints,
overload,
@@ -32,7 +33,7 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.channels.named_barrier_value import NamedBarrierValue
from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN
from langgraph.constants import EMPTY_SEQ, NS_END, NS_SEP, SELF, TAG_HIDDEN
from langgraph.errors import ErrorCode, InvalidUpdateError, create_error_message
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
from langgraph.managed.base import (
@@ -46,10 +47,10 @@ 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, RetryPolicy
from langgraph.types import All, Checkpointer, Control, RetryPolicy
from langgraph.utils.fields import get_field_default
from langgraph.utils.pydantic import create_model
from langgraph.utils.runnable import coerce_to_runnable
from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
logger = logging.getLogger(__name__)
@@ -71,6 +72,7 @@ class StateNodeSpec(NamedTuple):
metadata: Optional[dict[str, Any]]
input: Type[Any]
retry_policy: Optional[RetryPolicy]
ends: Optional[tuple[str, ...]] = EMPTY_SEQ
class StateGraph(Graph):
@@ -338,8 +340,29 @@ class StateGraph(Graph):
f"'{character}' is a reserved character and is not allowed in the node names."
)
if input is None:
input = _get_input_schema_from_type_hint(action)
ends = EMPTY_SEQ
try:
if isfunction(action) and (
hints := get_type_hints(getattr(action, "__call__"))
or get_type_hints(action)
):
if input is None:
first_parameter_name = next(
iter(inspect.signature(action).parameters.keys())
)
if input_hint := hints.get(first_parameter_name):
if isinstance(input_hint, type) and get_type_hints(input_hint):
input = input_hint
if (
(rtn := hints.get("return"))
and get_origin(rtn) is Control
and (rargs := get_args(rtn))
and get_origin(rargs[0]) is Literal
and (vals := get_args(rargs[0]))
):
ends = vals
except (TypeError, StopIteration):
pass
if input is not None:
self._add_schema(input)
self.nodes[cast(str, node)] = StateNodeSpec(
@@ -347,6 +370,7 @@ class StateGraph(Graph):
metadata,
input=input or self.schema,
retry_policy=retry,
ends=ends,
)
return self
@@ -468,6 +492,9 @@ class StateGraph(Graph):
for key, node in self.nodes.items():
compiled.attach_node(key, node)
for key, node in self.nodes.items():
compiled.attach_branch(key, SELF, CONTROL_BRANCH, with_reader=False)
for start, end in self.edges:
compiled.attach_edge(start, end)
@@ -518,11 +545,23 @@ class CompiledStateGraph(CompiledGraph):
if is_writable_managed_value(v)
]
def _get_root(input: Any) -> Any:
if isinstance(input, Control):
return input.update_state
else:
return input
def _get_state_key(input: Union[None, dict, Any], *, key: str) -> Any:
if input is None:
return SKIP_WRITE
elif isinstance(input, dict):
if all(k not in output_keys for k in input):
raise InvalidUpdateError(
f"Expected node {key} to update at least one of {output_keys}, got {input}"
)
return input.get(key, SKIP_WRITE)
elif isinstance(input, Control):
return _get_state_key(input.update_state, key=key)
elif get_type_hints(type(input)):
value = getattr(input, key, SKIP_WRITE)
return value if value is not None else SKIP_WRITE
@@ -535,7 +574,7 @@ class CompiledStateGraph(CompiledGraph):
# state updaters
write_entries = (
[ChannelWriteEntry("__root__", skip_none=True)]
[ChannelWriteEntry("__root__", skip_none=True, mapper=_get_root)]
if output_keys == ["__root__"]
else [
ChannelWriteEntry(key, mapper=partial(_get_state_key, key=key))
@@ -578,7 +617,6 @@ class CompiledStateGraph(CompiledGraph):
ChannelWrite(
[ChannelWriteEntry(key, key)] + write_entries,
tags=[TAG_HIDDEN],
require_at_least_one_of=output_keys,
),
],
metadata=node.metadata,
@@ -615,7 +653,9 @@ class CompiledStateGraph(CompiledGraph):
[ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN]
)
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def attach_branch(
self, start: str, name: str, branch: Branch, *, with_reader: bool = True
) -> None:
def branch_writer(
packets: Sequence[Union[str, Send]], config: RunnableConfig
) -> None:
@@ -648,7 +688,8 @@ class CompiledStateGraph(CompiledGraph):
else self.builder.schema
)
self.nodes[start] |= branch.run(
branch_writer, _get_state_reader(self.builder, schema)
branch_writer,
_get_state_reader(self.builder, schema) if with_reader else None,
)
# attach branch subscribers
@@ -697,6 +738,42 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
return schema(**input)
def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
if not isinstance(value, Control):
return EMPTY_SEQ
rtn: list[Union[str, Send]] = []
if isinstance(value.trigger, str):
rtn.append(value.trigger)
else:
rtn.extend(value.trigger)
if isinstance(value.send, Send):
rtn.append(value.send)
else:
rtn.extend(value.send)
return rtn
async def _acontrol_branch(value: Any) -> None:
if not isinstance(value, Control):
return EMPTY_SEQ
rtn: list[Union[str, Send]] = []
if isinstance(value.trigger, str):
rtn.append(value.trigger)
else:
rtn.extend(value.trigger)
if isinstance(value.send, Send):
rtn.append(value.send)
else:
rtn.extend(value.send)
return rtn
CONTROL_BRANCH_PATH = RunnableCallable(
_control_branch, _acontrol_branch, tags=[TAG_HIDDEN], trace=False, recurse=False
)
CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None)
def _get_channels(
schema: Type[dict],
) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]:
@@ -823,21 +900,3 @@ def _get_schema(
if k in channels and isinstance(channels[k], BaseChannel)
},
)
def _get_input_schema_from_type_hint(
action: Optional[RunnableLike],
) -> Optional[Type[Any]]:
if not isfunction(action) and not ismethod(getattr(action, "__call__", None)):
return None
action = cast(Callable, action)
try:
hints = get_type_hints(getattr(action, "__call__")) or get_type_hints(action)
first_parameter_name = next(iter(inspect.signature(action).parameters.keys()))
input_hint = hints.get(first_parameter_name)
if isinstance(input_hint, type) and get_type_hints(input_hint):
return input_hint
except (TypeError, StopIteration):
pass
return None
libs/langgraph/langgraph/types.py
+18
View File
@@ -4,11 +4,13 @@ from typing import (
TYPE_CHECKING,
Any,
Callable,
Generic,
Literal,
NamedTuple,
Optional,
Sequence,
Type,
TypeVar,
Union,
cast,
)
@@ -221,6 +223,22 @@ class Send:
)
N = TypeVar("N")
class Control(Generic[N]):
def __init__(
self,
*,
update_state: Optional[dict[str, Any]] = None,
trigger: Union[str, Sequence[str]] = (),
send: Union[Send, Sequence[Send]] = (),
) -> None:
self.update_state = update_state
self.trigger = trigger
self.send = send
StreamChunk = tuple[tuple[str, ...], str, Any]
libs/langgraph/tests/test_pregel.py
+1 -1
View File
@@ -6458,7 +6458,7 @@ def test_root_graph(
),
AIMessage(content="answer", id="ai2"),
AIMessage(
content="an extra message", id="00000000-0000-4000-8000-000000000091"
content="an extra message", id="00000000-0000-4000-8000-000000000092"
),
HumanMessage(content="what is weather in la"),
],
libs/langgraph/tests/test_pregel_async.py
+60 -1
View File
@@ -62,7 +62,7 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from langgraph.pregel.retry import RetryPolicy
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore
from langgraph.types import Interrupt, PregelTask, Send, StreamWriter
from langgraph.types import Control, Interrupt, PregelTask, Send, StreamWriter
from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence
from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
@@ -1969,6 +1969,65 @@ async def test_max_concurrency() -> None:
assert node2.currently == 0
async def test_max_concurrency_control() -> None:
async def node1(state) -> Control[Literal["2"]]:
return Control(update_state=["1"], send=[Send("2", state)] * 100)
node2_currently = 0
node2_max_currently = 0
async def node2(state) -> Control[Literal["3"]]:
nonlocal node2_currently, node2_max_currently
node2_currently += 1
if node2_currently > node2_max_currently:
node2_max_currently = node2_currently
await asyncio.sleep(0.1)
node2_currently -= 1
return Control(update_state=["2"], trigger="3")
async def node3(state) -> Literal["3"]:
return ["3"]
builder = StateGraph(Annotated[list, operator.add])
builder.add_node("1", node1)
builder.add_node("2", node2)
builder.add_node("3", node3)
builder.add_edge(START, "1")
graph = builder.compile()
assert (
graph.get_graph().draw_mermaid()
== """%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
1(1)
2(2)
3([3]):::last
__start__ --> 1;
1 -.-> 2;
2 -.-> 3;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
"""
)
assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *(["2"] * 100), "3"]
assert node2_max_currently == 100
assert node2_currently == 0
node2_max_currently = 0
assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [
"0",
"1",
*(["2"] * 100),
"3",
]
assert node2_max_currently == 10
assert node2_currently == 0
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_invoke_checkpoint_three(
mocker: MockerFixture, checkpointer_name: str