diff --git a/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb b/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb index c41653c50..02f8782eb 100644 --- a/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb +++ b/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb @@ -235,7 +235,7 @@ " # Call the chat bot\n", " chat_bot_response = my_chat_bot(messages)\n", " # Respond with an AI Message\n", - " return {\"messages\":[AIMessage(content=chat_bot_response[\"content\"])]}" + " return {\"messages\": [AIMessage(content=chat_bot_response[\"content\"])]}" ] }, { @@ -270,7 +270,7 @@ " # Call the simulated user\n", " response = simulated_user.invoke({\"messages\": new_messages})\n", " # This response is an AI message - we need to flip this to be a human message\n", - " return {\"messages\":[HumanMessage(content=response.content)]}" + " return {\"messages\": [HumanMessage(content=response.content)]}" ] }, { @@ -331,6 +331,7 @@ "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "graph_builder = StateGraph(State)\n", "graph_builder.add_node(\"user\", simulated_user_node)\n", "graph_builder.add_node(\"chat_bot\", chat_bot_node)\n", diff --git a/examples/chatbots/information-gather-prompting.ipynb b/examples/chatbots/information-gather-prompting.ipynb index ad4cdfe35..ad375f02f 100644 --- a/examples/chatbots/information-gather-prompting.ipynb +++ b/examples/chatbots/information-gather-prompting.ipynb @@ -79,7 +79,7 @@ "\n", "\n", "def info_chain(state):\n", - " messages = get_messages_info(state['messages'])\n", + " messages = get_messages_info(state[\"messages\"])\n", " response = llm_with_tool.invoke(messages)\n", " return {\"messages\": [response]}" ] @@ -126,7 +126,7 @@ "\n", "\n", "def prompt_gen_chain(state):\n", - " messages = get_prompt_messages(state['messages'])\n", + " messages = get_prompt_messages(state[\"messages\"])\n", " response = llm.invoke(messages)\n", " return {\"messages\": [response]}" ] @@ -158,7 +158,7 @@ "\n", "\n", "def get_state(state) -> Literal[\"add_tool_message\", \"info\", \"__end__\"]:\n", - " messages = state['messages']\n", + " messages = state[\"messages\"]\n", " if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:\n", " return \"add_tool_message\"\n", " elif not isinstance(messages[-1], HumanMessage):\n", @@ -190,9 +190,11 @@ "from typing import Annotated\n", "from typing_extensions import TypedDict\n", "\n", + "\n", "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "memory = MemorySaver()\n", "workflow = StateGraph(State)\n", "workflow.add_node(\"info\", info_chain)\n", @@ -201,9 +203,14 @@ "\n", "@workflow.add_node\n", "def add_tool_message(state: State):\n", - " return {\"messages\": [ToolMessage(\n", - " content=\"Prompt generated!\", tool_call_id=state['messages'][-1].tool_calls[0][\"id\"]\n", - " )]}\n", + " return {\n", + " \"messages\": [\n", + " ToolMessage(\n", + " content=\"Prompt generated!\",\n", + " tool_call_id=state[\"messages\"][-1].tool_calls[0][\"id\"],\n", + " )\n", + " ]\n", + " }\n", "\n", "\n", "workflow.add_conditional_edges(\"info\", get_state)\n", @@ -364,7 +371,7 @@ " for output in graph.stream(\n", " {\"messages\": [HumanMessage(content=user)]}, config=config, stream_mode=\"updates\"\n", " ):\n", - " last_message = next(iter(output.values()))['messages'][-1]\n", + " last_message = next(iter(output.values()))[\"messages\"][-1]\n", " last_message.pretty_print()\n", "\n", " if output and \"prompt\" in output:\n", diff --git a/examples/create-react-agent-hitl.ipynb b/examples/create-react-agent-hitl.ipynb index 39392f205..2bc6d306f 100644 --- a/examples/create-react-agent-hitl.ipynb +++ b/examples/create-react-agent-hitl.ipynb @@ -239,7 +239,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.1" + "version": "3.12.2" } }, "nbformat": 4, diff --git a/examples/customer-support/customer-support.ipynb b/examples/customer-support/customer-support.ipynb index cf451188f..07927338c 100644 --- a/examples/customer-support/customer-support.ipynb +++ b/examples/customer-support/customer-support.ipynb @@ -225,7 +225,14 @@ "\n", "Define the (`fetch_user_flight_information`) tool to let the agent see the current user's flight information. Then define tools to search for flights and manage the passenger's bookings stored in the SQL database.\n", "\n", - "We use `ensure_config` to pass in the `passenger_id` in via configurable parameters. The LLM never has to provide these explicitly, they are provided for a given invocation of the graph so that each user cannot access other passengers' booking information." + "We the can [access the RunnableConfig](https://python.langchain.com/v0.2/docs/how_to/tool_configure/#inferring-by-parameter-type) for a given run to check the `passenger_id` of the user accessing this application. The LLM never has to provide these explicitly, they are provided for a given invocation of the graph so that each user cannot access other passengers' booking information.\n", + "\n", + "
\n", + "

Compatibility

\n", + "

\n", + " This tutorial expects `langchain-core>=0.2.16` to use the injected RunnableConfig. Prior to that, you'd use `ensure_config` to collect the config from context.\n", + "

\n", + "
\n" ] }, { @@ -240,18 +247,17 @@ "from typing import Optional\n", "\n", "import pytz\n", - "from langchain_core.runnables import ensure_config\n", + "from langchain_core.runnables import RunnableConfig\n", "\n", "\n", "@tool\n", - "def fetch_user_flight_information() -> list[dict]:\n", + "def fetch_user_flight_information(config: RunnableConfig) -> list[dict]:\n", " \"\"\"Fetch all tickets for the user along with corresponding flight information and seat assignments.\n", "\n", " Returns:\n", " A list of dictionaries where each dictionary contains the ticket details,\n", " associated flight details, and the seat assignments for each ticket belonging to the user.\n", " \"\"\"\n", - " config = ensure_config() # Fetch from the context\n", " configuration = config.get(\"configurable\", {})\n", " passenger_id = configuration.get(\"passenger_id\", None)\n", " if not passenger_id:\n", @@ -328,9 +334,10 @@ "\n", "\n", "@tool\n", - "def update_ticket_to_new_flight(ticket_no: str, new_flight_id: int) -> str:\n", + "def update_ticket_to_new_flight(\n", + " ticket_no: str, new_flight_id: int, *, config: RunnableConfig\n", + ") -> str:\n", " \"\"\"Update the user's ticket to a new valid flight.\"\"\"\n", - " config = ensure_config()\n", " configuration = config.get(\"configurable\", {})\n", " passenger_id = configuration.get(\"passenger_id\", None)\n", " if not passenger_id:\n", @@ -396,9 +403,8 @@ "\n", "\n", "@tool\n", - "def cancel_ticket(ticket_no: str) -> str:\n", + "def cancel_ticket(ticket_no: str, *, config: RunnableConfig) -> str:\n", " \"\"\"Cancel the user's ticket and remove it from the database.\"\"\"\n", - " config = ensure_config()\n", " configuration = config.get(\"configurable\", {})\n", " passenger_id = configuration.get(\"passenger_id\", None)\n", " if not passenger_id:\n", @@ -4407,7 +4413,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.2" + "version": "3.12.2" } }, "nbformat": 4, diff --git a/examples/human_in_the_loop/review-tool-calls.ipynb b/examples/human_in_the_loop/review-tool-calls.ipynb index 5a4aeaf56..84373b90b 100644 --- a/examples/human_in_the_loop/review-tool-calls.ipynb +++ b/examples/human_in_the_loop/review-tool-calls.ipynb @@ -141,15 +141,18 @@ " print(\"----\")\n", " return \"Sunny!\"\n", "\n", - "model = ChatAnthropic(model_name=\"claude-3-5-sonnet-20240620\").bind_tools([weather_search])\n", + "\n", + "model = ChatAnthropic(model_name=\"claude-3-5-sonnet-20240620\").bind_tools(\n", + " [weather_search]\n", + ")\n", + "\n", "\n", "class State(MessagesState):\n", " \"\"\"Simple state.\"\"\"\n", "\n", + "\n", "def call_llm(state):\n", - " return {\n", - " \"messages\": [model.invoke(state['messages'])]\n", - " }\n", + " return {\"messages\": [model.invoke(state[\"messages\"])]}\n", "\n", "\n", "def human_review_node(state):\n", @@ -159,28 +162,30 @@ "def run_tool(state):\n", " new_messages = []\n", " tools = {\"weather_search\": weather_search}\n", - " tool_calls = state['messages'][-1].tool_calls\n", + " tool_calls = state[\"messages\"][-1].tool_calls\n", " for tool_call in tool_calls:\n", - " tool = tools[tool_call['name']]\n", - " result = tool.invoke(tool_call['args'])\n", - " new_messages.append({\n", - " \"role\": \"tool\",\n", - " \"name\": tool_call['name'],\n", - " \"content\": result,\n", - " \"tool_call_id\": tool_call['id']\n", - " })\n", + " tool = tools[tool_call[\"name\"]]\n", + " result = tool.invoke(tool_call[\"args\"])\n", + " new_messages.append(\n", + " {\n", + " \"role\": \"tool\",\n", + " \"name\": tool_call[\"name\"],\n", + " \"content\": result,\n", + " \"tool_call_id\": tool_call[\"id\"],\n", + " }\n", + " )\n", " return {\"messages\": new_messages}\n", "\n", "\n", "def route_after_llm(state) -> Literal[END, \"human_review_node\"]:\n", - " if len(state['messages'][-1].tool_calls) == 0:\n", + " if len(state[\"messages\"][-1].tool_calls) == 0:\n", " return END\n", " else:\n", " return \"human_review_node\"\n", "\n", "\n", "def route_after_human(state) -> Literal[\"run_tool\", \"call_llm\"]:\n", - " if isinstance(state['messages'][-1], AIMessage):\n", + " if isinstance(state[\"messages\"][-1], AIMessage):\n", " return \"run_tool\"\n", " else:\n", " return \"call_llm\"\n", @@ -460,35 +465,35 @@ "print(\"Current State:\")\n", "print(state.values)\n", "print(\"\\nCurrent Tool Call ID:\")\n", - "current_content = state.values['messages'][-1].content\n", - "current_id = state.values['messages'][-1].id\n", - "tool_call_id = state.values['messages'][-1].tool_calls[0]['id']\n", + "current_content = state.values[\"messages\"][-1].content\n", + "current_id = state.values[\"messages\"][-1].id\n", + "tool_call_id = state.values[\"messages\"][-1].tool_calls[0][\"id\"]\n", "print(tool_call_id)\n", "\n", "# We now need to construct a replacement tool call.\n", "# We will change the argument to be `San Francisco, USA`\n", "# Note that we could change any number of arguments or tool names - it just has to be a valid one\n", "new_message = {\n", - " \"role\": \"assistant\", \n", + " \"role\": \"assistant\",\n", " \"content\": current_content,\n", " \"tool_calls\": [\n", " {\n", " \"id\": tool_call_id,\n", " \"name\": \"weather_search\",\n", - " \"args\": {\"city\": \"San Francisco, USA\"}\n", + " \"args\": {\"city\": \"San Francisco, USA\"},\n", " }\n", " ],\n", " # This is important - this needs to be the same as the message you replacing!\n", " # Otherwise, it will show up as a separate message\n", - " \"id\": current_id\n", + " \"id\": current_id,\n", "}\n", "graph.update_state(\n", " # This is the config which represents this thread\n", - " thread, \n", + " thread,\n", " # This is the updated value we want to push\n", - " {\"messages\": [new_message]}, \n", + " {\"messages\": [new_message]},\n", " # We push this update acting as our human_review_node\n", - " as_node=\"human_review_node\"\n", + " as_node=\"human_review_node\",\n", ")\n", "\n", "# Let's now continue executing from here\n", @@ -595,26 +600,26 @@ "print(\"Current State:\")\n", "print(state.values)\n", "print(\"\\nCurrent Tool Call ID:\")\n", - "tool_call_id = state.values['messages'][-1].tool_calls[0]['id']\n", + "tool_call_id = state.values[\"messages\"][-1].tool_calls[0][\"id\"]\n", "print(tool_call_id)\n", "\n", "# We now need to construct a replacement tool call.\n", "# We will change the argument to be `San Francisco, USA`\n", "# Note that we could change any number of arguments or tool names - it just has to be a valid one\n", "new_message = {\n", - " \"role\": \"tool\", \n", + " \"role\": \"tool\",\n", " # This is our natural language feedback\n", " \"content\": \"User requested changes: pass in the country as well\",\n", " \"name\": \"weather_search\",\n", - " \"tool_call_id\": tool_call_id\n", + " \"tool_call_id\": tool_call_id,\n", "}\n", "graph.update_state(\n", " # This is the config which represents this thread\n", - " thread, \n", + " thread,\n", " # This is the updated value we want to push\n", - " {\"messages\": [new_message]}, \n", + " {\"messages\": [new_message]},\n", " # We push this update acting as our human_review_node\n", - " as_node=\"human_review_node\"\n", + " as_node=\"human_review_node\",\n", ")\n", "\n", "# Let's now continue executing from here\n", diff --git a/examples/input_output_schema.ipynb b/examples/input_output_schema.ipynb index 16779ba6e..4837b40a8 100644 --- a/examples/input_output_schema.ipynb +++ b/examples/input_output_schema.ipynb @@ -33,15 +33,19 @@ "from langgraph.graph import StateGraph, START, END\n", "from typing import TypedDict\n", "\n", + "\n", "class InputState(TypedDict):\n", " question: str\n", "\n", + "\n", "class OutputState(TypedDict):\n", " answer: str\n", "\n", + "\n", "def answer_node(state: InputState):\n", " return {\"answer\": \"bye\"}\n", "\n", + "\n", "graph = StateGraph(input=InputState, output=OutputState)\n", "graph.add_node(answer_node)\n", "graph.add_edge(START, \"answer_node\")\n", diff --git a/examples/llm-compiler/LLMCompiler.ipynb b/examples/llm-compiler/LLMCompiler.ipynb index bfe5736d8..2f2aa81ed 100644 --- a/examples/llm-compiler/LLMCompiler.ipynb +++ b/examples/llm-compiler/LLMCompiler.ipynb @@ -526,7 +526,7 @@ " \"tasks\": tasks,\n", " }\n", " )\n", - " return {\"messages\":[scheduled_tasks]}" + " return {\"messages\": [scheduled_tasks]}" ] }, { @@ -653,7 +653,7 @@ " )\n", " ]\n", " else:\n", - " return {\"messages\":response + [AIMessage(content=decision.action.response)]}\n", + " return {\"messages\": response + [AIMessage(content=decision.action.response)]}\n", "\n", "\n", "def select_recent_messages(state) -> dict:\n", @@ -726,9 +726,11 @@ "from langgraph.graph.message import add_messages\n", "from typing import Annotated\n", "\n", + "\n", "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "graph_builder = StateGraph(State)\n", "\n", "# 1. Define vertices\n", @@ -794,7 +796,9 @@ } ], "source": [ - "for step in chain.stream({\"messages\":[HumanMessage(content=\"What's the GDP of New York?\")]}):\n", + "for step in chain.stream(\n", + " {\"messages\": [HumanMessage(content=\"What's the GDP of New York?\")]}\n", + "):\n", " print(step)\n", " print(\"---\")" ] diff --git a/examples/many-tools.ipynb b/examples/many-tools.ipynb index b415f5065..107c9921f 100644 --- a/examples/many-tools.ipynb +++ b/examples/many-tools.ipynb @@ -328,9 +328,9 @@ " \"set more_information_needed False and populate a blank string for the query.\"\n", " )\n", " input_messages = [system] + state[\"messages\"]\n", - " response = llm.bind_tools(\n", - " [QueryForTools], tool_choice=True\n", - " ).invoke(input_messages)\n", + " response = llm.bind_tools([QueryForTools], tool_choice=True).invoke(\n", + " input_messages\n", + " )\n", " query = response.tool_calls[0][\"args\"][\"query\"]\n", " tool_documents = vector_store.similarity_search(query)\n", " if hack_remove_tool_condition:\n", diff --git a/examples/pass-run-time-values-to-tools.ipynb b/examples/pass-run-time-values-to-tools.ipynb index 3c74bd06d..ab6680230 100644 --- a/examples/pass-run-time-values-to-tools.ipynb +++ b/examples/pass-run-time-values-to-tools.ipynb @@ -329,6 +329,7 @@ "\n", "tools = [get_context, cite_context_sources]\n", "\n", + "\n", "# Define the function that calls the model\n", "def call_model(state, config):\n", " messages = state[\"messages\"]\n", diff --git a/examples/pass_private_state.ipynb b/examples/pass_private_state.ipynb index 2e60d805e..d34ecc7e1 100644 --- a/examples/pass_private_state.ipynb +++ b/examples/pass_private_state.ipynb @@ -72,12 +72,12 @@ "# Node to retrieve documents\n", "def retrieve_documents(state: QueryOutputState) -> DocumentOutputState:\n", " # Replace this with real logic\n", - " return {\"docs\": [state['query']] * 2}\n", + " return {\"docs\": [state[\"query\"]] * 2}\n", "\n", "\n", "# Node to generate answer\n", "def generate(state: GenerateInputState) -> OverallState:\n", - " return {\"answer\": \"\\n\\n\".join(state['docs'] + [state['question']])}\n", + " return {\"answer\": \"\\n\\n\".join(state[\"docs\"] + [state[\"question\"]])}\n", "\n", "\n", "graph = StateGraph(OverallState)\n", diff --git a/examples/persistence.ipynb b/examples/persistence.ipynb index 869043cf1..715f7c0b5 100644 --- a/examples/persistence.ipynb +++ b/examples/persistence.ipynb @@ -587,7 +587,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.9" + "version": "3.12.2" } }, "nbformat": 4, diff --git a/examples/persistence_mongodb.ipynb b/examples/persistence_mongodb.ipynb index 99ce73057..54dfdc638 100644 --- a/examples/persistence_mongodb.ipynb +++ b/examples/persistence_mongodb.ipynb @@ -630,7 +630,7 @@ " upsert=True,\n", " )\n", " )\n", - " await self.db[\"checkpoint_writes\"].bulk_write(operations)\n" + " await self.db[\"checkpoint_writes\"].bulk_write(operations)" ] }, { @@ -685,7 +685,9 @@ "metadata": {}, "outputs": [], "source": [ - "with MongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n", + "with MongoDBSaver.from_conn_info(\n", + " host=\"localhost\", port=27017, db_name=\"checkpoints\"\n", + ") as checkpointer:\n", " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", " config = {\"configurable\": {\"thread_id\": \"1\"}}\n", " res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)\n", @@ -796,10 +798,14 @@ "metadata": {}, "outputs": [], "source": [ - "async with AsyncMongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n", + "async with AsyncMongoDBSaver.from_conn_info(\n", + " host=\"localhost\", port=27017, db_name=\"checkpoints\"\n", + ") as checkpointer:\n", " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", " config = {\"configurable\": {\"thread_id\": \"2\"}}\n", - " res = await graph.ainvoke({\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config)\n", + " res = await graph.ainvoke(\n", + " {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n", + " )\n", "\n", " latest_checkpoint = await checkpointer.aget(config)\n", " latest_checkpoint_tuple = await checkpointer.aget_tuple(config)\n", diff --git a/examples/persistence_postgres.ipynb b/examples/persistence_postgres.ipynb index ef8be17cb..5ea838cee 100644 --- a/examples/persistence_postgres.ipynb +++ b/examples/persistence_postgres.ipynb @@ -122,7 +122,7 @@ "metadata": {}, "outputs": [], "source": [ - "DB_URI = \"postgresql://postgres:postgres@localhost:5441/postgres?sslmode=disable\"" + "DB_URI = \"postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable\"" ] }, { @@ -134,10 +134,9 @@ "source": [ "from psycopg.rows import dict_row\n", "\n", - "connection_kwargs ={\n", + "connection_kwargs = {\n", " \"autocommit\": True,\n", " \"prepare_threshold\": 0,\n", - " \"row_factory\": dict_row,\n", "}" ] }, @@ -166,7 +165,7 @@ " # Example configuration\n", " conninfo=DB_URI,\n", " max_size=20,\n", - " kwargs=connection_kwargs\n", + " kwargs=connection_kwargs,\n", ")\n", "\n", "with pool.connection() as conn:\n", @@ -394,7 +393,7 @@ " # Example configuration\n", " conninfo=DB_URI,\n", " max_size=20,\n", - " kwargs=connection_kwargs\n", + " kwargs=connection_kwargs,\n", ") as pool, pool.connection() as conn:\n", " checkpointer = AsyncPostgresSaver(conn)\n", "\n", @@ -551,9 +550,9 @@ ], "metadata": { "kernelspec": { - "display_name": "langgraph-postgres", + "display_name": "langgraph", "language": "python", - "name": "langgraph-postgres" + "name": "langgraph" }, "language_info": { "codemirror_mode": { diff --git a/examples/persistence_redis.ipynb b/examples/persistence_redis.ipynb index 3b2ad1170..34d78a53f 100644 --- a/examples/persistence_redis.ipynb +++ b/examples/persistence_redis.ipynb @@ -530,7 +530,9 @@ "\n", " @classmethod\n", " @asynccontextmanager\n", - " async def from_conn_info(cls, *, host: str, port: int, db: int) -> AsyncIterator[\"AsyncRedisSaver\"]:\n", + " async def from_conn_info(\n", + " cls, *, host: str, port: int, db: int\n", + " ) -> AsyncIterator[\"AsyncRedisSaver\"]:\n", " conn = None\n", " try:\n", " conn = AsyncRedis(host=host, port=port, db=db)\n", @@ -887,10 +889,14 @@ "metadata": {}, "outputs": [], "source": [ - "async with AsyncRedisSaver.from_conn_info(host=\"localhost\", port=6379, db=0) as checkpointer:\n", + "async with AsyncRedisSaver.from_conn_info(\n", + " host=\"localhost\", port=6379, db=0\n", + ") as checkpointer:\n", " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", " config = {\"configurable\": {\"thread_id\": \"2\"}}\n", - " res = await graph.ainvoke({\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config)\n", + " res = await graph.ainvoke(\n", + " {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n", + " )\n", "\n", " latest_checkpoint = await checkpointer.aget(config)\n", " latest_checkpoint_tuple = await checkpointer.aget_tuple(config)\n", diff --git a/examples/reflection/reflection.ipynb b/examples/reflection/reflection.ipynb index 810990e38..ca42160bd 100644 --- a/examples/reflection/reflection.ipynb +++ b/examples/reflection/reflection.ipynb @@ -269,7 +269,7 @@ "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", - " \n", + "\n", "async def generation_node(state: Sequence[BaseMessage]):\n", " return await generate.ainvoke({\"messages\": state})\n", "\n", diff --git a/examples/reflexion/reflexion.ipynb b/examples/reflexion/reflexion.ipynb index 670e6eb5b..8c183c817 100644 --- a/examples/reflexion/reflexion.ipynb +++ b/examples/reflexion/reflexion.ipynb @@ -392,6 +392,7 @@ "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "MAX_ITERATIONS = 5\n", "builder = StateGraph(State)\n", "builder.add_node(\"draft\", first_responder.respond)\n", diff --git a/examples/streaming-content.ipynb b/examples/streaming-content.ipynb index ff182ca21..8c2c7ab36 100644 --- a/examples/streaming-content.ipynb +++ b/examples/streaming-content.ipynb @@ -68,7 +68,9 @@ " # It's completely optional, but useful if you have many functions with similar names\n", " gen = RunnableGenerator(my_generator).with_config(\n", " tags=[\"should_stream\"],\n", - " callbacks=config.get(\"callbacks\", []) # <-- Propagate callbacks (Python <= 3.10)\n", + " callbacks=config.get(\n", + " \"callbacks\", []\n", + " ), # <-- Propagate callbacks (Python <= 3.10)\n", " )\n", " async for message in gen.astream(state):\n", " messages.append(message)\n", diff --git a/examples/tutorials/rag-agent-testing-local.ipynb b/examples/tutorials/rag-agent-testing-local.ipynb index f9e89e56d..3105d3342 100644 --- a/examples/tutorials/rag-agent-testing-local.ipynb +++ b/examples/tutorials/rag-agent-testing-local.ipynb @@ -169,9 +169,7 @@ "from langchain_core.output_parsers import JsonOutputParser\n", "\n", "# JSON\n", - "llm = ChatOllama(model=\"llama3.1\", \n", - " format=\"json\", \n", - " temperature=0)\n", + "llm = ChatOllama(model=\"llama3.1\", format=\"json\", temperature=0)\n", "\n", "\n", "prompt = PromptTemplate(\n", @@ -210,6 +208,7 @@ "from IPython.display import Image, display\n", "from langgraph.graph import START, END, StateGraph\n", "\n", + "\n", "class GraphState(TypedDict):\n", " \"\"\"\n", " Represents the state of our graph.\n", @@ -381,21 +380,22 @@ "metadata": {}, "outputs": [], "source": [ - "import uuid \n", + "import uuid\n", + "\n", "\n", "def predict_custom_agent_answer(example: dict):\n", - " \n", " config = {\"configurable\": {\"thread_id\": str(uuid.uuid4())}}\n", - " \n", + "\n", " state_dict = custom_graph.invoke(\n", " {\"question\": example[\"input\"], \"steps\": []}, config\n", " )\n", - " \n", + "\n", " return {\"response\": state_dict[\"generation\"], \"steps\": state_dict[\"steps\"]}\n", "\n", + "\n", "example = {\"input\": \"What are the types of agent memory?\"}\n", - "#response = predict_custom_agent_answer(example)\n", - "#response" + "# response = predict_custom_agent_answer(example)\n", + "# response" ] }, { @@ -544,6 +544,7 @@ " \"generate_answer\",\n", "]\n", "\n", + "\n", "def check_trajectory_custom(root_run: Run, example: Example) -> dict:\n", " \"\"\"\n", " Check if all expected tools are called in exact order and without any additional tool calls.\n", diff --git a/examples/tutorials/tool-calling-agent-local.ipynb b/examples/tutorials/tool-calling-agent-local.ipynb index 5c60336e8..12c43038f 100644 --- a/examples/tutorials/tool-calling-agent-local.ipynb +++ b/examples/tutorials/tool-calling-agent-local.ipynb @@ -134,6 +134,7 @@ " for d in web_results\n", " ]\n", "\n", + "\n", "# Tool list\n", "tools = [retrieve_documents, web_search]" ] @@ -152,9 +153,11 @@ "from langgraph.graph.message import AnyMessage, add_messages\n", "from typing_extensions import TypedDict\n", "\n", + "\n", "class State(TypedDict):\n", " messages: Annotated[list[AnyMessage], add_messages]\n", "\n", + "\n", "class Assistant:\n", " def __init__(self, runnable: Runnable):\n", " \"\"\"\n", @@ -291,6 +294,7 @@ "source": [ "import uuid\n", "\n", + "\n", "def predict_react_agent_answer(example: dict):\n", " \"\"\"Use this for answer evaluation\"\"\"\n", "\n", diff --git a/libs/checkpoint-postgres/README.md b/libs/checkpoint-postgres/README.md index 24a77673a..24652a2b2 100644 --- a/libs/checkpoint-postgres/README.md +++ b/libs/checkpoint-postgres/README.md @@ -44,7 +44,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint @@ -87,7 +86,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 433bc231d..cb4a79c35 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -66,7 +66,7 @@ class PostgresSaver(BasePostgresSaver): the first time checkpointer is used. """ with self.lock: - with self.conn.cursor(binary=True) as cur: + with self.conn.cursor(binary=True, row_factory=dict_row) as cur: try: version = cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" @@ -127,31 +127,33 @@ class PostgresSaver(BasePostgresSaver): if limit: query += f" LIMIT {limit}" # if we change this to use .stream() we need to make sure to close the cursor - for value in self.conn.execute(query, args, binary=True): - yield CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], + with self._cursor() as cur: + cur.execute(query, args, binary=True) + for value in cur: + yield CheckpointTuple( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["checkpoint_id"], + } + }, + { + **self._load_checkpoint(value["checkpoint"]), + "channel_values": self._load_blobs(value["channel_values"]), + }, + self._load_metadata(value["metadata"]), + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": self._load_blobs(value["channel_values"]), - }, - self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None, - self._load_writes(value["pending_writes"]), - ) + if value["parent_checkpoint_id"] + else None, + self._load_writes(value["pending_writes"]), + ) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. @@ -199,7 +201,7 @@ class PostgresSaver(BasePostgresSaver): where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1" with self._cursor() as cur: - cur = self.conn.execute( + cur.execute( self.SELECT_SQL + where, args, binary=True, @@ -336,7 +338,7 @@ class PostgresSaver(BasePostgresSaver): # in multiple threads/coroutines, but only one cursor can be # used at a time try: - with self.conn.cursor(binary=True) as cur: + with self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur finally: if pipeline: @@ -344,8 +346,10 @@ class PostgresSaver(BasePostgresSaver): elif pipeline: # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock - with self.lock, self.conn.pipeline(), self.conn.cursor(binary=True) as cur: + with self.lock, self.conn.pipeline(), self.conn.cursor( + binary=True, row_factory=dict_row + ) as cur: yield cur else: - with self.lock, self.conn.cursor(binary=True) as cur: + with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 79ae1ddf7..569159d91 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -64,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver): the first time checkpointer is used. """ async with self.lock: - async with self.conn.cursor(binary=True) as cur: + async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: try: results = await cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" @@ -110,33 +110,35 @@ class AsyncPostgresSaver(BasePostgresSaver): if limit: query += f" LIMIT {limit}" # if we change this to use .stream() we need to make sure to close the cursor - async for value in await self.conn.execute(query, args, binary=True): - yield CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], + async with self._cursor() as cur: + await cur.execute(query, args, binary=True) + async for value in cur: + yield CheckpointTuple( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["checkpoint_id"], + } + }, + { + **self._load_checkpoint(value["checkpoint"]), + "channel_values": await asyncio.to_thread( + self._load_blobs, value["channel_values"] + ), + }, + self._load_metadata(value["metadata"]), + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": await asyncio.to_thread( - self._load_blobs, value["channel_values"] - ), - }, - self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None, - await asyncio.to_thread(self._load_writes, value["pending_writes"]), - ) + if value["parent_checkpoint_id"] + else None, + await asyncio.to_thread(self._load_writes, value["pending_writes"]), + ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database asynchronously. @@ -163,7 +165,7 @@ class AsyncPostgresSaver(BasePostgresSaver): where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1" async with self._cursor() as cur: - cur = await self.conn.execute( + await cur.execute( self.SELECT_SQL + where, args, binary=True, @@ -293,7 +295,7 @@ class AsyncPostgresSaver(BasePostgresSaver): # in multiple threads/coroutines, but only one cursor can be # used at a time try: - async with self.conn.cursor(binary=True) as cur: + async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur finally: if pipeline: @@ -302,9 +304,11 @@ class AsyncPostgresSaver(BasePostgresSaver): # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock async with self.lock, self.conn.pipeline(), self.conn.cursor( - binary=True + binary=True, row_factory=dict_row ) as cur: yield cur else: - async with self.lock, self.conn.cursor(binary=True) as cur: + async with self.lock, self.conn.cursor( + binary=True, row_factory=dict_row + ) as cur: yield cur diff --git a/libs/checkpoint-sqlite/README.md b/libs/checkpoint-sqlite/README.md index 86ca9cafc..73fe94333 100644 --- a/libs/checkpoint-sqlite/README.md +++ b/libs/checkpoint-sqlite/README.md @@ -35,7 +35,6 @@ with SqliteSaver.from_conn_string(":memory:") as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint @@ -78,7 +77,6 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint @@ -89,4 +87,4 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: # list checkpoints [c async for c in checkpointer.alist(read_config)] -``` \ No newline at end of file +``` diff --git a/libs/checkpoint/README.md b/libs/checkpoint/README.md index 7f6b26d6e..19c7d3807 100644 --- a/libs/checkpoint/README.md +++ b/libs/checkpoint/README.md @@ -74,7 +74,6 @@ checkpoint = { } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 21d021ad5..a4f76f87b 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -97,9 +97,6 @@ class Checkpoint(TypedDict): pending_sends: List[SendProtocol] """List of packets sent to nodes but not yet processed. Cleared by the next checkpoint.""" - current_tasks: Dict[str, TaskInfo] - """Map from task ID to task info.""" - # TODO remove this def empty_checkpoint() -> Checkpoint: @@ -111,7 +108,6 @@ def empty_checkpoint() -> Checkpoint: channel_versions={}, versions_seen={}, pending_sends=[], - current_tasks={}, ) @@ -124,7 +120,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: channel_versions=checkpoint["channel_versions"].copy(), versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()}, pending_sends=checkpoint.get("pending_sends", []).copy(), - current_tasks=checkpoint.get("current_tasks", {}).copy(), ) @@ -156,7 +151,6 @@ def create_checkpoint( channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], pending_sends=checkpoint.get("pending_sends", []), - current_tasks={}, ) @@ -451,3 +445,4 @@ saving regular writes. Each Checkpointer implementation should use this mapping in put_writes. """ WRITES_IDX_MAP = {ERROR: -1} +# TODO To store scheduled status of tasks, add a special channel here diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 206ec08e1..45f0ed360 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -7,6 +7,7 @@ import re from collections import deque from datetime import date, datetime, time, timedelta, timezone from enum import Enum +from inspect import isclass from ipaddress import ( IPv4Address, IPv4Interface, @@ -111,7 +112,7 @@ class JsonPlusSerializer(SerializerProtocol): obj.__class__, method="fromhex", args=[obj.hex()] ) elif isinstance(obj, BaseException): - return self._encode_constructor_args(obj.__class__, args=obj.args) + return repr(obj) else: raise TypeError( f"Object of type {obj.__class__.__name__} is not JSON serializable" @@ -135,6 +136,8 @@ class JsonPlusSerializer(SerializerProtocol): method = getattr(cls, value["method"]) else: method = cls + if isclass(method) and issubclass(method, BaseException): + return None if value["args"] and value["kwargs"]: return method(*value["args"], **value["kwargs"]) elif value["args"]: diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index b17a77238..7ceea436e 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint" -version = "1.0.3" +version = "1.0.4" description = "Library with base interfaces for LangGraph checkpoint savers." authors = [] license = "MIT" diff --git a/libs/cli/pyproject.toml b/libs/cli/pyproject.toml index 508571f06..02c373e95 100644 --- a/libs/cli/pyproject.toml +++ b/libs/cli/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-cli" -version = "0.1.50" +version = "0.1.51" description = "CLI for interacting with LangGraph API" authors = [] license = "MIT" diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 603de9395..da3862616 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -108,5 +108,5 @@ class Send: @dataclass class Interrupt: - when: Literal["before", "during", "after"] - value: Any = None + value: Any + when: Literal["during"] = "during" diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 27a7689fa..25ba93a72 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -40,7 +40,7 @@ class NodeInterrupt(GraphInterrupt): """Raised by a node to interrupt execution.""" def __init__(self, value: Any) -> None: - super().__init__([Interrupt("during", value)]) + super().__init__([Interrupt(value)]) class EmptyInputError(Exception): diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 5da78a26c..3e4243c61 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -120,14 +120,14 @@ class AsyncBackgroundExecutor(AsyncContextManager): def done(self, task: asyncio.Task) -> None: try: - task.result() - except GraphInterrupt: - # This exception is an interruption signal, not an error - # so we don't want to re-raise it on exit - self.tasks.pop(task) - except BaseException: - pass - else: + if exc := task.exception(): + # This exception is an interruption signal, not an error + # so we don't want to re-raise it on exit + if isinstance(exc, GraphInterrupt): + self.tasks.pop(task) + else: + self.tasks.pop(task) + except asyncio.CancelledError: self.tasks.pop(task) async def __aenter__(self) -> Submit: @@ -146,12 +146,13 @@ class AsyncBackgroundExecutor(AsyncContextManager): # wait for all tasks to finish if self.tasks: await asyncio.wait(self.tasks) - # re-raise the first exception that occurred in a task + # if there's already an exception being raised, don't raise another one if exc_type is None: - # if there's already an exception being raised, don't raise another one + # re-raise the first exception that occurred in a task for task in self.tasks: try: - task.result() + if exc := task.exception(): + raise exc except asyncio.CancelledError: pass diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index a501232a9..b4e6d4118 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -155,10 +155,6 @@ class PregelLoop: self.stream_keys = stream_keys self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) - def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None: - """Mark tasks as scheduled, to be used by queue-based executors.""" - raise NotImplementedError - def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: """Put writes for a task, to be read by the next tick.""" if not writes: diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 1c3b20e9b..bc0b67fb4 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.11" +version = "0.2.12" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index c92bc4688..4835ec611 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -33,6 +33,142 @@ ''' # --- +# name: test_branch_then[memory] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[memory].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_branch_then[postgres] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[postgres].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_branch_then[postgres_pipe] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[postgres_pipe].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_branch_then[sqlite] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[sqlite].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_conditional_entrypoint_graph '{"title": "LangGraphInput"}' # --- @@ -601,6 +737,1386 @@ ''' # --- +# name: test_conditional_graph[memory] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[memory].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[memory].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[memory].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[memory].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[memory].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[memory].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[postgres].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres_pipe] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pipe].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pipe].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres_pipe].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pipe].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pipe].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[postgres_pipe].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[sqlite] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[sqlite].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[sqlite].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[sqlite].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[sqlite].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[sqlite].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[sqlite].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_conditional_state_graph '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- @@ -682,6 +2198,330 @@ ''' # --- +# name: test_conditional_state_graph[memory] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[memory].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[memory].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[memory].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_state_graph[postgres] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[postgres].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_state_graph[postgres_pipe] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres_pipe].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres_pipe].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[postgres_pipe].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_state_graph[sqlite] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[sqlite].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[sqlite].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[sqlite].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- # name: test_dynamic_interrupt ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% @@ -713,6 +2553,58 @@ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class ''' graph TD; @@ -796,6 +2688,286 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2 ''' graph TD; @@ -866,6 +3038,286 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_via_branch ''' graph TD; @@ -879,6 +3331,58 @@ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- # name: test_message_graph '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' # --- @@ -960,6 +3464,330 @@ ''' # --- +# name: test_message_graph[memory] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[memory].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[memory].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[memory].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  end  .-> __end__; + + ''' +# --- +# name: test_message_graph[postgres] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[postgres].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  end  .-> __end__; + + ''' +# --- +# name: test_message_graph[postgres_pipe] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres_pipe].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres_pipe].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[postgres_pipe].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  end  .-> __end__; + + ''' +# --- +# name: test_message_graph[sqlite] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[sqlite].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[sqlite].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[sqlite].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  end  .-> __end__; + + ''' +# --- # name: test_nested_graph ''' graph TD; @@ -1347,6 +4175,78 @@ ''' # --- +# name: test_start_branch_then[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[postgres] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[postgres_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[sqlite] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_state_graph_w_config '{"title": "LangGraphConfig", "type": "object", "properties": {"configurable": {"$ref": "#/definitions/Configurable"}}, "definitions": {"Configurable": {"title": "Configurable", "type": "object", "properties": {"tools": {"title": "Tools", "type": "array", "items": {"type": "string"}}}}}}' # --- diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py index a98962cdc..836cf9371 100644 --- a/libs/langgraph/tests/any_str.py +++ b/libs/langgraph/tests/any_str.py @@ -23,24 +23,6 @@ class AnyVersion: return hash(str(self)) -class ExceptionLike: - def __init__(self, exc: Exception) -> None: - self.exc = exc - - def __eq__(self, value: object) -> bool: - return ( - isinstance(value, Exception) - and self.exc.__class__ == value.__class__ - and str(self.exc) == str(value) - ) - - def __hash__(self) -> int: - return hash((self.exc.__class__, str(self.exc))) - - def __repr__(self) -> str: - return str(self.exc) - - class UnsortedSequence: def __init__(self, *values: Any) -> None: self.seq = values diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index b02cbb65b..624a711c7 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -24,8 +24,6 @@ class NoopSerializer(SerializerProtocol): class MemorySaverAssertImmutable(MemorySaver): - serde = NoopSerializer() - storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]] def __init__( @@ -74,15 +72,6 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver): should produce a side effect that can be asserted. """ - serde = NoopSerializer() - - def __init__( - self, - *, - serde: Optional[SerializerProtocol] = None, - ) -> None: - super().__init__(serde=serde) - def put( self, config: RunnableConfig, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index df6ecf8ec..6df2403a4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -49,9 +49,6 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.checkpoint.serde.base import SerializerProtocol -from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.constants import ERROR, Interrupt, Send from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.graph import END, Graph @@ -72,13 +69,11 @@ from langgraph.pregel import ( from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore -from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence +from tests.any_str import AnyStr, AnyVersion, UnsortedSequence from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, - MemorySaverAssertImmutable, MemorySaverNoPending, - NoopSerializer, ) from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -1253,7 +1248,16 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] -def test_invoke_checkpoint(mocker: MockerFixture) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_invoke_checkpoint_two( + mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) errored_once = False @@ -1276,8 +1280,6 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = MemorySaverAssertImmutable() - app = Pregel( nodes={"one": one}, channels={ @@ -1287,37 +1289,37 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: }, input_channels="input", output_channels="output", - checkpointer=memory, + checkpointer=checkpointer, retry_policy=RetryPolicy(), ) # total starts out as 0, so output is 0+2=2 assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) + checkpoint = checkpointer.get({"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 app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 assert errored_once, "errored and retried" - checkpoint_tup = memory.get_tuple({"configurable": {"thread_id": "1"}}) + checkpoint_tup = checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) assert checkpoint_tup is not None assert checkpoint_tup.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): app.invoke(4, {"configurable": {"thread_id": "1"}}) # checkpoint is not updated, error is recorded - checkpoint_tup = memory.get_tuple({"configurable": {"thread_id": "1"}}) + checkpoint_tup = checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) assert checkpoint_tup is not None assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 assert checkpoint_tup.pending_writes == [ - (AnyStr(), ERROR, ExceptionLike(ValueError("Input is too large"))) + (AnyStr(), ERROR, "ValueError('Input is too large')") ] # on a new thread, total starts out as 0, so output is 0+5=5 assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) + checkpoint = checkpointer.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 - checkpoint = memory.get({"configurable": {"thread_id": "2"}}) + checkpoint = checkpointer.get({"configurable": {"thread_id": "2"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 5 @@ -1377,7 +1379,7 @@ def test_pending_writes_resume( assert state.next == ("one", "two") assert state.tasks == ( PregelTask(AnyStr(), "one"), - PregelTask(AnyStr(), "two", ExceptionLike(ConnectionError("I'm not good"))), + PregelTask(AnyStr(), "two", 'ConnectionError("I\'m not good")'), ) assert state.metadata == {"source": "loop", "step": 0, "writes": None} # should contain pending write of "one" @@ -1387,7 +1389,7 @@ def test_pending_writes_resume( expected_writes = [ (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), ERROR, ExceptionLike(ConnectionError("I'm not good"))), + (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), ] assert len(checkpoint.pending_writes) == 3 assert all(w in expected_writes for w in checkpoint.pending_writes) @@ -1433,7 +1435,6 @@ def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "one": { @@ -1492,7 +1493,6 @@ def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "__input__": {}, @@ -1523,7 +1523,7 @@ def test_pending_writes_resume( pending_writes=UnsortedSequence( (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), "__error__", ExceptionLike(ConnectionError("I'm not good"))), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), (AnyStr(), "two", "two"), (AnyStr(), "value", 3), ), @@ -1540,7 +1540,6 @@ def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": {"__input__": {}}, "channel_versions": { @@ -1607,7 +1606,14 @@ async def test_checkpointer_null_pending_writes() -> None: ] * 4 -def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_invoke_checkpoint_three( + mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -1622,121 +1628,119 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: | raise_if_above_10 ) - with SqliteSaver.from_conn_string(":memory:") as memory: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=memory, - ) + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + ) - thread_1 = {"configurable": {"thread_id": "1"}} - # total starts out as 0, so output is 0+2=2 - assert app.invoke(2, thread_1, debug=1) == 2 - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 2 - assert state.next == () - assert ( - state.config["configurable"]["checkpoint_id"] == memory.get(thread_1)["id"] - ) - # total is now 2, so output is 2+3=5 - assert app.invoke(3, thread_1) == 5 - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert ( - state.config["configurable"]["checkpoint_id"] == memory.get(thread_1)["id"] - ) - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - app.invoke(4, thread_1) - # checkpoint is updated with new input - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert state.next == ("one",) - """we checkpoint inputs and it failed on "one", so the next node is one""" - # we can recover from error by sending new inputs - assert app.invoke(2, thread_1) == 9 - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 16, "total is now 7+9=16" - assert state.next == () + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert app.invoke(2, thread_1, debug=1) == 2 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert state.next == () + assert ( + state.config["configurable"]["checkpoint_id"] + == checkpointer.get(thread_1)["id"] + ) + # total is now 2, so output is 2+3=5 + assert app.invoke(3, thread_1) == 5 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == checkpointer.get(thread_1)["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + app.invoke(4, thread_1) + # checkpoint is updated with new input + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert app.invoke(2, thread_1) == 9 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () - thread_2 = {"configurable": {"thread_id": "2"}} - # on a new thread, total starts out as 0, so output is 0+5=5 - assert app.invoke(5, thread_2, debug=True) == 5 - state = app.get_state({"configurable": {"thread_id": "1"}}) - assert state is not None - assert state.values.get("total") == 16 - assert state.next == (), "checkpoint of other thread not touched" - state = app.get_state(thread_2) - assert state is not None - assert state.values.get("total") == 5 - assert state.next == () + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert app.invoke(5, thread_2, debug=True) == 5 + state = app.get_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == (), "checkpoint of other thread not touched" + state = app.get_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () - assert len(list(app.get_state_history(thread_1, limit=1))) == 1 - # list all checkpoints for thread 1 - thread_1_history = [c for c in app.get_state_history(thread_1)] - # there are 7 checkpoints - assert len(thread_1_history) == 7 - assert Counter(c.metadata["source"] for c in thread_1_history) == { - "input": 4, - "loop": 3, - } - # sorted descending - assert ( - thread_1_history[0].config["configurable"]["checkpoint_id"] - > thread_1_history[1].config["configurable"]["checkpoint_id"] - ) - # cursor pagination - cursored = list( - app.get_state_history(thread_1, limit=1, before=thread_1_history[0].config) - ) - assert len(cursored) == 1 - assert cursored[0].config == thread_1_history[1].config - # the last checkpoint - assert thread_1_history[0].values["total"] == 16 - # the first "loop" checkpoint - assert thread_1_history[-2].values["total"] == 2 - # can get each checkpoint using aget with config - assert ( - memory.get(thread_1_history[0].config)["id"] - == thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - assert ( - memory.get(thread_1_history[1].config)["id"] - == thread_1_history[1].config["configurable"]["checkpoint_id"] - ) + assert len(list(app.get_state_history(thread_1, limit=1))) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c for c in app.get_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = list( + app.get_state_history(thread_1, limit=1, before=thread_1_history[0].config) + ) + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert ( + checkpointer.get(thread_1_history[0].config)["id"] + == thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + assert ( + checkpointer.get(thread_1_history[1].config)["id"] + == thread_1_history[1].config["configurable"]["checkpoint_id"] + ) - thread_1_next_config = app.update_state(thread_1_history[1].config, 10) - # update creates a new checkpoint - assert ( - thread_1_next_config["configurable"]["checkpoint_id"] - > thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - # update makes new checkpoint child of the previous one - assert ( - app.get_state(thread_1_next_config).parent_config - == thread_1_history[1].config - ) - # 1 more checkpoint in history - assert len(list(app.get_state_history(thread_1))) == 8 - assert Counter( - c.metadata["source"] for c in app.get_state_history(thread_1) - ) == { - "update": 1, - "input": 4, - "loop": 3, - } - # the latest checkpoint is the updated one - assert app.get_state(thread_1) == app.get_state(thread_1_next_config) + thread_1_next_config = app.update_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # update makes new checkpoint child of the previous one + assert ( + app.get_state(thread_1_next_config).parent_config == thread_1_history[1].config + ) + # 1 more checkpoint in history + assert len(list(app.get_state_history(thread_1))) == 8 + assert Counter(c.metadata["source"] for c in app.get_state_history(thread_1)) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert app.get_state(thread_1) == app.get_state(thread_1_next_config) def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -1935,7 +1939,13 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup.call_count == 1, "Expected cleanup to be called once" -def test_conditional_graph(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_conditional_graph( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: from copy import deepcopy from langchain_core.agents import AgentAction, AgentFinish @@ -1944,6 +1954,10 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: from langchain_core.runnables import RunnablePassthrough from langchain_core.tools import tool + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + # Assemble the tools @tool() def search_api(query: str) -> str: @@ -1981,7 +1995,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ) if data.get("intermediate_steps") is None: data["intermediate_steps"] = [] - data["intermediate_steps"].append((agent_action, observation)) + data["intermediate_steps"].append([agent_action, observation]) return data # Define decision-making logic @@ -2017,22 +2031,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: 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" @@ -2053,14 +2067,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2068,14 +2082,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", @@ -2088,22 +2102,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", - ), + ], ], } }, @@ -2111,22 +2125,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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" @@ -2138,7 +2152,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2246,14 +2260,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -2261,14 +2275,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentAction( tool="search_api", @@ -2284,14 +2298,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2305,14 +2319,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2331,14 +2345,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2353,7 +2367,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -2453,14 +2467,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -2468,14 +2482,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentAction( tool="search_api", @@ -2491,14 +2505,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2512,14 +2526,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2538,14 +2552,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2560,10 +2574,10 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) - config = {"configurable": {"thread_id": "2"}} + config = {"configurable": {"thread_id": "3"}} llm.i = 0 # reset the llm assert [ @@ -2616,14 +2630,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2631,14 +2645,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", @@ -2654,22 +2668,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", - ), + ], ], } }, @@ -2677,22 +2691,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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" @@ -2788,14 +2802,24 @@ def test_conditional_entrypoint_to_multiple_state_graph( } +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_conditional_state_graph( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, + mocker: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> 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 + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) setup = mocker.Mock() teardown = mocker.Mock() @@ -2878,7 +2902,7 @@ def test_conditional_state_graph( observation = {t.name: t for t in tools}[agent_action.tool].invoke( agent_action.tool_input ) - return {"intermediate_steps": [(agent_action, observation)]} + return {"intermediate_steps": [[agent_action, observation]]} # Define decision-making logic def should_continue(data: AgentState) -> str: @@ -2915,22 +2939,22 @@ def test_conditional_state_graph( 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" @@ -2951,14 +2975,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2974,14 +2998,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="another", log="tool:search_api:another", ), "result for another", - ), + ], ], } }, @@ -2997,7 +3021,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -3091,14 +3115,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3131,14 +3155,14 @@ def test_conditional_state_graph( log="finish:a really nice answer", ), "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], }, tasks=(), @@ -3163,7 +3187,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], debug=True, ) @@ -3253,14 +3277,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3292,14 +3316,14 @@ def test_conditional_state_graph( log="finish:a really nice answer", ), "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], }, tasks=(), @@ -3323,7 +3347,7 @@ def test_conditional_state_graph( # test w interrupt before all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before="*", debug=True, ) @@ -3387,14 +3411,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3406,14 +3430,14 @@ def test_conditional_state_graph( tool="search_api", tool_input="query", log="tool:search_api:query" ), "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], }, tasks=(PregelTask(AnyStr(), "agent"),), @@ -3426,14 +3450,14 @@ def test_conditional_state_graph( "writes": { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3455,7 +3479,7 @@ def test_conditional_state_graph( # test w interrupt after all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after="*", ) config = {"configurable": {"thread_id": "4"}} @@ -3504,14 +3528,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3523,14 +3547,14 @@ def test_conditional_state_graph( tool="search_api", tool_input="query", log="tool:search_api:query" ), "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], }, tasks=(PregelTask(AnyStr(), "agent"),), @@ -3543,14 +3567,14 @@ def test_conditional_state_graph( "writes": { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -4018,8 +4042,13 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize("serde", [NoopSerializer(), JsonPlusSerializer()]) -def test_state_graph_packets(serde: SerializerProtocol) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_state_graph_packets( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -4032,6 +4061,10 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ) from langchain_core.tools import tool + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] @@ -4252,7 +4285,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(serde=serde), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -4529,9 +4562,15 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_message_graph( snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from copy import deepcopy @@ -4548,6 +4587,10 @@ def test_message_graph( from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.tools import tool + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @@ -4751,7 +4794,7 @@ def test_message_graph( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -4981,7 +5024,7 @@ def test_message_graph( ) app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -5244,15 +5287,20 @@ def test_message_graph( metadata={ "source": "update", "step": 6, - "writes": {"tools": ("ai", "an extra message")}, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_root_graph( - snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from copy import deepcopy @@ -5269,6 +5317,10 @@ def test_root_graph( from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.tools import tool + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @@ -5470,7 +5522,7 @@ def test_root_graph( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -5700,7 +5752,7 @@ def test_root_graph( ) app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -5963,7 +6015,7 @@ def test_root_graph( metadata={ "source": "update", "step": 6, - "writes": {"tools": ("ai", "an extra message")}, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5999,7 +6051,7 @@ def test_root_graph( }, ) new_workflow.add_edge("tools", "agent") - new_app = new_workflow.compile(checkpointer=app_w_interrupt.checkpointer) + new_app = new_workflow.compile(checkpointer=checkpointer) model.i = 0 # reset the llm # previous state is converted to new schema @@ -6035,7 +6087,7 @@ def test_root_graph( metadata={ "source": "update", "step": 6, - "writes": {"tools": ("ai", "an extra message")}, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6318,7 +6370,15 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] -def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_dynamic_interrupt( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -6356,49 +6416,56 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: "market": "US", } - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile(checkpointer=saver) + tool_two = tool_two_graph.compile(checkpointer=checkpointer) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { - "my_key": "value ⛰️", - "market": "DE", - } - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "source": "loop", - "step": 0, - "writes": None, - }, - { - "source": "input", - "step": -1, - "writes": {"my_key": "value ⛰️", "market": "DE"}, - }, - ] - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - interrupts=(Interrupt("during", "Just because..."),), - ), + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + interrupts=(Interrupt("Just because..."),), ), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) -def test_start_branch_then(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_start_branch_then( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -6443,148 +6510,155 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: "market": "US", } - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - store=MemoryStore(), - checkpointer=saver, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) + tool_two = tool_two_graph.compile( + store=MemoryStore(), + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { - "my_key": "value ⛰️", - "market": "DE", - } - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "source": "loop", - "step": 0, - "writes": None, - }, - { - "source": "input", - "step": -1, - "writes": {"my_key": "value ⛰️", "market": "DE"}, - }, - ] - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread1, debug=1) == { - "my_key": "value ⛰️ slow", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ slow", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value ⛰️ slow", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ slow", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) + thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) - # update state - tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) + thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == { + "my_key": "value", + "market": "US", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + # update state + tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "valuekey", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread3, debug=1) == { + "my_key": "valuekey fast", + "market": "US", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) -def test_branch_then(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_branch_then( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -6614,504 +6688,509 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "market": "US", } - with SqliteSaver.from_conn_string(":memory:") as saver: - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=saver) - thread10 = {"configurable": {"thread_id": "10"}} - assert [ - *tool_two.stream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - "next": [], - "tasks": [], - }, - }, - ] - - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + assert [ + *tool_two.stream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 0, + "writes": None, + }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + "next": ["tool_two_slow"], + "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "name": "tool_two_slow", + "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + "next": ["finish"], + "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition::then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "name": "finish", + "result": [("my_key", " finished")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + "next": [], + "tasks": [], + }, + }, + ] - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_before=["finish"] - ) + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) - thread1 = {"configurable": {"thread_id": "1"}} + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_before=["finish"] + ) + + thread1 = {"configurable": {"thread_id": "11"}} + + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared slow", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={ "my_key": "value prepared slow", "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={ - "my_key": "value prepared slow", - "market": "DE", - }, - tasks=(PregelTask(AnyStr(), "finish"),), - next=("finish",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + }, + tasks=(PregelTask(AnyStr(), "finish"),), + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) - # update state - tool_two.update_state(thread1, {"my_key": "er"}) - assert tool_two.get_state(thread1) == StateSnapshot( - values={ - "my_key": "value prepared slower", - "market": "DE", - }, - tasks=(PregelTask(AnyStr(), "finish"),), - next=("finish",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "update", - "step": 3, - "writes": {"tool_two_slow": {"my_key": "er"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_after=["prepare"] - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", + # update state + tool_two.update_state(thread1, {"my_key": "er"}) + assert tool_two.get_state(thread1) == StateSnapshot( + values={ + "my_key": "value prepared slower", "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + }, + tasks=(PregelTask(AnyStr(), "finish"),), + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "update", + "step": 3, + "writes": {"tool_two_slow": {"my_key": "er"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_after=["prepare"] + ) - thread3 = {"configurable": {"thread_id": "3"}} - # update an empty thread before first run - uconfig = tool_two.update_state(thread3, {"my_key": "key", "market": "DE"}) - # check current state - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare"),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - }, - parent_config=None, - ) - # run from this point - assert tool_two.invoke(None, thread3) == { - "my_key": "key prepared", - "market": "DE", - } - # get state after first node - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "key prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=uconfig, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread3, debug=1) == { - "my_key": "key prepared slow finished", - "market": "DE", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "key prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "21"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "22"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + + thread3 = {"configurable": {"thread_id": "23"}} + # update an empty thread before first run + uconfig = tool_two.update_state(thread3, {"my_key": "key", "market": "DE"}) + # check current state + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare"),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + }, + parent_config=None, + ) + # run from this point + assert tool_two.invoke(None, thread3) == { + "my_key": "key prepared", + "market": "DE", + } + # get state after first node + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=uconfig, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread3, debug=1) == { + "my_key": "key prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) -def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_in_one_fan_out_state_graph_waiting_edge( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -7176,7 +7255,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7195,10 +7274,10 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["qa"], ) - config = {"configurable": {"thread_id": "1"}} + config = {"configurable": {"thread_id": "2"}} assert [ c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) @@ -7232,9 +7311,17 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> ] +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_in_one_fan_out_state_graph_waiting_edge_via_branch( - snapshot: SnapshotAssertion, + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -7302,7 +7389,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7321,11 +7408,19 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, + mocker: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() teardown = mocker.Mock() @@ -7450,7 +7545,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7485,11 +7580,19 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( } +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, + mocker: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from pydantic import BaseModel, ConfigDict, ValidationError + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() teardown = mocker.Mock() @@ -7612,7 +7715,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7647,7 +7750,17 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( } -def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -7716,7 +7829,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0d6f21bf7..6938071de 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -46,7 +46,6 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.constants import ERROR, Interrupt, Send from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.graph import END, Graph, StateGraph @@ -67,7 +66,7 @@ from langgraph.pregel import ( from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore -from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence +from tests.any_str import AnyStr, AnyVersion, UnsortedSequence from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -290,7 +289,7 @@ async def test_dynamic_interrupt( PregelTask( AnyStr(), "tool_two", - interrupts=(Interrupt("during", "Just because..."),), + interrupts=(Interrupt("Just because..."),), ), ), config=tup.config, @@ -1599,7 +1598,7 @@ async def test_pending_writes_resume( assert state.next == ("one", "two") assert state.tasks == ( PregelTask(AnyStr(), "one"), - PregelTask(AnyStr(), "two", ExceptionLike(ValueError("I'm not good"))), + PregelTask(AnyStr(), "two", 'ValueError("I\'m not good")'), ) assert state.metadata == {"source": "loop", "step": 0, "writes": None} # should contain pending write of "one" @@ -1609,7 +1608,7 @@ async def test_pending_writes_resume( expected_writes = [ (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), ERROR, ExceptionLike(ValueError("I'm not good"))), + (AnyStr(), ERROR, 'ValueError("I\'m not good")'), ] assert len(checkpoint.pending_writes) == 3 assert all(w in expected_writes for w in checkpoint.pending_writes) @@ -1659,7 +1658,6 @@ async def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "one": { @@ -1718,7 +1716,6 @@ async def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "__input__": {}, @@ -1749,7 +1746,7 @@ async def test_pending_writes_resume( pending_writes=UnsortedSequence( (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), "__error__", ExceptionLike(ValueError("I'm not good"))), + (AnyStr(), "__error__", 'ValueError("I\'m not good")'), (AnyStr(), "two", "two"), (AnyStr(), "value", 3), ), @@ -1766,7 +1763,6 @@ async def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": {"__input__": {}}, "channel_versions": { @@ -1811,7 +1807,14 @@ async def test_cond_edge_after_send() -> None: assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"] -async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_invoke_checkpoint_three( + mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -1826,121 +1829,118 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: | raise_if_above_10 ) - async with AsyncSqliteSaver.from_conn_string(":memory:") as memory: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=memory, - debug=True, - ) + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + debug=True, + ) - thread_1 = {"configurable": {"thread_id": "1"}} - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, thread_1) == 2 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 2 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await memory.aget(thread_1))["id"] - ) - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, thread_1) == 5 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await memory.aget(thread_1))["id"] - ) - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, thread_1) - # checkpoint is not updated - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert state.next == ("one",) - """we checkpoint inputs and it failed on "one", so the next node is one""" - # we can recover from error by sending new inputs - assert await app.ainvoke(2, thread_1) == 9 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 16, "total is now 7+9=16" - assert state.next == () + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, thread_1) == 2 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await checkpointer.aget(thread_1))["id"] + ) + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, thread_1) == 5 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await checkpointer.aget(thread_1))["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, thread_1) + # checkpoint is not updated + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert await app.ainvoke(2, thread_1) == 9 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () - thread_2 = {"configurable": {"thread_id": "2"}} - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, thread_2) == 5 - state = await app.aget_state({"configurable": {"thread_id": "1"}}) - assert state is not None - assert state.values.get("total") == 16 - assert state.next == () - state = await app.aget_state(thread_2) - assert state is not None - assert state.values.get("total") == 5 - assert state.next == () + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, thread_2) == 5 + state = await app.aget_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == () + state = await app.aget_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () - assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 - # list all checkpoints for thread 1 - thread_1_history = [c async for c in app.aget_state_history(thread_1)] - # there are 7 checkpoints - assert len(thread_1_history) == 7 - assert Counter(c.metadata["source"] for c in thread_1_history) == { - "input": 4, - "loop": 3, - } - # sorted descending - assert ( - thread_1_history[0].config["configurable"]["checkpoint_id"] - > thread_1_history[1].config["configurable"]["checkpoint_id"] + assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c async for c in app.aget_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = [ + c + async for c in app.aget_state_history( + thread_1, limit=1, before=thread_1_history[0].config ) - # cursor pagination - cursored = [ - c - async for c in app.aget_state_history( - thread_1, limit=1, before=thread_1_history[0].config - ) - ] - assert len(cursored) == 1 - assert cursored[0].config == thread_1_history[1].config - # the last checkpoint - assert thread_1_history[0].values["total"] == 16 - # the first "loop" checkpoint - assert thread_1_history[-2].values["total"] == 2 - # can get each checkpoint using aget with config - assert (await memory.aget(thread_1_history[0].config))[ - "id" - ] == thread_1_history[0].config["configurable"]["checkpoint_id"] - assert (await memory.aget(thread_1_history[1].config))[ - "id" - ] == thread_1_history[1].config["configurable"]["checkpoint_id"] + ] + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert (await checkpointer.aget(thread_1_history[0].config))[ + "id" + ] == thread_1_history[0].config["configurable"]["checkpoint_id"] + assert (await checkpointer.aget(thread_1_history[1].config))[ + "id" + ] == thread_1_history[1].config["configurable"]["checkpoint_id"] - thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) - # update creates a new checkpoint - assert ( - thread_1_next_config["configurable"]["checkpoint_id"] - > thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - # 1 more checkpoint in history - assert len([c async for c in app.aget_state_history(thread_1)]) == 8 - assert Counter( - [c.metadata["source"] async for c in app.aget_state_history(thread_1)] - ) == { - "update": 1, - "input": 4, - "loop": 3, - } - # the latest checkpoint is the updated one - assert await app.aget_state(thread_1) == await app.aget_state( - thread_1_next_config - ) + thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # 1 more checkpoint in history + assert len([c async for c in app.aget_state_history(thread_1)]) == 8 + assert Counter( + [c.metadata["source"] async for c in app.aget_state_history(thread_1)] + ) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert await app.aget_state(thread_1) == await app.aget_state(thread_1_next_config) async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -2152,7 +2152,13 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup_async.call_count == 1, "Expected cleanup to be called once" -async def test_conditional_graph() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_conditional_graph( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: from copy import deepcopy from langchain_core.agents import AgentAction, AgentFinish @@ -2161,6 +2167,8 @@ async def test_conditional_graph() -> None: from langchain_core.runnables import RunnablePassthrough from langchain_core.tools import tool + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + # Assemble the tools @tool() def search_api(query: str) -> str: @@ -2198,7 +2206,7 @@ async def test_conditional_graph() -> None: ) if data.get("intermediate_steps") is None: data["intermediate_steps"] = [] - data["intermediate_steps"].append((agent_action, observation)) + data["intermediate_steps"].append([agent_action, observation]) return data # Define decision-making logic @@ -2228,22 +2236,22 @@ async def test_conditional_graph() -> None: 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" @@ -2266,14 +2274,14 @@ async def test_conditional_graph() -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2281,14 +2289,14 @@ async def test_conditional_graph() -> None: "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", @@ -2301,22 +2309,22 @@ async def test_conditional_graph() -> None: "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", - ), + ], ], } }, @@ -2324,22 +2332,22 @@ async def test_conditional_graph() -> None: "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" @@ -2376,14 +2384,14 @@ async def test_conditional_graph() -> None: { "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", @@ -2394,22 +2402,22 @@ async def test_conditional_graph() -> None: { "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" @@ -2420,7 +2428,7 @@ async def test_conditional_graph() -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2530,14 +2538,14 @@ async def test_conditional_graph() -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -2545,14 +2553,14 @@ async def test_conditional_graph() -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentAction( tool="search_api", @@ -2568,14 +2576,14 @@ async def test_conditional_graph() -> None: { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2589,14 +2597,14 @@ async def test_conditional_graph() -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2617,14 +2625,14 @@ async def test_conditional_graph() -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2641,7 +2649,7 @@ async def test_conditional_graph() -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -2752,14 +2760,14 @@ async def test_conditional_graph() -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -2767,14 +2775,14 @@ async def test_conditional_graph() -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentAction( tool="search_api", @@ -2790,14 +2798,14 @@ async def test_conditional_graph() -> None: { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2811,14 +2819,14 @@ async def test_conditional_graph() -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2839,14 +2847,14 @@ async def test_conditional_graph() -> None: "agent": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, @@ -2863,10 +2871,10 @@ async def test_conditional_graph() -> None: # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) - config = {"configurable": {"thread_id": "2"}} + config = {"configurable": {"thread_id": "3"}} llm.i = 0 # reset the llm assert [ @@ -2926,14 +2934,14 @@ async def test_conditional_graph() -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2941,14 +2949,14 @@ async def test_conditional_graph() -> None: "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", @@ -2964,22 +2972,22 @@ async def test_conditional_graph() -> None: "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", - ), + ], ], } }, @@ -2987,22 +2995,22 @@ async def test_conditional_graph() -> None: "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" @@ -3105,7 +3113,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: observation = {t.name: t for t in tools}[agent_action.tool].invoke( agent_action.tool_input ) - return {"intermediate_steps": [(agent_action, observation)]} + return {"intermediate_steps": [[agent_action, observation]]} # Define decision-making logic def should_continue(data: AgentState) -> str: @@ -3137,22 +3145,22 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: 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" @@ -3173,14 +3181,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3196,14 +3204,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="another", log="tool:search_api:another", ), "result for another", - ), + ], ], } }, @@ -3358,14 +3366,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3398,14 +3406,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: log="finish:a really nice answer", ), "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], }, tasks=(), @@ -3534,14 +3542,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3573,14 +3581,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: log="finish:a really nice answer", ), "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], }, tasks=(), @@ -5013,7 +5021,15 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] -async def test_start_branch_then() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_start_branch_then( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5058,176 +5074,169 @@ async def test_start_branch_then() -> None: "market": "US", } - async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - store=MemoryStore(), - checkpointer=saver, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) + tool_two = tool_two_graph.compile( + store=MemoryStore(), + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value", - "market": "DE", - } - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "source": "loop", - "step": 0, - "writes": None, - }, - { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value slow", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value slow", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) + thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value slow", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value slow", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) + thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) - # update state - await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) + thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + -1 + ].config, + ) + # update state + await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "valuekey", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], + metadata={ + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + -1 + ].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "valuekey fast", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + -1 + ].config, + ) -async def test_branch_then() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_branch_then( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5255,606 +5264,578 @@ async def test_branch_then() -> None: "market": "US", } - async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=saver) - thread10 = {"configurable": {"thread_id": "10"}} - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, - "values": {"my_key": ""}, - "metadata": { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + "values": {"my_key": ""}, + "metadata": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} - ], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], + "values": { + "my_key": "value", + "market": "DE", }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], + "metadata": { + "source": "loop", + "step": 0, + "writes": None, }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, - "values": { - "my_key": "value prepared slow", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], + "values": { + "my_key": "value prepared", + "market": "DE", }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], + "metadata": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, }, + "next": ["tool_two_slow"], + "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "name": "tool_two_slow", + "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - "next": [], - "tasks": [], }, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + "next": ["finish"], + "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, - ] + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition::then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "name": "finish", + "result": [("my_key", " finished")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + "next": [], + "tasks": [], + }, + }, + ] - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "11"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "1"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, - "values": {"my_key": ""}, - "metadata": { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "1"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + "values": {"my_key": ""}, + "metadata": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70", - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "1"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} - ], }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 0, + "writes": None, + }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + "next": ["tool_two_slow"], + "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + }, + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) + thread2 = {"configurable": {"thread_id": "12"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) - async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_after=["prepare"] - ) + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_after=["prepare"] + ) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) + thread1 = {"configurable": {"thread_id": "21"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) + thread2 = {"configurable": {"thread_id": "22"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) - thread3 = {"configurable": {"thread_id": "3"}} - # update an empty thread before first run - uconfig = await tool_two.aupdate_state( - thread3, {"my_key": "key", "market": "DE"} - ) - # check current state - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare"),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - }, - parent_config=None, - ) - # run from this point - assert await tool_two.ainvoke(None, thread3) == { - "my_key": "key prepared", - "market": "DE", - } - # get state after first node - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=uconfig, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "key prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) + thread3 = {"configurable": {"thread_id": "23"}} + # update an empty thread before first run + uconfig = await tool_two.aupdate_state(thread3, {"my_key": "key", "market": "DE"}) + # check current state + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare"),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + }, + parent_config=None, + ) + # run from this point + assert await tool_two.ainvoke(None, thread3) == { + "my_key": "key prepared", + "market": "DE", + } + # get state after first node + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=uconfig, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "key prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + -1 + ].config, + ) async def test_in_one_fan_out_state_graph_waiting_edge() -> None: diff --git a/libs/sdk-js/package.json b/libs/sdk-js/package.json index db509e9c2..a10bcdab8 100644 --- a/libs/sdk-js/package.json +++ b/libs/sdk-js/package.json @@ -1,6 +1,6 @@ { "name": "@langchain/langgraph-sdk", - "version": "0.0.5", + "version": "0.0.6", "description": "Client library for interacting with the LangGraph API", "type": "module", "packageManager": "yarn@1.22.19", diff --git a/libs/sdk-js/src/client.mts b/libs/sdk-js/src/client.mts index bf98dccb4..fbf2dd12c 100644 --- a/libs/sdk-js/src/client.mts +++ b/libs/sdk-js/src/client.mts @@ -571,6 +571,7 @@ export class RunsClient extends BaseClient { ); let parser: EventSourceParser; + let onEndEvent: () => void; const textDecoder = new TextDecoder(); const stream: ReadableStream<{ event: string; data: any }> = ( @@ -594,9 +595,17 @@ export class RunsClient extends BaseClient { }); } }); + onEndEvent = () => { + ctrl.enqueue({ event: "end", data: undefined }); + }; }, async transform(chunk) { - parser.feed(textDecoder.decode(chunk)); + const payload = textDecoder.decode(chunk); + parser.feed(payload); + + // eventsource-parser will ignore events + // that are not terminated by a newline + if (payload.trim() === "event: end") onEndEvent(); }, }), ); diff --git a/libs/sdk-js/src/schema.ts b/libs/sdk-js/src/schema.ts index 9c2df19a9..1d772b7eb 100644 --- a/libs/sdk-js/src/schema.ts +++ b/libs/sdk-js/src/schema.ts @@ -2,6 +2,16 @@ import type { JSONSchema7 } from "json-schema"; type Optional = T | null | undefined; +type RunStatus = + | "pending" + | "running" + | "error" + | "success" + | "timeout" + | "interrupted"; + +type ThreadStatus = "idle" | "busy" | "interrupted"; + export interface Config { /** * Tags for this call and any sub-calls (eg. a Chain calling an LLM). @@ -80,6 +90,7 @@ export interface Thread { created_at: string; updated_at: string; metadata: Metadata; + status: ThreadStatus; } export interface Cron { @@ -112,12 +123,6 @@ export interface Run { assistant_id: string; created_at: string; updated_at: string; - status: - | "pending" - | "running" - | "error" - | "success" - | "timeout" - | "interrupted"; + status: RunStatus; metadata: Metadata; } diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index ba8aa2f42..24b2ef3b0 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -25,6 +25,7 @@ from langgraph_sdk.schema import ( Assistant, Config, Cron, + DisconnectMode, GraphSchema, Metadata, MultitaskStrategy, @@ -963,6 +964,8 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, feedback_keys: Optional[list[str]] = None, + on_disconnect: Optional[DisconnectMode] = None, + webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> AsyncIterator[StreamPart]: ... @@ -980,6 +983,8 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, feedback_keys: Optional[list[str]] = None, + on_disconnect: Optional[DisconnectMode] = None, + webhook: Optional[str] = None, ) -> AsyncIterator[StreamPart]: ... @@ -996,6 +1001,7 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, feedback_keys: Optional[list[str]] = None, + on_disconnect: Optional[DisconnectMode] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> AsyncIterator[StreamPart]: @@ -1019,6 +1025,8 @@ class RunsClient: webhook: Webhook to call after LangGraph API call is done. multitask_strategy: Multitask strategy to use. Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'. + on_disconnect: The disconnect mode to use. + Must be one of 'cancel' or 'continue'. Returns: AsyncIterator[StreamPart]: Asynchronous iterator of stream results. @@ -1061,6 +1069,7 @@ class RunsClient: "webhook": webhook, "checkpoint_id": checkpoint_id, "multitask_strategy": multitask_strategy, + "on_disconnect": on_disconnect, } endpoint = ( f"/threads/{thread_id}/runs/stream" @@ -1129,9 +1138,7 @@ class RunsClient: config: The configuration for the assistant. checkpoint_id: The checkpoint to start streaming from. interrupt_before: Nodes to interrupt immediately before they get executed. - interrupt_after: Nodes to Nodes to interrupt immediately after they get executed. - webhook: Webhook to call after LangGraph API call is done. multitask_strategy: Multitask strategy to use. Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'. @@ -1242,6 +1249,8 @@ class RunsClient: checkpoint_id: Optional[str] = None, interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, + webhook: Optional[str] = None, + on_disconnect: Optional[DisconnectMode] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> Union[list[dict], dict[str, Any]]: ... @@ -1257,6 +1266,8 @@ class RunsClient: config: Optional[Config] = None, interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, + webhook: Optional[str] = None, + on_disconnect: Optional[DisconnectMode] = None, ) -> Union[list[dict], dict[str, Any]]: ... @@ -1272,6 +1283,7 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, webhook: Optional[str] = None, + on_disconnect: Optional[DisconnectMode] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> Union[list[dict], dict[str, Any]]: """Create a run, wait until it finishes and return the final state. @@ -1286,12 +1298,12 @@ class RunsClient: config: The configuration for the assistant. checkpoint_id: The checkpoint to start streaming from. interrupt_before: Nodes to interrupt immediately before they get executed. - interrupt_after: Nodes to Nodes to interrupt immediately after they get executed. - webhook: Webhook to call after LangGraph API call is done. multitask_strategy: Multitask strategy to use. Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'. + on_disconnect: The disconnect mode to use. + Must be one of 'cancel' or 'continue'. Returns: Union[list[dict], dict[str, Any]]: The output of the run. @@ -1351,6 +1363,7 @@ class RunsClient: "webhook": webhook, "checkpoint_id": checkpoint_id, "multitask_strategy": multitask_strategy, + "on_disconnect": on_disconnect, } endpoint = ( f"/threads/{thread_id}/runs/wait" if thread_id is not None else "/runs/wait" @@ -1451,6 +1464,28 @@ class RunsClient: """ # noqa: E501 return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join") + def join_stream(self, thread_id: str, run_id: str) -> AsyncIterator[StreamPart]: + """Stream output from a run in real-time, until the run is done. + Output is not buffered, so any output produced before this call will + not be received here. + + Args: + thread_id: The thread ID to join. + run_id: The run ID to join. + + Returns: + None + + Example Usage: + + await client.runs.join( + thread_id="thread_id_to_join", + run_id="run_id_to_join" + ) + + """ # noqa: E501 + return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET") + async def delete(self, thread_id: str, run_id: str) -> None: """Delete a run. diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index a0fac51a0..c3232c88e 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -9,6 +9,8 @@ ThreadStatus = Literal["idle", "busy", "interrupted"] StreamMode = Literal["values", "messages", "updates", "events", "debug"] +DisconnectMode = Literal["cancel", "continue"] + MultitaskStrategy = Literal["reject", "interrupt", "rollback", "enqueue"] OnConflictBehavior = Literal["raise", "do_nothing"]