mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 08:32:24 +02:00
lib: Add metadata arg to add_node() (#990)
* lib: Add metadata arg to add_node() - use metadata when drawing graph - use metadata for tracing * Update core * Lock
This commit is contained in:
@@ -20,9 +20,8 @@ from typing import (
|
||||
from langchain_core.runnables import Runnable
|
||||
from langchain_core.runnables.base import RunnableLike
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.runnables.graph import (
|
||||
Node as RunnableGraphNode,
|
||||
)
|
||||
from langchain_core.runnables.graph import Graph as DrawableGraph
|
||||
from langchain_core.runnables.graph import Node as DrawableNode
|
||||
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.checkpoint import BaseCheckpointSaver
|
||||
@@ -32,11 +31,16 @@ from langgraph.pregel import Channel, Pregel
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.types import All
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.utils import DrawableGraph, RunnableCallable, coerce_to_runnable
|
||||
from langgraph.utils import RunnableCallable, coerce_to_runnable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NodeSpec(NamedTuple):
|
||||
runnable: Runnable
|
||||
metadata: Optional[dict[str, Any]] = None
|
||||
|
||||
|
||||
class Branch(NamedTuple):
|
||||
path: Runnable[Any, Union[Hashable, list[Hashable]]]
|
||||
ends: Optional[dict[Hashable, str]]
|
||||
@@ -114,7 +118,7 @@ class Branch(NamedTuple):
|
||||
|
||||
class Graph:
|
||||
def __init__(self) -> None:
|
||||
self.nodes: dict[str, Runnable] = {}
|
||||
self.nodes: dict[str, NodeSpec] = {}
|
||||
self.edges = set[tuple[str, str]]()
|
||||
self.branches: defaultdict[str, dict[str, Branch]] = defaultdict(dict)
|
||||
self.support_multiple_edges = False
|
||||
@@ -125,15 +129,30 @@ class Graph:
|
||||
return self.edges
|
||||
|
||||
@overload
|
||||
def add_node(self, node: RunnableLike) -> None:
|
||||
def add_node(
|
||||
self,
|
||||
node: RunnableLike,
|
||||
*,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
...
|
||||
|
||||
@overload
|
||||
def add_node(self, node: str, action: RunnableLike) -> None:
|
||||
def add_node(
|
||||
self,
|
||||
node: str,
|
||||
action: RunnableLike,
|
||||
*,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
...
|
||||
|
||||
def add_node(
|
||||
self, node: Union[str, RunnableLike], action: Optional[RunnableLike] = None
|
||||
self,
|
||||
node: Union[str, RunnableLike],
|
||||
action: Optional[RunnableLike] = None,
|
||||
*,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
if self.compiled:
|
||||
logger.warning(
|
||||
@@ -148,7 +167,9 @@ class Graph:
|
||||
if node == END or node == START:
|
||||
raise ValueError(f"Node `{node}` is reserved.")
|
||||
|
||||
self.nodes[node] = coerce_to_runnable(action, name=node, trace=False)
|
||||
self.nodes[node] = NodeSpec(
|
||||
coerce_to_runnable(action, name=node, trace=False), metadata
|
||||
)
|
||||
|
||||
def add_edge(self, start_key: str, end_key: str) -> None:
|
||||
if self.compiled:
|
||||
@@ -385,11 +406,11 @@ class Graph:
|
||||
class CompiledGraph(Pregel):
|
||||
builder: Graph
|
||||
|
||||
def attach_node(self, key: str, node: Runnable) -> None:
|
||||
def attach_node(self, key: str, node: NodeSpec) -> None:
|
||||
self.channels[key] = EphemeralValue(Any)
|
||||
self.nodes[key] = (
|
||||
PregelNode(channels=[], triggers=[])
|
||||
| node
|
||||
PregelNode(channels=[], triggers=[], metadata=node.metadata)
|
||||
| node.runnable
|
||||
| ChannelWrite([ChannelWriteEntry(key)], tags=[TAG_HIDDEN])
|
||||
)
|
||||
cast(list[str], self.stream_channels).append(key)
|
||||
@@ -441,14 +462,14 @@ class CompiledGraph(Pregel):
|
||||
) -> DrawableGraph:
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
graph = DrawableGraph()
|
||||
start_nodes: dict[str, RunnableGraphNode] = {
|
||||
start_nodes: dict[str, DrawableNode] = {
|
||||
START: graph.add_node(self.get_input_schema(config), START)
|
||||
}
|
||||
end_nodes: dict[str, RunnableGraphNode] = {
|
||||
end_nodes: dict[str, DrawableNode] = {
|
||||
END: graph.add_node(self.get_output_schema(config), END)
|
||||
}
|
||||
|
||||
for key, node in self.builder.nodes.items():
|
||||
for key, (node, metadata) in self.builder.nodes.items():
|
||||
if xray:
|
||||
subgraph = (
|
||||
node.get_graph(
|
||||
@@ -469,7 +490,7 @@ class CompiledGraph(Pregel):
|
||||
start_nodes[key] = n
|
||||
end_nodes[key] = n
|
||||
else:
|
||||
n = graph.add_node(node, key)
|
||||
n = graph.add_node(node, key, metadata=metadata)
|
||||
start_nodes[key] = n
|
||||
end_nodes[key] = n
|
||||
for start, end in sorted(self.builder._all_edges):
|
||||
|
||||
@@ -28,7 +28,15 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue
|
||||
from langgraph.checkpoint import BaseCheckpointSaver
|
||||
from langgraph.constants import TAG_HIDDEN
|
||||
from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
|
||||
from langgraph.graph.graph import (
|
||||
END,
|
||||
START,
|
||||
Branch,
|
||||
CompiledGraph,
|
||||
Graph,
|
||||
NodeSpec,
|
||||
Send,
|
||||
)
|
||||
from langgraph.managed.base import ManagedValue, is_managed_value
|
||||
from langgraph.pregel.read import ChannelRead, PregelNode
|
||||
from langgraph.pregel.types import All
|
||||
@@ -332,7 +340,7 @@ class CompiledStateGraph(CompiledGraph):
|
||||
) -> type[BaseModel]:
|
||||
return self.get_output_schema(config)
|
||||
|
||||
def attach_node(self, key: str, node: Optional[Runnable]) -> None:
|
||||
def attach_node(self, key: str, node: Optional[NodeSpec]) -> None:
|
||||
state_keys = list(self.builder.channels)
|
||||
|
||||
def _get_state_key(input: dict, config: RunnableConfig, *, key: str) -> Any:
|
||||
@@ -399,7 +407,8 @@ class CompiledStateGraph(CompiledGraph):
|
||||
require_at_least_one_of=state_keys,
|
||||
),
|
||||
],
|
||||
).pipe(node)
|
||||
metadata=node.metadata,
|
||||
).pipe(node.runnable)
|
||||
|
||||
def attach_edge(self, starts: Union[str, Sequence[str]], end: str) -> None:
|
||||
if isinstance(starts, str):
|
||||
|
||||
@@ -151,6 +151,7 @@ class PregelNode(RunnableBindingBase):
|
||||
mapper: Optional[Callable[[Any], Any]] = None,
|
||||
writers: Optional[list[Runnable]] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
metadata: Optional[Mapping[str, Any]] = None,
|
||||
bound: Optional[Runnable[Any, Any]] = None,
|
||||
kwargs: Optional[Mapping[str, Any]] = None,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
@@ -163,7 +164,9 @@ class PregelNode(RunnableBindingBase):
|
||||
writers=writers or [],
|
||||
bound=bound or DEFAULT_BOUND,
|
||||
kwargs=kwargs or {},
|
||||
config=merge_configs(config, {"tags": tags or []}),
|
||||
config=merge_configs(
|
||||
config, {"tags": tags or [], "metadata": metadata or {}}
|
||||
),
|
||||
**other_kwargs,
|
||||
)
|
||||
|
||||
|
||||
@@ -18,7 +18,6 @@ from langchain_core.runnables.config import (
|
||||
run_in_executor,
|
||||
var_child_runnable_config,
|
||||
)
|
||||
from langchain_core.runnables.graph import Edge, Graph, Node, is_uuid
|
||||
from langchain_core.runnables.utils import accepts_config
|
||||
from typing_extensions import TypeGuard
|
||||
|
||||
@@ -132,37 +131,6 @@ class RunnableCallable(Runnable):
|
||||
return ret
|
||||
|
||||
|
||||
class DrawableGraph(Graph):
|
||||
def extend(
|
||||
self, graph: Graph, prefix: str = ""
|
||||
) -> tuple[Optional[Node], Optional[Node]]:
|
||||
if all(is_uuid(node.id) for node in graph.nodes.values()):
|
||||
super().extend(graph)
|
||||
return graph.first_node(), graph.last_node()
|
||||
|
||||
new_nodes = {
|
||||
f"{prefix}:{k}": Node(f"{prefix}:{k}", v.data)
|
||||
for k, v in graph.nodes.items()
|
||||
}
|
||||
new_edges = [
|
||||
Edge(
|
||||
f"{prefix}:{edge.source}",
|
||||
f"{prefix}:{edge.target}",
|
||||
edge.data,
|
||||
edge.conditional,
|
||||
)
|
||||
for edge in graph.edges
|
||||
]
|
||||
self.nodes.update(new_nodes)
|
||||
self.edges.extend(new_edges)
|
||||
first = graph.first_node()
|
||||
last = graph.last_node()
|
||||
return (
|
||||
Node(f"{prefix}:{first.id}", first.data) if first else None,
|
||||
Node(f"{prefix}:{last.id}", last.data) if last else None,
|
||||
)
|
||||
|
||||
|
||||
def is_async_callable(
|
||||
func: Any,
|
||||
) -> TypeGuard[Callable[..., Awaitable]]:
|
||||
|
||||
Generated
+5
-5
@@ -1,4 +1,4 @@
|
||||
# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
|
||||
# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
|
||||
|
||||
[[package]]
|
||||
name = "aiohttp"
|
||||
@@ -1746,13 +1746,13 @@ langchain-core = ">=0.2.2rc1,<0.3"
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "0.2.11"
|
||||
version = "0.2.15"
|
||||
description = "Building applications with LLMs through composability"
|
||||
optional = false
|
||||
python-versions = "<4.0,>=3.8.1"
|
||||
files = [
|
||||
{file = "langchain_core-0.2.11-py3-none-any.whl", hash = "sha256:c7ca4dc4d88e3c69fd7916c95a7027c2b1a11c2db5a51141c3ceb8afac212208"},
|
||||
{file = "langchain_core-0.2.11.tar.gz", hash = "sha256:7a4661b50604eeb20c3373fbfd8a4f1b74482a6ab4e0f9df11e96821ead8ef0c"},
|
||||
{file = "langchain_core-0.2.15-py3-none-any.whl", hash = "sha256:3bf7afaef96d7c1af0d9d223833bdee5fafc46755dc10f9c7576a85d4f6c5240"},
|
||||
{file = "langchain_core-0.2.15.tar.gz", hash = "sha256:ce03ab0a5c45b4ebfe5475eb07bf081cd21218421ff4cf26b8d2e5573ae2bd42"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -4130,4 +4130,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = ">=3.9.0,<4.0"
|
||||
content-hash = "1ec1c06495c9a1a564c5082600a4fb335278341a19d11993b981b72684679948"
|
||||
content-hash = "19250230952cb11ee6b5c820ae0a2b589ce520a59e0723204dcee0c46c3b739e"
|
||||
|
||||
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.9.0,<4.0"
|
||||
langchain-core = ">=0.2.11,<0.3"
|
||||
langchain-core = ">=0.2.15,<0.3"
|
||||
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -1460,7 +1460,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
workflow = Graph()
|
||||
|
||||
workflow.add_node("agent", agent)
|
||||
workflow.add_node("tools", execute_tools)
|
||||
workflow.add_node("tools", execute_tools, metadata={"version": 2, "variant": "b"})
|
||||
|
||||
workflow.set_entry_point("agent")
|
||||
|
||||
@@ -1474,6 +1474,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_graph().draw_mermaid() == snapshot
|
||||
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
|
||||
assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
|
||||
Reference in New Issue
Block a user