From 3d58aff93edb3f940bb152bed641f2e7ac642605 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 13 Jan 2024 10:44:42 -0800 Subject: [PATCH] Add async test --- tests/test_pregel_async.py | 194 ++++++++++++++++++++++++++++++++++++- 1 file changed, 192 insertions(+), 2 deletions(-) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 3ee58d257..96adce78f 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1,7 +1,7 @@ import asyncio import operator from contextlib import asynccontextmanager, contextmanager -from typing import Any, AsyncGenerator, AsyncIterator, Generator +from typing import Annotated, Any, AsyncGenerator, AsyncIterator, Generator, TypedDict import pytest from langchain_core.runnables import RunnablePassthrough @@ -13,7 +13,7 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.memory import MemorySaver -from langgraph.graph import END, Graph +from langgraph.graph import END, Graph, StateGraph from langgraph.pregel import Channel, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -817,3 +817,193 @@ async def test_conditional_graph() -> None: # 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 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 + + class AgentState(TypedDict): + input: str + agent_outcome: AgentAction | AgentFinish | None + 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) -> AgentFinish | AgentAction: + 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 [ + deepcopy(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" + ), + } + }, + ]