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

Merge pull request #1309 from langchain-ai/nc/11aug/skip-context-get-state

Browse Source 
Skip initializing Context channels when calling get_state/get_state_history
This commit is contained in:
Nuno Campos
2024-08-11 10:56:54 -07:00
committed by GitHub
parent 00c1b6a2b0 9634468260
commit 4e643baaa6
6 changed files with 648 additions and 492 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+29 -8
View File
@@ -52,6 +52,7 @@ from langgraph.channels.base import (
BaseChannel,
)
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.manager import (
AsyncChannelsManager,
ChannelsManager,
@@ -360,7 +361,12 @@ class Pregel(
checkpoint = saved.checkpoint if saved else empty_checkpoint()
config = saved.config if saved else config
with ChannelsManager(
self.channels, checkpoint, config
{
k: LastValue(None) if isinstance(c, Context) else c
for k, c in self.channels.items()
},
checkpoint,
config,
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
@@ -375,7 +381,7 @@ class Pregel(
)
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
tuple(t.name for t in next_tasks),
saved.config if saved else config,
saved.metadata if saved else None,
saved.checkpoint["ts"] if saved else None,
@@ -392,7 +398,12 @@ class Pregel(
config = saved.config if saved else config
async with AsyncChannelsManager(
self.channels, checkpoint, config
{
k: LastValue(None) if isinstance(c, Context) else c
for k, c in self.channels.items()
},
checkpoint,
config,
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
@@ -407,7 +418,7 @@ class Pregel(
)
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
tuple(t.name for t in next_tasks),
saved.config if saved else config,
saved.metadata if saved else None,
saved.checkpoint["ts"] if saved else None,
@@ -434,7 +445,12 @@ class Pregel(
config, before=before, limit=limit, filter=filter
):
with ChannelsManager(
self.channels, checkpoint, config
{
k: LastValue(None) if isinstance(c, Context) else c
for k, c in self.channels.items()
},
checkpoint,
config,
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
@@ -449,7 +465,7 @@ class Pregel(
)
yield StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
tuple(t.name for t in next_tasks),
config,
metadata,
checkpoint["ts"],
@@ -480,7 +496,12 @@ class Pregel(
_,
) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
async with AsyncChannelsManager(
self.channels, checkpoint, config
{
k: LastValue(None) if isinstance(c, Context) else c
for k, c in self.channels.items()
},
checkpoint,
config,
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
@@ -495,7 +516,7 @@ class Pregel(
)
yield StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
tuple(t.name for t in next_tasks),
config,
metadata,
checkpoint["ts"],
libs/langgraph/langgraph/pregel/algo.py
+11 -5
View File
@@ -328,7 +328,7 @@ def prepare_next_tasks(
)
)
else:
tasks.append(PregelTaskDescription(packet.node, packet.arg))
tasks.append(PregelTaskDescription(packet.node))
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
@@ -348,7 +348,11 @@ def prepare_next_tasks(
> seen.get(chan, null_version)
):
try:
val = next(_proc_input(step, name, proc, managed, channels))
val = next(
_proc_input(
step, name, proc, managed, channels, for_execution=for_execution
)
)
except StopIteration:
continue
@@ -415,7 +419,7 @@ def prepare_next_tasks(
)
)
else:
tasks.append(PregelTaskDescription(name, val))
tasks.append(PregelTaskDescription(name))
return tasks
@@ -425,6 +429,8 @@ def _proc_input(
proc: PregelNode,
managed: ManagedValueMapping,
channels: Mapping[str, BaseChannel],
*,
for_execution: bool,
) -> Iterator[Any]:
# If all trigger channels subscribed by this process are not empty
# then invoke the process with the values of all non-empty channels
@@ -444,7 +450,7 @@ def _proc_input(
for key, chan in proc.channels.items():
if is_managed_value(chan):
managed_values[key] = managed[key](
step, PregelTaskDescription(name, val)
step, PregelTaskDescription(name)
)
val.update(managed_values)
@@ -465,7 +471,7 @@ def _proc_input(
)
# If the process has a mapper, apply it to the value
if proc.mapper is not None:
if for_execution and proc.mapper is not None:
val = proc.mapper(val)
yield val
libs/langgraph/langgraph/pregel/types.py
-1
View File
@@ -58,7 +58,6 @@ class RetryPolicy(NamedTuple):
class PregelTaskDescription(NamedTuple):
name: str
input: Any
class PregelExecutableTask(NamedTuple):
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+7 -93
View File
@@ -739,17 +739,6 @@
}),
}),
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/definitions/InnerObject',
}),
@@ -761,29 +750,13 @@
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1.2
dict({
'definitions': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'title': 'Answer',
@@ -796,20 +769,12 @@
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/definitions/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'answer',
'docs',
]),
'title': 'State',
'title': 'Output',
'type': 'object',
})
# ---
@@ -844,25 +809,6 @@
}),
}),
'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',
}),
@@ -874,41 +820,17 @@
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'title': 'Input',
'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',
'type': 'string',
}),
'docs': dict({
'items': dict({
@@ -917,20 +839,12 @@
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'answer',
'docs',
]),
'title': 'State',
'title': 'Output',
'type': 'object',
})
# ---
libs/langgraph/tests/test_pregel.py
+330 -211
View File
@@ -10,6 +10,7 @@ from typing import (
Any,
Dict,
Generator,
Iterator,
List,
Literal,
Optional,
@@ -2555,21 +2556,47 @@ def test_conditional_entrypoint_to_multiple_state_graph(
}
def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
def test_conditional_state_graph(
snapshot: SnapshotAssertion, mocker: MockerFixture
) -> None:
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.language_models.fake import FakeStreamingListLLM
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
setup = mocker.Mock()
teardown = mocker.Mock()
@contextmanager
def assert_ctx_once() -> Iterator[None]:
assert setup.call_count == 0
assert teardown.call_count == 0
try:
yield
finally:
assert setup.call_count == 1
assert teardown.call_count == 1
setup.reset_mock()
teardown.reset_mock()
@contextmanager
def make_httpx_client() -> Iterator[httpx.Client]:
setup()
with httpx.Client() as client:
try:
yield client
finally:
teardown()
class AgentState(TypedDict, total=False):
input: Annotated[str, UntrackedValue]
agent_outcome: Optional[Union[AgentAction, AgentFinish]]
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
session: Annotated[httpx.Client, Context(httpx.Client)]
session: Annotated[httpx.Client, Context(make_httpx_client)]
class ToolState(TypedDict, total=False):
agent_outcome: Union[AgentAction, AgentFinish]
session: Annotated[httpx.Client, Context(httpx.Client)]
session: Annotated[httpx.Client, Context(make_httpx_client)]
# Assemble the tools
@tool()
@@ -2652,84 +2679,88 @@ def test_conditional_state_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.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
assert [*app.stream({"input": "what is weather in sf"})] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
with assert_ctx_once():
assert app.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
],
}
},
{
"agent": {
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
"result for query",
),
}
},
]
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
with assert_ctx_once():
assert [*app.stream({"input": "what is weather in sf"})] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
}
},
{
"agent": {
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
},
]
# test state get/update methods with interrupt_after
@@ -2739,17 +2770,21 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
with assert_ctx_once():
assert [
c
for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
@@ -2777,16 +2812,17 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config,
)
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
with assert_ctx_once():
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
@@ -2816,41 +2852,43 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
with assert_ctx_once():
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
with assert_ctx_once():
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
@@ -6866,10 +6904,34 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
snapshot: SnapshotAssertion,
snapshot: SnapshotAssertion, mocker: MockerFixture
) -> None:
from langchain_core.pydantic_v1 import BaseModel, ValidationError
setup = mocker.Mock()
teardown = mocker.Mock()
@contextmanager
def assert_ctx_once() -> Iterator[None]:
assert setup.call_count == 0
assert teardown.call_count == 0
try:
yield
finally:
assert setup.call_count == 1
assert teardown.call_count == 1
setup.reset_mock()
teardown.reset_mock()
@contextmanager
def make_httpx_client() -> Iterator[httpx.Client]:
setup()
with httpx.Client() as client:
try:
yield client
finally:
teardown()
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
@@ -6883,10 +6945,22 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
yo: int
class State(BaseModel):
class Config:
arbitrary_types_allowed = True
query: str
inner: InnerObject
answer: Optional[str] = None
docs: Annotated[list[str], sorted_add]
client: Annotated[httpx.Client, Context(make_httpx_client)]
class Input(BaseModel):
query: str
inner: InnerObject
class Output(BaseModel):
answer: str
docs: list[str]
class StateUpdate(BaseModel):
query: Optional[str] = None
@@ -6914,7 +6988,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
assert isinstance(data, State)
return "retriever_two"
workflow = StateGraph(State)
workflow = StateGraph(State, input=Input, output=Output)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
@@ -6937,23 +7011,25 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
assert app.get_input_schema().schema() == snapshot
assert app.get_output_schema().schema() == snapshot
with pytest.raises(ValidationError):
with pytest.raises(ValidationError), assert_ctx_once():
app.invoke({"query": {}})
assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == {
"query": "analyzed: query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
"inner": {"yo": 1},
}
with assert_ctx_once():
assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == {
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
}
assert [*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
with assert_ctx_once():
assert [
*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
@@ -6961,37 +7037,64 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
for c in app_w_interrupt.stream(
{"query": "what is weather in sf", "inner": {"yo": 1}}, config
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
with assert_ctx_once():
assert [
c
for c in app_w_interrupt.stream(
{"query": "what is weather in sf", "inner": {"yo": 1}}, config
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
with assert_ctx_once():
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert app_w_interrupt.update_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
with assert_ctx_once():
assert app_w_interrupt.update_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
}
}
}
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
snapshot: SnapshotAssertion,
snapshot: SnapshotAssertion, mocker: MockerFixture
) -> None:
from pydantic import BaseModel, ValidationError
from pydantic import BaseModel, ConfigDict, ValidationError
setup = mocker.Mock()
teardown = mocker.Mock()
@contextmanager
def assert_ctx_once() -> Iterator[None]:
assert setup.call_count == 0
assert teardown.call_count == 0
try:
yield
finally:
assert setup.call_count == 1
assert teardown.call_count == 1
setup.reset_mock()
teardown.reset_mock()
@contextmanager
def make_httpx_client() -> Iterator[httpx.Client]:
setup()
with httpx.Client() as client:
try:
yield client
finally:
teardown()
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
@@ -7006,16 +7109,27 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
yo: int
class State(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
query: str
inner: InnerObject
answer: Optional[str] = None
docs: Annotated[list[str], sorted_add]
client: Annotated[httpx.Client, Context(make_httpx_client)]
class StateUpdate(BaseModel):
query: Optional[str] = None
answer: Optional[str] = None
docs: Optional[list[str]] = None
class Input(BaseModel):
query: str
inner: InnerObject
class Output(BaseModel):
answer: str
docs: list[str]
def rewrite_query(data: State) -> State:
return {"query": f"query: {data.query}"}
@@ -7036,7 +7150,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
assert isinstance(data, State)
return "retriever_two"
workflow = StateGraph(State)
workflow = StateGraph(State, input=Input, output=Output)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
@@ -7059,23 +7173,25 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
assert app.get_input_schema().schema() == snapshot
assert app.get_output_schema().schema() == snapshot
with pytest.raises(ValidationError):
with pytest.raises(ValidationError), assert_ctx_once():
app.invoke({"query": {}})
assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == {
"query": "analyzed: query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
"inner": {"yo": 1},
}
with assert_ctx_once():
assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == {
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
}
assert [*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
with assert_ctx_once():
assert [
*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
@@ -7083,31 +7199,34 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
for c in app_w_interrupt.stream(
{"query": "what is weather in sf", "inner": {"yo": 1}}, config
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
with assert_ctx_once():
assert [
c
for c in app_w_interrupt.stream(
{"query": "what is weather in sf", "inner": {"yo": 1}}, config
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
with assert_ctx_once():
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert app_w_interrupt.update_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
with assert_ctx_once():
assert app_w_interrupt.update_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
}
}
}
def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
libs/langgraph/tests/test_pregel_async.py
+271 -174
View File
@@ -2664,12 +2664,27 @@ async def test_conditional_graph() -> None:
]
async def test_conditional_graph_state() -> None:
async def test_conditional_graph_state(mocker: MockerFixture) -> None:
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.language_models.fake import FakeStreamingListLLM
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
setup = mocker.Mock()
teardown = mocker.Mock()
@asynccontextmanager
async def assert_ctx_once() -> AsyncIterator[None]:
assert setup.call_count == 0
assert teardown.call_count == 0
try:
yield
finally:
assert setup.call_count == 1
assert teardown.call_count == 1
setup.reset_mock()
teardown.reset_mock()
class MyPydanticContextModel(BaseModel):
class Config:
arbitrary_types_allowed = True
@@ -2682,11 +2697,13 @@ async def test_conditional_graph_state() -> None:
config: RunnableConfig,
) -> AsyncIterator[MyPydanticContextModel]:
assert isinstance(config, dict)
setup()
session = httpx.AsyncClient()
try:
yield MyPydanticContextModel(session=session, something_else="hello")
finally:
await session.aclose()
teardown()
class AgentState(TypedDict):
input: Annotated[str, UntrackedValue]
@@ -2768,86 +2785,91 @@ async def test_conditional_graph_state() -> None:
app = workflow.compile()
assert await app.ainvoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
async with assert_ctx_once():
assert await app.ainvoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
],
}
},
{
"agent": {
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
"result for query",
),
}
},
]
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
async with assert_ctx_once():
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
}
},
{
"agent": {
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
},
]
async with assert_ctx_once():
patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
patch_paths = {op["path"] for log in patches for op in log.ops}
# Check that agent (one of the nodes) has its output streamed to the logs
@@ -2887,20 +2909,23 @@ async def test_conditional_graph_state() -> None:
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
async with assert_ctx_once():
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
},
]
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
@@ -2934,16 +2959,17 @@ async def test_conditional_graph_state() -> None:
][-1].config,
)
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
async with assert_ctx_once():
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
@@ -2977,41 +3003,43 @@ async def test_conditional_graph_state() -> None:
][-1].config,
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
async with assert_ctx_once():
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
async with assert_ctx_once():
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
@@ -5462,10 +5490,34 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
snapshot: SnapshotAssertion,
snapshot: SnapshotAssertion, mocker: MockerFixture
) -> None:
from langchain_core.pydantic_v1 import BaseModel, ValidationError
setup = mocker.Mock()
teardown = mocker.Mock()
@asynccontextmanager
async def assert_ctx_once() -> AsyncIterator[None]:
assert setup.call_count == 0
assert teardown.call_count == 0
try:
yield
finally:
assert setup.call_count == 1
assert teardown.call_count == 1
setup.reset_mock()
teardown.reset_mock()
@asynccontextmanager
async def make_httpx_client() -> AsyncIterator[httpx.AsyncClient]:
setup()
async with httpx.AsyncClient() as client:
try:
yield client
finally:
teardown()
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
@@ -5476,9 +5528,20 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
return sorted(operator.add(x, y))
class State(BaseModel):
class Config:
arbitrary_types_allowed = True
query: str
answer: Optional[str] = None
docs: Annotated[list[str], sorted_add]
client: Annotated[httpx.AsyncClient, Context(make_httpx_client)]
class Input(BaseModel):
query: str
class Output(BaseModel):
answer: str
docs: list[str]
class StateUpdate(BaseModel):
query: Optional[str] = None
@@ -5505,7 +5568,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
assert isinstance(data, State)
return "retriever_two"
workflow = StateGraph(State)
workflow = StateGraph(State, input=Input, output=Output)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
@@ -5526,22 +5589,24 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
assert app.get_graph().draw_ascii() == snapshot
with pytest.raises(ValidationError):
await app.ainvoke({"query": {}})
async with assert_ctx_once():
with pytest.raises(ValidationError):
await app.ainvoke({"query": {}})
assert await app.ainvoke({"query": "what is weather in sf"}) == {
"query": "analyzed: query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
}
async with assert_ctx_once():
assert await app.ainvoke({"query": "what is weather in sf"}) == {
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
}
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
async with assert_ctx_once():
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
@@ -5549,31 +5614,63 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"query": "what is weather in sf"}, config
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
async with assert_ctx_once():
assert [
c
async for c in app_w_interrupt.astream(
{"query": "what is weather in sf"}, config
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
async with assert_ctx_once():
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert await app_w_interrupt.aupdate_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"query": "analyzed: query: what is weather in sf",
"answer": "doc1,doc2,doc3,doc4",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
)
async with assert_ctx_once():
assert await app_w_interrupt.aupdate_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
}
}
}
async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(