Remove implementation of get_graph for @entrypoint

This commit is contained in:
Nuno Campos
2025-01-23 15:39:36 -08:00
parent 7b552ebf4b
commit f014d96d1c
5 changed files with 2 additions and 4665 deletions
+1 -1
View File
@@ -60,7 +60,7 @@ MAXFAIL ?=
MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),)
test_watch:
make start-postgres && poetry run ptw . -- --ff -vv -x $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \
make start-postgres && poetry run ptw . -- --ff -vv -x --snapshot-update --tb short $(TEST); \
EXIT_CODE=$$?; \
make stop-postgres; \
exit $$EXIT_CODE
+1 -101
View File
@@ -3,7 +3,6 @@ import concurrent.futures
import functools
import inspect
import types
from collections.abc import Iterator
from dataclasses import dataclass
from typing import (
Any,
@@ -18,17 +17,12 @@ from typing import (
overload,
)
from langchain_core.runnables.base import Runnable
from langchain_core.runnables.config import RunnableConfig
from langchain_core.runnables.graph import Graph, Node
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import END, PREVIOUS, START, TAG_HIDDEN
from langgraph.pregel import Pregel
from langgraph.pregel.call import P, T, call, get_runnable_for_entrypoint
from langgraph.pregel.protocol import PregelProtocol
from langgraph.pregel.read import PregelNode
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
@@ -523,7 +517,7 @@ class entrypoint:
else:
output_type = save_type = sig.return_annotation
return EntrypointPregel(
return Pregel(
nodes={
func.__name__: PregelNode(
bound=bound,
@@ -554,97 +548,3 @@ class entrypoint:
store=self.store,
config_type=self.config_schema,
)
class EntrypointPregel(Pregel):
def get_graph(
self,
config: Optional[RunnableConfig] = None,
*,
xray: Union[int, bool] = False,
) -> Graph:
name, entrypoint = next(iter(self.nodes.items()))
graph = Graph()
node = Node(f"__{name}", name, entrypoint.bound, None)
graph.nodes[node.id] = node
candidates: list[tuple[Node, Union[Callable, PregelProtocol]]] = [
*_find_children(entrypoint.bound, node)
]
seen: set[Union[Callable, PregelProtocol]] = set()
for parent, child in candidates:
if child in seen:
continue
else:
seen.add(child)
if callable(child):
node = Node(f"__{child.__name__}", child.__name__, child, None) # type: ignore[arg-type]
graph.nodes[node.id] = node
graph.add_edge(parent, node, conditional=True)
graph.add_edge(node, parent)
candidates.extend(_find_children(child, node))
elif isinstance(child, Runnable):
if xray > 0:
graph = child.get_graph(config, xray=xray - 1 if xray else 0)
graph.trim_first_node()
graph.trim_last_node()
s, e = graph.extend(graph, prefix=child.name or "")
if s is None:
raise ValueError(
f"Could not extend subgraph '{child.name}' due to missing entrypoint"
)
else:
graph.add_edge(parent, s, conditional=True)
if e is not None:
graph.add_edge(e, parent)
else:
node = graph.add_node(child, child.name)
graph.add_edge(parent, node, conditional=True)
graph.add_edge(node, parent)
return graph
def _find_children(
candidate: Union[Callable, Runnable], parent: Node
) -> Iterator[tuple[Node, Union[Callable, PregelProtocol]]]:
from langchain_core.runnables.utils import get_function_nonlocals
from langgraph.utils.runnable import (
RunnableCallable,
RunnableLambda,
RunnableSeq,
RunnableSequence,
)
candidates: list[Union[Callable, Runnable]] = []
if callable(candidate) and getattr(candidate, "_is_pregel_task", False) is True:
candidates.extend(
nl.__self__ if hasattr(nl, "__self__") else nl
for nl in get_function_nonlocals(
candidate.__wrapped__
if hasattr(candidate, "__wrapped__") and callable(candidate.__wrapped__)
else candidate
)
)
else:
candidates.append(candidate)
for c in candidates:
if callable(c) and getattr(c, "_is_pregel_task", False) is True:
yield (parent, c)
elif isinstance(c, PregelProtocol):
yield (parent, c)
elif isinstance(c, RunnableSequence) or isinstance(c, RunnableSeq):
candidates.extend(c.steps)
elif isinstance(c, RunnableLambda):
candidates.extend(c.deps)
elif isinstance(c, RunnableCallable):
if c.func is not None:
candidates.extend(
nl.__self__ if hasattr(nl, "__self__") else nl
for nl in get_function_nonlocals(c.func)
)
elif c.afunc is not None:
candidates.extend(
nl.__self__ if hasattr(nl, "__self__") else nl
for nl in get_function_nonlocals(c.afunc)
)
File diff suppressed because one or more lines are too long
@@ -1,334 +1,4 @@
# serializer version: 1
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[memory]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio_pipe]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio_pool]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[sqlite_aio]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2.1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2.2
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory]
'''
graph TD;
@@ -1055,253 +725,6 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio_pipe]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio_pool]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite_aio]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_nested_graph
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| inner |
+-------+
*
*
*
+------+
| side |
+------+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_send_react_interrupt_control[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
@@ -1392,128 +815,3 @@
'''
# ---
# name: test_weather_subgraph[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_aio_pipe]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_aio_pool]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[sqlite_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
-12
View File
@@ -1568,8 +1568,6 @@ def test_imp_nested(
"title": "LangGraphOutput",
}
assert graph.get_graph().draw_mermaid() == snapshot
thread1 = {"configurable": {"thread_id": "1"}}
assert [*graph.stream([0, 1], thread1)] == [
{"submapper": "0"},
@@ -1619,8 +1617,6 @@ def test_imp_stream_order(
fut_baz = baz(fut_bar.result())
return fut_baz.result()
assert graph.get_graph().draw_mermaid() == snapshot
thread1 = {"configurable": {"thread_id": "1"}}
assert [c for c in graph.stream({"a": "0"}, thread1)] == [
{
@@ -5048,8 +5044,6 @@ def test_interrupt_functional(
fut_bar = bar(bar_input)
return fut_bar.result()
assert graph.get_graph().draw_mermaid() == snapshot
config = {"configurable": {"thread_id": "1"}}
# First run, interrupted at bar
graph.invoke({"a": ""}, config)
@@ -5081,8 +5075,6 @@ def test_interrupt_task_functional(
fut_bar = bar(fut_foo.result())
return fut_bar.result()
assert graph.get_graph().draw_mermaid() == snapshot
config = {"configurable": {"thread_id": "1"}}
# First run, interrupted at bar
assert not graph.invoke({"a": ""}, config)
@@ -5572,8 +5564,6 @@ def test_falsy_return_from_task(
falsy_task().result()
interrupt("test")
assert graph.get_graph().draw_mermaid() == snapshot
configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
graph.invoke({"a": 5}, configurable)
graph.invoke(Command(resume="123"), configurable)
@@ -5606,8 +5596,6 @@ def test_multiple_interrupts_imperative(
return {"values": values}
assert graph.get_graph().draw_mermaid() == snapshot
configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
graph.invoke({}, configurable)
graph.invoke(Command(resume="a"), configurable)