import asyncio import json import operator from contextlib import asynccontextmanager, contextmanager from typing import ( Annotated, Any, AsyncGenerator, AsyncIterator, Generator, Optional, TypedDict, Union, ) import pytest from langchain_core.runnables import RunnablePassthrough from pytest_mock import MockerFixture from langgraph.channels.base import InvalidUpdateError from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph from langgraph.prebuilt.chat_agent_executor import ( create_function_calling_executor, create_tool_calling_executor, ) from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( nodes={ "one": chain, }, channels={ "input": LastValue(int), "output": LastValue(int), }, input="input", output="output", ) graph = Graph() graph.add_node("add_one", add_one) graph.set_entry_point("add_one") graph.set_finish_point("add_one") gapp = graph.compile() assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"} assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"} assert await app.ainvoke(2) == 3 assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3} assert await gapp.ainvoke(2) == 3 async def test_invoke_single_process_in_out_implicit_channels( mocker: MockerFixture, ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel(nodes={"one": chain}) assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == {"title": "LangGraphOutput"} assert await app.ainvoke(2) == 3 async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = ( Channel.subscribe_to("input") | add_one | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) ) app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"]) assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", "properties": { "output": {"title": "Output"}, "fixed": {"title": "Fixed"}, "output_plus_one": {"title": "Output Plus One"}, }, } assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} async def test_invoke_single_process_in_out_reserved_is_last( mocker: MockerFixture, ) -> None: add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1}) chain = ( Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step]) | add_one | Channel.write_to("output") ) app = Pregel(nodes={"one": chain}) assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == {"title": "LangGraphOutput"} assert await app.ainvoke(2) == {"input": 3, "is_last_step": False} assert await app.ainvoke(2, {"recursion_limit": 1}) == { "input": 3, "is_last_step": True, } async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( nodes={"one": chain}, output=["output"], ) assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", "properties": {"output": {"title": "Output"}}, } assert await app.ainvoke(2) == {"output": 3} async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( nodes={ "one": chain, }, input=["input"], output=["output"], ) assert app.input_schema.schema() == { "title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input"}}, } assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", "properties": {"output": {"title": "Output"}}, } assert await app.ainvoke({"input": 2}) == {"output": 3} async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") app = Pregel(nodes={"one": one, "two": two}) assert await app.ainvoke(2) == 4 assert await app.ainvoke(2, input_keys="inbox") == 3 with pytest.raises(GraphRecursionError): await app.ainvoke(2, {"recursion_limit": 1}) step = 0 async for values in app.astream(2): step += 1 if step == 1: assert values == { "inbox": 3, } elif step == 2: assert values == { "output": 4, } assert step == 2 step = 0 async for values in app.astream(2): step += 1 if step == 1: assert values == { "inbox": 3, } # modify inbox value values["inbox"] = 5 elif step == 2: # output is different now assert values == { "output": 6, } assert step == 2 graph = Graph() graph.add_node("add_one", add_one) graph.add_node("add_one_more", add_one) graph.set_entry_point("add_one") graph.set_finish_point("add_one_more") graph.add_edge("add_one", "add_one_more") gapp = graph.compile() assert await gapp.ainvoke(2) == 4 step = 0 async for values in gapp.astream(2): step += 1 if step == 1: assert values == { "add_one": 3, } elif step == 2: assert values == { "add_one_more": 4, } elif step == 3: assert values == { "__end__": 4, } assert step == 3 step = 0 async for values in gapp.astream(2): step += 1 if step == 1: assert values == { "add_one": 3, } # modify value before running next step values["add_one"] = 5 elif step == 2: # output is different now assert values == { "add_one_more": 6, } elif step == 3: assert values == { "__end__": 6, } assert step == 3 async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") memory = MemorySaver() app = Pregel( nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"] ) # start execution, stop at inbox assert await app.ainvoke(2, {"configurable": {"thread_id": 1}}) is None # inbox == 3 checkpoint = await memory.aget({"configurable": {"thread_id": 1}}) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 3 # resume execution, finish assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 4 # start execution again, stop at inbox assert await app.ainvoke(20, {"configurable": {"thread_id": 1}}) is None # inbox == 21 checkpoint = await memory.aget({"configurable": {"thread_id": 1}}) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 21 # send a new value in, interrupting the previous execution assert await app.ainvoke(3, {"configurable": {"thread_id": 1}}) is None assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 5 async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") pubsub = Pregel( nodes={"one": one, "two": two}, channels={"inbox": Topic(int)}, input=["input", "inbox"], ) # [12 + 1, 2 + 1 + 1] assert [ c async for c in pubsub.astream({"input": 2, "inbox": 12}, output_keys="output") ] == [13, 4] assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [ {"inbox": [3], "output": 13}, {"output": 4}, ] async def test_batch_two_processes_in_out() -> None: async def add_one_with_delay(inp: int) -> int: await asyncio.sleep(inp / 10) return inp + 1 one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") app = Pregel( nodes={"one": one, "two": two}, channels={"one": LastValue(int)}, ) assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] assert await app.abatch([3, 2, 1, 3, 5], output_keys=["output"]) == [ {"output": 5}, {"output": 4}, {"output": 3}, {"output": 5}, {"output": 7}, ] graph = Graph() graph.add_node("add_one", add_one_with_delay) graph.add_node("add_one_more", add_one_with_delay) graph.set_entry_point("add_one") graph.set_finish_point("add_one_more") graph.add_edge("add_one", "add_one_more") gapp = graph.compile() assert await gapp.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): nodes[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") app = Pregel(nodes=nodes) # No state is left over from previous invocations for _ in range(10): assert await app.ainvoke(2, {"recursion_limit": test_size}) == 2 + test_size # Concurrent invocations do not interfere with each other assert await asyncio.gather( *(app.ainvoke(2, {"recursion_limit": test_size}) for _ in range(10)) ) == [2 + test_size for _ in range(10)] async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): nodes[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") app = Pregel(nodes=nodes) # No state is left over from previous invocations for _ in range(3): # Then invoke pubsub assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ 2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size, ] # Concurrent invocations do not interfere with each other assert await asyncio.gather( *(app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) for _ in range(3)) ) == [ [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size] for _ in range(3) ] async def test_invoke_two_processes_two_in_two_out_invalid( mocker: MockerFixture, ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel(nodes={"one": one, "two": two}) with pytest.raises(InvalidUpdateError): # LastValue channels can only be updated once per iteration await app.ainvoke(2) async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( nodes={"one": one, "two": two}, channels={"output": Topic(int)}, ) # An Topic channel accumulates updates into a sequence assert await app.ainvoke(2) == [3, 3] async def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: if input > 10: raise ValueError("Input is too large") return input one = ( Channel.subscribe_to(["input"]).join(["total"]) | add_one | Channel.write_to("output", "total") | raise_if_above_10 ) memory = MemorySaver() app = Pregel( nodes={"one": one}, channels={"total": BinaryOperatorAggregate(int, operator.add)}, checkpointer=memory, ) # total starts out as 0, so output is 0+2=2 assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 2 # total is now 2, so output is 2+3=5 assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) # checkpoint is not updated checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 # on a new thread, total starts out as 0, so output is 0+5=5 assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 checkpoint = await memory.aget({"configurable": {"thread_id": "2"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 5 async def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: if input > 10: raise ValueError("Input is too large") return input one = ( Channel.subscribe_to(["input"]).join(["total"]) | add_one | Channel.write_to("output", "total") | raise_if_above_10 ) memory = AsyncSqliteSaver.from_conn_string(":memory:") app = Pregel( nodes={"one": one}, channels={"total": BinaryOperatorAggregate(int, operator.add)}, checkpointer=memory, debug=True, ) # total starts out as 0, so output is 0+2=2 assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 2 # total is now 2, so output is 2+3=5 assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) # checkpoint is not updated checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 # on a new thread, total starts out as 0, so output is 0+5=5 assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 checkpoint = await memory.aget({"configurable": {"thread_id": "2"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 5 await memory.conn.close() async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") chain_four = ( Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output") ) app = Pregel( nodes={ "one": one, "chain_three": chain_three, "chain_four": chain_four, }, channels={"inbox": Topic(int)}, ) # Then invoke app # We get a single array result as chain_four waits for all publishers to finish # before operating on all elements published to topic_two as an array for _ in range(100): assert await app.ainvoke(2) == [13, 13] assert await asyncio.gather(*(app.ainvoke(2) for _ in range(100))) == [ [13, 13] for _ in range(100) ] async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) inner_app = Pregel( nodes={ "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") } ) one = ( Channel.subscribe_to("input") | add_10_each | Channel.write_to("inbox_one").map() ) two = ( Channel.subscribe_to("inbox_one") | inner_app.map() | sorted | Channel.write_to("outbox_one") ) chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") app = Pregel( nodes={ "one": one, "two": two, "chain_three": chain_three, }, channels={ "inbox_one": Topic(int), "outbox_one": LastValue(int), }, ) # Then invoke pubsub for _ in range(10): assert await app.ainvoke([2, 3]) == 27 assert await asyncio.gather(*(app.ainvoke([2, 3]) for _ in range(10))) == [ 27 for _ in range(10) ] async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = ( Channel.subscribe_to("input") | add_one | Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") app = Pregel(nodes={"one": one, "two": two}) # Then invoke pubsub assert [c async for c in app.astream(2)] == [ {"between": 3, "output": 3}, {"output": 4}, ] async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") two = Channel.subscribe_to("between") | add_one app = Pregel(nodes={"one": one, "two": two}) # It finishes executing (once no more messages being published) # but returns nothing, as nothing was published to "output" topic assert await app.ainvoke(2) is None async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: setup_sync = mocker.Mock() cleanup_sync = mocker.Mock() setup_async = mocker.Mock() cleanup_async = mocker.Mock() @contextmanager def an_int() -> Generator[int, None, None]: setup_sync() try: yield 5 finally: cleanup_sync() @asynccontextmanager async def an_int_async() -> AsyncGenerator[int, None]: setup_async() try: yield 5 finally: cleanup_async() add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") app = Pregel( nodes={"one": one, "two": two}, channels={ "inbox": Topic(int), "ctx": Context(an_int, an_int_async, typ=int), }, output=["inbox", "output"], ) async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]: i = 0 async for chunk in aiter: yield i, chunk i += 1 assert setup_sync.call_count == 0 assert cleanup_sync.call_count == 0 assert setup_async.call_count == 0 assert cleanup_async.call_count == 0 async for i, chunk in aenumerate(app.astream(2)): assert setup_sync.call_count == 0, "Sync context manager should not be used" assert cleanup_sync.call_count == 0, "Sync context manager should not be used" assert setup_async.call_count == 1, "Expected setup to be called once" assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet" if i == 0: assert chunk == {"inbox": [3]} elif i == 1: assert chunk == {"output": 4} else: assert False, "Expected only two chunks" assert setup_sync.call_count == 0 assert cleanup_sync.call_count == 0 assert setup_async.call_count == 1, "Expected setup to be called once" assert cleanup_async.call_count == 1, "Expected cleanup to be called once" async def test_conditional_graph() -> None: from copy import deepcopy from langchain.llms.fake import FakeStreamingListLLM from langchain_community.tools import tool from langchain_core.agents import AgentAction, AgentFinish from langchain_core.prompts import PromptTemplate from langchain_core.runnables import RunnablePassthrough # Assemble the tools @tool() def search_api(query: str) -> str: """Searches the API for the query.""" return f"result for {query}" tools = [search_api] # Construct the agent prompt = PromptTemplate.from_template("Hello!") llm = FakeStreamingListLLM( responses=[ "tool:search_api:query", "tool:search_api:another", "finish:answer", ] ) async def agent_parser(input: str) -> Union[AgentAction, AgentFinish]: if input.startswith("finish"): _, answer = input.split(":") return AgentFinish(return_values={"answer": answer}, log=input) else: _, tool_name, tool_input = input.split(":") return AgentAction(tool=tool_name, tool_input=tool_input, log=input) agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser) # Define tool execution logic async def execute_tools(data: dict) -> dict: agent_action: AgentAction = data.pop("agent_outcome") observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke( agent_action.tool_input ) if data.get("intermediate_steps") is None: data["intermediate_steps"] = [] data["intermediate_steps"].append((agent_action, observation)) return data # Define decision-making logic def should_continue(data: dict) -> str: # Logic to decide whether to continue in the loop or exit if isinstance(data["agent_outcome"], AgentFinish): return "exit" else: return "continue" # Define a new graph workflow = Graph() workflow.add_node("agent", agent) workflow.add_node("tools", execute_tools) workflow.set_entry_point("agent") workflow.add_conditional_edges( "agent", should_continue, {"continue": "tools", "exit": END} ) workflow.add_edge("tools", "agent") 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 [ deepcopy(c) async for c in app.astream({"input": "what is weather in sf"}) ] == [ { "agent": { "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), } }, { "tools": { "input": "what is weather in sf", "intermediate_steps": [ ( AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", ) ], } }, { "agent": { "input": "what is weather in sf", "intermediate_steps": [ ( AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", ) ], "agent_outcome": AgentAction( tool="search_api", tool_input="another", log="tool:search_api:another", ), } }, { "tools": { "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": { "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" ), } }, { "__end__": { "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" ), } }, ] 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 assert "/logs/agent/streamed_output/-" in patch_paths async def test_conditional_graph_state() -> None: from langchain.llms.fake import FakeStreamingListLLM from langchain_community.tools import tool from langchain_core.agents import AgentAction, AgentFinish from langchain_core.prompts import PromptTemplate class AgentState(TypedDict): input: str agent_outcome: Optional[Union[AgentAction, AgentFinish]] intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] # Assemble the tools @tool() def search_api(query: str) -> str: """Searches the API for the query.""" return f"result for {query}" tools = [search_api] # Construct the agent prompt = PromptTemplate.from_template("Hello!") llm = FakeStreamingListLLM( responses=[ "tool:search_api:query", "tool:search_api:another", "finish:answer", ] ) def agent_parser(input: str) -> dict[str, Union[AgentAction, AgentFinish]]: if input.startswith("finish"): _, answer = input.split(":") return { "agent_outcome": AgentFinish( return_values={"answer": answer}, log=input ) } else: _, tool_name, tool_input = input.split(":") return { "agent_outcome": AgentAction( tool=tool_name, tool_input=tool_input, log=input ) } agent = prompt | llm | agent_parser # Define tool execution logic def execute_tools(data: AgentState) -> dict: agent_action: AgentAction = data.pop("agent_outcome") observation = {t.name: t for t in tools}[agent_action.tool].invoke( agent_action.tool_input ) return {"intermediate_steps": [(agent_action, observation)]} # Define decision-making logic def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit if isinstance(data["agent_outcome"], AgentFinish): return "exit" else: return "continue" # Define a new graph workflow = StateGraph(AgentState) workflow.add_node("agent", agent) workflow.add_node("tools", execute_tools) workflow.set_entry_point("agent") workflow.add_conditional_edges( "agent", should_continue, {"continue": "tools", "exit": END} ) workflow.add_edge("tools", "agent") 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", ), ], } }, { "agent": { "agent_outcome": AgentFinish( return_values={"answer": "answer"}, log="finish:answer" ), } }, { "__end__": { "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" ), } }, ] async def test_prebuilt_tool_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.messages import AIMessage, HumanMessage, ToolMessage class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @tool() def search_api(query: str) -> str: """Searches the API for the query.""" return f"result for {query}" tools = [search_api] app = create_tool_calling_executor( FakeFuntionChatModel( responses=[ AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call123", "type": "function", "function": { "name": "search_api", "arguments": json.dumps("query"), }, } ] }, ), AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call234", "type": "function", "function": { "name": "search_api", "arguments": json.dumps("another"), }, }, { "id": "tool_call567", "type": "function", "function": { "name": "search_api", "arguments": '"a third one"', }, }, ] }, ), AIMessage(content="answer"), ] ), tools, ) assert await app.ainvoke( {"messages": [HumanMessage(content="what is weather in sf")]} ) == { "messages": [ HumanMessage(content="what is weather in sf"), AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call123", "type": "function", "function": { "name": "search_api", "arguments": '"query"', }, } ] }, ), ToolMessage(content="result for query", tool_call_id="tool_call123"), AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call234", "type": "function", "function": { "name": "search_api", "arguments": '"another"', }, }, { "id": "tool_call567", "type": "function", "function": { "name": "search_api", "arguments": '"a third one"', }, }, ] }, ), ToolMessage(content="result for another", tool_call_id="tool_call234"), ToolMessage(content="result for a third one", tool_call_id="tool_call567"), AIMessage(content="answer"), ] } assert [ c async for c in app.astream( {"messages": [HumanMessage(content="what is weather in sf")]} ) ] == [ { "agent": { "messages": [ AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call123", "type": "function", "function": { "name": "search_api", "arguments": '"query"', }, } ] }, ) ] } }, { "action": { "messages": [ ToolMessage(content="result for query", tool_call_id="tool_call123") ] } }, { "agent": { "messages": [ AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call234", "type": "function", "function": { "name": "search_api", "arguments": '"another"', }, }, { "id": "tool_call567", "type": "function", "function": { "name": "search_api", "arguments": '"a third one"', }, }, ] }, ) ] } }, { "action": { "messages": [ ToolMessage( content="result for another", tool_call_id="tool_call234" ), ToolMessage( content="result for a third one", tool_call_id="tool_call567" ), ] } }, {"agent": {"messages": [AIMessage(content="answer")]}}, { "__end__": { "messages": [ HumanMessage(content="what is weather in sf"), AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call123", "type": "function", "function": { "name": "search_api", "arguments": '"query"', }, } ] }, ), ToolMessage( content="result for query", tool_call_id="tool_call123" ), AIMessage( content="", additional_kwargs={ "tool_calls": [ { "id": "tool_call234", "type": "function", "function": { "name": "search_api", "arguments": '"another"', }, }, { "id": "tool_call567", "type": "function", "function": { "name": "search_api", "arguments": '"a third one"', }, }, ] }, ), ToolMessage( content="result for another", tool_call_id="tool_call234" ), ToolMessage( content="result for a third one", tool_call_id="tool_call567" ), AIMessage(content="answer"), ] } }, ] async def test_prebuilt_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @tool() def search_api(query: str) -> str: """Searches the API for the query.""" return f"result for {query}" tools = [search_api] app = create_function_calling_executor( FakeFuntionChatModel( responses=[ AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": json.dumps("query"), } }, ), AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": json.dumps("another"), } }, ), AIMessage(content="answer"), ] ), tools, ) assert await app.ainvoke( {"messages": [HumanMessage(content="what is weather in sf")]} ) == { "messages": [ HumanMessage(content="what is weather in sf"), AIMessage( content="", additional_kwargs={ "function_call": {"name": "search_api", "arguments": '"query"'} }, ), FunctionMessage(content="result for query", name="search_api"), AIMessage( content="", additional_kwargs={ "function_call": {"name": "search_api", "arguments": '"another"'} }, ), FunctionMessage(content="result for another", name="search_api"), AIMessage(content="answer"), ] } assert [ c async for c in app.astream( {"messages": [HumanMessage(content="what is weather in sf")]} ) ] == [ { "agent": { "messages": [ AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": '"query"', } }, ) ] } }, { "action": { "messages": [ FunctionMessage(content="result for query", name="search_api") ] } }, { "agent": { "messages": [ AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": '"another"', } }, ) ] } }, { "action": { "messages": [ FunctionMessage(content="result for another", name="search_api") ] } }, {"agent": {"messages": [AIMessage(content="answer")]}}, { "__end__": { "messages": [ HumanMessage(content="what is weather in sf"), AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": '"query"', } }, ), FunctionMessage(content="result for query", name="search_api"), AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": '"another"', } }, ), FunctionMessage(content="result for another", name="search_api"), AIMessage(content="answer"), ] } }, ] async def test_message_graph() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.agents import AgentAction from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @tool() def search_api(query: str) -> str: """Searches the API for the query.""" return f"result for {query}" tools = [search_api] model = FakeFuntionChatModel( responses=[ AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": json.dumps("query"), } }, ), AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": json.dumps("another"), } }, ), AIMessage(content="answer"), ] ) tool_executor = ToolExecutor(tools) # Define the function that determines whether to continue or not def should_continue(messages): last_message = messages[-1] # If there is no function call, then we finish if "function_call" not in last_message.additional_kwargs: return "end" # Otherwise if there is, we continue else: return "continue" async def call_tool(messages): # Based on the continue condition # we know the last message involves a function call last_message = messages[-1] # We construct an AgentAction from the function_call action = AgentAction( tool=last_message.additional_kwargs["function_call"]["name"], tool_input=json.loads( last_message.additional_kwargs["function_call"]["arguments"] ), log="", ) # We call the tool_executor and get back a response response = await tool_executor.ainvoke(action) # We use the response to create a FunctionMessage return FunctionMessage(content=str(response), name=action.tool) # Define a new graph workflow = MessageGraph() # Define the two nodes we will cycle between workflow.add_node("agent", model) workflow.add_node("action", call_tool) # Set the entrypoint as `agent` # This means that this node is the first one called workflow.set_entry_point("agent") # We now add a conditional edge workflow.add_conditional_edges( # First, we define the start node. We use `agent`. # This means these are the edges taken after the `agent` node is called. "agent", # Next, we pass in the function that will determine which node is called next. should_continue, # Finally we pass in a mapping. # The keys are strings, and the values are other nodes. # END is a special node marking that the graph should finish. # What will happen is we will call `should_continue`, and then the output of that # will be matched against the keys in this mapping. # Based on which one it matches, that node will then be called. { # If `tools`, then we call the tool node. "continue": "action", # Otherwise we finish. "end": END, }, ) # We now add a normal edge from `tools` to `agent`. # This means that after `tools` is called, `agent` node is called next. workflow.add_edge("action", "agent") # Finally, we compile it! # This compiles it into a LangChain Runnable, # meaning you can use it as you would any other runnable app = workflow.compile() assert await app.ainvoke(HumanMessage(content="what is weather in sf")) == [ HumanMessage(content="what is weather in sf"), AIMessage( content="", additional_kwargs={ "function_call": {"name": "search_api", "arguments": '"query"'} }, ), FunctionMessage(content="result for query", name="search_api"), AIMessage( content="", additional_kwargs={ "function_call": {"name": "search_api", "arguments": '"another"'} }, ), FunctionMessage(content="result for another", name="search_api"), AIMessage(content="answer"), ] assert [ c async for c in app.astream([HumanMessage(content="what is weather in sf")]) ] == [ { "agent": AIMessage( content="", additional_kwargs={ "function_call": {"name": "search_api", "arguments": '"query"'} }, ) }, {"action": FunctionMessage(content="result for query", name="search_api")}, { "agent": AIMessage( content="", additional_kwargs={ "function_call": {"name": "search_api", "arguments": '"another"'} }, ) }, {"action": FunctionMessage(content="result for another", name="search_api")}, {"agent": AIMessage(content="answer")}, { "__end__": [ HumanMessage(content="what is weather in sf"), AIMessage( content="", additional_kwargs={ "function_call": {"name": "search_api", "arguments": '"query"'} }, ), FunctionMessage(content="result for query", name="search_api"), AIMessage( content="", additional_kwargs={ "function_call": { "name": "search_api", "arguments": '"another"', } }, ), FunctionMessage(content="result for another", name="search_api"), AIMessage(content="answer"), ] }, ] async def test_in_one_fan_out_out_one_graph_state() -> None: def sorted_add(x: list[str], y: list[str]) -> list[str]: return sorted(operator.add(x, y)) class State(TypedDict, total=False): query: str answer: str docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: return {"query": f'query: {data["query"]}'} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: return {"docs": ["doc3", "doc4"]} async def qa(data: State) -> State: return {"answer": ",".join(data["docs"])} workflow = StateGraph(State) workflow.add_node("rewrite_query", rewrite_query) workflow.add_node("retriever_one", retriever_one) workflow.add_node("retriever_two", retriever_two) workflow.add_node("qa", qa) workflow.set_entry_point("rewrite_query") workflow.add_edge("rewrite_query", "retriever_one") workflow.add_edge("rewrite_query", "retriever_two") workflow.add_edge("retriever_one", "qa") workflow.add_edge("retriever_two", "qa") workflow.set_finish_point("qa") app = workflow.compile() assert await app.ainvoke({"query": "what is weather in sf"}) == { "query": "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"}}, { "retriever_two": {"docs": ["doc3", "doc4"]}, "retriever_one": {"docs": ["doc1", "doc2"]}, }, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, { "__end__": { "query": "query: what is weather in sf", "answer": "doc1,doc2,doc3,doc4", "docs": ["doc1", "doc2", "doc3", "doc4"], } }, ]