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:
Nuno Campos
2024-07-11 17:03:22 -07:00
committed by GitHub
parent bd100fdfef
commit 71d9d02d26
8 changed files with 156 additions and 130 deletions
+37 -16
View File
@@ -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):
+12 -3
View File
@@ -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):
+4 -1
View File
@@ -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,
)
-32
View File
@@ -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]]:
+5 -5
View File
@@ -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"
+1 -1
View File
@@ -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
+2 -1
View File
@@ -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