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

Merge branch 'main' into vb/update-get-state

Browse Source
This commit is contained in:
vbarda
2024-08-22 15:34:14 -04:00
parent 426125cf96 a261e1a497
commit 5654d8fc86
32 changed files with 1265 additions and 540 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/chatbots/information-gather-prompting.ipynb
+97 -47
View File
File diff suppressed because one or more lines are too long
examples/memory/manage-conversation-history.ipynb
+5 -5
View File
@@ -268,8 +268,8 @@
"\n",
"\n",
"def filter_messages(messages: list):\n",
" # This is very simple helper function which only ever uses the last two messages\n",
" return messages[-2:]\n",
" # This is very simple helper function which only ever uses the last message\n",
" return messages[-1:]\n",
"\n",
"\n",
"# Define the function that calls the model\n",
@@ -372,9 +372,9 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"display_name": "langgraph-example-dev",
"language": "python",
"name": "python3"
"name": "langgraph-example-dev"
},
"language_info": {
"codemirror_mode": {
@@ -386,7 +386,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.1"
"version": "3.11.9"
}
},
"nbformat": 4,
examples/rag/langgraph_agentic_rag.ipynb
+306 -10
View File
@@ -20,7 +20,10 @@
"id": "969fb438",
"metadata": {},
"outputs": [],
"source": ["%%capture --no-stderr\n%pip install -U --quiet langchain-community tiktoken langchain-openai langchainhub chromadb langchain langgraph langchain-text-splitters"]
"source": [
"%%capture --no-stderr\n",
"%pip install -U --quiet langchain-community tiktoken langchain-openai langchainhub chromadb langchain langgraph langchain-text-splitters"
]
},
{
"cell_type": "code",
@@ -28,7 +31,22 @@
"id": "e4958a8c",
"metadata": {},
"outputs": [],
"source": ["import getpass\nimport os\n\n\ndef _set_env(key: str):\n if key not in os.environ:\n os.environ[key] = getpass.getpass(f\"{key}:\")\n\n\n_set_env(\"OPENAI_API_KEY\")\n\n# (Optional) For tracing\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n_set_env(\"LANGCHAIN_API_KEY\")"]
"source": [
"import getpass\n",
"import os\n",
"\n",
"\n",
"def _set_env(key: str):\n",
" if key not in os.environ:\n",
" os.environ[key] = getpass.getpass(f\"{key}:\")\n",
"\n",
"\n",
"_set_env(\"OPENAI_API_KEY\")\n",
"\n",
"# (Optional) For tracing\n",
"os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
"_set_env(\"LANGCHAIN_API_KEY\")"
]
},
{
"cell_type": "markdown",
@@ -46,7 +64,34 @@
"id": "e50c9efe-4abe-42fa-b35a-05eeeede9ec6",
"metadata": {},
"outputs": [],
"source": ["from langchain_community.document_loaders import WebBaseLoader\nfrom langchain_community.vectorstores import Chroma\nfrom langchain_openai import OpenAIEmbeddings\nfrom langchain_text_splitters import RecursiveCharacterTextSplitter\n\nurls = [\n \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n]\n\ndocs = [WebBaseLoader(url).load() for url in urls]\ndocs_list = [item for sublist in docs for item in sublist]\n\ntext_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n chunk_size=100, chunk_overlap=50\n)\ndoc_splits = text_splitter.split_documents(docs_list)\n\n# Add to vectorDB\nvectorstore = Chroma.from_documents(\n documents=doc_splits,\n collection_name=\"rag-chroma\",\n embedding=OpenAIEmbeddings(),\n)\nretriever = vectorstore.as_retriever()"]
"source": [
"from langchain_community.document_loaders import WebBaseLoader\n",
"from langchain_community.vectorstores import Chroma\n",
"from langchain_openai import OpenAIEmbeddings\n",
"from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
"\n",
"urls = [\n",
" \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n",
" \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n",
" \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n",
"]\n",
"\n",
"docs = [WebBaseLoader(url).load() for url in urls]\n",
"docs_list = [item for sublist in docs for item in sublist]\n",
"\n",
"text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
" chunk_size=100, chunk_overlap=50\n",
")\n",
"doc_splits = text_splitter.split_documents(docs_list)\n",
"\n",
"# Add to vectorDB\n",
"vectorstore = Chroma.from_documents(\n",
" documents=doc_splits,\n",
" collection_name=\"rag-chroma\",\n",
" embedding=OpenAIEmbeddings(),\n",
")\n",
"retriever = vectorstore.as_retriever()"
]
},
{
"cell_type": "markdown",
@@ -62,7 +107,17 @@
"id": "0b97bdd8-d7e3-444d-ac96-5ef4725f9048",
"metadata": {},
"outputs": [],
"source": ["from langchain.tools.retriever import create_retriever_tool\n\nretriever_tool = create_retriever_tool(\n retriever,\n \"retrieve_blog_posts\",\n \"Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.\",\n)\n\ntools = [retriever_tool]"]
"source": [
"from langchain.tools.retriever import create_retriever_tool\n",
"\n",
"retriever_tool = create_retriever_tool(\n",
" retriever,\n",
" \"retrieve_blog_posts\",\n",
" \"Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.\",\n",
")\n",
"\n",
"tools = [retriever_tool]"
]
},
{
"cell_type": "markdown",
@@ -86,7 +141,19 @@
"id": "0e378706-47d5-425a-8ba0-57b9acffbd0c",
"metadata": {},
"outputs": [],
"source": ["from typing import Annotated, Sequence, TypedDict\n\nfrom langchain_core.messages import BaseMessage\n\nfrom langgraph.graph.message import add_messages\n\n\nclass AgentState(TypedDict):\n # The add_messages function defines how an update should be processed\n # Default is to replace. add_messages says \"append\"\n messages: Annotated[Sequence[BaseMessage], add_messages]"]
"source": [
"from typing import Annotated, Sequence, TypedDict\n",
"\n",
"from langchain_core.messages import BaseMessage\n",
"\n",
"from langgraph.graph.message import add_messages\n",
"\n",
"\n",
"class AgentState(TypedDict):\n",
" # The add_messages function defines how an update should be processed\n",
" # Default is to replace. add_messages says \"append\"\n",
" messages: Annotated[Sequence[BaseMessage], add_messages]"
]
},
{
"attachments": {
@@ -129,7 +196,173 @@
]
}
],
"source": ["from typing import Annotated, Literal, Sequence, TypedDict\n\nfrom langchain import hub\nfrom langchain_core.messages import BaseMessage, HumanMessage\nfrom langchain_core.output_parsers import StrOutputParser\nfrom langchain_core.prompts import PromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\nfrom langchain_openai import ChatOpenAI\n\nfrom langgraph.prebuilt import tools_condition\n\n### Edges\n\n\ndef grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n \"\"\"\n Determines whether the retrieved documents are relevant to the question.\n\n Args:\n state (messages): The current state\n\n Returns:\n str: A decision for whether the documents are relevant or not\n \"\"\"\n\n print(\"---CHECK RELEVANCE---\")\n\n # Data model\n class grade(BaseModel):\n \"\"\"Binary score for relevance check.\"\"\"\n\n binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n\n # LLM\n model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n\n # LLM with tool and validation\n llm_with_tool = model.with_structured_output(grade)\n\n # Prompt\n prompt = PromptTemplate(\n template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n Here is the retrieved document: \\n\\n {context} \\n\\n\n Here is the user question: {question} \\n\n If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n input_variables=[\"context\", \"question\"],\n )\n\n # Chain\n chain = prompt | llm_with_tool\n\n messages = state[\"messages\"]\n last_message = messages[-1]\n\n question = messages[0].content\n docs = last_message.content\n\n scored_result = chain.invoke({\"question\": question, \"context\": docs})\n\n score = scored_result.binary_score\n\n if score == \"yes\":\n print(\"---DECISION: DOCS RELEVANT---\")\n return \"generate\"\n\n else:\n print(\"---DECISION: DOCS NOT RELEVANT---\")\n print(score)\n return \"rewrite\"\n\n\n### Nodes\n\n\ndef agent(state):\n \"\"\"\n Invokes the agent model to generate a response based on the current state. Given\n the question, it will decide to retrieve using the retriever tool, or simply end.\n\n Args:\n state (messages): The current state\n\n Returns:\n dict: The updated state with the agent response appended to messages\n \"\"\"\n print(\"---CALL AGENT---\")\n messages = state[\"messages\"]\n model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4-turbo\")\n model = model.bind_tools(tools)\n response = model.invoke(messages)\n # We return a list, because this will get added to the existing list\n return {\"messages\": [response]}\n\n\ndef rewrite(state):\n \"\"\"\n Transform the query to produce a better question.\n\n Args:\n state (messages): The current state\n\n Returns:\n dict: The updated state with re-phrased question\n \"\"\"\n\n print(\"---TRANSFORM QUERY---\")\n messages = state[\"messages\"]\n question = messages[0].content\n\n msg = [\n HumanMessage(\n content=f\"\"\" \\n \n Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n Here is the initial question:\n \\n ------- \\n\n {question} \n \\n ------- \\n\n Formulate an improved question: \"\"\",\n )\n ]\n\n # Grader\n model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n response = model.invoke(msg)\n return {\"messages\": [response]}\n\n\ndef generate(state):\n \"\"\"\n Generate answer\n\n Args:\n state (messages): The current state\n\n Returns:\n dict: The updated state with re-phrased question\n \"\"\"\n print(\"---GENERATE---\")\n messages = state[\"messages\"]\n question = messages[0].content\n last_message = messages[-1]\n\n question = messages[0].content\n docs = last_message.content\n\n # Prompt\n prompt = hub.pull(\"rlm/rag-prompt\")\n\n # LLM\n llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0, streaming=True)\n\n # Post-processing\n def format_docs(docs):\n return \"\\n\\n\".join(doc.page_content for doc in docs)\n\n # Chain\n rag_chain = prompt | llm | StrOutputParser()\n\n # Run\n response = rag_chain.invoke({\"context\": docs, \"question\": question})\n return {\"messages\": [response]}\n\n\nprint(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\nprompt = hub.pull(\"rlm/rag-prompt\").pretty_print() # Show what the prompt looks like"]
"source": [
"from typing import Annotated, Literal, Sequence, TypedDict\n",
"\n",
"from langchain import hub\n",
"from langchain_core.messages import BaseMessage, HumanMessage\n",
"from langchain_core.output_parsers import StrOutputParser\n",
"from langchain_core.prompts import PromptTemplate\n",
"from langchain_core.pydantic_v1 import BaseModel, Field\n",
"from langchain_openai import ChatOpenAI\n",
"\n",
"from langgraph.prebuilt import tools_condition\n",
"\n",
"### Edges\n",
"\n",
"\n",
"def grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n",
" \"\"\"\n",
" Determines whether the retrieved documents are relevant to the question.\n",
"\n",
" Args:\n",
" state (messages): The current state\n",
"\n",
" Returns:\n",
" str: A decision for whether the documents are relevant or not\n",
" \"\"\"\n",
"\n",
" print(\"---CHECK RELEVANCE---\")\n",
"\n",
" # Data model\n",
" class grade(BaseModel):\n",
" \"\"\"Binary score for relevance check.\"\"\"\n",
"\n",
" binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n",
"\n",
" # LLM\n",
" model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
"\n",
" # LLM with tool and validation\n",
" llm_with_tool = model.with_structured_output(grade)\n",
"\n",
" # Prompt\n",
" prompt = PromptTemplate(\n",
" template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
" Here is the retrieved document: \\n\\n {context} \\n\\n\n",
" Here is the user question: {question} \\n\n",
" If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
" Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n",
" input_variables=[\"context\", \"question\"],\n",
" )\n",
"\n",
" # Chain\n",
" chain = prompt | llm_with_tool\n",
"\n",
" messages = state[\"messages\"]\n",
" last_message = messages[-1]\n",
"\n",
" question = messages[0].content\n",
" docs = last_message.content\n",
"\n",
" scored_result = chain.invoke({\"question\": question, \"context\": docs})\n",
"\n",
" score = scored_result.binary_score\n",
"\n",
" if score == \"yes\":\n",
" print(\"---DECISION: DOCS RELEVANT---\")\n",
" return \"generate\"\n",
"\n",
" else:\n",
" print(\"---DECISION: DOCS NOT RELEVANT---\")\n",
" print(score)\n",
" return \"rewrite\"\n",
"\n",
"\n",
"### Nodes\n",
"\n",
"\n",
"def agent(state):\n",
" \"\"\"\n",
" Invokes the agent model to generate a response based on the current state. Given\n",
" the question, it will decide to retrieve using the retriever tool, or simply end.\n",
"\n",
" Args:\n",
" state (messages): The current state\n",
"\n",
" Returns:\n",
" dict: The updated state with the agent response appended to messages\n",
" \"\"\"\n",
" print(\"---CALL AGENT---\")\n",
" messages = state[\"messages\"]\n",
" model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4-turbo\")\n",
" model = model.bind_tools(tools)\n",
" response = model.invoke(messages)\n",
" # We return a list, because this will get added to the existing list\n",
" return {\"messages\": [response]}\n",
"\n",
"\n",
"def rewrite(state):\n",
" \"\"\"\n",
" Transform the query to produce a better question.\n",
"\n",
" Args:\n",
" state (messages): The current state\n",
"\n",
" Returns:\n",
" dict: The updated state with re-phrased question\n",
" \"\"\"\n",
"\n",
" print(\"---TRANSFORM QUERY---\")\n",
" messages = state[\"messages\"]\n",
" question = messages[0].content\n",
"\n",
" msg = [\n",
" HumanMessage(\n",
" content=f\"\"\" \\n \n",
" Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n",
" Here is the initial question:\n",
" \\n ------- \\n\n",
" {question} \n",
" \\n ------- \\n\n",
" Formulate an improved question: \"\"\",\n",
" )\n",
" ]\n",
"\n",
" # Grader\n",
" model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
" response = model.invoke(msg)\n",
" return {\"messages\": [response]}\n",
"\n",
"\n",
"def generate(state):\n",
" \"\"\"\n",
" Generate answer\n",
"\n",
" Args:\n",
" state (messages): The current state\n",
"\n",
" Returns:\n",
" dict: The updated state with re-phrased question\n",
" \"\"\"\n",
" print(\"---GENERATE---\")\n",
" messages = state[\"messages\"]\n",
" question = messages[0].content\n",
" last_message = messages[-1]\n",
"\n",
" docs = last_message.content\n",
"\n",
" # Prompt\n",
" prompt = hub.pull(\"rlm/rag-prompt\")\n",
"\n",
" # LLM\n",
" llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0, streaming=True)\n",
"\n",
" # Post-processing\n",
" def format_docs(docs):\n",
" return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
"\n",
" # Chain\n",
" rag_chain = prompt | llm | StrOutputParser()\n",
"\n",
" # Run\n",
" response = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
" return {\"messages\": [response]}\n",
"\n",
"\n",
"print(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\n",
"prompt = hub.pull(\"rlm/rag-prompt\").pretty_print() # Show what the prompt looks like"
]
},
{
"cell_type": "markdown",
@@ -150,7 +383,48 @@
"id": "8718a37f-83c2-4f16-9850-e61e0f49c3d4",
"metadata": {},
"outputs": [],
"source": ["from langgraph.graph import END, StateGraph, START\nfrom langgraph.prebuilt import ToolNode\n\n# Define a new graph\nworkflow = StateGraph(AgentState)\n\n# Define the nodes we will cycle between\nworkflow.add_node(\"agent\", agent) # agent\nretrieve = ToolNode([retriever_tool])\nworkflow.add_node(\"retrieve\", retrieve) # retrieval\nworkflow.add_node(\"rewrite\", rewrite) # Re-writing the question\nworkflow.add_node(\n \"generate\", generate\n) # Generating a response after we know the documents are relevant\n# Call agent node to decide to retrieve or not\nworkflow.add_edge(START, \"agent\")\n\n# Decide whether to retrieve\nworkflow.add_conditional_edges(\n \"agent\",\n # Assess agent decision\n tools_condition,\n {\n # Translate the condition outputs to nodes in our graph\n \"tools\": \"retrieve\",\n END: END,\n },\n)\n\n# Edges taken after the `action` node is called.\nworkflow.add_conditional_edges(\n \"retrieve\",\n # Assess agent decision\n grade_documents,\n)\nworkflow.add_edge(\"generate\", END)\nworkflow.add_edge(\"rewrite\", \"agent\")\n\n# Compile\ngraph = workflow.compile()"]
"source": [
"from langgraph.graph import END, StateGraph, START\n",
"from langgraph.prebuilt import ToolNode\n",
"\n",
"# Define a new graph\n",
"workflow = StateGraph(AgentState)\n",
"\n",
"# Define the nodes we will cycle between\n",
"workflow.add_node(\"agent\", agent) # agent\n",
"retrieve = ToolNode([retriever_tool])\n",
"workflow.add_node(\"retrieve\", retrieve) # retrieval\n",
"workflow.add_node(\"rewrite\", rewrite) # Re-writing the question\n",
"workflow.add_node(\n",
" \"generate\", generate\n",
") # Generating a response after we know the documents are relevant\n",
"# Call agent node to decide to retrieve or not\n",
"workflow.add_edge(START, \"agent\")\n",
"\n",
"# Decide whether to retrieve\n",
"workflow.add_conditional_edges(\n",
" \"agent\",\n",
" # Assess agent decision\n",
" tools_condition,\n",
" {\n",
" # Translate the condition outputs to nodes in our graph\n",
" \"tools\": \"retrieve\",\n",
" END: END,\n",
" },\n",
")\n",
"\n",
"# Edges taken after the `action` node is called.\n",
"workflow.add_conditional_edges(\n",
" \"retrieve\",\n",
" # Assess agent decision\n",
" grade_documents,\n",
")\n",
"workflow.add_edge(\"generate\", END)\n",
"workflow.add_edge(\"rewrite\", \"agent\")\n",
"\n",
"# Compile\n",
"graph = workflow.compile()"
]
},
{
"cell_type": "code",
@@ -169,7 +443,15 @@
"output_type": "display_data"
}
],
"source": ["from IPython.display import Image, display\n\ntry:\n display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\nexcept Exception:\n # This requires some extra dependencies and is optional\n pass"]
"source": [
"from IPython.display import Image, display\n",
"\n",
"try:\n",
" display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\n",
"except Exception:\n",
" # This requires some extra dependencies and is optional\n",
" pass"
]
},
{
"cell_type": "code",
@@ -203,7 +485,21 @@
]
}
],
"source": ["import pprint\n\ninputs = {\n \"messages\": [\n (\"user\", \"What does Lilian Weng say about the types of agent memory?\"),\n ]\n}\nfor output in graph.stream(inputs):\n for key, value in output.items():\n pprint.pprint(f\"Output from node '{key}':\")\n pprint.pprint(\"---\")\n pprint.pprint(value, indent=2, width=80, depth=None)\n pprint.pprint(\"\\n---\\n\")"]
"source": [
"import pprint\n",
"\n",
"inputs = {\n",
" \"messages\": [\n",
" (\"user\", \"What does Lilian Weng say about the types of agent memory?\"),\n",
" ]\n",
"}\n",
"for output in graph.stream(inputs):\n",
" for key, value in output.items():\n",
" pprint.pprint(f\"Output from node '{key}':\")\n",
" pprint.pprint(\"---\")\n",
" pprint.pprint(value, indent=2, width=80, depth=None)\n",
" pprint.pprint(\"\\n---\\n\")"
]
},
{
"cell_type": "code",
@@ -211,7 +507,7 @@
"id": "189333cc-5d34-4869-9f9b-741210e1096f",
"metadata": {},
"outputs": [],
"source": [""]
"source": []
}
],
"metadata": {
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
-10
View File
@@ -318,16 +318,6 @@ class PostgresSaver(BasePostgresSaver):
task_id (str): Identifier for the task creating the writes.
"""
with self._cursor(pipeline=True) as cur:
cur.execute(
self.DELETE_WRITES_SQL,
(
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
len(writes),
),
)
cur.executemany(
self.UPSERT_CHECKPOINT_WRITES_SQL,
self._dump_writes(
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
-10
View File
@@ -274,16 +274,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
task_id (str): Identifier for the task creating the writes.
"""
async with self._cursor(pipeline=True) as cur:
await cur.execute(
self.DELETE_WRITES_SQL,
(
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
len(writes),
),
)
await cur.executemany(
self.UPSERT_CHECKPOINT_WRITES_SQL,
await asyncio.to_thread(
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+2 -11
View File
@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
from psycopg.types.json import Jsonb
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
Checkpoint,
EmptyChannelError,
@@ -105,15 +106,6 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
"""
DELETE_WRITES_SQL = """
DELETE FROM checkpoint_writes
WHERE thread_id = %s
AND checkpoint_ns = %s
AND checkpoint_id = %s
AND task_id = %s
AND idx >= %s
"""
class BasePostgresSaver(BaseCheckpointSaver):
SELECT_SQL = SELECT_SQL
@@ -121,7 +113,6 @@ class BasePostgresSaver(BaseCheckpointSaver):
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
DELETE_WRITES_SQL = DELETE_WRITES_SQL
jsonplus_serde = JsonPlusSerializer()
@@ -210,7 +201,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
checkpoint_ns,
checkpoint_id,
task_id,
idx,
WRITES_IDX_MAP.get(channel, idx),
channel,
*self.serde.dumps_typed(value),
)
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+8 -23
View File
@@ -1,12 +1,13 @@
import sqlite3
import threading
from contextlib import contextmanager
from contextlib import closing, contextmanager
from hashlib import md5
from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
@@ -318,9 +319,7 @@ class SqliteSaver(BaseCheckpointSaver):
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
with self.cursor(transaction=False) as cur, self.cursor(
transaction=False
) as writes_cur:
with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
cur.execute(query, param_values)
for (
thread_id,
@@ -331,13 +330,9 @@ class SqliteSaver(BaseCheckpointSaver):
checkpoint,
metadata,
) in cur:
writes_cur.execute(
wcur.execute(
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
thread_id,
checkpoint_ns,
checkpoint_id,
),
(thread_id, checkpoint_ns, checkpoint_id),
)
yield CheckpointTuple(
{
@@ -362,7 +357,7 @@ class SqliteSaver(BaseCheckpointSaver):
),
[
(task_id, channel, self.serde.loads_typed((type, value)))
for task_id, channel, type, value in writes_cur
for task_id, channel, type, value in wcur
],
)
@@ -438,25 +433,15 @@ class SqliteSaver(BaseCheckpointSaver):
task_id (str): Identifier for the task creating the writes.
"""
with self.lock, self.cursor() as cur:
cur.execute(
"DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
len(writes),
),
)
cur.executemany(
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
"INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
[
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
idx,
WRITES_IDX_MAP.get(channel, idx),
channel,
*self.serde.dumps_typed(value),
)
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+9 -22
View File
@@ -16,6 +16,7 @@ import aiosqlite
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
@@ -329,14 +330,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
"""
await self.setup()
where, param_values = search_where(config, filter, before)
where, params = search_where(config, filter, before)
query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
FROM checkpoints
{where}
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
async with self.conn.execute(query, param_values) as cursor:
async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur:
async for (
thread_id,
checkpoint_ns,
@@ -345,14 +346,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
type,
checkpoint,
metadata,
) in cursor:
writes_cur = await self.conn.execute(
) in cur:
await wcur.execute(
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
thread_id,
checkpoint_ns,
checkpoint_id,
),
(thread_id, checkpoint_ns, checkpoint_id),
)
yield CheckpointTuple(
{
@@ -377,7 +374,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
),
[
(task_id, channel, self.serde.loads_typed((type, value)))
async for task_id, channel, type, value in writes_cur
async for task_id, channel, type, value in wcur
],
)
@@ -445,25 +442,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
"""
await self.setup()
async with self.conn.cursor() as cur:
await cur.execute(
"DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
len(writes),
),
)
await cur.executemany(
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
"INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
[
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
idx,
WRITES_IDX_MAP.get(channel, idx),
channel,
*self.serde.dumps_typed(value),
)
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+14
View File
@@ -22,6 +22,7 @@ from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import (
ERROR,
ChannelProtocol,
SendProtocol,
)
@@ -98,6 +99,7 @@ class Checkpoint(TypedDict):
Cleared by the next checkpoint."""
current_tasks: Dict[str, TaskInfo]
"""Map from task ID to task info."""
# TODO remove this
def empty_checkpoint() -> Checkpoint:
@@ -140,6 +142,8 @@ def create_checkpoint(
else:
values: dict[str, Any] = {}
for k, v in channels.items():
if k not in checkpoint["channel_versions"]:
continue
try:
values[k] = v.checkpoint()
except EmptyChannelError:
@@ -437,3 +441,13 @@ def get_checkpoint_id(config: RunnableConfig) -> Optional[str]:
return config["configurable"].get(
"checkpoint_id", config["configurable"].get("thread_ts")
)
"""
Mapping from error type to error index.
Regular writes just map to their index in the list of writes being saved.
Special writes (e.g. errors) map to negative indices, to avoid those writes from
saving regular writes.
Each Checkpointer implementation should use this mapping in put_writes.
"""
WRITES_IDX_MAP = {ERROR: -1}
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+16 -12
View File
@@ -8,6 +8,7 @@ from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
@@ -52,6 +53,9 @@ class MemorySaver(
# thread ID -> checkpoint NS -> checkpoint ID -> checkpoint mapping
storage: defaultdict[str, dict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]]
writes: defaultdict[
tuple[str, str, str], dict[tuple[str, int], tuple[str, str, bytes]]
]
def __init__(
self,
@@ -60,7 +64,7 @@ class MemorySaver(
) -> None:
super().__init__(serde=serde)
self.storage = defaultdict(lambda: defaultdict(dict))
self.writes = defaultdict(list)
self.writes = defaultdict(dict)
def __enter__(self) -> "MemorySaver":
return self
@@ -103,7 +107,7 @@ class MemorySaver(
if checkpoint_id := get_checkpoint_id(config):
if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id):
checkpoint, metadata, parent_checkpoint_id = saved
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
return CheckpointTuple(
config=config,
checkpoint=self.serde.loads_typed(checkpoint),
@@ -125,7 +129,7 @@ class MemorySaver(
if checkpoints := self.storage[thread_id][checkpoint_ns]:
checkpoint_id = max(checkpoints.keys())
checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
return CheckpointTuple(
config={
"configurable": {
@@ -206,7 +210,8 @@ class MemorySaver(
elif limit is not None:
limit -= 1
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
yield CheckpointTuple(
config={
"configurable": {
@@ -216,9 +221,6 @@ class MemorySaver(
}
},
checkpoint=self.serde.loads_typed(checkpoint),
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
],
metadata=metadata,
parent_config={
"configurable": {
@@ -229,6 +231,9 @@ class MemorySaver(
}
if parent_checkpoint_id
else None,
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
],
)
def put(
@@ -293,11 +298,10 @@ class MemorySaver(
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
checkpoint_id = config["configurable"]["checkpoint_id"]
key = (thread_id, checkpoint_ns, checkpoint_id)
self.writes[key] = [w for w in self.writes[key] if w[0] != task_id]
self.writes[key].extend(
[(task_id, c, self.serde.dumps_typed(v)) for c, v in writes]
)
outer_key = (thread_id, checkpoint_ns, checkpoint_id)
for idx, (c, v) in enumerate(writes):
inner_key = (task_id, WRITES_IDX_MAP.get(c, idx))
self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v))
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Asynchronous version of get_tuple.
libs/checkpoint/langgraph/checkpoint/serde/types.py
+2
View File
@@ -12,6 +12,8 @@ from typing import (
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
ERROR = "__error__"
Value = TypeVar("Value")
Update = TypeVar("Update")
C = TypeVar("C")
libs/langgraph/langgraph/channels/base.py
+22 -4
View File
@@ -2,9 +2,9 @@ from abc import ABC, abstractmethod
from contextlib import asynccontextmanager, contextmanager
from typing import (
Any,
AsyncGenerator,
Generator,
AsyncIterator,
Generic,
Iterator,
Optional,
Sequence,
TypeVar,
@@ -21,6 +21,8 @@ C = TypeVar("C")
class BaseChannel(Generic[Value, Update, C], ABC):
key: str = ""
@property
@abstractmethod
def ValueType(self) -> Any:
@@ -43,19 +45,35 @@ class BaseChannel(Generic[Value, Update, C], ABC):
@abstractmethod
def from_checkpoint(
self, checkpoint: Optional[C], config: RunnableConfig
) -> Generator[Self, None, None]:
) -> Iterator[Self]:
"""Return a new identical channel, optionally initialized from a checkpoint.
If the checkpoint contains complex data structures, they should be copied."""
@contextmanager
def from_checkpoint_named(
self, checkpoint: Optional[C], config: RunnableConfig
) -> Iterator[Self]:
with self.from_checkpoint(checkpoint, config) as value:
value.key = self.key
yield value
@asynccontextmanager
async def afrom_checkpoint(
self, checkpoint: Optional[C], config: RunnableConfig
) -> AsyncGenerator[Self, None]:
) -> AsyncIterator[Self]:
"""Return a new identical channel, optionally initialized from a checkpoint.
If the checkpoint contains complex data structures, they should be copied."""
with self.from_checkpoint(checkpoint, config) as value:
yield value
@asynccontextmanager
async def afrom_checkpoint_named(
self, checkpoint: Optional[C], config: RunnableConfig
) -> AsyncIterator[Self]:
async with self.afrom_checkpoint(checkpoint, config) as value:
value.key = self.key
yield value
# state methods
@abstractmethod
libs/langgraph/langgraph/channels/context.py
+3 -1
View File
@@ -112,7 +112,9 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
def update(self, values: Sequence[None]) -> bool:
if values:
raise InvalidUpdateError("Context channel does not accept writes.")
raise InvalidUpdateError(
f"At key '{self.key}': Context channel does not accept writes."
)
return False
def get(self) -> Value:
libs/langgraph/langgraph/channels/dynamic_barrier_value.py
+1 -1
View File
@@ -69,7 +69,7 @@ class DynamicBarrierValue(
if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]:
if len(wait_for_names) > 1:
raise InvalidUpdateError(
"Received multiple WaitForNames updates in the same step."
f"At key '{self.key}': Received multiple WaitForNames updates in the same step."
)
self.names = wait_for_names[0].names
return True
libs/langgraph/langgraph/channels/ephemeral_value.py
+1 -1
View File
@@ -58,7 +58,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
return False
if len(values) != 1 and self.guard:
raise InvalidUpdateError(
"EphemeralValue can only receive one value per step."
f"At key '{self.key}': EphemeralValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
)
self.value = values[-1]
libs/langgraph/langgraph/channels/last_value.py
+3 -1
View File
@@ -52,7 +52,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
if len(values) == 0:
return False
if len(values) != 1:
raise InvalidUpdateError("LastValue can only receive one value per step.")
raise InvalidUpdateError(
f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values."
)
self.value = values[-1]
return True
libs/langgraph/langgraph/channels/named_barrier_value.py
+3 -1
View File
@@ -53,7 +53,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
self.seen.add(value)
updated = True
else:
raise InvalidUpdateError(f"Value {value} not in {self.names}")
raise InvalidUpdateError(
f"At key '{self.key}': Value {value} not in {self.names}"
)
return updated
def get(self) -> Value:
libs/langgraph/langgraph/channels/untracked_value.py
+1 -1
View File
@@ -49,7 +49,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
return False
if len(values) != 1 and self.guard:
raise InvalidUpdateError(
"UntrackedValue can only receive one value per step."
f"At key '{self.key}': UntrackedValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
)
self.value = values[-1]
libs/langgraph/langgraph/errors.py
+2 -2
View File
@@ -1,4 +1,4 @@
from typing import Any
from typing import Any, Sequence
from langgraph.checkpoint.base import EmptyChannelError
from langgraph.constants import Interrupt
@@ -32,7 +32,7 @@ class InvalidUpdateError(Exception):
class GraphInterrupt(Exception):
"""Raised when a subgraph is interrupted."""
def __init__(self, interrupts: list[Interrupt]) -> None:
def __init__(self, interrupts: Sequence[Interrupt] = ()) -> None:
super().__init__(interrupts)
libs/langgraph/langgraph/graph/graph.py
+4 -2
View File
@@ -198,12 +198,14 @@ class Graph:
raise ValueError("END cannot be a start node")
if end_key == START:
raise ValueError("START cannot be an end node")
if not self.support_multiple_edges and start_key in set(
# run this validation only for non-StateGraph graphs
if not hasattr(self, "channels") and start_key in set(
start for start, _ in self.edges
):
raise ValueError(
f"Already found path for node '{start_key}'.\n"
"For multiple edges, use StateGraph with an annotated state key."
"For multiple edges, use StateGraph with an Annotated state key."
)
self.edges.add((start_key, end_key))
libs/langgraph/langgraph/graph/state.py
+14 -9
View File
@@ -1,3 +1,4 @@
import inspect
import logging
import typing
import warnings
@@ -196,10 +197,6 @@ class StateGraph(Graph):
)
else:
self.managed[key] = managed
if any(
isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()
):
self.support_multiple_edges = True
@overload
def add_node(
@@ -338,10 +335,13 @@ class StateGraph(Graph):
hints := get_type_hints(action.__call__) or get_type_hints(action)
):
if input is None:
input_hint = hints[list(hints.keys())[0]]
if isinstance(input_hint, type) and get_type_hints(input_hint):
input = input_hint
except TypeError:
first_parameter_name = next(
iter(inspect.signature(action).parameters.keys())
)
if input_hint := hints.get(first_parameter_name):
if isinstance(input_hint, type) and get_type_hints(input_hint):
input = input_hint
except (TypeError, StopIteration):
pass
if input is not None:
self._add_schema(input)
@@ -727,10 +727,15 @@ def _get_channel(
else:
raise ValueError(f"This {annotation} not allowed in this position")
elif channel := _is_field_channel(annotation):
channel.key = name
return channel
elif channel := _is_field_binop(annotation):
channel.key = name
return channel
return LastValue(annotation)
fallback = LastValue(annotation)
fallback.key = name
return fallback
def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]:
libs/langgraph/langgraph/pregel/__init__.py
+54 -135
View File
@@ -71,23 +71,14 @@ from langgraph.constants import (
)
from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
from langgraph.managed.base import ManagedValueSpec
from langgraph.pregel.algo import (
apply_writes,
local_read,
prepare_next_tasks,
should_interrupt,
)
from langgraph.pregel.algo import apply_writes, local_read, prepare_next_tasks
from langgraph.pregel.debug import (
map_debug_task_results,
print_step_checkpoint,
print_step_tasks,
print_step_writes,
tasks_w_writes,
)
from langgraph.pregel.io import (
map_output_updates,
read_channels,
)
from langgraph.pregel.io import read_channels
from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
@@ -503,6 +494,7 @@ class Pregel(
config=config,
metadata=None,
created_at=None,
parent_config=None,
tasks=(),
)
@@ -578,6 +570,7 @@ class Pregel(
config=config,
metadata=None,
created_at=None,
parent_config=None,
tasks=(),
)
@@ -723,7 +716,7 @@ class Pregel(
# update channels
with ChannelsManager(self.channels, checkpoint, config) as (
channels,
managed,
_,
):
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
@@ -759,31 +752,6 @@ class Pregel(
checkpoint, channels, [task], self.checkpointer.get_next_version
), "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# check interrupt before
if tasks := should_interrupt(
checkpoint,
self.interrupt_before_nodes,
prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
step + 2,
for_execution=False,
),
):
for t in tasks:
self.checkpointer.put_writes(
{
"configurable": {
**checkpoint_config["configurable"],
"checkpoint_id": checkpoint["id"],
}
},
[(INTERRUPT, Interrupt("before"))],
t.id,
)
return self.checkpointer.put(
checkpoint_config,
checkpoint,
@@ -866,7 +834,7 @@ class Pregel(
# update channels, acting as the chosen node
async with AsyncChannelsManager(self.channels, checkpoint, config) as (
channels,
managed,
_,
):
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
@@ -902,35 +870,6 @@ class Pregel(
checkpoint, channels, [task], self.checkpointer.get_next_version
), "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# check interrupt before
if tasks := should_interrupt(
checkpoint,
self.interrupt_before_nodes,
prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
step + 2,
for_execution=False,
),
):
await asyncio.gather(
*(
self.checkpointer.aput_writes(
{
"configurable": {
**checkpoint_config["configurable"],
"checkpoint_id": checkpoint["id"],
}
},
[(INTERRUPT, Interrupt("before"))],
t.id,
)
for t in tasks
)
)
return await self.checkpointer.aput(
checkpoint_config,
checkpoint,
@@ -1111,6 +1050,7 @@ class Pregel(
nodes=self.nodes,
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
) as loop:
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
@@ -1119,7 +1059,6 @@ class Pregel(
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
manager=run_manager,
@@ -1188,26 +1127,17 @@ class Pregel(
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
# yield updates output for the finished task
if "updates" in stream_modes:
yield from _with_mode(
"updates",
isinstance(stream_mode, list),
map_output_updates(output_keys, [task]),
)
if "debug" in stream_modes:
yield from _with_mode(
"debug",
isinstance(stream_mode, list),
map_debug_task_results(
loop.step,
[task],
self.stream_channels_list,
),
)
else:
# remove references to loop vars
del fut, task
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
if _should_stop_others(done):
break
@@ -1222,21 +1152,21 @@ class Pregel(
[w for t in loop.tasks for w in t.writes],
self.stream_channels_list,
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# set final channel values as run output
run_manager.on_chain_end(loop.output)
except BaseException as e:
@@ -1368,6 +1298,7 @@ class Pregel(
nodes=self.nodes,
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
) as loop:
aioloop = asyncio.get_event_loop()
# Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1377,7 +1308,6 @@ class Pregel(
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
manager=run_manager,
@@ -1448,28 +1378,17 @@ class Pregel(
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
# yield updates output for the finished task
if "updates" in stream_modes:
for chunk in _with_mode(
"updates",
isinstance(stream_mode, list),
map_output_updates(output_keys, [task]),
):
yield chunk
if "debug" in stream_modes:
for chunk in _with_mode(
"debug",
isinstance(stream_mode, list),
map_debug_task_results(
loop.step,
[task],
self.stream_channels_list,
),
):
yield chunk
else:
# remove references to loop vars
del fut, task
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
if _should_stop_others(done):
break
@@ -1484,21 +1403,21 @@ class Pregel(
[w for t in loop.tasks for w in t.writes],
self.stream_channels_list,
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# set final channel values as run output
await run_manager.on_chain_end(loop.output)
except BaseException as e:
libs/langgraph/langgraph/pregel/algo.py
+1 -6
View File
@@ -198,12 +198,7 @@ def apply_writes(
updated_channels: set[str] = set()
for chan, vals in pending_writes_by_channel.items():
if chan in channels:
try:
updated = channels[chan].update(vals)
except InvalidUpdateError as e:
raise InvalidUpdateError(
f"Invalid update for channel {chan} with values {vals}"
) from e
updated = channels[chan].update(vals)
if updated and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, channels[chan]
libs/langgraph/langgraph/pregel/debug.py
+12 -4
View File
@@ -1,5 +1,6 @@
import json
from collections import defaultdict
from dataclasses import asdict
from datetime import datetime, timezone
from pprint import pformat
from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypedDict, Union
@@ -25,6 +26,8 @@ class TaskPayload(TypedDict):
class TaskResultPayload(TypedDict):
id: str
name: str
error: Optional[str]
interrupts: list[dict]
result: list[tuple[str, Any]]
@@ -97,11 +100,14 @@ def map_debug_tasks(
def map_debug_task_results(
step: int,
tasks: list[PregelExecutableTask],
stream_channels_list: Sequence[str],
tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
stream_keys: Union[str, Sequence[str]],
) -> Iterator[DebugOutputTaskResult]:
stream_channels_list = (
[stream_keys] if isinstance(stream_keys, str) else stream_keys
)
ts = datetime.now(timezone.utc).isoformat()
for name, _, _, writes, config, _, _, _ in tasks:
for (name, _, _, _, config, _, _, _), writes in tasks:
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
@@ -116,7 +122,9 @@ def map_debug_task_results(
"payload": {
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
"name": name,
"error": next((w[1] for w in writes if w[0] == ERROR), None),
"result": [w for w in writes if w[0] in stream_channels_list],
"interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
},
}
@@ -150,7 +158,7 @@ def map_debug_checkpoint(
else {
"id": t.id,
"name": t.name,
"interrupts": t.interrupts,
"interrupts": tuple(asdict(i) for i in t.interrupts),
}
for t in tasks_w_writes(tasks, pending_writes)
],
libs/langgraph/langgraph/pregel/io.py
+11 -8
View File
@@ -3,7 +3,7 @@ from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union
from langchain_core.runnables.utils import AddableDict
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.constants import TAG_HIDDEN
from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
from langgraph.pregel.log import logger
from langgraph.pregel.types import PregelExecutableTask
@@ -95,19 +95,22 @@ class AddableUpdatesDict(AddableDict):
def map_output_updates(
output_channels: Union[str, Sequence[str]],
tasks: list[PregelExecutableTask],
tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
) -> Iterator[dict[str, Union[Any, dict[str, Any]]]]:
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
output_tasks = [
t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags")
(t, ww)
for t, ww in tasks
if (not t.config or TAG_HIDDEN not in t.config.get("tags"))
and all(k not in (ERROR, INTERRUPT) for k, _ in ww)
]
if not output_tasks:
return
if isinstance(output_channels, str):
updated = [
(task.name, value)
for task in output_tasks
for chan, value in task.writes
for task, writes in output_tasks
for chan, value in writes
if chan == output_channels
]
else:
@@ -116,10 +119,10 @@ def map_output_updates(
task.name,
{chan: value for chan, value in task.writes if chan in output_channels},
)
for task in output_tasks
if any(chan in output_channels for chan, _ in task.writes)
for task, writes in output_tasks
if any(chan in output_channels for chan, _ in writes)
]
grouped = {t.name: [] for t in output_tasks}
grouped = {t.name: [] for t, _ in output_tasks}
for node, value in updated:
grouped[node].append(value)
for node, value in grouped.items():
libs/langgraph/langgraph/pregel/loop.py
+35 -16
View File
@@ -41,7 +41,6 @@ from langgraph.constants import (
ERROR,
INPUT,
INTERRUPT,
Interrupt,
)
from langgraph.errors import EmptyInputError, GraphInterrupt
from langgraph.managed.base import (
@@ -56,7 +55,11 @@ from langgraph.pregel.algo import (
prepare_next_tasks,
should_interrupt,
)
from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
from langgraph.pregel.debug import (
map_debug_checkpoint,
map_debug_task_results,
map_debug_tasks,
)
from langgraph.pregel.executor import (
AsyncBackgroundExecutor,
BackgroundExecutor,
@@ -90,6 +93,7 @@ class PregelLoop:
nodes: Mapping[str, PregelNode]
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
output_keys: Union[str, Sequence[str]]
stream_keys: Union[str, Sequence[str]]
is_nested: bool
checkpointer_get_next_version: Callable[[Optional[V]], V]
@@ -138,6 +142,7 @@ class PregelLoop:
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]],
stream_keys: Union[str, Sequence[str]],
) -> None:
self.stream = deque()
self.input = input
@@ -147,6 +152,7 @@ class PregelLoop:
self.nodes = nodes
self.specs = specs
self.output_keys = output_keys
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:
@@ -155,6 +161,8 @@ class PregelLoop:
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:
return
self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
if self.checkpointer_put_writes is not None:
self.submit(
@@ -172,12 +180,22 @@ class PregelLoop:
writes,
task_id,
)
if task := next((t for t in self.tasks if t.id == task_id), None):
self.stream.extend(
("updates", v)
for v in map_output_updates(self.output_keys, [(task, writes)])
)
self.stream.extend(
("debug", v)
for v in map_debug_task_results(
self.step, [(task, writes)], self.stream_keys
)
)
def tick(
self,
*,
input_keys: Union[str, Sequence[str]],
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
interrupt_after: Sequence[str] = EMPTY_SEQ,
interrupt_before: Sequence[str] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
@@ -213,17 +231,18 @@ class PregelLoop:
self._put_checkpoint(
{
"source": "loop",
"writes": single(map_output_updates(self.output_keys, self.tasks)),
"writes": single(
map_output_updates(
self.output_keys, [(t, t.writes) for t in self.tasks]
)
),
}
)
# after execution, check if we should interrupt
if tasks := should_interrupt(self.checkpoint, interrupt_after, self.tasks):
if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
self.status = "interrupt_after"
interrupts = [(t.id, Interrupt("after")) for t in tasks]
for tid, interrupt in interrupts:
self.put_writes(tid, [(INTERRUPT, interrupt)])
if self.is_nested:
raise GraphInterrupt([i[1] for i in interrupts])
raise GraphInterrupt()
else:
return False
else:
@@ -256,7 +275,7 @@ class PregelLoop:
self.step - 1, # printing checkpoint for previous step
self.checkpoint_config,
self.channels,
stream_keys,
self.stream_keys,
self.checkpoint_metadata,
self.checkpoint,
self.tasks,
@@ -281,20 +300,16 @@ class PregelLoop:
if all(task.writes for task in self.tasks):
return self.tick(
input_keys=input_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
manager=manager,
)
# before execution, check if we should interrupt
if tasks := should_interrupt(self.checkpoint, interrupt_before, self.tasks):
if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
self.status = "interrupt_before"
interrupts = [(t.id, Interrupt("before")) for t in tasks]
for tid, interrupt in interrupts:
self.put_writes(tid, [(INTERRUPT, interrupt)])
if self.is_nested:
raise GraphInterrupt([i[1] for i in interrupts])
raise GraphInterrupt()
else:
return False
@@ -435,6 +450,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(
input,
@@ -444,6 +460,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
nodes=nodes,
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
)
self.stack = ExitStack()
if checkpointer:
@@ -522,6 +539,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(
input,
@@ -531,6 +549,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
nodes=nodes,
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
)
self.store = AsyncBatchedStore(self.store) if self.store else None
self.stack = AsyncExitStack()
libs/langgraph/langgraph/pregel/manager.py
+4 -2
View File
@@ -41,7 +41,7 @@ def ChannelsManager(
yield (
{
k: stack.enter_context(
v.from_checkpoint(checkpoint["channel_values"].get(k), config)
v.from_checkpoint_named(checkpoint["channel_values"].get(k), config)
)
for k, v in channel_specs.items()
},
@@ -95,7 +95,9 @@ async def AsyncChannelsManager(
# channels: enter each channel with checkpoint
{
k: await stack.enter_async_context(
v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
v.afrom_checkpoint_named(
checkpoint["channel_values"].get(k), config
)
)
for k, v in channel_specs.items()
},
libs/langgraph/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.2.9"
version = "0.2.11"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
license = "MIT"
libs/langgraph/tests/any_str.py
+35
View File
@@ -1,3 +1,6 @@
from typing import Any, Sequence
class AnyStr(str):
def __init__(self) -> None:
super().__init__()
@@ -9,6 +12,17 @@ class AnyStr(str):
return hash(str(self))
class AnyVersion:
def __init__(self) -> None:
super().__init__()
def __eq__(self, other: object) -> bool:
return isinstance(other, (str, int, float))
def __hash__(self) -> int:
return hash(str(self))
class ExceptionLike:
def __init__(self, exc: Exception) -> None:
self.exc = exc
@@ -22,3 +36,24 @@ class ExceptionLike:
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
def __eq__(self, value: object) -> bool:
return (
isinstance(value, Sequence)
and len(self.seq) == len(value)
and all(a in value for a in self.seq)
)
def __hash__(self) -> int:
return hash(frozenset(self.seq))
def __repr__(self) -> str:
return repr(self.seq)
libs/langgraph/tests/test_pregel.py
+216 -74
View File
@@ -72,7 +72,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, ExceptionLike
from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence
from tests.fake_tracer import FakeTracer
from tests.memory_assert import (
MemorySaverAssertCheckpointMetadata,
@@ -203,6 +203,21 @@ def test_graph_validation() -> None:
with pytest.raises(ValueError, match="Invalid reducer"):
StateGraph(BadReducerState)
def node_b(state: State) -> State:
return {"hello": "world"}
builder = StateGraph(State)
builder.add_node("a", node_b)
builder.add_node("b", node_b)
builder.add_node("c", node_b)
builder.set_entry_point("a")
builder.add_edge("a", "b")
builder.add_edge("a", "c")
graph = builder.compile()
with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
graph.invoke({"hello": "there"})
def test_checkpoint_errors() -> None:
class FaultyGetCheckpointer(MemorySaver):
@@ -774,7 +789,7 @@ def test_invoke_two_processes_in_out_interrupt(
),
]
# forking from any previous checkpoint w/out forking should do nothing
# re-running from any previous checkpoint w/out forking should do nothing
assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == []
assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == []
assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == []
@@ -1038,6 +1053,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
"name": "one",
"result": [("inbox", 3)],
"error": None,
"interrupts": [],
},
},
{
@@ -1048,6 +1065,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
"name": "two",
"result": [("output", 13)],
"error": None,
"interrupts": [],
},
},
{
@@ -1069,6 +1088,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
"name": "two",
"result": [("output", 4)],
"error": None,
"interrupts": [],
},
},
]
@@ -1196,6 +1217,21 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
# LastValue channels can only be updated once per iteration
app.invoke(2)
class State(TypedDict):
hello: str
def my_node(input: State) -> State:
return {"hello": "world"}
builder = StateGraph(State)
builder.add_node("one", my_node)
builder.add_node("two", my_node)
builder.set_conditional_entry_point(lambda _: ["one", "two"])
graph = builder.compile()
with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
graph.invoke({"hello": "there"}, debug=True)
def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -1380,6 +1416,147 @@ def test_pending_writes_resume(
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
assert graph.invoke(None, thread1) == {"value": 6}
# check all final checkpoints
checkpoints = [c for c in checkpointer.list(thread1)]
# we should have 3
assert len(checkpoints) == 3
# the last one not too interesting for this test
assert checkpoints[0] == CheckpointTuple(
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
checkpoint={
"v": 1,
"id": AnyStr(),
"ts": AnyStr(),
"current_tasks": {},
"pending_sends": [],
"versions_seen": {
"one": {
"start:one": AnyVersion(),
},
"two": {
"start:two": AnyVersion(),
},
"__input__": {},
"__start__": {
"__start__": AnyVersion(),
},
"__interrupt__": {
"value": AnyVersion(),
"__start__": AnyVersion(),
"start:one": AnyVersion(),
"start:two": AnyVersion(),
},
},
"channel_versions": {
"one": AnyVersion(),
"two": AnyVersion(),
"value": AnyVersion(),
"__start__": AnyVersion(),
"start:one": AnyVersion(),
"start:two": AnyVersion(),
},
"channel_values": {"one": "one", "two": "two", "value": 6},
},
metadata={
"step": 1,
"source": "loop",
"writes": {"one": {"value": 2}, "two": {"value": 3}},
},
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"],
}
},
pending_writes=[],
)
# the previous one we assert that pending writes contains both
# - original error
# - successful writes from resuming after preventing error
assert checkpoints[1] == CheckpointTuple(
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
checkpoint={
"v": 1,
"id": AnyStr(),
"ts": AnyStr(),
"current_tasks": {},
"pending_sends": [],
"versions_seen": {
"__input__": {},
"__start__": {
"__start__": AnyVersion(),
},
},
"channel_versions": {
"value": AnyVersion(),
"__start__": AnyVersion(),
"start:one": AnyVersion(),
"start:two": AnyVersion(),
},
"channel_values": {
"value": 1,
"start:one": "__start__",
"start:two": "__start__",
},
},
metadata={"step": 0, "source": "loop", "writes": None},
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
}
},
pending_writes=UnsortedSequence(
(AnyStr(), "one", "one"),
(AnyStr(), "value", 2),
(AnyStr(), "__error__", ExceptionLike(ConnectionError("I'm not good"))),
(AnyStr(), "two", "two"),
(AnyStr(), "value", 3),
),
)
assert checkpoints[2] == CheckpointTuple(
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
checkpoint={
"v": 1,
"id": AnyStr(),
"ts": AnyStr(),
"current_tasks": {},
"pending_sends": [],
"versions_seen": {"__input__": {}},
"channel_versions": {
"__start__": AnyVersion(),
},
"channel_values": {"__start__": {"value": 1}},
},
metadata={"step": -1, "source": "input", "writes": {"value": 1}},
parent_config=None,
pending_writes=UnsortedSequence(
(AnyStr(), "value", 1),
(AnyStr(), "start:one", "__start__"),
(AnyStr(), "start:two", "__start__"),
),
)
def test_cond_edge_after_send() -> None:
class Node:
@@ -1749,7 +1926,6 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
assert cleanup.call_count == 0
for i, chunk in enumerate(app.stream(2)):
assert setup.call_count == 1, "Expected setup to be called once"
assert cleanup.call_count == 0, "Expected cleanup to not be called yet"
if i == 0:
assert chunk == {"inbox": [3]}
elif i == 1:
@@ -2205,7 +2381,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
),
},
},
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2251,7 +2427,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
"input": "what is weather in sf",
},
},
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2412,7 +2588,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
),
},
},
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3013,7 +3189,7 @@ def test_conditional_state_graph(
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3053,7 +3229,7 @@ def test_conditional_state_graph(
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3162,7 +3338,7 @@ def test_conditional_state_graph(
values={
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "agent"),),
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3187,7 +3363,7 @@ def test_conditional_state_graph(
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3240,7 +3416,7 @@ def test_conditional_state_graph(
)
],
},
tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "agent"),),
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4842,13 +5018,7 @@ def test_message_graph(
id="ai1",
),
],
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4893,7 +5063,7 @@ def test_message_graph(
],
),
],
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4975,7 +5145,7 @@ def test_message_graph(
id="ai2",
),
],
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5567,13 +5737,7 @@ def test_root_graph(
id="ai1",
),
],
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5618,7 +5782,7 @@ def test_root_graph(
],
),
],
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5700,7 +5864,7 @@ def test_root_graph(
id="ai2",
),
],
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6022,6 +6186,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
"name": "rewrite_query",
"result": [("query", "query: what is weather in sf")],
"error": None,
"interrupts": [],
},
},
),
@@ -6076,6 +6242,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
"name": "retriever_two",
"result": [("docs", ["doc3", "doc4"])],
"error": None,
"interrupts": [],
},
},
),
@@ -6093,6 +6261,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
"name": "retriever_one",
"result": [("docs", ["doc1", "doc2"])],
"error": None,
"interrupts": [],
},
},
),
@@ -6132,6 +6302,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
"name": "qa",
"result": [("answer", "doc1,doc2,doc3,doc4")],
"error": None,
"interrupts": [],
},
},
),
@@ -6302,13 +6474,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_slow",
interrupts=(Interrupt("before"),),
),
),
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"],
@@ -6342,13 +6508,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
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"],
@@ -6382,13 +6542,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread3) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
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"],
@@ -6399,13 +6553,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
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",
interrupts=(Interrupt("before"),),
),
),
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"],
@@ -6549,6 +6697,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
"name": "prepare",
"result": [("my_key", " prepared")],
"error": None,
"interrupts": [],
},
},
{
@@ -6601,6 +6751,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
"error": None,
"interrupts": [],
},
},
{
@@ -6651,6 +6803,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
"name": "finish",
"result": [("my_key", " finished")],
"error": None,
"interrupts": [],
},
},
{
@@ -6700,13 +6854,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value prepared", "market": "DE"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_slow",
interrupts=(Interrupt("before"),),
),
),
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"],
@@ -6744,13 +6892,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value prepared", "market": "US"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
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"],
@@ -6797,7 +6939,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"my_key": "value prepared slow",
"market": "DE",
},
tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "finish"),),
next=("finish",),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6816,7 +6958,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"my_key": "value prepared slower",
"market": "DE",
},
tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "finish"),),
next=("finish",),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -7073,7 +7215,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
"query": "analyzed: query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
},
tasks=(PregelTask(AnyStr(), "qa", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "qa"),),
next=("qa",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
libs/langgraph/tests/test_pregel_async.py
+342 -120
View File
@@ -67,7 +67,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, ExceptionLike
from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence
from tests.fake_tracer import FakeTracer
from tests.memory_assert import (
MemorySaverAssertCheckpointMetadata,
@@ -212,7 +212,15 @@ async def test_node_cancellation_on_other_node_exception() -> None:
assert inner_task_cancelled
async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
)
async def test_dynamic_interrupt(
checkpointer_name: str, snapshot: SnapshotAssertion, request: pytest.FixtureRequest
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
@@ -250,51 +258,49 @@ async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
"market": "US",
}
async with AsyncSqliteSaver.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"):
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 ⛰️",
"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"},
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_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 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"},
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
"tool_two",
interrupts=(Interrupt("during", "Just because..."),),
),
config=tup.config,
created_at=tup.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,
)
),
config=tup.config,
created_at=tup.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,
)
# TODO use aget_state_history
@pytest.mark.parametrize(
@@ -1283,6 +1289,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
"name": "one",
"result": [("inbox", 3)],
"error": None,
"interrupts": [],
},
},
{
@@ -1293,6 +1301,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
"name": "two",
"result": [("output", 13)],
"error": None,
"interrupts": [],
},
},
{
@@ -1314,6 +1324,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
"name": "two",
"result": [("output", 4)],
"error": None,
"interrupts": [],
},
},
]
@@ -1608,6 +1620,10 @@ async def test_pending_writes_resume(
error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR)
assert error_write[0] != non_error_writes[0][0]
# TODO arguably this shouldn't even run the failed task again,
# and should require empty update_state (ie new checkpoint_id)
# in order to try again
# resume execution
with pytest.raises(ValueError, match="I'm not good"):
await graph.ainvoke(None, thread1)
@@ -1626,6 +1642,147 @@ async def test_pending_writes_resume(
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
assert await graph.ainvoke(None, thread1) == {"value": 6}
# check all final checkpoints
checkpoints = [c async for c in checkpointer.alist(thread1)]
# we should have 3
assert len(checkpoints) == 3
# the last one not too interesting for this test
assert checkpoints[0] == CheckpointTuple(
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
checkpoint={
"v": 1,
"id": AnyStr(),
"ts": AnyStr(),
"current_tasks": {},
"pending_sends": [],
"versions_seen": {
"one": {
"start:one": AnyVersion(),
},
"two": {
"start:two": AnyVersion(),
},
"__input__": {},
"__start__": {
"__start__": AnyVersion(),
},
"__interrupt__": {
"value": AnyVersion(),
"__start__": AnyVersion(),
"start:one": AnyVersion(),
"start:two": AnyVersion(),
},
},
"channel_versions": {
"one": AnyVersion(),
"two": AnyVersion(),
"value": AnyVersion(),
"__start__": AnyVersion(),
"start:one": AnyVersion(),
"start:two": AnyVersion(),
},
"channel_values": {"one": "one", "two": "two", "value": 6},
},
metadata={
"step": 1,
"source": "loop",
"writes": {"one": {"value": 2}, "two": {"value": 3}},
},
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"],
}
},
pending_writes=[],
)
# the previous one we assert that pending writes contains both
# - original error
# - successful writes from resuming after preventing error
assert checkpoints[1] == CheckpointTuple(
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
checkpoint={
"v": 1,
"id": AnyStr(),
"ts": AnyStr(),
"current_tasks": {},
"pending_sends": [],
"versions_seen": {
"__input__": {},
"__start__": {
"__start__": AnyVersion(),
},
},
"channel_versions": {
"value": AnyVersion(),
"__start__": AnyVersion(),
"start:one": AnyVersion(),
"start:two": AnyVersion(),
},
"channel_values": {
"value": 1,
"start:one": "__start__",
"start:two": "__start__",
},
},
metadata={"step": 0, "source": "loop", "writes": None},
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
}
},
pending_writes=UnsortedSequence(
(AnyStr(), "one", "one"),
(AnyStr(), "value", 2),
(AnyStr(), "__error__", ExceptionLike(ValueError("I'm not good"))),
(AnyStr(), "two", "two"),
(AnyStr(), "value", 3),
),
)
assert checkpoints[2] == CheckpointTuple(
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
checkpoint={
"v": 1,
"id": AnyStr(),
"ts": AnyStr(),
"current_tasks": {},
"pending_sends": [],
"versions_seen": {"__input__": {}},
"channel_versions": {
"__start__": AnyVersion(),
},
"channel_values": {"__start__": {"value": 1}},
},
metadata={"step": -1, "source": "input", "writes": {"value": 1}},
parent_config=None,
pending_writes=UnsortedSequence(
(AnyStr(), "value", 1),
(AnyStr(), "start:one", "__start__"),
(AnyStr(), "start:two", "__start__"),
),
)
async def test_cond_edge_after_send() -> None:
class Node:
@@ -1983,7 +2140,6 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
assert setup_sync.call_count == 0, "Sync context manager should not be used"
assert cleanup_sync.call_count == 0, "Sync context manager should not be used"
assert setup_async.call_count == 1, "Expected setup to be called once"
assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet"
if i == 0:
assert chunk == {"inbox": [3]}
elif i == 1:
@@ -2516,13 +2672,7 @@ async def test_conditional_graph() -> None:
),
},
},
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2572,13 +2722,7 @@ async def test_conditional_graph() -> None:
"input": "what is weather in sf",
},
},
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2750,13 +2894,7 @@ async def test_conditional_graph() -> None:
),
},
},
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -3324,13 +3462,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
),
"intermediate_steps": [],
},
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -3374,7 +3506,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
tasks=(PregelTask(AnyStr(), "tools"),),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -4749,6 +4881,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
"name": "rewrite_query",
"result": [("query", "query: what is weather in sf")],
"error": None,
"interrupts": [],
},
},
),
@@ -4803,6 +4937,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
"name": "retriever_two",
"result": [("docs", ["doc3", "doc4"])],
"error": None,
"interrupts": [],
},
},
),
@@ -4820,6 +4956,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
"name": "retriever_one",
"result": [("docs", ["doc1", "doc2"])],
"error": None,
"interrupts": [],
},
},
),
@@ -4859,6 +4997,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
"name": "qa",
"result": [("answer", "doc1,doc2,doc3,doc4")],
"error": None,
"interrupts": [],
},
},
),
@@ -4949,13 +5089,7 @@ async def test_start_branch_then() -> None:
]
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value", "market": "DE"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_slow",
interrupts=(Interrupt("before"),),
),
),
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[
@@ -4997,13 +5131,7 @@ async def test_start_branch_then() -> None:
}
assert await tool_two.aget_state(thread2) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
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[
@@ -5045,13 +5173,7 @@ async def test_start_branch_then() -> None:
}
assert await tool_two.aget_state(thread3) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
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[
@@ -5066,13 +5188,7 @@ async def test_start_branch_then() -> None:
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",
interrupts=(Interrupt("before"),),
),
),
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[
@@ -5223,6 +5339,8 @@ async def test_branch_then() -> None:
"id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
"name": "prepare",
"result": [("my_key", " prepared")],
"error": None,
"interrupts": [],
},
},
{
@@ -5275,6 +5393,8 @@ async def test_branch_then() -> None:
"id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
"error": None,
"interrupts": [],
},
},
{
@@ -5325,6 +5445,8 @@ async def test_branch_then() -> None:
"id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
"name": "finish",
"result": [("my_key", " finished")],
"error": None,
"interrupts": [],
},
},
{
@@ -5368,19 +5490,125 @@ async def test_branch_then() -> None:
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 [
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(),
},
},
"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": ()}],
},
},
{
"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(),
},
},
"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",
interrupts=(Interrupt("before"),),
),
),
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[
@@ -5426,13 +5654,7 @@ async def test_branch_then() -> None:
}
assert await tool_two.aget_state(thread2) == StateSnapshot(
values={"my_key": "value prepared", "market": "US"},
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
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[
libs/langgraph/tests/test_state.py
+41 -1
View File
@@ -2,10 +2,11 @@ from typing import Annotated as Annotated2
from typing import Any
import pytest
from langchain_core.runnables import RunnableConfig
from pydantic.v1 import BaseModel
from typing_extensions import Annotated, TypedDict
from langgraph.graph.state import _warn_invalid_state_schema
from langgraph.graph.state import StateGraph, _warn_invalid_state_schema
class State(BaseModel):
@@ -46,3 +47,42 @@ def test_doesnt_warn_valid_schema(schema: Any):
# Assert the function does not raise a warning
with pytest.warns(None):
_warn_invalid_state_schema(schema)
def test_state_schema_with_type_hint():
class InputState(TypedDict):
question: str
class OutputState(TypedDict):
input_state: InputState
def complete_hint(state: InputState) -> OutputState:
return {"input_state": state}
def miss_first_hint(state, config: RunnableConfig) -> OutputState:
return {"input_state": state}
def only_return_hint(state, config) -> OutputState:
return {"input_state": state}
def miss_all_hint(state, config):
return {"input_state": state}
graph = StateGraph(input=InputState, output=OutputState)
actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
for action in actions:
graph.add_node(action)
graph.set_entry_point(actions[0].__name__)
for i in range(len(actions) - 1):
graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
graph.set_finish_point(actions[-1].__name__)
graph = graph.compile()
input_state = InputState(question="Hello World!")
output_state = OutputState(input_state=input_state)
for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
node_name = actions[i].__name__
assert c[node_name] == output_state