diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 8f9583964..939cbc256 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -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"], diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index e30cec4aa..2ddbcd873 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -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 diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 19d3c3301..7784cf023 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -58,7 +58,6 @@ class RetryPolicy(NamedTuple): class PregelTaskDescription(NamedTuple): name: str - input: Any class PregelExecutableTask(NamedTuple): diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index f1ef888ec..004daa1e1 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -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', }) # --- diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 71f2f8c59..af7f0920a 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e5c305dd8..67b62941b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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(