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

Merge branch 'main' into patch-1

Browse Source
This commit is contained in:
LEE KYU WON
2024-07-21 18:12:38 +09:00
committed by GitHub
parent 58af2c7d9f 610b6cc78c
commit e8b05faff4
47 changed files with 2163 additions and 1663 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/.gitignore
+1
View File
@@ -1,2 +1,3 @@
*.ipynb
site/
docs/tutorials/**/*.png
docs/_scripts/copy_notebooks.py
+1
View File
@@ -56,6 +56,7 @@ _MANUAL = {
"human_in_the_loop/time-travel.ipynb",
"human_in_the_loop/edit-graph-state.ipynb",
"human_in_the_loop/wait-user-input.ipynb",
"node-retries.ipynb",
],
"tutorials": [
"introduction.ipynb",
docs/docs/cloud/how-tos/cloud_examples/img/webhook_results.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 322 KiB

docs/docs/cloud/index.md
+6
View File
@@ -1,5 +1,11 @@
# LangGraph Cloud (beta)
!!! tip
- LangGraph is an MIT-licensed open-source library, which we are committed to maintaining and growing for the community.
- LangGraph Cloud is an optional managed hosting service for LangGraph, which provides additional features geared towards production deployments.
- We are actively contributing improvements back to LangGraph informed by our work on LangGraph Cloud.
- You can always deploy LangGraph applications on your own infrastructure using the open-source LangGraph project.
!!! danger "Important"
LangGraph Cloud is a closed source, paid product in an invite-only stage. We are currently focused on providing high bandwidth support to make our select early customers successful. If you are interested in applying for access, please fill out [this form](https://www.langchain.com/langgraph-cloud-beta).
docs/docs/how-tos/index.md
+1
View File
@@ -68,6 +68,7 @@ These guides show how to use different streaming modes.
- [How to add runtime configuration to your graph](configuration.ipynb)
- [How to use a Pydantic model as your state](state-model.ipynb)
- [How to use a context object in state](state-context-key.ipynb)
- [How to add node retries](node-retries.ipynb)
## Prebuilt ReAct Agent
docs/docs/reference/checkpoints.md
-4
View File
@@ -18,12 +18,10 @@ You can [compile][langgraph.graph.MessageGraph.compile] any LangGraph workflow w
### BaseCheckpointSaver
::: langgraph.checkpoint.base.BaseCheckpointSaver
handler: python
### SerializerProtocol
::: langgraph.checkpoint.SerializerProtocol
handler: python
## Implementations
@@ -32,12 +30,10 @@ LangGraph also natively provides the following checkpoint implementations.
### MemorySaver
::: langgraph.checkpoint.memory.MemorySaver
handler: python
### AsyncSqliteSaver
::: langgraph.checkpoint.aiosqlite.AsyncSqliteSaver
handler: python
### SqliteSaver
docs/docs/reference/graphs.md
+5 -1
View File
@@ -65,4 +65,8 @@ builder.add_conditional_edges("my_node", my_condition)
## Send
::: langgraph.constants.Send
::: langgraph.constants.Send
## RetryPolicy
::: langgraph.pregel.types.RetryPolicy
docs/docs/reference/prebuilt.md
+10 -1
View File
@@ -55,4 +55,13 @@ from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ValidationNode
```
::: langgraph.prebuilt.ValidationNode
::: langgraph.prebuilt.ValidationNode
## InjectedState
```python
from langgraph.prebuilt import InjectedState
```
::: langgraph.prebuilt.InjectedState
handler: python
docs/mkdocs.yml
+1
View File
@@ -163,6 +163,7 @@ nav:
- Add runtime configuration: how-tos/configuration.ipynb
- Use Pydantic model as state: how-tos/state-model.ipynb
- Use a context object in state: how-tos/state-context-key.ipynb
- Add node retries: how-tos/node-retries.ipynb
- Prebuilt ReAct Agent:
- Create a ReAct agent: how-tos/create-react-agent.ipynb
- Add memory to a ReAct agent: how-tos/create-react-agent-memory.ipynb
examples/chatbots/customer-support.ipynb
+24 -262
View File
@@ -32,10 +32,7 @@
"scrolled": true
},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph langchain-community langchain-openai scikit-learn"
]
"source": ["%%capture --no-stderr\n%pip install -U langgraph langchain-community langchain-openai scikit-learn"]
},
{
"cell_type": "markdown",
@@ -51,15 +48,7 @@
"id": "3d1ef253-6b0c-4481-868c-e1fe84f2c8ff",
"metadata": {},
"outputs": [],
"source": [
"import requests\n",
"\n",
"url = \"https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db\"\n",
"response = requests.get(url)\n",
"\n",
"with open(\"Chinook.db\", \"wb\") as file:\n",
" file.write(response.content)"
]
"source": ["import requests\n\nurl = \"https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db\"\nresponse = requests.get(url)\n\nwith open(\"Chinook.db\", \"wb\") as file:\n file.write(response.content)"]
},
{
"cell_type": "code",
@@ -88,12 +77,7 @@
"output_type": "execute_result"
}
],
"source": [
"from langchain_community.utilities import SQLDatabase\n",
"\n",
"db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")\n",
"db.get_usable_table_names()"
]
"source": ["from langchain_community.utilities import SQLDatabase\n\ndb = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")\ndb.get_usable_table_names()"]
},
{
"cell_type": "markdown",
@@ -112,11 +96,7 @@
"id": "d9ea4e80-30e6-4d46-b480-35f0be2fb055",
"metadata": {},
"outputs": [],
"source": [
"from langchain_openai import ChatOpenAI\n",
"\n",
"model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4o\")"
]
"source": ["from langchain_openai import ChatOpenAI\n\nmodel = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4o\")"]
},
{
"cell_type": "markdown",
@@ -138,9 +118,7 @@
"id": "ea958e9f-ab1f-49b5-bd85-16332055297c",
"metadata": {},
"outputs": [],
"source": [
"from langchain_core.messages import HumanMessage, SystemMessage"
]
"source": ["from langchain_core.messages import HumanMessage, SystemMessage"]
},
{
"cell_type": "markdown",
@@ -159,12 +137,7 @@
"id": "975b039a",
"metadata": {},
"outputs": [],
"source": [
"# This tool is given to the agent to look up information about a customer\n",
"def get_customer_info(customer_id: int):\n",
" \"\"\"Look up customer info given their ID. ALWAYS make sure you have the customer ID before invoking this.\"\"\"\n",
" return db.run(f\"SELECT * FROM Customer WHERE CustomerID = {customer_id};\")"
]
"source": ["# This tool is given to the agent to look up information about a customer\ndef get_customer_info(customer_id: int):\n \"\"\"Look up customer info given their ID. ALWAYS make sure you have the customer ID before invoking this.\"\"\"\n return db.run(f\"SELECT * FROM Customer WHERE CustomerID = {customer_id};\")"]
},
{
"cell_type": "code",
@@ -172,20 +145,7 @@
"id": "1d5fa446",
"metadata": {},
"outputs": [],
"source": [
"customer_prompt = \"\"\"Your job is to help a user update their profile.\n",
"\n",
"You only have certain tools you can use. These tools require specific input. If you don't know the required input, then ask the user for it.\n",
"\n",
"If you are unable to help the user, you can \"\"\"\n",
"\n",
"\n",
"def get_customer_messages(messages):\n",
" return [SystemMessage(content=customer_prompt)] + messages\n",
"\n",
"\n",
"customer_chain = get_customer_messages | model.bind_tools([get_customer_info])"
]
"source": ["customer_prompt = \"\"\"Your job is to help a user update their profile.\n\nYou only have certain tools you can use. These tools require specific input. If you don't know the required input, then ask the user for it.\n\nIf you are unable to help the user, you can \"\"\"\n\n\ndef get_customer_messages(messages):\n return [SystemMessage(content=customer_prompt)] + messages\n\n\ncustomer_chain = get_customer_messages | model.bind_tools([get_customer_info])"]
},
{
"cell_type": "markdown",
@@ -206,19 +166,7 @@
"id": "a8604a3b-b484-4b2b-a914-4236cb98c524",
"metadata": {},
"outputs": [],
"source": [
"from langchain_community.vectorstores import SKLearnVectorStore\n",
"from langchain_openai import OpenAIEmbeddings\n",
"\n",
"artists = db._execute(\"select * from Artist\")\n",
"songs = db._execute(\"select * from Track\")\n",
"artist_retriever = SKLearnVectorStore.from_texts(\n",
" [a[\"Name\"] for a in artists], OpenAIEmbeddings(), metadatas=artists\n",
").as_retriever()\n",
"song_retriever = SKLearnVectorStore.from_texts(\n",
" [a[\"Name\"] for a in songs], OpenAIEmbeddings(), metadatas=songs\n",
").as_retriever()"
]
"source": ["from langchain_community.vectorstores import SKLearnVectorStore\nfrom langchain_openai import OpenAIEmbeddings\n\nartists = db._execute(\"select * from Artist\")\nsongs = db._execute(\"select * from Track\")\nartist_retriever = SKLearnVectorStore.from_texts(\n [a[\"Name\"] for a in artists], OpenAIEmbeddings(), metadatas=artists\n).as_retriever()\nsong_retriever = SKLearnVectorStore.from_texts(\n [a[\"Name\"] for a in songs], OpenAIEmbeddings(), metadatas=songs\n).as_retriever()"]
},
{
"cell_type": "markdown",
@@ -234,16 +182,7 @@
"id": "0a2a2b74",
"metadata": {},
"outputs": [],
"source": [
"def get_albums_by_artist(artist):\n",
" \"\"\"Get albums by an artist (or similar artists).\"\"\"\n",
" docs = artist_retriever.get_relevant_documents(artist)\n",
" artist_ids = \", \".join([str(d.metadata[\"ArtistId\"]) for d in docs])\n",
" return db.run(\n",
" f\"SELECT Title, Name FROM Album LEFT JOIN Artist ON Album.ArtistId = Artist.ArtistId WHERE Album.ArtistId in ({artist_ids});\",\n",
" include_columns=True,\n",
" )"
]
"source": ["def get_albums_by_artist(artist):\n \"\"\"Get albums by an artist (or similar artists).\"\"\"\n docs = artist_retriever.get_relevant_documents(artist)\n artist_ids = \", \".join([str(d.metadata[\"ArtistId\"]) for d in docs])\n return db.run(\n f\"SELECT Title, Name FROM Album LEFT JOIN Artist ON Album.ArtistId = Artist.ArtistId WHERE Album.ArtistId in ({artist_ids});\",\n include_columns=True,\n )"]
},
{
"cell_type": "markdown",
@@ -259,16 +198,7 @@
"id": "da533f50",
"metadata": {},
"outputs": [],
"source": [
"def get_tracks_by_artist(artist):\n",
" \"\"\"Get songs by an artist (or similar artists).\"\"\"\n",
" docs = artist_retriever.invoke(artist)\n",
" artist_ids = \", \".join([str(d.metadata[\"ArtistId\"]) for d in docs])\n",
" return db.run(\n",
" f\"SELECT Track.Name as SongName, Artist.Name as ArtistName FROM Album LEFT JOIN Artist ON Album.ArtistId = Artist.ArtistId LEFT JOIN Track ON Track.AlbumId = Album.AlbumId WHERE Album.ArtistId in ({artist_ids});\",\n",
" include_columns=True,\n",
" )"
]
"source": ["def get_tracks_by_artist(artist):\n \"\"\"Get songs by an artist (or similar artists).\"\"\"\n docs = artist_retriever.invoke(artist)\n artist_ids = \", \".join([str(d.metadata[\"ArtistId\"]) for d in docs])\n return db.run(\n f\"SELECT Track.Name as SongName, Artist.Name as ArtistName FROM Album LEFT JOIN Artist ON Album.ArtistId = Artist.ArtistId LEFT JOIN Track ON Track.AlbumId = Album.AlbumId WHERE Album.ArtistId in ({artist_ids});\",\n include_columns=True,\n )"]
},
{
"cell_type": "markdown",
@@ -284,11 +214,7 @@
"id": "b3c07010",
"metadata": {},
"outputs": [],
"source": [
"def check_for_songs(song_title):\n",
" \"\"\"Check if a song exists by its name.\"\"\"\n",
" return song_retriever.invoke(song_title)"
]
"source": ["def check_for_songs(song_title):\n \"\"\"Check if a song exists by its name.\"\"\"\n return song_retriever.invoke(song_title)"]
},
{
"cell_type": "markdown",
@@ -304,23 +230,7 @@
"id": "72a14d5c",
"metadata": {},
"outputs": [],
"source": [
"song_system_message = \"\"\"Your job is to help a customer find any songs they are looking for. \n",
"\n",
"You only have certain tools you can use. If a customer asks you to look something up that you don't know how, politely tell them what you can help with.\n",
"\n",
"When looking up artists and songs, sometimes the artist/song will not be found. In that case, the tools will return information \\\n",
"on similar songs and artists. This is intentional, it is not the tool messing up.\"\"\"\n",
"\n",
"\n",
"def get_song_messages(messages):\n",
" return [SystemMessage(content=song_system_message)] + messages\n",
"\n",
"\n",
"song_recc_chain = get_song_messages | model.bind_tools(\n",
" [get_albums_by_artist, get_tracks_by_artist, check_for_songs]\n",
")"
]
"source": ["song_system_message = \"\"\"Your job is to help a customer find any songs they are looking for. \n\nYou only have certain tools you can use. If a customer asks you to look something up that you don't know how, politely tell them what you can help with.\n\nWhen looking up artists and songs, sometimes the artist/song will not be found. In that case, the tools will return information \\\non similar songs and artists. This is intentional, it is not the tool messing up.\"\"\"\n\n\ndef get_song_messages(messages):\n return [SystemMessage(content=song_system_message)] + messages\n\n\nsong_recc_chain = get_song_messages | model.bind_tools(\n [get_albums_by_artist, get_tracks_by_artist, check_for_songs]\n)"]
},
{
"cell_type": "code",
@@ -339,10 +249,7 @@
"output_type": "execute_result"
}
],
"source": [
"msgs = [HumanMessage(content=\"hi! can you help me find songs by amy whinehouse?\")]\n",
"song_recc_chain.invoke(msgs)"
]
"source": ["msgs = [HumanMessage(content=\"hi! can you help me find songs by amy whinehouse?\")]\nsong_recc_chain.invoke(msgs)"]
},
{
"cell_type": "markdown",
@@ -360,32 +267,7 @@
"id": "73e74268",
"metadata": {},
"outputs": [],
"source": [
"from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
"from langchain_core.pydantic_v1 import BaseModel, Field\n",
"\n",
"\n",
"class Router(BaseModel):\n",
" \"\"\"Call this if you are able to route the user to the appropriate representative.\"\"\"\n",
"\n",
" choice: str = Field(description=\"should be one of: music, customer\")\n",
"\n",
"\n",
"system_message = \"\"\"Your job is to help as a customer service representative for a music store.\n",
"\n",
"You should interact politely with customers to try to figure out how you can help. You can help in a few ways:\n",
"\n",
"- Updating user information: if a customer wants to update the information in the user database. Call the router with `customer`\n",
"- Recommending music: if a customer wants to find some music or information about music. Call the router with `music`\n",
"\n",
"If the user is asking or wants to ask about updating or accessing their information, send them to that route.\n",
"If the user is asking or wants to ask about music, send them to that route.\n",
"Otherwise, respond.\"\"\"\n",
"\n",
"\n",
"def get_messages(messages):\n",
" return [SystemMessage(content=system_message)] + messages"
]
"source": ["from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\nfrom langchain_core.pydantic_v1 import BaseModel, Field\n\n\nclass Router(BaseModel):\n \"\"\"Call this if you are able to route the user to the appropriate representative.\"\"\"\n\n choice: str = Field(description=\"should be one of: music, customer\")\n\n\nsystem_message = \"\"\"Your job is to help as a customer service representative for a music store.\n\nYou should interact politely with customers to try to figure out how you can help. You can help in a few ways:\n\n- Updating user information: if a customer wants to update the information in the user database. Call the router with `customer`\n- Recommending music: if a customer wants to find some music or information about music. Call the router with `music`\n\nIf the user is asking or wants to ask about updating or accessing their information, send them to that route.\nIf the user is asking or wants to ask about music, send them to that route.\nOtherwise, respond.\"\"\"\n\n\ndef get_messages(messages):\n return [SystemMessage(content=system_message)] + messages"]
},
{
"cell_type": "code",
@@ -393,9 +275,7 @@
"id": "ddf27314",
"metadata": {},
"outputs": [],
"source": [
"chain = get_messages | model.bind_tools([Router])"
]
"source": ["chain = get_messages | model.bind_tools([Router])"]
},
{
"cell_type": "code",
@@ -414,10 +294,7 @@
"output_type": "execute_result"
}
],
"source": [
"msgs = [HumanMessage(content=\"hi! can you help me find a good song?\")]\n",
"chain.invoke(msgs)"
]
"source": ["msgs = [HumanMessage(content=\"hi! can you help me find a good song?\")]\nchain.invoke(msgs)"]
},
{
"cell_type": "code",
@@ -436,10 +313,7 @@
"output_type": "execute_result"
}
],
"source": [
"msgs = [HumanMessage(content=\"hi! what's the email you have for me?\")]\n",
"chain.invoke(msgs)"
]
"source": ["msgs = [HumanMessage(content=\"hi! what's the email you have for me?\")]\nchain.invoke(msgs)"]
},
{
"cell_type": "code",
@@ -447,15 +321,7 @@
"id": "bd6ddd8b-7500-46a7-811d-3bcb937bda51",
"metadata": {},
"outputs": [],
"source": [
"from langchain_core.messages import AIMessage\n",
"\n",
"\n",
"def add_name(message, name):\n",
" _dict = message.dict()\n",
" _dict[\"name\"] = name\n",
" return AIMessage(**_dict)"
]
"source": ["from langchain_core.messages import AIMessage\n\n\ndef add_name(message, name):\n _dict = message.dict()\n _dict[\"name\"] = name\n return AIMessage(**_dict)"]
},
{
"cell_type": "code",
@@ -463,45 +329,7 @@
"id": "27494de5-8345-4c23-bc0e-81e0dd5d47d8",
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"\n",
"from langgraph.graph import END\n",
"\n",
"\n",
"def _get_last_ai_message(messages):\n",
" for m in messages[::-1]:\n",
" if isinstance(m, AIMessage):\n",
" return m\n",
" return None\n",
"\n",
"\n",
"def _is_tool_call(msg):\n",
" return hasattr(msg, \"additional_kwargs\") and \"tool_calls\" in msg.additional_kwargs\n",
"\n",
"\n",
"def _route(messages):\n",
" last_message = messages[-1]\n",
" if isinstance(last_message, AIMessage):\n",
" if not last_message.tool_calls:\n",
" return END\n",
" else:\n",
" if last_message.name == \"general\":\n",
" if len(last_message.tool_calls) > 1:\n",
" raise ValueError(\"Too many tools\")\n",
" return last_message.tool_calls[0][\"args\"][\"choice\"]\n",
" else:\n",
" return \"tools\"\n",
" last_m = _get_last_ai_message(messages)\n",
" if last_m is None:\n",
" return \"general\"\n",
" if last_m.name == \"music\":\n",
" return \"music\"\n",
" elif last_m.name == \"customer\":\n",
" return \"customer\"\n",
" else:\n",
" return \"general\""
]
"source": ["import json\n\nfrom langgraph.graph import END, START\n\n\ndef _get_last_ai_message(messages):\n for m in messages[::-1]:\n if isinstance(m, AIMessage):\n return m\n return None\n\n\ndef _is_tool_call(msg):\n return hasattr(msg, \"additional_kwargs\") and \"tool_calls\" in msg.additional_kwargs\n\n\ndef _route(messages):\n last_message = messages[-1]\n if isinstance(last_message, AIMessage):\n if not last_message.tool_calls:\n return END\n else:\n if last_message.name == \"general\":\n if len(last_message.tool_calls) > 1:\n raise ValueError(\"Too many tools\")\n return last_message.tool_calls[0][\"args\"][\"choice\"]\n else:\n return \"tools\"\n last_m = _get_last_ai_message(messages)\n if last_m is None:\n return \"general\"\n if last_m.name == \"music\":\n return \"music\"\n elif last_m.name == \"customer\":\n return \"customer\"\n else:\n return \"general\""]
},
{
"cell_type": "code",
@@ -509,12 +337,7 @@
"id": "8aec704a-46fe-4fb3-bdee-11c3bbffc370",
"metadata": {},
"outputs": [],
"source": [
"from langgraph.prebuilt import ToolNode\n",
"\n",
"tools = [get_albums_by_artist, get_tracks_by_artist, check_for_songs, get_customer_info]\n",
"tool_node = ToolNode(tools)"
]
"source": ["from langgraph.prebuilt import ToolNode\n\ntools = [get_albums_by_artist, get_tracks_by_artist, check_for_songs, get_customer_info]\ntool_node = ToolNode(tools)"]
},
{
"cell_type": "code",
@@ -522,16 +345,7 @@
"id": "4d5b75c6-73e0-4922-a765-a15be63f869e",
"metadata": {},
"outputs": [],
"source": [
"def _filter_out_routes(messages):\n",
" ms = []\n",
" for m in messages:\n",
" if _is_tool_call(m):\n",
" if m.name == \"general\":\n",
" continue\n",
" ms.append(m)\n",
" return ms"
]
"source": ["def _filter_out_routes(messages):\n ms = []\n for m in messages:\n if _is_tool_call(m):\n if m.name == \"general\":\n continue\n ms.append(m)\n return ms"]
},
{
"cell_type": "code",
@@ -539,13 +353,7 @@
"id": "fd4dbf98-dbb3-411a-bad6-2bb334072aaf",
"metadata": {},
"outputs": [],
"source": [
"from functools import partial\n",
"\n",
"general_node = _filter_out_routes | chain | partial(add_name, name=\"general\")\n",
"music_node = _filter_out_routes | song_recc_chain | partial(add_name, name=\"music\")\n",
"customer_node = _filter_out_routes | customer_chain | partial(add_name, name=\"customer\")"
]
"source": ["from functools import partial\n\ngeneral_node = _filter_out_routes | chain | partial(add_name, name=\"general\")\nmusic_node = _filter_out_routes | song_recc_chain | partial(add_name, name=\"music\")\ncustomer_node = _filter_out_routes | customer_chain | partial(add_name, name=\"customer\")"]
},
{
"cell_type": "code",
@@ -553,33 +361,7 @@
"id": "dcade924",
"metadata": {},
"outputs": [],
"source": [
"from langgraph.checkpoint.sqlite import SqliteSaver\n",
"\n",
"from langgraph.graph import MessageGraph\n",
"\n",
"memory = SqliteSaver.from_conn_string(\":memory:\")\n",
"graph = MessageGraph()\n",
"nodes = {\n",
" \"general\": \"general\",\n",
" \"music\": \"music\",\n",
" END: END,\n",
" \"tools\": \"tools\",\n",
" \"customer\": \"customer\",\n",
"}\n",
"# Define a new graph\n",
"workflow = MessageGraph()\n",
"workflow.add_node(\"general\", general_node)\n",
"workflow.add_node(\"music\", music_node)\n",
"workflow.add_node(\"customer\", customer_node)\n",
"workflow.add_node(\"tools\", tool_node)\n",
"workflow.add_conditional_edges(\"general\", _route, nodes)\n",
"workflow.add_conditional_edges(\"tools\", _route, nodes)\n",
"workflow.add_conditional_edges(\"music\", _route, nodes)\n",
"workflow.add_conditional_edges(\"customer\", _route, nodes)\n",
"workflow.set_conditional_entry_point(_route, nodes)\n",
"graph = workflow.compile()"
]
"source": ["from langgraph.checkpoint.sqlite import SqliteSaver\n\nfrom langgraph.graph import MessageGraph\n\nmemory = SqliteSaver.from_conn_string(\":memory:\")\ngraph = MessageGraph()\nnodes = {\n \"general\": \"general\",\n \"music\": \"music\",\n END: END,\n \"tools\": \"tools\",\n \"customer\": \"customer\",\n}\n# Define a new graph\nworkflow = MessageGraph()\nworkflow.add_node(\"general\", general_node)\nworkflow.add_node(\"music\", music_node)\nworkflow.add_node(\"customer\", customer_node)\nworkflow.add_node(\"tools\", tool_node)\nworkflow.add_conditional_edges(\"general\", _route, nodes)\nworkflow.add_conditional_edges(\"tools\", _route, nodes)\nworkflow.add_conditional_edges(\"music\", _route, nodes)\nworkflow.add_conditional_edges(\"customer\", _route, nodes)\nworkflow.add_conditional_edges(START, _route, nodes)\ngraph = workflow.compile()"]
},
{
"cell_type": "code",
@@ -715,27 +497,7 @@
]
}
],
"source": [
"import uuid\n",
"\n",
"from langchain_core.messages import HumanMessage\n",
"\n",
"from langgraph.graph.graph import START\n",
"\n",
"history = []\n",
"while True:\n",
" user = input(\"User (q/Q to quit): \")\n",
" if user in {\"q\", \"Q\"}:\n",
" print(\"AI: Byebye\")\n",
" break\n",
" history.append(HumanMessage(content=user))\n",
" async for output in graph.astream(history):\n",
" for key, value in output.items():\n",
" print(f\"Output from node '{key}':\")\n",
" print(\"---\")\n",
" print(value)\n",
" print(\"\\n---\\n\")"
]
"source": ["import uuid\n\nfrom langchain_core.messages import HumanMessage\n\nfrom langgraph.graph.graph import START\n\nhistory = []\nwhile True:\n user = input(\"User (q/Q to quit): \")\n if user in {\"q\", \"Q\"}:\n print(\"AI: Byebye\")\n break\n history.append(HumanMessage(content=user))\n async for output in graph.astream(history):\n for key, value in output.items():\n print(f\"Output from node '{key}':\")\n print(\"---\")\n print(value)\n print(\"\\n---\\n\")"]
}
],
"metadata": {
examples/human_in_the_loop/breakpoints.ipynb
+7 -3
View File
@@ -133,21 +133,26 @@
"from langgraph.checkpoint.memory import MemorySaver\n",
"from IPython.display import Image, display\n",
"\n",
"\n",
"class State(TypedDict):\n",
" input: str\n",
"\n",
"\n",
"def step_1(state):\n",
" print(\"---Step 1---\")\n",
" pass\n",
"\n",
"\n",
"def step_2(state):\n",
" print(\"---Step 2---\")\n",
" pass\n",
"\n",
"\n",
"def step_3(state):\n",
" print(\"---Step 3---\")\n",
" pass\n",
"\n",
"\n",
"builder = StateGraph(State)\n",
"builder.add_node(\"step_1\", step_1)\n",
"builder.add_node(\"step_2\", step_2)\n",
@@ -160,7 +165,7 @@
"# Set up memory\n",
"memory = MemorySaver()\n",
"\n",
"# Add \n",
"# Add\n",
"graph = builder.compile(checkpointer=memory, interrupt_before=[\"step_3\"])\n",
"\n",
"# View\n",
@@ -222,8 +227,7 @@
"\n",
"user_approval = input(\"Do you want to go to Step 3? (yes/no): \")\n",
"\n",
"if user_approval.lower() == 'yes':\n",
" \n",
"if user_approval.lower() == \"yes\":\n",
" # If approved, continue the graph execution\n",
" for event in graph.stream(None, thread, stream_mode=\"values\"):\n",
" print(event)\n",
examples/human_in_the_loop/edit-graph-state.ipynb
+6 -1
View File
@@ -135,21 +135,26 @@
"from langgraph.checkpoint.memory import MemorySaver\n",
"from IPython.display import Image, display\n",
"\n",
"\n",
"class State(TypedDict):\n",
" input: str\n",
"\n",
"\n",
"def step_1(state):\n",
" print(\"---Step 1---\")\n",
" pass\n",
"\n",
"\n",
"def step_2(state):\n",
" print(\"---Step 2---\")\n",
" pass\n",
"\n",
"\n",
"def step_3(state):\n",
" print(\"---Step 3---\")\n",
" pass\n",
"\n",
"\n",
"builder = StateGraph(State)\n",
"builder.add_node(\"step_1\", step_1)\n",
"builder.add_node(\"step_2\", step_2)\n",
@@ -162,7 +167,7 @@
"# Set up memory\n",
"memory = MemorySaver()\n",
"\n",
"# Add \n",
"# Add\n",
"graph = builder.compile(checkpointer=memory, interrupt_before=[\"step_2\"])\n",
"\n",
"# View\n",
examples/human_in_the_loop/wait-user-input.ipynb
+8 -3
View File
@@ -137,25 +137,30 @@
"from langgraph.checkpoint.memory import MemorySaver\n",
"from IPython.display import Image, display\n",
"\n",
"\n",
"class State(TypedDict):\n",
" input: str\n",
" user_feedback: str\n",
"\n",
"\n",
"def step_1(state):\n",
" print(\"---Step 1---\")\n",
" pass\n",
"\n",
"\n",
"def human_feedback(state):\n",
" print(\"---human_feedback---\")\n",
" pass\n",
"\n",
"\n",
"def step_3(state):\n",
" print(\"---Step 3---\")\n",
" pass\n",
"\n",
"\n",
"builder = StateGraph(State)\n",
"builder.add_node(\"step_1\", step_1)\n",
"builder.add_node(\"human_feedback\", step_2)\n",
"builder.add_node(\"human_feedback\", human_feedback)\n",
"builder.add_node(\"step_3\", step_3)\n",
"builder.add_edge(START, \"step_1\")\n",
"builder.add_edge(\"step_1\", \"human_feedback\")\n",
@@ -165,7 +170,7 @@
"# Set up memory\n",
"memory = MemorySaver()\n",
"\n",
"# Add \n",
"# Add\n",
"graph = builder.compile(checkpointer=memory, interrupt_before=[\"human_feedback\"])\n",
"\n",
"# View\n",
@@ -253,7 +258,7 @@
"\n",
"# We now update the state as if we are the human_feedback node\n",
"graph.update_state(thread, {\"user_feedback\": user_input}, as_node=\"human_feedback\")\n",
" \n",
"\n",
"# We can check the state\n",
"print(\"--State after update--\")\n",
"print(graph.get_state(thread))\n",
examples/memory/add-summary-conversation-history.ipynb
+14 -15
View File
@@ -110,23 +110,26 @@
"\n",
"memory = SqliteSaver.from_conn_string(\":memory:\")\n",
"\n",
"\n",
"# We will add a `summary` attribute (in addition to `messages` key,\n",
"# which MessagesState already has)\n",
"class State(MessagesState):\n",
" summary: str\n",
"\n",
"\n",
"# We will use this model for both the conversation and the summarization\n",
"model = ChatAnthropic(model_name=\"claude-3-haiku-20240307\")\n",
"\n",
"\n",
"# Define the logic to call the model\n",
"def call_model(state: State):\n",
" # If a summary exists, we add this in as a system message\n",
" summary = state.get('summary', '')\n",
" summary = state.get(\"summary\", \"\")\n",
" if summary:\n",
" system_message = f\"Summary of conversation earlier: {summary}\"\n",
" messages = [SystemMessage(content=system_message)] + state['messages']\n",
" messages = [SystemMessage(content=system_message)] + state[\"messages\"]\n",
" else:\n",
" messages = state['messages']\n",
" messages = state[\"messages\"]\n",
" response = model.invoke(messages)\n",
" # We return a list, because this will get added to the existing list\n",
" return {\"messages\": [response]}\n",
@@ -145,7 +148,7 @@
"\n",
"def summarize_conversation(state: State):\n",
" # First, we summarize the conversation\n",
" summary = state.get('summary', '')\n",
" summary = state.get(\"summary\", \"\")\n",
" if summary:\n",
" # If a summary already exists, we use a different system prompt\n",
" # to summarize it than if one didn't\n",
@@ -155,17 +158,13 @@
" )\n",
" else:\n",
" summary_message = \"Create a summary of the conversation above:\"\n",
" \n",
" messages = state['messages'] + [HumanMessage(content=summary_message)]\n",
"\n",
" messages = state[\"messages\"] + [HumanMessage(content=summary_message)]\n",
" response = model.invoke(messages)\n",
" # We now need to delete messages that we no longer want to show up\n",
" # I will delete all but the last two messages, but you can change this\n",
" delete_messages = [RemoveMessage(id=m.id) for m in state['messages'][:-2]]\n",
" return {\n",
" \"summary\": response.content,\n",
" \"messages\": delete_messages\n",
" }\n",
" \n",
" delete_messages = [RemoveMessage(id=m.id) for m in state[\"messages\"][:-2]]\n",
" return {\"summary\": response.content, \"messages\": delete_messages}\n",
"\n",
"\n",
"# Define a new graph\n",
@@ -212,10 +211,10 @@
"source": [
"def print_update(update):\n",
" for k, v in update.items():\n",
" for m in v['messages']:\n",
" for m in v[\"messages\"]:\n",
" m.pretty_print()\n",
" if 'summary' in v:\n",
" print(v['summary'])"
" if \"summary\" in v:\n",
" print(v[\"summary\"])"
]
},
{
examples/memory/delete-messages.ipynb
+10 -5
View File
@@ -252,7 +252,7 @@
}
],
"source": [
"messages = app.get_state(config).values['messages']\n",
"messages = app.get_state(config).values[\"messages\"]\n",
"messages"
]
},
@@ -292,6 +292,7 @@
],
"source": [
"from langchain_core.messages import RemoveMessage\n",
"\n",
"app.update_state(config, {\"messages\": RemoveMessage(id=messages[0].id)})"
]
},
@@ -323,7 +324,7 @@
}
],
"source": [
"messages = app.get_state(config).values['messages']\n",
"messages = app.get_state(config).values[\"messages\"]\n",
"messages"
]
},
@@ -349,10 +350,11 @@
"\n",
"\n",
"def delete_messages(state):\n",
" messages = state['messages']\n",
" messages = state[\"messages\"]\n",
" if len(messages) > 3:\n",
" return {\"messages\": [RemoveMessage(id=m.id) for m in messages[:-3]]}\n",
"\n",
"\n",
"# We need to modify the logic to call delete_messages rather than end right away\n",
"def should_continue(state: MessagesState) -> Literal[\"action\", \"delete_messages\"]:\n",
" \"\"\"Return the next node to execute.\"\"\"\n",
@@ -374,7 +376,10 @@
"\n",
"\n",
"workflow.add_edge(START, \"agent\")\n",
"workflow.add_conditional_edges(\"agent\", should_continue,)\n",
"workflow.add_conditional_edges(\n",
" \"agent\",\n",
" should_continue,\n",
")\n",
"workflow.add_edge(\"action\", \"agent\")\n",
"\n",
"# This is the new edge we're adding: after we delete messages, we finish\n",
@@ -450,7 +455,7 @@
}
],
"source": [
"messages = app.get_state(config).values['messages']\n",
"messages = app.get_state(config).values[\"messages\"]\n",
"messages"
]
},
examples/node-retries.ipynb
+120
View File
@@ -0,0 +1,120 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# How to add node retry policies\n",
"\n",
"There are many use cases where you may wish for your node to have a custom retry policy, for example if you are calling an API, querying a database, or calling an LLM, etc. \n",
"\n",
"In order to configure the retry policty, you have to pass the `retry` parameter to the `add_node` function. The `retry` parameter takes in a `RetryPolicy` named tuple object. Below we instantiate a `RetryPolicy` object with the default parameters:"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"RetryPolicy(initial_interval=0.5, backoff_factor=2.0, max_interval=128.0, max_attempts=3, jitter=True, retry_on=<function default_retry_on at 0x1157419e0>)"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from langgraph.pregel import RetryPolicy\n",
"\n",
"RetryPolicy()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If you want more information on what each of the parameters does, be sure to read the [reference](https://langchain-ai.github.io/langgraph/reference/graphs/#retrypolicy).\n",
"\n",
"## Passing a retry policy to a node\n",
"\n",
"Lastly, we can pass `RetryPolicy` objects when we call the `add_node` function. In the example below we pass two different retry policies to each of our nodes:"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [],
"source": [
"import operator\n",
"import sqlite3\n",
"from typing import Annotated, Sequence, TypedDict\n",
"\n",
"from langchain_anthropic import ChatAnthropic\n",
"from langchain_core.messages import BaseMessage\n",
"\n",
"from langgraph.graph import END, StateGraph, START\n",
"from langchain_community.utilities import SQLDatabase\n",
"from langchain_core.messages import AIMessage\n",
"\n",
"db = SQLDatabase.from_uri(\"sqlite:///:memory:\")\n",
"\n",
"model = ChatAnthropic(model_name=\"claude-2.1\")\n",
"\n",
"\n",
"class AgentState(TypedDict):\n",
" messages: Annotated[Sequence[BaseMessage], operator.add]\n",
"\n",
"\n",
"def query_database(state):\n",
" query_result = db.run(\"SELECT * FROM Artist LIMIT 10;\")\n",
" return {\"messages\": [AIMessage(content=query_result)]}\n",
"\n",
"\n",
"def call_model(state):\n",
" response = model.invoke(state[\"messages\"])\n",
" return {\"messages\": [response]}\n",
"\n",
"\n",
"# Define a new graph\n",
"workflow = StateGraph(AgentState)\n",
"workflow.add_node(\n",
" \"query_database\",\n",
" query_database,\n",
" retry=RetryPolicy(retry_on=sqlite3.OperationalError),\n",
")\n",
"workflow.add_node(\"model\", call_model, retry=RetryPolicy(max_attempts=5))\n",
"workflow.add_edge(START, \"model\")\n",
"workflow.add_edge(\"model\", \"query_database\")\n",
"workflow.add_edge(\"query_database\", END)\n",
"\n",
"app = workflow.compile()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "env",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
examples/pass-config-to-tools.ipynb
+2 -1
View File
@@ -91,7 +91,8 @@
"@tool(parse_docstring=True)\n",
"def update_favorite_pets(\n",
" # NOTE: config arg does not need to be added to docstring, as we don't want it to be included in the function signature attached to the LLM\n",
" pets: List[str], config: RunnableConfig\n",
" pets: List[str],\n",
" config: RunnableConfig,\n",
") -> None:\n",
" \"\"\"Add the list of favorite pets.\n",
"\n",
examples/pass-run-time-values-to-tools.ipynb
+42 -63
View File
@@ -86,30 +86,30 @@
"source": [
"## Defining the tools\n",
"\n",
"We'll want our tool to take graph state as an input, but we don't want the model to try to generate this input when calling the tool. We can use the `InjectedToolArg` annotation to mark `state` as being injected at runtime. Any argument annotated with `InjectedToolArg` will not be generated by the model.\n",
"We'll want our tool to take graph state as an input, but we don't want the model to try to generate this input when calling the tool. We can use the `InjectedState` annotation to mark arguments as required graph state (or some field of graph state. These arguments will not be generated by the model. When using `ToolNode`, graph state will automatically be passed in to the relevant tools and arguments.\n",
"\n",
"In this example we'll create a tool that returns Documents and then another tool that actually cites the Documents that justify a claim."
]
},
{
"cell_type": "code",
"execution_count": 63,
"execution_count": 6,
"id": "1d36e782-80f4-4334-b7d7-ee4c79864480",
"metadata": {},
"outputs": [],
"source": [
"from typing import List, Tuple\n",
"from typing_extensions import Annotated\n",
"\n",
"from langchain_core.documents import Document\n",
"from langchain_core.pydantic_v1 import BaseModel\n",
"from langchain_core.tools import InjectedToolArg, tool\n",
"from typing_extensions import Annotated\n",
"from langchain_core.tools import tool\n",
"\n",
"from langgraph.prebuilt import InjectedState\n",
"\n",
"\n",
"@tool(parse_docstring=True, response_format=\"content_and_artifact\")\n",
"def get_context(\n",
" question: List[str], state: Annotated[dict, InjectedToolArg]\n",
") -> Tuple[str, List[Document]]:\n",
"def get_context(question: List[str]) -> Tuple[str, List[Document]]:\n",
" \"\"\"Get context on the question.\n",
"\n",
" Args:\n",
@@ -136,7 +136,7 @@
"\n",
"@tool(parse_docstring=True, response_format=\"content_and_artifact\")\n",
"def cite_context_sources(\n",
" claim: str, state: Annotated[dict, InjectedToolArg]\n",
" claim: str, state: Annotated[dict, InjectedState]\n",
") -> Tuple[str, List[Document]]:\n",
" \"\"\"Cite which source a claim was based on.\n",
"\n",
@@ -175,31 +175,30 @@
},
{
"cell_type": "code",
"execution_count": 64,
"execution_count": 9,
"id": "1092929b-c939-4b2a-9f9c-e725b0e34af2",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'title': 'get_contextSchema',\n",
" 'description': 'Get context on the question.',\n",
"{'title': 'cite_context_sourcesSchema',\n",
" 'description': 'Cite which source a claim was based on.',\n",
" 'type': 'object',\n",
" 'properties': {'question': {'title': 'Question',\n",
" 'description': 'The user question',\n",
" 'type': 'array',\n",
" 'items': {'type': 'string'}},\n",
" 'properties': {'claim': {'title': 'Claim',\n",
" 'description': 'The claim that was made.',\n",
" 'type': 'string'},\n",
" 'state': {'title': 'State', 'type': 'object'}},\n",
" 'required': ['question', 'state']}"
" 'required': ['claim', 'state']}"
]
},
"execution_count": 64,
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"get_context.get_input_schema().schema()"
"cite_context_sources.get_input_schema().schema()"
]
},
{
@@ -212,30 +211,29 @@
},
{
"cell_type": "code",
"execution_count": 65,
"execution_count": 11,
"id": "3912bb51-3107-4335-a659-021c5d89fb37",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'title': 'get_context',\n",
" 'description': 'Get context on the question.',\n",
"{'title': 'cite_context_sources',\n",
" 'description': 'Cite which source a claim was based on.',\n",
" 'type': 'object',\n",
" 'properties': {'question': {'title': 'Question',\n",
" 'description': 'The user question',\n",
" 'type': 'array',\n",
" 'items': {'type': 'string'}}},\n",
" 'required': ['question']}"
" 'properties': {'claim': {'title': 'Claim',\n",
" 'description': 'The claim that was made.',\n",
" 'type': 'string'}},\n",
" 'required': ['claim']}"
]
},
"execution_count": 65,
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"get_context.tool_call_schema.schema()"
"cite_context_sources.tool_call_schema.schema()"
]
},
{
@@ -258,7 +256,7 @@
},
{
"cell_type": "code",
"execution_count": 66,
"execution_count": 12,
"id": "ea793afa-2eab-4901-910d-6eed90cd6564",
"metadata": {},
"outputs": [],
@@ -302,7 +300,7 @@
},
{
"cell_type": "code",
"execution_count": 67,
"execution_count": 18,
"id": "3b541bb9-900c-40d0-964d-7b5dfee30667",
"metadata": {},
"outputs": [],
@@ -312,7 +310,7 @@
"from langchain_core.messages import ToolMessage\n",
"from langchain_openai import ChatOpenAI\n",
"\n",
"from langgraph.prebuilt import ToolExecutor, ToolInvocation\n",
"from langgraph.prebuilt import ToolNode\n",
"\n",
"model = ChatOpenAI(model=\"gpt-4o\", temperature=0)\n",
"\n",
@@ -330,8 +328,6 @@
"\n",
"\n",
"tools = [get_context, cite_context_sources]\n",
"tool_map = {tool_.name: tool_ for tool_ in tools}\n",
"\n",
"\n",
"# Define the function that calls the model\n",
"def call_model(state, config):\n",
@@ -342,25 +338,8 @@
" return {\"messages\": [response]}\n",
"\n",
"\n",
"# Helper function for adding state to each tool call's arguments\n",
"def inject_state(message, state):\n",
" tool_calls = []\n",
" for tool_call in message.tool_calls:\n",
" tool_call_copy = deepcopy(tool_call)\n",
" tool_call_copy[\"args\"][\"state\"] = state\n",
" tool_calls.append(tool_call_copy)\n",
" return tool_calls\n",
"\n",
"\n",
"# Define the function to execute tools\n",
"def call_tool(state, config):\n",
" messages = state[\"messages\"]\n",
" last_message = messages[-1]\n",
" tool_messages = []\n",
" for tool_call in inject_state(last_message, state):\n",
" tool_messages.append(tool_map[tool_call[\"name\"]].invoke(tool_call, config))\n",
" # We return a list, because this will get added to the existing list\n",
" return {\"messages\": tool_messages}"
"# ToolNode will automatically take care of injecting state into tools\n",
"tool_node = ToolNode(tools)"
]
},
{
@@ -375,7 +354,7 @@
},
{
"cell_type": "code",
"execution_count": 68,
"execution_count": 19,
"id": "813ae66c-3b58-4283-a02a-36da72a2ab90",
"metadata": {},
"outputs": [],
@@ -387,7 +366,7 @@
"\n",
"# Define the two nodes we will cycle between\n",
"workflow.add_node(\"agent\", call_model)\n",
"workflow.add_node(\"action\", call_tool)\n",
"workflow.add_node(\"action\", tool_node)\n",
"\n",
"# Set the entrypoint as `agent`\n",
"# This means that this node is the first one called\n",
@@ -426,7 +405,7 @@
},
{
"cell_type": "code",
"execution_count": 69,
"execution_count": 20,
"id": "a8afd6ef",
"metadata": {},
"outputs": [
@@ -464,7 +443,7 @@
},
{
"cell_type": "code",
"execution_count": 70,
"execution_count": 21,
"id": "8edb04b9-40b6-46f1-a7a8-4b2d8aba7752",
"metadata": {},
"outputs": [
@@ -474,19 +453,19 @@
"text": [
"Output from node 'agent':\n",
"---\n",
"{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_aFUFt3TdazRnmD3FTZfxFAgL', 'function': {'arguments': '{\"question\":[\"what\\'s the latest news about FooBar\"]}', 'name': 'get_context'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 22, 'prompt_tokens': 87, 'total_tokens': 109}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-adf99f00-a903-49f2-b0c3-37b84b9b801f-0', tool_calls=[{'name': 'get_context', 'args': {'question': [\"what's the latest news about FooBar\"]}, 'id': 'call_aFUFt3TdazRnmD3FTZfxFAgL', 'type': 'tool_call'}], usage_metadata={'input_tokens': 87, 'output_tokens': 22, 'total_tokens': 109})]}\n",
"{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_BidVTw5NiW2wp8Ez7m8dDoHI', 'function': {'arguments': '{\"question\":[\"latest news about FooBar\"]}', 'name': 'get_context'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 19, 'prompt_tokens': 87, 'total_tokens': 106}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-fcac1b73-563e-4f4c-b1b0-626f55d377be-0', tool_calls=[{'name': 'get_context', 'args': {'question': ['latest news about FooBar']}, 'id': 'call_BidVTw5NiW2wp8Ez7m8dDoHI', 'type': 'tool_call'}], usage_metadata={'input_tokens': 87, 'output_tokens': 19, 'total_tokens': 106})]}\n",
"\n",
"---\n",
"\n",
"Output from node 'action':\n",
"---\n",
"{'messages': [ToolMessage(content=\"FooBar company just raised 1 Billion dollars!\\n\\nFooBar company is now only hiring AI's\\n\\nFooBar company was founded in 2019\\n\\nFooBar company makes friendly robots\", name='get_context', tool_call_id='call_aFUFt3TdazRnmD3FTZfxFAgL', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!'), Document(metadata={'source': 'twitter'}, page_content=\"FooBar company is now only hiring AI's\"), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company was founded in 2019'), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company makes friendly robots')])]}\n",
"{'messages': [ToolMessage(content=\"FooBar company just raised 1 Billion dollars!\\n\\nFooBar company is now only hiring AI's\\n\\nFooBar company was founded in 2019\\n\\nFooBar company makes friendly robots\", name='get_context', tool_call_id='call_BidVTw5NiW2wp8Ez7m8dDoHI', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!'), Document(metadata={'source': 'twitter'}, page_content=\"FooBar company is now only hiring AI's\"), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company was founded in 2019'), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company makes friendly robots')])]}\n",
"\n",
"---\n",
"\n",
"Output from node 'agent':\n",
"---\n",
"{'messages': [AIMessage(content='The latest news about FooBar is that the company just raised 1 billion dollars!', response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 153, 'total_tokens': 171}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'stop', 'logprobs': None}, id='run-c229a397-fda3-415b-a188-1416fd5f21b7-0', usage_metadata={'input_tokens': 153, 'output_tokens': 18, 'total_tokens': 171})]}\n",
"{'messages': [AIMessage(content='The latest news about FooBar is that the company has just raised 1 billion dollars!', response_metadata={'token_usage': {'completion_tokens': 19, 'prompt_tokens': 150, 'total_tokens': 169}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'stop', 'logprobs': None}, id='run-a8407471-7715-4c16-bd46-c29e5751e882-0', usage_metadata={'input_tokens': 150, 'output_tokens': 19, 'total_tokens': 169})]}\n",
"\n",
"---\n",
"\n"
@@ -509,7 +488,7 @@
},
{
"cell_type": "code",
"execution_count": 71,
"execution_count": 22,
"id": "4a2128ed-e23f-4f25-a026-0c6590f01a1c",
"metadata": {},
"outputs": [
@@ -519,19 +498,19 @@
"text": [
"Output from node 'agent':\n",
"---\n",
"{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_qqB4kucZnVhrZ5mJSH1dF8Lb', 'function': {'arguments': '{\"claim\":\"The latest news about FooBar is that the company just raised 1 billion dollars!\"}', 'name': 'cite_context_sources'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 32, 'prompt_tokens': 185, 'total_tokens': 217}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-686d4706-81c9-4ca0-8f09-d9af02f4ad7f-0', tool_calls=[{'name': 'cite_context_sources', 'args': {'claim': 'The latest news about FooBar is that the company just raised 1 billion dollars!'}, 'id': 'call_qqB4kucZnVhrZ5mJSH1dF8Lb', 'type': 'tool_call'}], usage_metadata={'input_tokens': 185, 'output_tokens': 32, 'total_tokens': 217})]}\n",
"{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_EB0zaQypXMqEUzaqwflUr0zH', 'function': {'arguments': '{\"claim\":\"FooBar company just raised 1 Billion dollars!\"}', 'name': 'cite_context_sources'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 25, 'prompt_tokens': 183, 'total_tokens': 208}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b4952777-e2b3-4448-be87-200e6e80981b-0', tool_calls=[{'name': 'cite_context_sources', 'args': {'claim': 'FooBar company just raised 1 Billion dollars!'}, 'id': 'call_EB0zaQypXMqEUzaqwflUr0zH', 'type': 'tool_call'}], usage_metadata={'input_tokens': 183, 'output_tokens': 25, 'total_tokens': 208})]}\n",
"\n",
"---\n",
"\n",
"Output from node 'action':\n",
"---\n",
"{'messages': [ToolMessage(content='twitter', name='cite_context_sources', tool_call_id='call_qqB4kucZnVhrZ5mJSH1dF8Lb', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!')])]}\n",
"{'messages': [ToolMessage(content='twitter', name='cite_context_sources', tool_call_id='call_EB0zaQypXMqEUzaqwflUr0zH', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!')])]}\n",
"\n",
"---\n",
"\n",
"Output from node 'agent':\n",
"---\n",
"{'messages': [AIMessage(content='The information about FooBar raising 1 billion dollars came from Twitter.', response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 227, 'total_tokens': 242}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_18cc0f1fa0', 'finish_reason': 'stop', 'logprobs': None}, id='run-343ad465-9a62-4d72-91bf-ab29c4fe8781-0', usage_metadata={'input_tokens': 227, 'output_tokens': 15, 'total_tokens': 242})]}\n",
"{'messages': [AIMessage(content='The information that FooBar company just raised 1 billion dollars comes from Twitter.', response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 218, 'total_tokens': 235}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_400f27fa1f', 'finish_reason': 'stop', 'logprobs': None}, id='run-a0dede05-dadd-46f6-8654-746520d4cef8-0', usage_metadata={'input_tokens': 218, 'output_tokens': 17, 'total_tokens': 235})]}\n",
"\n",
"---\n",
"\n"
examples/persistence_mongodb.ipynb
+22 -13
View File
@@ -82,6 +82,7 @@
" return pickle.loads(data)\n",
" return super().loads(data)\n",
"\n",
"\n",
"class MongoDBSaver(AbstractContextManager, BaseCheckpointSaver):\n",
" \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database.\n",
"\n",
@@ -309,17 +310,19 @@
}
],
"source": [
"from langgraph.graph import StateGraph\n",
"from langgraph.graph import StateGraph, START, END\n",
"\n",
"checkpointer = MongoDBSaver(MongoClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\")\n",
"checkpointer = MongoDBSaver(\n",
" MongoClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\"\n",
")\n",
"builder = StateGraph(int)\n",
"builder.add_node(\"add_one\", lambda x: x + 1)\n",
"builder.set_entry_point(\"add_one\")\n",
"builder.set_finish_point(\"add_one\")\n",
"builder.add_edge(START, \"add_one\")\n",
"builder.add_edge(\"add_one\", END)\n",
"graph = builder.compile(checkpointer=checkpointer)\n",
"config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
"graph.get_state(config)\n",
"result = graph.invoke(3,config)\n",
"result = graph.invoke(3, config)\n",
"graph.get_state(config)"
]
},
@@ -572,7 +575,7 @@
"for doc in collection.find():\n",
" print(doc)\n",
"\n",
"#The checkpoints from both the examples have been saved in the database."
"# The checkpoints from both the examples have been saved in the database."
]
},
{
@@ -588,7 +591,7 @@
"metadata": {},
"outputs": [],
"source": [
"#Async package for MongoDB\n",
"# Async package for MongoDB\n",
"%pip install motor"
]
},
@@ -601,7 +604,7 @@
"import pickle\n",
"from contextlib import AbstractContextManager\n",
"from types import TracebackType\n",
"from typing import Any, Dict,Optional,AsyncIterator\n",
"from typing import Any, Dict, Optional, AsyncIterator\n",
"\n",
"from langchain_core.runnables import RunnableConfig\n",
"from typing_extensions import Self\n",
@@ -616,6 +619,7 @@
"from langgraph.serde.jsonplus import JsonPlusSerializer\n",
"from motor.motor_asyncio import AsyncIOMotorClient\n",
"\n",
"\n",
"class JsonPlusSerializerCompat(JsonPlusSerializer):\n",
" \"\"\"A serializer that supports loading pickled checkpoints for backwards compatibility.\n",
"\n",
@@ -643,6 +647,7 @@
" return pickle.loads(data)\n",
" return super().loads(data)\n",
"\n",
"\n",
"class MongoDBSaver(AbstractContextManager, BaseCheckpointSaver):\n",
" \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database.\n",
"\n",
@@ -842,15 +847,18 @@
"metadata": {},
"outputs": [],
"source": [
"from langgraph.graph import StateGraph\n",
"checkpointer = MongoDBSaver(AsyncIOMotorClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\")\n",
"from langgraph.graph import StateGraph, START\n",
"\n",
"checkpointer = MongoDBSaver(\n",
" AsyncIOMotorClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\"\n",
")\n",
"builder = StateGraph(int)\n",
"builder.add_node(\"add_one\", lambda x: x + 1)\n",
"builder.set_entry_point(\"add_one\")\n",
"builder.set_finish_point(\"add_one\")\n",
"builder.add_edge(START, \"add_one\")\n",
"builder.add_edge(\"add_one\", END)\n",
"graph = builder.compile(checkpointer=checkpointer)\n",
"config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
"res = await graph.ainvoke(3,config)"
"res = await graph.ainvoke(3, config)"
]
},
{
@@ -971,6 +979,7 @@
],
"source": [
"from pymongo import MongoClient\n",
"\n",
"client = MongoClient(MONGO_URI)\n",
"database = client[\"checkpoints_db\"]\n",
"collection = database[\"checkpoints_collection\"]\n",
examples/persistence_postgres.ipynb
+3 -4
View File
@@ -51,7 +51,7 @@
" Union,\n",
" Tuple,\n",
" List,\n",
" Sequence\n",
" Sequence,\n",
")\n",
"\n",
"import psycopg\n",
@@ -341,7 +341,6 @@
" writes: Sequence[Tuple[str, Any]],\n",
" task_id: str,\n",
" ) -> None:\n",
"\n",
" async with self._get_async_connection() as conn:\n",
" async with conn.cursor() as cur:\n",
" await cur.executemany(\n",
@@ -524,7 +523,7 @@
" pending_writes=[\n",
" (task_id, channel, self.serde.loads(value))\n",
" for task_id, channel, value in cur\n",
" ]\n",
" ],\n",
" )\n",
"\n",
" async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n",
@@ -594,7 +593,7 @@
" pending_writes=[\n",
" (task_id, channel, self.serde.loads(value))\n",
" async for task_id, channel, value in cur\n",
" ]\n",
" ],\n",
" )\n",
"\n",
" def _search_where(\n",
examples/persistence_redis.ipynb
+163 -43
View File
@@ -63,10 +63,13 @@
"logging.basicConfig(level=logging.INFO)\n",
"logger = logging.getLogger(__name__)\n",
"\n",
"\n",
"class JsonAndBinarySerializer(JsonPlusSerializer):\n",
" def _default(self, obj: Any) -> Any:\n",
" if isinstance(obj, (bytes, bytearray)):\n",
" return self._encode_constructor_args(obj.__class__, method=\"fromhex\", args=[obj.hex()])\n",
" return self._encode_constructor_args(\n",
" obj.__class__, method=\"fromhex\", args=[obj.hex()]\n",
" )\n",
" return super()._default(obj)\n",
"\n",
" def dumps(self, obj: Any) -> str:\n",
@@ -87,17 +90,25 @@
" logger.error(f\"Deserialization error: {e}\")\n",
" raise\n",
"\n",
"def initialize_sync_pool(host: str = 'localhost', port: int = 6379, db: int = 0, **kwargs) -> redis.ConnectionPool:\n",
"\n",
"def initialize_sync_pool(\n",
" host: str = \"localhost\", port: int = 6379, db: int = 0, **kwargs\n",
") -> redis.ConnectionPool:\n",
" \"\"\"Initialize a synchronous Redis connection pool.\"\"\"\n",
" try:\n",
" pool = redis.ConnectionPool(host=host, port=port, db=db, **kwargs)\n",
" logger.info(f\"Synchronous Redis pool initialized with host={host}, port={port}, db={db}\")\n",
" logger.info(\n",
" f\"Synchronous Redis pool initialized with host={host}, port={port}, db={db}\"\n",
" )\n",
" return pool\n",
" except Exception as e:\n",
" logger.error(f\"Error initializing sync pool: {e}\")\n",
" raise\n",
"\n",
"def initialize_async_pool(url: str = \"redis://localhost\", **kwargs) -> AsyncConnectionPool:\n",
"\n",
"def initialize_async_pool(\n",
" url: str = \"redis://localhost\", **kwargs\n",
") -> AsyncConnectionPool:\n",
" \"\"\"Initialize an asynchronous Redis connection pool.\"\"\"\n",
" try:\n",
" pool = AsyncConnectionPool.from_url(url, **kwargs)\n",
@@ -107,8 +118,11 @@
" logger.error(f\"Error initializing async pool: {e}\")\n",
" raise\n",
"\n",
"\n",
"@contextmanager\n",
"def _get_sync_connection(connection: Union[redis.Redis, redis.ConnectionPool, None]) -> Generator[redis.Redis, None, None]:\n",
"def _get_sync_connection(\n",
" connection: Union[redis.Redis, redis.ConnectionPool, None]\n",
") -> Generator[redis.Redis, None, None]:\n",
" conn = None\n",
" try:\n",
" if isinstance(connection, redis.Redis):\n",
@@ -125,8 +139,11 @@
" if conn:\n",
" conn.close()\n",
"\n",
"\n",
"@asynccontextmanager\n",
"async def _get_async_connection(connection: Union[AsyncRedis, AsyncConnectionPool, None]) -> AsyncGenerator[AsyncRedis, None]:\n",
"async def _get_async_connection(\n",
" connection: Union[AsyncRedis, AsyncConnectionPool, None]\n",
") -> AsyncGenerator[AsyncRedis, None]:\n",
" conn = None\n",
" try:\n",
" if isinstance(connection, AsyncRedis):\n",
@@ -143,27 +160,42 @@
" if conn:\n",
" await conn.aclose()\n",
"\n",
"\n",
"class RedisSaver(BaseCheckpointSaver):\n",
" sync_connection: Optional[Union[redis.Redis, redis.ConnectionPool]] = None\n",
" async_connection: Optional[Union[AsyncRedis, AsyncConnectionPool]] = None\n",
"\n",
" def __init__(self, sync_connection: Optional[Union[redis.Redis, redis.ConnectionPool]] = None, async_connection: Optional[Union[AsyncRedis, AsyncConnectionPool]] = None):\n",
" def __init__(\n",
" self,\n",
" sync_connection: Optional[Union[redis.Redis, redis.ConnectionPool]] = None,\n",
" async_connection: Optional[Union[AsyncRedis, AsyncConnectionPool]] = None,\n",
" ):\n",
" super().__init__(serde=JsonAndBinarySerializer())\n",
" self.sync_connection = sync_connection\n",
" self.async_connection = async_connection\n",
"\n",
" def put(self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata) -> RunnableConfig:\n",
" def put(\n",
" self,\n",
" config: RunnableConfig,\n",
" checkpoint: Checkpoint,\n",
" metadata: CheckpointMetadata,\n",
" ) -> RunnableConfig:\n",
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
" parent_ts = config[\"configurable\"].get(\"thread_ts\")\n",
" key = f\"checkpoint:{thread_id}:{checkpoint['ts']}\"\n",
" try:\n",
" with _get_sync_connection(self.sync_connection) as conn:\n",
" conn.hset(key, mapping={\n",
" \"checkpoint\": self.serde.dumps(checkpoint),\n",
" \"metadata\": self.serde.dumps(metadata),\n",
" \"parent_ts\": parent_ts if parent_ts else \"\"\n",
" })\n",
" logger.info(f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\")\n",
" conn.hset(\n",
" key,\n",
" mapping={\n",
" \"checkpoint\": self.serde.dumps(checkpoint),\n",
" \"metadata\": self.serde.dumps(metadata),\n",
" \"parent_ts\": parent_ts if parent_ts else \"\",\n",
" },\n",
" )\n",
" logger.info(\n",
" f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\"\n",
" )\n",
" except Exception as e:\n",
" logger.error(f\"Failed to put checkpoint: {e}\")\n",
" raise\n",
@@ -174,18 +206,28 @@
" },\n",
" }\n",
"\n",
" async def aput(self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata) -> RunnableConfig:\n",
" async def aput(\n",
" self,\n",
" config: RunnableConfig,\n",
" checkpoint: Checkpoint,\n",
" metadata: CheckpointMetadata,\n",
" ) -> RunnableConfig:\n",
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
" parent_ts = config[\"configurable\"].get(\"thread_ts\")\n",
" key = f\"checkpoint:{thread_id}:{checkpoint['ts']}\"\n",
" try:\n",
" async with _get_async_connection(self.async_connection) as conn:\n",
" await conn.hset(key, mapping={\n",
" \"checkpoint\": self.serde.dumps(checkpoint),\n",
" \"metadata\": self.serde.dumps(metadata),\n",
" \"parent_ts\": parent_ts if parent_ts else \"\"\n",
" })\n",
" logger.info(f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\")\n",
" await conn.hset(\n",
" key,\n",
" mapping={\n",
" \"checkpoint\": self.serde.dumps(checkpoint),\n",
" \"metadata\": self.serde.dumps(metadata),\n",
" \"parent_ts\": parent_ts if parent_ts else \"\",\n",
" },\n",
" )\n",
" logger.info(\n",
" f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\"\n",
" )\n",
" except Exception as e:\n",
" logger.error(f\"Failed to aput checkpoint: {e}\")\n",
" raise\n",
@@ -217,9 +259,20 @@
" checkpoint = self.serde.loads(checkpoint_data[b\"checkpoint\"].decode())\n",
" metadata = self.serde.loads(checkpoint_data[b\"metadata\"].decode())\n",
" parent_ts = checkpoint_data.get(b\"parent_ts\", b\"\").decode()\n",
" parent_config = {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}} if parent_ts else None\n",
" logger.info(f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\")\n",
" return CheckpointTuple(config=config, checkpoint=checkpoint, metadata=metadata, parent_config=parent_config)\n",
" parent_config = (\n",
" {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}}\n",
" if parent_ts\n",
" else None\n",
" )\n",
" logger.info(\n",
" f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\"\n",
" )\n",
" return CheckpointTuple(\n",
" config=config,\n",
" checkpoint=checkpoint,\n",
" metadata=metadata,\n",
" parent_config=parent_config,\n",
" )\n",
" except Exception as e:\n",
" logger.error(f\"Failed to get checkpoint tuple: {e}\")\n",
" raise\n",
@@ -245,22 +298,47 @@
" checkpoint = self.serde.loads(checkpoint_data[b\"checkpoint\"].decode())\n",
" metadata = self.serde.loads(checkpoint_data[b\"metadata\"].decode())\n",
" parent_ts = checkpoint_data.get(b\"parent_ts\", b\"\").decode()\n",
" parent_config = {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}} if parent_ts else None\n",
" logger.info(f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\")\n",
" return CheckpointTuple(config=config, checkpoint=checkpoint, metadata=metadata, parent_config=parent_config)\n",
" parent_config = (\n",
" {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}}\n",
" if parent_ts\n",
" else None\n",
" )\n",
" logger.info(\n",
" f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\"\n",
" )\n",
" return CheckpointTuple(\n",
" config=config,\n",
" checkpoint=checkpoint,\n",
" metadata=metadata,\n",
" parent_config=parent_config,\n",
" )\n",
" except Exception as e:\n",
" logger.error(f\"Failed to get checkpoint tuple: {e}\")\n",
" raise\n",
"\n",
" def list(self, config: Optional[RunnableConfig], *, filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None) -> Generator[CheckpointTuple, None, None]:\n",
" def list(\n",
" self,\n",
" config: Optional[RunnableConfig],\n",
" *,\n",
" filter: Optional[dict[str, Any]] = None,\n",
" before: Optional[RunnableConfig] = None,\n",
" limit: Optional[int] = None,\n",
" ) -> Generator[CheckpointTuple, None, None]:\n",
" thread_id = config[\"configurable\"][\"thread_id\"] if config else \"*\"\n",
" pattern = f\"checkpoint:{thread_id}:*\"\n",
" try:\n",
" with _get_sync_connection(self.sync_connection) as conn:\n",
" keys = conn.keys(pattern)\n",
" if before:\n",
" keys = [k for k in keys if k.decode().split(\":\")[-1] < before[\"configurable\"][\"thread_ts\"]]\n",
" keys = sorted(keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True)\n",
" keys = [\n",
" k\n",
" for k in keys\n",
" if k.decode().split(\":\")[-1]\n",
" < before[\"configurable\"][\"thread_ts\"]\n",
" ]\n",
" keys = sorted(\n",
" keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True\n",
" )\n",
" if limit:\n",
" keys = keys[:limit]\n",
" for key in keys:\n",
@@ -268,25 +346,53 @@
" if data and \"checkpoint\" in data and \"metadata\" in data:\n",
" thread_ts = key.decode().split(\":\")[-1]\n",
" yield CheckpointTuple(\n",
" config={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": thread_ts}},\n",
" config={\n",
" \"configurable\": {\n",
" \"thread_id\": thread_id,\n",
" \"thread_ts\": thread_ts,\n",
" }\n",
" },\n",
" checkpoint=self.serde.loads(data[\"checkpoint\"].decode()),\n",
" metadata=self.serde.loads(data[\"metadata\"].decode()),\n",
" parent_config={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": data.get(\"parent_ts\", b\"\").decode()}} if data.get(\"parent_ts\") else None,\n",
" parent_config={\n",
" \"configurable\": {\n",
" \"thread_id\": thread_id,\n",
" \"thread_ts\": data.get(\"parent_ts\", b\"\").decode(),\n",
" }\n",
" }\n",
" if data.get(\"parent_ts\")\n",
" else None,\n",
" )\n",
" logger.info(\n",
" f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\"\n",
" )\n",
" logger.info(f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\")\n",
" except Exception as e:\n",
" logger.error(f\"Failed to list checkpoints: {e}\")\n",
" raise\n",
"\n",
" async def alist(self, config: Optional[RunnableConfig], *, filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None) -> AsyncGenerator[CheckpointTuple, None]:\n",
" async def alist(\n",
" self,\n",
" config: Optional[RunnableConfig],\n",
" *,\n",
" filter: Optional[dict[str, Any]] = None,\n",
" before: Optional[RunnableConfig] = None,\n",
" limit: Optional[int] = None,\n",
" ) -> AsyncGenerator[CheckpointTuple, None]:\n",
" thread_id = config[\"configurable\"][\"thread_id\"] if config else \"*\"\n",
" pattern = f\"checkpoint:{thread_id}:*\"\n",
" try:\n",
" async with _get_async_connection(self.async_connection) as conn:\n",
" keys = await conn.keys(pattern)\n",
" if before:\n",
" keys = [k for k in keys if k.decode().split(\":\")[-1] < before[\"configurable\"][\"thread_ts\"]]\n",
" keys = sorted(keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True)\n",
" keys = [\n",
" k\n",
" for k in keys\n",
" if k.decode().split(\":\")[-1]\n",
" < before[\"configurable\"][\"thread_ts\"]\n",
" ]\n",
" keys = sorted(\n",
" keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True\n",
" )\n",
" if limit:\n",
" keys = keys[:limit]\n",
" for key in keys:\n",
@@ -294,15 +400,29 @@
" if data and \"checkpoint\" in data and \"metadata\" in data:\n",
" thread_ts = key.decode().split(\":\")[-1]\n",
" yield CheckpointTuple(\n",
" config={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": thread_ts}},\n",
" config={\n",
" \"configurable\": {\n",
" \"thread_id\": thread_id,\n",
" \"thread_ts\": thread_ts,\n",
" }\n",
" },\n",
" checkpoint=self.serde.loads(data[\"checkpoint\"].decode()),\n",
" metadata=self.serde.loads(data[\"metadata\"].decode()),\n",
" parent_config={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": data.get(\"parent_ts\", b\"\").decode()}} if data.get(\"parent_ts\") else None,\n",
" parent_config={\n",
" \"configurable\": {\n",
" \"thread_id\": thread_id,\n",
" \"thread_ts\": data.get(\"parent_ts\", b\"\").decode(),\n",
" }\n",
" }\n",
" if data.get(\"parent_ts\")\n",
" else None,\n",
" )\n",
" logger.info(\n",
" f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\"\n",
" )\n",
" logger.info(f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\")\n",
" except Exception as e:\n",
" logger.error(f\"Failed to list checkpoints: {e}\")\n",
" raise\n"
" raise"
]
},
{
@@ -538,7 +658,7 @@
"import redis\n",
"\n",
"# Initialize the Redis synchronous direct connection\n",
"sync_redis_direct = redis.Redis(host='172.25.0.4', port=6379, db=0)\n",
"sync_redis_direct = redis.Redis(host=\"172.25.0.4\", port=6379, db=0)\n",
"\n",
"# Initialize the RedisSaver with the synchronous direct connection\n",
"checkpointer = RedisSaver(sync_connection=sync_redis_direct)\n",
@@ -582,7 +702,7 @@
],
"source": [
"# Initialize a synchronous Redis connection pool\n",
"async_pool = initialize_async_pool(url='redis://172.25.0.4:6379/0')\n",
"async_pool = initialize_async_pool(url=\"redis://172.25.0.4:6379/0\")\n",
"\n",
"checkpointer = RedisSaver(async_connection=async_pool)"
]
@@ -687,7 +807,7 @@
"source": [
"from redis.asyncio import Redis as AsyncRedis\n",
"\n",
"async with await AsyncRedis(host='172.25.0.4', port=6379, db=0) as conn:\n",
"async with await AsyncRedis(host=\"172.25.0.4\", port=6379, db=0) as conn:\n",
" checkpointer = RedisSaver(async_connection=conn)\n",
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
" config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
examples/streaming-events-from-within-tools-without-langchain.ipynb
+49 -38
View File
@@ -88,40 +88,36 @@
"from openai import AsyncOpenAI\n",
"from langchain_core.language_models.chat_models import ChatGenerationChunk\n",
"from langchain_core.messages import AIMessageChunk\n",
"from langchain_core.runnables.config import ensure_config, get_callback_manager_for_config\n",
"from langchain_core.runnables.config import (\n",
" ensure_config,\n",
" get_callback_manager_for_config,\n",
")\n",
"\n",
"openai_client = AsyncOpenAI()\n",
"# define tool schema for openai tool calling\n",
"\n",
"tool = {\n",
" \"type\": \"function\",\n",
" \"function\": {\n",
" \"name\": \"get_items\",\n",
" \"description\": \"Use this tool to look up which items are in the given place.\",\n",
" \"parameters\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"place\": {\n",
" \"type\": \"string\"\n",
" }\n",
" },\n",
" \"required\": [\n",
" \"place\"\n",
" ]\n",
" }\n",
" }\n",
" \"type\": \"function\",\n",
" \"function\": {\n",
" \"name\": \"get_items\",\n",
" \"description\": \"Use this tool to look up which items are in the given place.\",\n",
" \"parameters\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\"place\": {\"type\": \"string\"}},\n",
" \"required\": [\"place\"],\n",
" },\n",
" },\n",
"}\n",
"\n",
"\n",
"async def call_model(state, config=None):\n",
" config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n",
" callback_manager = get_callback_manager_for_config(config)\n",
" messages = state[\"messages\"]\n",
" \n",
"\n",
" llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
" response = await openai_client.chat.completions.create(\n",
" messages=messages,\n",
" model=\"gpt-3.5-turbo\",\n",
" tools=[tool],\n",
" stream=True\n",
" messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
" )\n",
"\n",
" response_content = \"\"\n",
@@ -147,7 +143,10 @@
"\n",
" # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n",
" tool_call_chunk = ChatGenerationChunk(\n",
" message=AIMessageChunk(content=\"\", additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]})\n",
" message=AIMessageChunk(\n",
" content=\"\",\n",
" additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n",
" )\n",
" )\n",
" llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n",
" tool_call_function_arguments += delta.tool_calls[0].function.arguments\n",
@@ -156,8 +155,11 @@
" tool_calls = [\n",
" {\n",
" \"id\": tool_call_id,\n",
" \"function\": {\"name\": tool_call_function_name, \"arguments\": tool_call_function_arguments},\n",
" \"type\": \"function\"\n",
" \"function\": {\n",
" \"name\": tool_call_function_name,\n",
" \"arguments\": tool_call_function_arguments,\n",
" },\n",
" \"type\": \"function\",\n",
" }\n",
" ]\n",
" else:\n",
@@ -166,7 +168,7 @@
" response_message = {\n",
" \"role\": role,\n",
" \"content\": response_content,\n",
" \"tool_calls\": tool_calls\n",
" \"tool_calls\": tool_calls,\n",
" }\n",
" return {\"messages\": [response_message]}"
]
@@ -189,8 +191,10 @@
"import json\n",
"from langchain_core.callbacks import adispatch_custom_event\n",
"\n",
"\n",
"async def get_items(place: str) -> str:\n",
" \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
"\n",
" # this can be replaced with any actual streaming logic that you might have\n",
" def stream(place: str):\n",
" if \"bed\" in place: # For under the bed\n",
@@ -205,18 +209,22 @@
" await adispatch_custom_event(\n",
" # this will allow you to filter events by name\n",
" \"tool_call_token_stream\",\n",
" {\"function_name\": \"get_items\", \"arguments\": {\"place\": place}, \"tool_output_token\": token},\n",
" {\n",
" \"function_name\": \"get_items\",\n",
" \"arguments\": {\"place\": place},\n",
" \"tool_output_token\": token,\n",
" },\n",
" # this will allow you to filter events by tags\n",
" config={\"tags\": [\"tool_call\"]}\n",
" config={\"tags\": [\"tool_call\"]},\n",
" )\n",
" tokens.append(token)\n",
"\n",
" return \", \".join(tokens)\n",
"\n",
"\n",
"# define mapping to look up functions when running tools\n",
"function_name_to_function = {\n",
" \"get_items\": get_items\n",
"}\n",
"function_name_to_function = {\"get_items\": get_items}\n",
"\n",
"\n",
"async def call_tools(state):\n",
" messages = state[\"messages\"]\n",
@@ -225,17 +233,15 @@
" function_name = tool_call[\"function\"][\"name\"]\n",
" function_arguments = tool_call[\"function\"][\"arguments\"]\n",
" arguments = json.loads(function_arguments)\n",
" \n",
" function_response = await function_name_to_function[function_name](**arguments) \n",
"\n",
" function_response = await function_name_to_function[function_name](**arguments)\n",
" tool_message = {\n",
" \"tool_call_id\": tool_call[\"id\"],\n",
" \"role\": \"tool\",\n",
" \"name\": function_name,\n",
" \"content\": function_response,\n",
" }\n",
" return {\n",
" \"messages\": [tool_message]\n",
" }"
" return {\"messages\": [tool_message]}"
]
},
{
@@ -258,16 +264,19 @@
"\n",
"from langgraph.graph import StateGraph, END\n",
"\n",
"\n",
"class State(TypedDict):\n",
" messages: Annotated[list, operator.add]\n",
"\n",
"\n",
"def should_continue(state) -> Literal[\"tools\", END]:\n",
" messages = state['messages']\n",
" messages = state[\"messages\"]\n",
" last_message = messages[-1]\n",
" if last_message[\"tool_calls\"]:\n",
" return \"tools\"\n",
" return END\n",
"\n",
"\n",
"workflow = StateGraph(State)\n",
"workflow.set_entry_point(\"model\")\n",
"workflow.add_node(\"model\", call_model) # i.e. our \"agent\"\n",
@@ -310,7 +319,9 @@
}
],
"source": [
"async for event in graph.astream_events({\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"):\n",
"async for event in graph.astream_events(\n",
" {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n",
"):\n",
" tags = event.get(\"tags\", [])\n",
" if event[\"event\"] == \"on_custom_event\" and \"tool_call\" in tags:\n",
" print(\"Tool token\", event[\"data\"][\"tool_output_token\"])"
examples/streaming-tokens-without-langchain.ipynb
+42 -36
View File
@@ -88,40 +88,36 @@
"from openai import AsyncOpenAI\n",
"from langchain_core.language_models.chat_models import ChatGenerationChunk\n",
"from langchain_core.messages import AIMessageChunk\n",
"from langchain_core.runnables.config import ensure_config, get_callback_manager_for_config\n",
"from langchain_core.runnables.config import (\n",
" ensure_config,\n",
" get_callback_manager_for_config,\n",
")\n",
"\n",
"openai_client = AsyncOpenAI()\n",
"# define tool schema for openai tool calling\n",
"\n",
"tool = {\n",
" \"type\": \"function\",\n",
" \"function\": {\n",
" \"name\": \"get_items\",\n",
" \"description\": \"Use this tool to look up which items are in the given place.\",\n",
" \"parameters\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"place\": {\n",
" \"type\": \"string\"\n",
" }\n",
" },\n",
" \"required\": [\n",
" \"place\"\n",
" ]\n",
" }\n",
" }\n",
" \"type\": \"function\",\n",
" \"function\": {\n",
" \"name\": \"get_items\",\n",
" \"description\": \"Use this tool to look up which items are in the given place.\",\n",
" \"parameters\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\"place\": {\"type\": \"string\"}},\n",
" \"required\": [\"place\"],\n",
" },\n",
" },\n",
"}\n",
"\n",
"\n",
"async def call_model(state, config=None):\n",
" config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n",
" callback_manager = get_callback_manager_for_config(config)\n",
" messages = state[\"messages\"]\n",
" \n",
"\n",
" llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
" response = await openai_client.chat.completions.create(\n",
" messages=messages,\n",
" model=\"gpt-3.5-turbo\",\n",
" tools=[tool],\n",
" stream=True\n",
" messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
" )\n",
"\n",
" response_content = \"\"\n",
@@ -147,7 +143,10 @@
"\n",
" # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n",
" tool_call_chunk = ChatGenerationChunk(\n",
" message=AIMessageChunk(content=\"\", additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]})\n",
" message=AIMessageChunk(\n",
" content=\"\",\n",
" additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n",
" )\n",
" )\n",
" llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n",
" tool_call_function_arguments += delta.tool_calls[0].function.arguments\n",
@@ -156,8 +155,11 @@
" tool_calls = [\n",
" {\n",
" \"id\": tool_call_id,\n",
" \"function\": {\"name\": tool_call_function_name, \"arguments\": tool_call_function_arguments},\n",
" \"type\": \"function\"\n",
" \"function\": {\n",
" \"name\": tool_call_function_name,\n",
" \"arguments\": tool_call_function_arguments,\n",
" },\n",
" \"type\": \"function\",\n",
" }\n",
" ]\n",
" else:\n",
@@ -166,7 +168,7 @@
" response_message = {\n",
" \"role\": role,\n",
" \"content\": response_content,\n",
" \"tool_calls\": tool_calls\n",
" \"tool_calls\": tool_calls,\n",
" }\n",
" return {\"messages\": [response_message]}"
]
@@ -188,6 +190,7 @@
"source": [
"import json\n",
"\n",
"\n",
"async def get_items(place: str) -> str:\n",
" \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
" if \"bed\" in place: # For under the bed\n",
@@ -197,10 +200,10 @@
" else: # if the agent decides to ask about a different place\n",
" return \"cat snacks\"\n",
"\n",
"\n",
"# define mapping to look up functions when running tools\n",
"function_name_to_function = {\n",
" \"get_items\": get_items\n",
"}\n",
"function_name_to_function = {\"get_items\": get_items}\n",
"\n",
"\n",
"async def call_tools(state):\n",
" messages = state[\"messages\"]\n",
@@ -209,17 +212,15 @@
" function_name = tool_call[\"function\"][\"name\"]\n",
" function_arguments = tool_call[\"function\"][\"arguments\"]\n",
" arguments = json.loads(function_arguments)\n",
" \n",
" function_response = await function_name_to_function[function_name](**arguments) \n",
"\n",
" function_response = await function_name_to_function[function_name](**arguments)\n",
" tool_message = {\n",
" \"tool_call_id\": tool_call[\"id\"],\n",
" \"role\": \"tool\",\n",
" \"name\": function_name,\n",
" \"content\": function_response,\n",
" }\n",
" return {\n",
" \"messages\": [tool_message]\n",
" }"
" return {\"messages\": [tool_message]}"
]
},
{
@@ -242,16 +243,19 @@
"\n",
"from langgraph.graph import StateGraph, END\n",
"\n",
"\n",
"class State(TypedDict):\n",
" messages: Annotated[list, operator.add]\n",
"\n",
"\n",
"def should_continue(state) -> Literal[\"tools\", END]:\n",
" messages = state['messages']\n",
" messages = state[\"messages\"]\n",
" last_message = messages[-1]\n",
" if last_message[\"tool_calls\"]:\n",
" return \"tools\"\n",
" return END\n",
"\n",
"\n",
"workflow = StateGraph(State)\n",
"workflow.set_entry_point(\"model\")\n",
"workflow.add_node(\"model\", call_model) # i.e. our \"agent\"\n",
@@ -325,7 +329,9 @@
}
],
"source": [
"async for event in graph.astream_events({\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"):\n",
"async for event in graph.astream_events(\n",
" {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n",
"):\n",
" tags = event.get(\"tags\", [])\n",
" if event[\"event\"] == \"on_chat_model_stream\" and \"agent_llm\" in tags:\n",
" print(\"LLM token\", event[\"data\"][\"chunk\"].dict())"
examples/subgraph.ipynb
+41 -94
View File
@@ -22,8 +22,7 @@
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph"
"%%capture --no-stderr\n%pip install -U langgraph"
]
},
{
@@ -55,15 +54,17 @@
"from langgraph.checkpoint.memory import MemorySaver\n",
"from langgraph.graph import StateGraph, START, END\n",
"\n",
"\n",
"# The structure of the logs\n",
"class Logs(TypedDict):\n",
" id: str\n",
" question: str\n",
" docs: Optional[List]\n",
" answer: str \n",
" answer: str\n",
" grade: Optional[int]\n",
" grader: Optional[str]\n",
" feedback: Optional[str] \n",
" feedback: Optional[str]\n",
"\n",
"\n",
"# Failure Analysis Sub-graph\n",
"class FailureAnalysisState(TypedDict):\n",
@@ -71,23 +72,27 @@
" failures: List[Logs]\n",
" fa_summary: str\n",
"\n",
"\n",
"def get_failures(state):\n",
" docs = state['docs']\n",
" docs = state[\"docs\"]\n",
" failures = [doc for doc in docs if \"grade\" in doc]\n",
" return {\"failures\": failures}\n",
"\n",
"\n",
"def generate_summary(state):\n",
" failures = state['failures']\n",
" failures = state[\"failures\"]\n",
" # Add fxn: fa_summary = summarize(failures)\n",
" fa_summary = \"Poor quality retrieval of Chroma documentation.\" \n",
" fa_summary = \"Poor quality retrieval of Chroma documentation.\"\n",
" return {\"fa_summary\": fa_summary}\n",
"\n",
"\n",
"fa_builder = StateGraph(FailureAnalysisState)\n",
"fa_builder.add_node(\"get_failures\", get_failures)\n",
"fa_builder.add_node(\"generate_summary\", generate_summary)\n",
"fa_builder.set_entry_point(\"get_failures\")\n",
"fa_builder.add_edge(\"get_failures\",\"generate_summary\")\n",
"fa_builder.set_finish_point(\"generate_summary\")\n",
"fa_builder.add_edge(START, \"get_failures\")\n",
"fa_builder.add_edge(\"get_failures\", \"generate_summary\")\n",
"fa_builder.add_edge(\"generate_summary\", END)\n",
"\n",
"\n",
"# Summarization subgraph\n",
"class QuestionSummarizationState(TypedDict):\n",
@@ -95,29 +100,33 @@
" qs_summary: str\n",
" report: str\n",
"\n",
"\n",
"def generate_summary(state):\n",
" docs = state['docs']\n",
" docs = state[\"docs\"]\n",
" # Add fxn: summary = summarize(docs)\n",
" summary = \"Questions focused on usage of ChatOllama and Chroma vector store.\" \n",
" summary = \"Questions focused on usage of ChatOllama and Chroma vector store.\"\n",
" return {\"qs_summary\": summary}\n",
"\n",
"\n",
"def send_to_slack(state):\n",
" qs_summary = state['qs_summary']\n",
" qs_summary = state[\"qs_summary\"]\n",
" # Add fxn: report = report_generation(qs_summary)\n",
" report = \"foo bar baz\"\n",
" return {\"report\": report}\n",
"\n",
"\n",
"def format_report_for_slack(state):\n",
" report = state['report']\n",
" report = state[\"report\"]\n",
" # Add fxn: formatted_report = report_format(report)\n",
" formatted_report = \"foo bar\"\n",
" return {\"report\": formatted_report}\n",
"\n",
"\n",
"qs_builder = StateGraph(QuestionSummarizationState)\n",
"qs_builder.add_node(\"generate_summary\", generate_summary)\n",
"qs_builder.add_node(\"send_to_slack\", send_to_slack)\n",
"qs_builder.add_node(\"format_report_for_slack\", format_report_for_slack)\n",
"qs_builder.set_entry_point(\"generate_summary\")\n",
"qs_builder.add_edge(START, \"generate_summary\")\n",
"qs_builder.add_edge(\"generate_summary\", \"send_to_slack\")\n",
"qs_builder.add_edge(\"send_to_slack\", \"format_report_for_slack\")\n",
"qs_builder.add_edge(\"format_report_for_slack\", END)"
@@ -165,25 +174,28 @@
" feedback=\"The retrieved documents discuss vector stores in general, but not Chroma specifically\",\n",
")\n",
"\n",
"\n",
"# Entry Graph\n",
"class EntryGraphState(TypedDict):\n",
" raw_logs: Annotated[List[Dict], add]\n",
" docs: Annotated[List[Logs], add] # This will be used in sub-graphs\n",
" fa_summary: str # This will be generated in the FA sub-graph\n",
" report: str # This will be generated in the QS sub-graph\n",
" docs: Annotated[List[Logs], add] # This will be used in sub-graphs\n",
" fa_summary: str # This will be generated in the FA sub-graph\n",
" report: str # This will be generated in the QS sub-graph\n",
"\n",
"\n",
"def convert_logs_to_docs(state):\n",
" # Get logs\n",
" raw_logs = state['raw_logs']\n",
" docs = [question_answer,question_answer_feedback]\n",
" return {\"docs\": docs} \n",
" raw_logs = state[\"raw_logs\"]\n",
" docs = [question_answer, question_answer_feedback]\n",
" return {\"docs\": docs}\n",
"\n",
"\n",
"entry_builder = StateGraph(EntryGraphState)\n",
"entry_builder.add_node(\"convert_logs_to_docs\", convert_logs_to_docs)\n",
"entry_builder.add_node(\"question_summarization\", qs_builder.compile())\n",
"entry_builder.add_node(\"failure_analysis\", fa_builder.compile())\n",
"\n",
"entry_builder.set_entry_point(\"convert_logs_to_docs\")\n",
"entry_builder.add_edge(START, \"convert_logs_to_docs\")\n",
"entry_builder.add_edge(\"convert_logs_to_docs\", \"failure_analysis\")\n",
"entry_builder.add_edge(\"convert_logs_to_docs\", \"question_summarization\")\n",
"entry_builder.add_edge(\"failure_analysis\", END)\n",
@@ -192,6 +204,7 @@
"graph = entry_builder.compile()\n",
"\n",
"from IPython.display import Image, display\n",
"\n",
"# Setting xray to 1 will show the internal structure of the nested graph\n",
"display(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
]
@@ -233,7 +246,7 @@
}
],
"source": [
"raw_logs = [{\"foo\":\"bar\"},{\"foo\":\"baz\"}]\n",
"raw_logs = [{\"foo\": \"bar\"}, {\"foo\": \"baz\"}]\n",
"graph.invoke({\"raw_logs\": raw_logs}, debug=False)"
]
},
@@ -260,6 +273,7 @@
"\n",
"from typing_extensions import TypedDict\n",
"\n",
"\n",
"def reduce_list(left: list | None, right: list | None) -> list:\n",
" if not left:\n",
" left = []\n",
@@ -324,10 +338,7 @@
}
],
"source": [
"from IPython.display import Image, display\n",
"\n",
"# Setting xray to 1 will show the internal structure of the nested graph\n",
"display(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
"from IPython.display import Image, display\n\n# Setting xray to 1 will show the internal structure of the nested graph\ndisplay(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
]
},
{
@@ -445,42 +456,7 @@
"metadata": {},
"outputs": [],
"source": [
"import uuid\n",
"\n",
"\n",
"def reduce_list(left: list | None, right: list | None) -> list:\n",
" \"\"\"Append the right-hand list, replacing any elements with the same id in the left-hand list.\"\"\"\n",
" if not left:\n",
" left = []\n",
" if not right:\n",
" right = []\n",
" left_, right_ = [], []\n",
" for orig, new in [(left, left_), (right, right_)]:\n",
" for val in orig:\n",
" if not isinstance(val, dict):\n",
" val = {\"val\": val}\n",
" if \"id\" not in val:\n",
" val[\"id\"] = str(uuid.uuid4())\n",
" new.append(val)\n",
" # Merge the two lists\n",
" left_idx_by_id = {val[\"id\"]: i for i, val in enumerate(left_)}\n",
" merged = left_.copy()\n",
" for val in right_:\n",
" if (existing_idx := left_idx_by_id.get(val[\"id\"])) is not None:\n",
" merged[existing_idx] = val\n",
" else:\n",
" merged.append(val)\n",
" return merged\n",
"\n",
"\n",
"class ChildState(TypedDict):\n",
" name: str\n",
" path: Annotated[list[str], reduce_list]\n",
"\n",
"\n",
"class ParentState(TypedDict):\n",
" name: str\n",
" path: Annotated[list[str], reduce_list]"
"import uuid\n\n\ndef reduce_list(left: list | None, right: list | None) -> list:\n \"\"\"Append the right-hand list, replacing any elements with the same id in the left-hand list.\"\"\"\n if not left:\n left = []\n if not right:\n right = []\n left_, right_ = [], []\n for orig, new in [(left, left_), (right, right_)]:\n for val in orig:\n if not isinstance(val, dict):\n val = {\"val\": val}\n if \"id\" not in val:\n val[\"id\"] = str(uuid.uuid4())\n new.append(val)\n # Merge the two lists\n left_idx_by_id = {val[\"id\"]: i for i, val in enumerate(left_)}\n merged = left_.copy()\n for val in right_:\n if (existing_idx := left_idx_by_id.get(val[\"id\"])) is not None:\n merged[existing_idx] = val\n else:\n merged.append(val)\n return merged\n\n\nclass ChildState(TypedDict):\n name: str\n path: Annotated[list[str], reduce_list]\n\n\nclass ParentState(TypedDict):\n name: str\n path: Annotated[list[str], reduce_list]"
]
},
{
@@ -489,33 +465,7 @@
"metadata": {},
"outputs": [],
"source": [
"child_builder = StateGraph(ChildState)\n",
"\n",
"child_builder.add_node(\"child_start\", lambda state: {\"path\": [\"child_start\"]})\n",
"child_builder.add_edge(START, \"child_start\")\n",
"child_builder.add_node(\"child_middle\", lambda state: {\"path\": [\"child_middle\"]})\n",
"child_builder.add_node(\"child_end\", lambda state: {\"path\": [\"child_end\"]})\n",
"child_builder.add_edge(\"child_start\", \"child_middle\")\n",
"child_builder.add_edge(\"child_middle\", \"child_end\")\n",
"child_builder.add_edge(\"child_end\", END)\n",
"\n",
"builder = StateGraph(ParentState)\n",
"\n",
"builder.add_node(\"grandparent\", lambda state: {\"path\": [\"grandparent\"]})\n",
"builder.add_edge(START, \"grandparent\")\n",
"builder.add_node(\"parent\", lambda state: {\"path\": [\"parent\"]})\n",
"builder.add_node(\"child\", child_builder.compile())\n",
"builder.add_node(\"sibling\", lambda state: {\"path\": [\"sibling\"]})\n",
"builder.add_node(\"fin\", lambda state: {\"path\": [\"fin\"]})\n",
"\n",
"# Add connections\n",
"builder.add_edge(\"grandparent\", \"parent\")\n",
"builder.add_edge(\"parent\", \"child\")\n",
"builder.add_edge(\"parent\", \"sibling\")\n",
"builder.add_edge(\"child\", \"fin\")\n",
"builder.add_edge(\"sibling\", \"fin\")\n",
"builder.add_edge(\"fin\", END)\n",
"graph = builder.compile()"
"child_builder = StateGraph(ChildState)\n\nchild_builder.add_node(\"child_start\", lambda state: {\"path\": [\"child_start\"]})\nchild_builder.add_edge(START, \"child_start\")\nchild_builder.add_node(\"child_middle\", lambda state: {\"path\": [\"child_middle\"]})\nchild_builder.add_node(\"child_end\", lambda state: {\"path\": [\"child_end\"]})\nchild_builder.add_edge(\"child_start\", \"child_middle\")\nchild_builder.add_edge(\"child_middle\", \"child_end\")\nchild_builder.add_edge(\"child_end\", END)\n\nbuilder = StateGraph(ParentState)\n\nbuilder.add_node(\"grandparent\", lambda state: {\"path\": [\"grandparent\"]})\nbuilder.add_edge(START, \"grandparent\")\nbuilder.add_node(\"parent\", lambda state: {\"path\": [\"parent\"]})\nbuilder.add_node(\"child\", child_builder.compile())\nbuilder.add_node(\"sibling\", lambda state: {\"path\": [\"sibling\"]})\nbuilder.add_node(\"fin\", lambda state: {\"path\": [\"fin\"]})\n\n# Add connections\nbuilder.add_edge(\"grandparent\", \"parent\")\nbuilder.add_edge(\"parent\", \"child\")\nbuilder.add_edge(\"parent\", \"sibling\")\nbuilder.add_edge(\"child\", \"fin\")\nbuilder.add_edge(\"sibling\", \"fin\")\nbuilder.add_edge(\"fin\", END)\ngraph = builder.compile()"
]
},
{
@@ -535,10 +485,7 @@
}
],
"source": [
"from IPython.display import Image, display\n",
"\n",
"# Setting xray to 1 will show the internal structure of the nested graph\n",
"display(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
"from IPython.display import Image, display\n\n# Setting xray to 1 will show the internal structure of the nested graph\ndisplay(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
]
},
{
examples/tool-calling-errors.ipynb
+2 -2
View File
@@ -547,8 +547,8 @@
],
"source": [
"stream = app.stream(\n",
" {\"messages\": [(\"human\", \"Write me an incredible haiku about water.\")]},\n",
" {\"recursion_limit\": 10},\n",
" {\"messages\": [(\"human\", \"Write me an incredible haiku about water.\")]},\n",
" {\"recursion_limit\": 10},\n",
")\n",
"\n",
"for chunk in stream:\n",
examples/tool-calling.ipynb
+23 -6
View File
@@ -139,7 +139,14 @@
"source": [
"message_with_single_tool_call = AIMessage(\n",
" content=\"\",\n",
" tool_calls=[{'name': 'get_weather', 'args': {'location': 'sf'}, 'id': 'tool_call_id', 'type': 'tool_call'}]\n",
" tool_calls=[\n",
" {\n",
" \"name\": \"get_weather\",\n",
" \"args\": {\"location\": \"sf\"},\n",
" \"id\": \"tool_call_id\",\n",
" \"type\": \"tool_call\",\n",
" }\n",
" ],\n",
")\n",
"\n",
"tool_node.invoke({\"messages\": [message_with_single_tool_call]})"
@@ -175,9 +182,19 @@
"message_with_multiple_tool_calls = AIMessage(\n",
" content=\"\",\n",
" tool_calls=[\n",
" {'name': 'get_coolest_cities', 'args': {}, 'id': 'tool_call_id_1', 'type': 'tool_call'},\n",
" {'name': 'get_weather', 'args': {'location': 'sf'}, 'id': 'tool_call_id_2', 'type': 'tool_call'}\n",
" ]\n",
" {\n",
" \"name\": \"get_coolest_cities\",\n",
" \"args\": {},\n",
" \"id\": \"tool_call_id_1\",\n",
" \"type\": \"tool_call\",\n",
" },\n",
" {\n",
" \"name\": \"get_weather\",\n",
" \"args\": {\"location\": \"sf\"},\n",
" \"id\": \"tool_call_id_2\",\n",
" \"type\": \"tool_call\",\n",
" },\n",
" ],\n",
")\n",
"\n",
"tool_node.invoke({\"messages\": [message_with_multiple_tool_calls]})"
@@ -210,7 +227,6 @@
"from langgraph.prebuilt import ToolNode\n",
"\n",
"\n",
"\n",
"model_with_tools = ChatAnthropic(\n",
" model=\"claude-3-haiku-20240307\", temperature=0\n",
").bind_tools(tools)"
@@ -454,7 +470,8 @@
"# example with a multiple tool calls in succession\n",
"\n",
"for chunk in app.stream(\n",
" {\"messages\": [(\"human\", \"what's the weather in the coolest cities?\")]}, stream_mode=\"values\"\n",
" {\"messages\": [(\"human\", \"what's the weather in the coolest cities?\")]},\n",
" stream_mode=\"values\",\n",
"):\n",
" chunk[\"messages\"][-1].pretty_print()"
]
libs/langgraph/langgraph/channels/manager.py
+2
View File
@@ -61,4 +61,6 @@ def create_checkpoint(
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
pending_sends=checkpoint.get("pending_sends", []),
# checkpoints are saved only at the end of a step, ie. when current tasks should be cleared
current_tasks={},
)
libs/langgraph/langgraph/checkpoint/aiosqlite.py
+32 -17
View File
@@ -46,22 +46,29 @@ def not_implemented_sync_method(func: T) -> T:
class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
"""An asynchronous checkpoint saver that stores checkpoints in a SQLite database.
This class provides an asynchronous interface for saving and retrieving checkpoints
using a SQLite database. It's designed for use in asynchronous environments and
offers better performance for I/O-bound operations compared to synchronous alternatives.
Attributes:
conn (aiosqlite.Connection): The asynchronous SQLite database connection.
serde (SerializerProtocol): The serializer used for encoding/decoding checkpoints.
Tip:
Requires the [aiosqlite](https://pypi.org/project/aiosqlite/) package.
Install it with `pip install aiosqlite`.
Note:
While this class does support asynchronous checkpointing, it is not recommended
for production workloads, due to limitations in SQLite's write performance. For
production workloads, consider using a more robust database like PostgreSQL.
Warning:
While this class supports asynchronous checkpointing, it is not recommended
for production workloads due to limitations in SQLite's write performance.
For production use, consider a more robust database like PostgreSQL.
!!! Important
Tip:
Remember to **close the database connection** after executing your code,
otherwise, you may see the graph "hang" after execution (since the program
will not exit until the connection is closed).
The easiest way to do this is to use the `async with` statement, as shown in the
examples below.
The easiest way is to use the `async with` statement as shown in the examples.
```python
async with AsyncSqliteSaver.from_conn_string("checkpoints.sqlite") as saver:
@@ -72,12 +79,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
print(event)
```
Args:
conn (aiosqlite.Connection): The asynchronous SQLite database connection.
serde (Optional[SerializerProtocol]): The serializer to use for serializing and deserializing checkpoints. Defaults to JsonPlusSerializerCompat.
Examples:
Usage within a StateGraph:
Usage within StateGraph:
```pycon
>>> import asyncio
>>> import aiosqlite
@@ -95,8 +99,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
>>> asyncio.run(coro)
Output: 2
```
Raw usage:
```pycon
>>> import asyncio
>>> import aiosqlite
@@ -309,12 +313,13 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
on the provided config. The checkpoints are ordered by timestamp in descending order.
Args:
config (RunnableConfig): The config to use for listing the checkpoints.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
before (Optional[RunnableConfig]): List checkpoints created before this configuration.
limit (Optional[int]): Maximum number of checkpoints to return.
Yields:
AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
"""
await self.setup()
where, param_values = search_where(config, filter, before)
@@ -356,6 +361,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
Args:
config (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
Returns:
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
@@ -385,6 +391,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
writes: Sequence[Tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
This method saves intermediate writes associated with a checkpoint to the database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
"""
await self.setup()
async with self.conn.executemany(
"INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)",
libs/langgraph/langgraph/checkpoint/base.py
+31 -11
View File
@@ -1,5 +1,4 @@
from abc import ABC
from collections import defaultdict
from datetime import datetime, timezone
from typing import (
Any,
@@ -25,10 +24,13 @@ from langgraph.serde.base import SerializerProtocol
from langgraph.serde.jsonplus import JsonPlusSerializer
V = TypeVar("V", int, float, str)
PendingWrite = Tuple[str, str, Any]
# Marked as total=False to allow for future expansion.
class CheckpointMetadata(TypedDict, total=False):
"""Metadata associated with a checkpoint."""
source: Literal["input", "loop", "update"]
"""The source of the checkpoint.
- "input": The checkpoint was created from an input to invoke/stream/batch.
@@ -53,6 +55,10 @@ class CheckpointMetadata(TypedDict, total=False):
"""
class TaskInfo(TypedDict):
status: Literal["scheduled", "success", "error"]
class Checkpoint(TypedDict):
"""State snapshot at a given point in time."""
@@ -74,7 +80,7 @@ class Checkpoint(TypedDict):
The keys are channel names and the values are the logical time step
at which the channel was last updated.
"""
versions_seen: defaultdict[str, dict[str, Union[str, int, float]]]
versions_seen: dict[str, dict[str, Union[str, int, float]]]
"""Map from node ID to map from channel name to version seen.
This keeps track of the versions of the channels that each node has seen.
@@ -84,6 +90,8 @@ class Checkpoint(TypedDict):
pending_sends: List[Send]
"""List of packets sent to nodes but not yet processed.
Cleared by the next checkpoint."""
current_tasks: Dict[str, TaskInfo]
"""Map from task ID to task info."""
def empty_checkpoint() -> Checkpoint:
@@ -93,8 +101,9 @@ def empty_checkpoint() -> Checkpoint:
ts=datetime.now(timezone.utc).isoformat(),
channel_values={},
channel_versions={},
versions_seen=defaultdict(dict),
versions_seen={},
pending_sends=[],
current_tasks={},
)
@@ -105,20 +114,20 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
id=checkpoint["id"],
channel_values=checkpoint["channel_values"].copy(),
channel_versions=checkpoint["channel_versions"].copy(),
versions_seen=defaultdict(
dict,
{k: v.copy() for k, v in checkpoint["versions_seen"].items()},
),
versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
pending_sends=checkpoint.get("pending_sends", []).copy(),
current_tasks=checkpoint.get("current_tasks", {}).copy(),
)
class CheckpointTuple(NamedTuple):
"""A tuple containing a checkpoint and its associated data."""
config: RunnableConfig
checkpoint: Checkpoint
metadata: CheckpointMetadata
parent_config: Optional[RunnableConfig] = None
pending_writes: Optional[List[Tuple[str, str, Any]]] = None
pending_writes: Optional[List[PendingWrite]] = None
CheckpointThreadId = ConfigurableFieldSpec(
@@ -264,11 +273,13 @@ class BaseCheckpointSaver(ABC):
)
async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
"""
Asynchronously fetch a checkpoint using the given configuration.
"""Asynchronously fetch a checkpoint using the given configuration.
Args:
config (RunnableConfig): Configuration specifying which checkpoint to retrieve.
Returns:
Optional[Checkpoint]: The requested checkpoint, or None if not found.
"""
if value := await self.aget_tuple(config):
return value.checkpoint
@@ -281,6 +292,9 @@ class BaseCheckpointSaver(ABC):
Returns:
Optional[CheckpointTuple]: The requested checkpoint tuple, or None if not found.
Raises:
NotImplementedError: Implement this method in your custom checkpoint saver.
"""
raise NotImplementedError
@@ -296,12 +310,15 @@ class BaseCheckpointSaver(ABC):
Args:
config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
before (Optional[RunnableConfig]): List checkpoints created before this configuration.
limit (Optional[int]): Maximum number of checkpoints to return.
Returns:
AsyncIterator[CheckpointTuple]: Async iterator of matching checkpoint tuples.
Raises:
NotImplementedError: Implement this method in your custom checkpoint saver.
"""
raise NotImplementedError
yield
@@ -321,6 +338,9 @@ class BaseCheckpointSaver(ABC):
Returns:
RunnableConfig: Updated configuration after storing the checkpoint.
Raises:
NotImplementedError: Implement this method in your custom checkpoint saver.
"""
raise NotImplementedError
libs/langgraph/langgraph/checkpoint/memory.py
+28 -5
View File
@@ -107,15 +107,16 @@ class MemorySaver(BaseCheckpointSaver):
"""List checkpoints from the in-memory storage.
This method retrieves a list of checkpoint tuples from the in-memory storage based
on the provided config. The checkpoints are ordered by timestamp in insertion order.
on the provided criteria.
Args:
config (RunnableConfig): The config to use for listing the checkpoints.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
before (Optional[RunnableConfig]): List checkpoints created before this configuration.
limit (Optional[int]): Maximum number of checkpoints to return.
Yields:
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
"""
thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
for thread_id in thread_ids:
@@ -158,6 +159,7 @@ class MemorySaver(BaseCheckpointSaver):
Args:
config (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
Returns:
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
@@ -191,6 +193,7 @@ class MemorySaver(BaseCheckpointSaver):
Args:
config (RunnableConfig): The config to associate with the writes.
writes (list[tuple[str, Any]]): The writes to save.
task_id (str): Identifier for the task creating the writes.
Returns:
RunnableConfig: The updated config containing the saved writes' timestamp.
@@ -254,6 +257,16 @@ class MemorySaver(BaseCheckpointSaver):
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
"""Asynchronous version of put.
Args:
config (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
Returns:
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
"""
return await asyncio.get_running_loop().run_in_executor(
None, self.put, config, checkpoint, metadata
)
@@ -264,6 +277,16 @@ class MemorySaver(BaseCheckpointSaver):
writes: List[Tuple[str, Any]],
task_id: str,
) -> RunnableConfig:
"""Asynchronous version of put_writes.
This method is an asynchronous wrapper around put_writes that runs the synchronous
method in a separate thread using asyncio.
Args:
config (RunnableConfig): The config to associate with the writes.
writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
"""
return await asyncio.get_running_loop().run_in_executor(
None, self.put_writes, config, writes, task_id
)
libs/langgraph/langgraph/checkpoint/sqlite.py
+21
View File
@@ -309,6 +309,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
Args:
config (RunnableConfig): The config to use for listing the checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
@@ -410,6 +411,15 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
writes: Sequence[Tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint.
This method saves intermediate writes associated with a checkpoint to the SQLite database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
"""
with self.lock, self.cursor() as cur:
cur.executemany(
"INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)",
@@ -467,6 +477,17 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
raise NotImplementedError(_AIO_ERROR_MSG)
def get_next_version(self, current: Optional[str], channel: BaseChannel) -> str:
"""Generate the next version ID for a channel.
This method creates a new version identifier for a channel based on its current version.
Args:
current (Optional[str]): The current version identifier of the channel.
channel (BaseChannel): The channel being versioned.
Returns:
str: The next version identifier, which is guaranteed to be monotonically increasing.
"""
if current is None:
current_v = 0
else:
libs/langgraph/langgraph/constants.py
+2 -1
View File
@@ -1,10 +1,11 @@
from typing import Any
INPUT = "__input__"
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
INTERRUPT = "__interrupt__"
TASKS = "__pregel_tasks"
RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ}
RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, INPUT}
TAG_HIDDEN = "langsmith:hidden"
START = "__start__"
libs/langgraph/langgraph/graph/state.py
+1
View File
@@ -254,6 +254,7 @@ class StateGraph(Graph):
action (Optional[RunnableLike]): The action associated with the node. (default: None)
metadata (Optional[dict[str, Any]]): The metadata associated with the node. (default: None)
input (Optional[Type[Any]]): The input schema for the node. (default: the graph's input schema)
retry (Optional[RetryPolicy]): The policy for retrying the node. (default: None)
Raises:
ValueError: If the key is already being used as a state key.
libs/langgraph/langgraph/prebuilt/__init__.py
+2 -1
View File
@@ -3,7 +3,7 @@ from langgraph.prebuilt import chat_agent_executor
from langgraph.prebuilt.agent_executor import create_agent_executor
from langgraph.prebuilt.chat_agent_executor import create_react_agent
from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
from langgraph.prebuilt.tool_node import ToolNode, tools_condition
from langgraph.prebuilt.tool_node import InjectedState, ToolNode, tools_condition
from langgraph.prebuilt.tool_validator import ValidationNode
__all__ = [
@@ -15,4 +15,5 @@ __all__ = [
"ToolNode",
"tools_condition",
"ValidationNode",
"InjectedState",
]
libs/langgraph/langgraph/prebuilt/tool_node.py
+189 -61
View File
@@ -1,11 +1,24 @@
import asyncio
from typing import Any, Callable, Dict, Literal, Optional, Sequence, Union
from copy import copy
from typing import (
Any,
Callable,
Dict,
List,
Literal,
Optional,
Sequence,
Tuple,
Union,
cast,
)
from langchain_core.messages import AIMessage, AnyMessage, ToolCall, ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.config import get_executor_for_config
from langchain_core.tools import BaseTool
from langchain_core.runnables.config import get_config_list, get_executor_for_config
from langchain_core.tools import BaseTool, InjectedToolArg
from langchain_core.tools import tool as create_tool
from typing_extensions import get_args
from langgraph.utils import RunnableCallable
@@ -60,47 +73,49 @@ class ToolNode(RunnableCallable):
def _func(
self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
) -> Any:
if isinstance(input, list):
output_type = "list"
message: AnyMessage = input[-1]
elif messages := input.get("messages", []):
output_type = "dict"
message = messages[-1]
else:
raise ValueError("No message found in input")
if not isinstance(message, AIMessage):
raise ValueError("Last message is not an AIMessage")
def run_one(call: ToolCall):
if (requested_tool := call["name"]) not in self.tools_by_name:
content = INVALID_TOOL_NAME_ERROR_TEMPLATE.format(
requested_tool=requested_tool,
available_tools=", ".join(self.tools_by_name.keys()),
)
return ToolMessage(
content, name=requested_tool, tool_call_id=call["id"]
)
try:
input = {**call, **{"type": "tool_call"}}
return self.tools_by_name[call["name"]].invoke(input, config)
except Exception as e:
if not self.handle_tool_errors:
raise e
content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
tool_calls, output_type = self._parse_input(input)
config_list = get_config_list(config, len(tool_calls))
with get_executor_for_config(config) as executor:
outputs = [*executor.map(run_one, message.tool_calls)]
if output_type == "list":
return outputs
else:
return {"messages": outputs}
outputs = [*executor.map(self._run_one, tool_calls, config_list)]
return outputs if output_type == "list" else {"messages": outputs}
async def _afunc(
self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
) -> Any:
tool_calls, output_type = self._parse_input(input)
outputs = await asyncio.gather(
*(self._arun_one(call, config) for call in tool_calls)
)
return outputs if output_type == "list" else {"messages": outputs}
def _run_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
if invalid_tool_message := self._validate_tool_call(call):
return invalid_tool_message
try:
input = {**call, **{"type": "tool_call"}}
return self.tools_by_name[call["name"]].invoke(input, config)
except Exception as e:
if not self.handle_tool_errors:
raise e
content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
async def _arun_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
if invalid_tool_message := self._validate_tool_call(call):
return invalid_tool_message
try:
input = {**call, **{"type": "tool_call"}}
return await self.tools_by_name[call["name"]].ainvoke(input, config)
except Exception as e:
if not self.handle_tool_errors:
raise e
content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
def _parse_input(
self, input: Union[list[AnyMessage], dict[str, Any]]
) -> Tuple[List[ToolCall], Literal["list", "dict"]]:
if isinstance(input, list):
output_type = "list"
message: AnyMessage = input[-1]
@@ -113,30 +128,54 @@ class ToolNode(RunnableCallable):
if not isinstance(message, AIMessage):
raise ValueError("Last message is not an AIMessage")
async def run_one(call: ToolCall):
if (requested_tool := call["name"]) not in self.tools_by_name:
content = INVALID_TOOL_NAME_ERROR_TEMPLATE.format(
requested_tool=requested_tool,
available_tools=", ".join(self.tools_by_name.keys()),
)
return ToolMessage(
content, name=requested_tool, tool_call_id=call["id"]
)
tool_calls = [
self._inject_state(call, input)
for call in cast(AIMessage, message).tool_calls
]
return tool_calls, output_type
try:
input = {**call, **{"type": "tool_call"}}
return await self.tools_by_name[call["name"]].ainvoke(input, config)
except Exception as e:
if not self.handle_tool_errors:
raise e
content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
outputs = await asyncio.gather(*(run_one(call) for call in message.tool_calls))
if output_type == "list":
return outputs
def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]:
if (requested_tool := call["name"]) not in self.tools_by_name:
content = INVALID_TOOL_NAME_ERROR_TEMPLATE.format(
requested_tool=requested_tool,
available_tools=", ".join(self.tools_by_name.keys()),
)
return ToolMessage(content, name=requested_tool, tool_call_id=call["id"])
else:
return {"messages": outputs}
return None
def _inject_state(
self, tool_call: ToolCall, input: Union[list[AnyMessage], dict[str, Any]]
) -> ToolCall:
if tool_call["name"] not in self.tools_by_name:
return tool_call
state_args = _get_state_args(self.tools_by_name[tool_call["name"]])
if state_args and not isinstance(input, dict):
required_fields = list(state_args.values())
if (
len(required_fields) == 1
and required_fields[0] == "messages"
or required_fields[0] is None
):
input = {"messages": input}
else:
err_msg = (
f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
f"graph state dict as input."
)
if any(state_field for state_field in state_args.values()):
required_fields_str = ", ".join(f for f in required_fields if f)
err_msg += f" State should contain fields {required_fields_str}."
raise ValueError(err_msg)
tool_call_copy: ToolCall = copy(tool_call)
tool_call_copy["args"] = {
**tool_call_copy["args"],
**{
tool_arg: cast(dict, input)[state_field] if state_field else input
for tool_arg, state_field in state_args.items()
},
}
return tool_call_copy
def tools_condition(
@@ -194,3 +233,92 @@ def tools_condition(
if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
return "tools"
return "__end__"
class InjectedState(InjectedToolArg):
"""Annotation for a Tool arg that is meant to be populated with the graph state.
Any Tool argument annotated with InjectedState will be hidden from a tool-calling
model, so that the model doesn't attempt to generate the argument. If using
ToolNode, the appropriate graph state field will be automatically injected into
the model-generated tool args.
Args:
field: The key from state to insert. If None, the entire state is expected to
be passed in.
Example:
```python
from typing import List
from typing_extensions import Annotated, TypedDict
from langchain_core.messages import BaseMessage, AIMessage
from langchain_core.tools import tool
from langgraph.prebuilt import InjectedState, ToolNode
class AgentState(TypedDict):
messages: List[BaseMessage]
foo: str
@tool
def state_tool(x: int, state: Annotated[dict, InjectedState]) -> str:
'''Do something with state.'''
if len(state["messages"]) > 2:
return state["foo"] + str(x)
else:
return "not enough messages"
@tool
def foo_tool(x: int, foo: Annotated[str, InjectedState("foo")]) -> str:
'''Do something else with state.'''
return foo + str(x + 1)
node = ToolNode([state_tool, foo_tool])
tool_call1 = {"name": "state_tool", "args": {"x": 1}, "id": "1", "type": "tool_call"}
tool_call2 = {"name": "foo_tool", "args": {"x": 1}, "id": "2", "type": "tool_call"}
state = {
"messages": [AIMessage("", tool_calls=[tool_call1, tool_call2])],
"foo": "bar",
}
node.invoke(state)
```
```pycon
[
ToolMessage(content='not enough messages', name='state_tool', tool_call_id='1'),
ToolMessage(content='bar2', name='foo_tool', tool_call_id='2')
]
```
""" # noqa: E501
def __init__(self, field: Optional[str] = None) -> None:
self.field = field
def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]:
full_schema = tool.get_input_schema()
tool_args_to_state_fields: Dict = {}
for name, type_ in full_schema.__annotations__.items():
injections = [
type_arg
for type_arg in get_args(type_)
if isinstance(type_arg, InjectedState)
or (isinstance(type_arg, type) and issubclass(type_arg, InjectedState))
]
if len(injections) > 1:
raise ValueError(
"A tool argument should not be annotated with InjectedState more than "
f"once. Received arg {name} with annotations {injections}."
)
elif len(injections) == 1:
injection = injections[0]
if isinstance(injection, InjectedState) and injection.field:
tool_args_to_state_fields[name] = injection.field
else:
tool_args_to_state_fields[name] = None
else:
pass
return tool_args_to_state_fields
libs/langgraph/langgraph/pregel/__init__.py
+135 -939
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/langgraph/pregel/algo.py
+427
View File
@@ -0,0 +1,427 @@
import json
from collections import defaultdict, deque
from functools import partial
from typing import (
Any,
Callable,
Iterator,
Literal,
Mapping,
NamedTuple,
Optional,
Protocol,
Sequence,
Union,
overload,
)
from uuid import UUID, uuid5
from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
from langchain_core.runnables.config import (
RunnableConfig,
merge_configs,
patch_config,
)
from langgraph.channels.base import BaseChannel
from langgraph.channels.context import Context
from langgraph.channels.manager import ChannelsManager, create_checkpoint
from langgraph.checkpoint.base import Checkpoint, copy_checkpoint
from langgraph.constants import (
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
INTERRUPT,
RESERVED,
TAG_HIDDEN,
TASKS,
Send,
)
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.managed.base import ManagedValueMapping, is_managed_value
from langgraph.pregel.io import read_channel, read_channels
from langgraph.pregel.log import logger
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescription
class WritesProtocol(Protocol):
name: str
writes: Sequence[tuple[str, Any]]
triggers: Sequence[str]
class PregelTaskWrites(NamedTuple):
name: str
writes: Sequence[tuple[str, Any]]
triggers: Sequence[str]
def should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Union[All, Sequence[str]],
tasks: list[PregelExecutableTask],
) -> bool:
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
null_version = version_type()
seen = checkpoint["versions_seen"].get(INTERRUPT, {})
return (
# interrupt if any channel has been updated since last interrupt
any(
version > seen.get(chan, null_version)
for chan, version in checkpoint["channel_versions"].items()
)
# and any triggered node is in interrupt_nodes list
and any(
task.name
for task in tasks
if (
(not task.config or TAG_HIDDEN not in task.config.get("tags"))
if interrupt_nodes == "*"
else task.name in interrupt_nodes
)
)
)
def local_read(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
task: WritesProtocol,
config: RunnableConfig,
select: Union[list[str], str],
fresh: bool = False,
) -> Union[dict[str, Any], Any]:
if fresh:
new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1)
context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
with ChannelsManager(
{k: v for k, v in channels.items() if k not in context_channels},
new_checkpoint,
config,
) as channels:
all_channels = {**channels, **context_channels}
apply_writes(new_checkpoint, all_channels, [task], None)
return read_channels(all_channels, select)
else:
return read_channels(channels, select)
def local_write(
commit: Callable[[Sequence[tuple[str, Any]]], None],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
writes: Sequence[tuple[str, Any]],
) -> None:
for chan, value in writes:
if chan == TASKS:
if not isinstance(value, Send):
raise InvalidUpdateError(
f"Invalid packet type, expected Packet, got {value}"
)
if value.node not in processes:
raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
elif chan not in channels:
logger.warning(f"Skipping write for channel '{chan}' which has no readers")
commit(writes)
def increment(current: Optional[int], channel: BaseChannel) -> int:
return current + 1 if current is not None else 1
def apply_writes(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
tasks: Sequence[WritesProtocol],
get_next_version: Optional[Callable[[int, BaseChannel], int]],
) -> None:
# update seen versions
for task in tasks:
checkpoint["versions_seen"].setdefault(task.name, {}).update(
{
chan: checkpoint["channel_versions"][chan]
for chan in task.triggers
if chan in checkpoint["channel_versions"]
}
)
# Find the highest version of all channels
if checkpoint["channel_versions"]:
max_version = max(checkpoint["channel_versions"].values())
else:
max_version = None
# Consume all channels that were read
for chan in {
chan for task in tasks for chan in task.triggers if chan not in RESERVED
}:
if channels[chan].consume():
if get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, channels[chan]
)
# clear pending sends
if checkpoint["pending_sends"]:
checkpoint["pending_sends"].clear()
# Group writes by channel
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
for task in tasks:
for chan, val in task.writes:
if chan == TASKS:
checkpoint["pending_sends"].append(val)
else:
pending_writes_by_channel[chan].append(val)
# Find the highest version of all channels
if checkpoint["channel_versions"]:
max_version = max(checkpoint["channel_versions"].values())
else:
max_version = None
# Apply writes to channels
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
if updated and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, channels[chan]
)
updated_channels.add(chan)
# Channels that weren't updated in this step are notified of a new step
for chan in channels:
if chan not in updated_channels:
if channels[chan].update([]) and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, channels[chan]
)
@overload
def prepare_next_tasks(
checkpoint: Checkpoint,
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
for_execution: Literal[False],
manager: Literal[None] = None,
) -> list[PregelTaskDescription]:
...
@overload
def prepare_next_tasks(
checkpoint: Checkpoint,
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
for_execution: Literal[True],
manager: Union[None, ParentRunManager, AsyncParentRunManager],
) -> list[PregelExecutableTask]:
...
def prepare_next_tasks(
checkpoint: Checkpoint,
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
*,
for_execution: bool,
manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]:
tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = []
# Consume pending packets
for packet in checkpoint["pending_sends"]:
if not isinstance(packet, Send):
logger.warn(f"Ignoring invalid packet type {type(packet)} in pending sends")
continue
if for_execution:
proc = processes[packet.node]
if node := proc.get_node():
triggers = [TASKS]
metadata = {
"langgraph_step": step,
"langgraph_node": packet.node,
"langgraph_triggers": triggers,
"langgraph_task_idx": len(tasks),
}
task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
writes = deque()
tasks.append(
PregelExecutableTask(
packet.node,
packet.arg,
node,
writes,
patch_config(
merge_configs(
config,
processes[packet.node].config,
{"metadata": metadata},
),
run_name=packet.node,
callbacks=(
manager.get_child(f"graph:step:{step}")
if manager
else None
),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write, writes.extend, processes, channels
),
CONFIG_KEY_READ: partial(
local_read,
checkpoint,
channels,
PregelTaskWrites(packet.node, writes, triggers),
config,
),
},
),
triggers,
proc.retry_policy,
task_id,
)
)
else:
tasks.append(PregelTaskDescription(packet.node, packet.arg))
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
null_version = version_type()
if null_version is None:
return tasks
for name, proc in processes.items():
seen = checkpoint["versions_seen"].get(name, {})
# If any of the channels read by this process were updated
if triggers := sorted(
chan
for chan in proc.triggers
if not isinstance(
read_channel(channels, chan, return_exception=True), EmptyChannelError
)
and checkpoint["channel_versions"].get(chan, null_version)
> seen.get(chan, null_version)
):
try:
val = next(_proc_input(step, name, proc, managed, channels))
except StopIteration:
continue
if for_execution:
if node := proc.get_node():
metadata = {
"langgraph_step": step,
"langgraph_node": name,
"langgraph_triggers": triggers,
"langgraph_task_idx": len(tasks),
}
task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
writes = deque()
tasks.append(
PregelExecutableTask(
name,
val,
node,
writes,
patch_config(
merge_configs(
config,
proc.config,
{"metadata": metadata},
),
run_name=name,
callbacks=(
manager.get_child(f"graph:step:{step}")
if manager
else None
),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write, writes.extend, processes, channels
),
CONFIG_KEY_READ: partial(
local_read,
checkpoint,
channels,
PregelTaskWrites(name, writes, triggers),
config,
),
},
),
triggers,
proc.retry_policy,
task_id,
)
)
else:
tasks.append(PregelTaskDescription(name, val))
return tasks
def _proc_input(
step: int,
name: str,
proc: PregelNode,
managed: ManagedValueMapping,
channels: Mapping[str, BaseChannel],
) -> Iterator[Any]:
# If all trigger channels subscribed by this process are not empty
# then invoke the process with the values of all non-empty channels
if isinstance(proc.channels, dict):
try:
val: dict = {
k: read_channel(
channels,
chan,
catch=chan not in proc.triggers,
)
for k, chan in proc.channels.items()
if isinstance(chan, str)
}
managed_values = {}
for key, chan in proc.channels.items():
if is_managed_value(chan):
managed_values[key] = managed[key](
step, PregelTaskDescription(name, val)
)
val.update(managed_values)
except EmptyChannelError:
return
elif isinstance(proc.channels, list):
for chan in proc.channels:
try:
val = read_channel(channels, chan, catch=False)
break
except EmptyChannelError:
pass
else:
return
else:
raise RuntimeError(
"Invalid channels type, expected list or dict, got {proc.channels}"
)
# If the process has a mapper, apply it to the value
if proc.mapper is not None:
val = proc.mapper(val)
yield val
libs/langgraph/langgraph/pregel/executor.py
+3 -2
View File
@@ -6,6 +6,7 @@ from contextvars import copy_context
from types import TracebackType
from typing import (
AsyncContextManager,
Awaitable,
Callable,
Iterator,
Optional,
@@ -78,7 +79,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
def submit(
self,
fn: Callable[P, T],
fn: Callable[P, Awaitable[T]],
*args: P.args,
__name__: Optional[str] = None,
__cancel_on_exit__: bool = False,
@@ -101,7 +102,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
else:
self.tasks.pop(task)
async def __aenter__(self) -> Submit:
async def __aenter__(self) -> "submit":
return self.submit
async def exit(self) -> None:
libs/langgraph/langgraph/pregel/loop.py
+423
View File
@@ -0,0 +1,423 @@
import asyncio
from collections import deque
from contextlib import AsyncExitStack, ExitStack
from types import TracebackType
from typing import (
TYPE_CHECKING,
Any,
AsyncContextManager,
Callable,
ContextManager,
List,
Literal,
Mapping,
Optional,
Sequence,
Tuple,
Type,
TypeVar,
Union,
)
from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel
from langgraph.channels.manager import (
AsyncChannelsManager,
ChannelsManager,
create_checkpoint,
)
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
PendingWrite,
copy_checkpoint,
empty_checkpoint,
)
from langgraph.constants import INPUT, INTERRUPT
from langgraph.managed.base import (
AsyncManagedValuesManager,
ManagedValueMapping,
ManagedValuesManager,
)
from langgraph.pregel.algo import (
PregelTaskWrites,
apply_writes,
increment,
prepare_next_tasks,
should_interrupt,
)
from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
from langgraph.pregel.executor import (
AsyncBackgroundExecutor,
BackgroundExecutor,
Submit,
)
from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
from langgraph.pregel.types import PregelExecutableTask
if TYPE_CHECKING:
from langgraph.pregel import Pregel
V = TypeVar("V")
INPUT_DONE = object()
class PregelLoop:
input: Optional[Any]
config: RunnableConfig
checkpointer: Optional[BaseCheckpointSaver]
checkpointer_get_next_version: Callable[[Optional[V]], V]
checkpointer_put_writes: Optional[
Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
]
checkpointer_put: Optional[
Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], Any]
]
graph: "Pregel"
submit: Submit
channels: Mapping[str, BaseChannel]
managed: ManagedValueMapping
checkpoint: Checkpoint
checkpoint_config: RunnableConfig
checkpoint_metadata: CheckpointMetadata
checkpoint_pending_writes: Optional[List[PendingWrite]]
step: int
status: Literal[
"pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
]
tasks: Sequence[PregelExecutableTask]
stream: deque[Tuple[str, Any]]
# public
def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
"""Mark tasks as scheduled, to be used by queue-based executors."""
raise NotImplementedError
def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
"""Put writes for a task, to be read by the next tick."""
self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
if self.checkpointer_put_writes is not None:
self.submit(
self.checkpointer_put_writes,
{
**self.checkpoint_config,
"configurable": {
**self.checkpoint_config["configurable"],
"thread_ts": self.checkpoint["id"],
},
},
writes,
task_id,
)
def tick(
self,
*,
output_keys: Union[str, Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Sequence[str]] = None,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
) -> bool:
"""Execute a single iteration of the Pregel loop.
Returns True if more iterations are needed."""
if self.status != "pending":
raise RuntimeError("Cannot tick when status is no longer 'pending'")
if self.input is not INPUT_DONE:
self._first()
elif all(task.writes for task in self.tasks):
writes = [w for t in self.tasks for w in t.writes]
# all tasks have finished
apply_writes(
self.checkpoint,
self.channels,
self.tasks,
self.checkpointer_get_next_version,
)
# produce values output
self.stream.extend(
("values", v)
for v in map_output_values(output_keys, writes, self.channels)
)
# clear pending writes
self.checkpoint_pending_writes.clear()
# save checkpoint
self._put_checkpoint(
{
"source": "loop",
"writes": single(
map_output_updates(output_keys, self.tasks)
if self.graph.stream_mode == "updates"
else map_output_values(output_keys, writes, self.channels)
),
}
)
# after execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
self.status = "interrupt_after"
return False
else:
return False
# check if iteration limit is reached
if self.step > self.config["recursion_limit"]:
self.status = "out_of_steps"
return False
# prepare next tasks
self.tasks = prepare_next_tasks(
self.checkpoint,
self.graph.nodes,
self.channels,
self.managed,
self.config,
self.step,
for_execution=True,
manager=manager,
)
# if no more tasks, we're done
if not self.tasks:
self.status = "done"
return False
# if there are pending writes from a previous loop, apply them
if self.checkpoint_pending_writes:
for tid, k, v in self.checkpoint_pending_writes:
if task := next((t for t in self.tasks if t.id == tid), None):
task.writes.append((k, v))
# if all tasks have finished, re-tick
if all(task.writes for task in self.tasks):
return self.tick()
# before execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
self.status = "interrupt_before"
return False
# produce debug output
self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks))
return True
# private
def _first(self) -> None:
# map inputs to channel updates
if input_writes := deque(map_input(self.graph.input_channels, self.input)):
# discard any unfinished tasks from previous checkpoint
discard_tasks = prepare_next_tasks(
self.checkpoint,
self.graph.nodes,
self.channels,
self.managed,
self.config,
self.step,
for_execution=True,
)
# apply input writes
apply_writes(
self.checkpoint,
self.channels,
discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
self.checkpointer_get_next_version,
)
# save input checkpoint
self._put_checkpoint({"source": "input", "writes": self.input})
else:
# no input is taken as signal to proceed past previous interrupt
self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
for k in self.channels:
if k in self.checkpoint["channel_versions"]:
version = self.checkpoint["channel_versions"][k]
self.checkpoint["versions_seen"][INTERRUPT][k] = version
# done with input
self.input = INPUT_DONE
def _put_checkpoint(
self,
metadata: CheckpointMetadata,
) -> None:
# assign step
metadata["step"] = self.step
# bail if no checkpointer
if self.checkpointer_put is not None:
# create new checkpoint
self.checkpoint_metadata = metadata
self.checkpoint = create_checkpoint(
self.checkpoint, self.channels, self.step
)
# save it, without blocking
self.submit(
self.checkpointer_put,
self.checkpoint_config,
copy_checkpoint(self.checkpoint),
self.checkpoint_metadata,
)
self.checkpoint_config = {
**self.checkpoint_config,
"configurable": {
**self.checkpoint_config["configurable"],
"thread_ts": self.checkpoint["id"],
},
}
# produce debug output
self.stream.extend(
("debug", v)
for v in map_debug_checkpoint(
self.step,
self.checkpoint_config,
self.channels,
self.graph.stream_channels_asis,
self.checkpoint_metadata,
)
)
# increment step
self.step += 1
class SyncPregelLoop(PregelLoop, ContextManager):
def __init__(
self,
input: Optional[Any],
*,
config: RunnableConfig,
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
) -> None:
self.stream = deque()
self.stack = ExitStack()
self.input = input
self.config = config
self.checkpointer = checkpointer
self.checkpointer_get_next_version = (
checkpointer.get_next_version if checkpointer else increment
)
self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None
self.checkpointer_put = checkpointer.put if checkpointer else None
self.graph = graph
# TODO if managed values no longer needs graph we can replace with
# managed_specs, channel_specs
# context manager
def __enter__(self) -> Self:
saved = (
self.checkpointer.get_tuple(self.config) if self.checkpointer else None
) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
self.checkpoint_config = {
**self.config,
**saved.config,
"configurable": {
**self.config.get("configurable", {}),
**saved.config.get("configurable", {}),
},
}
self.checkpoint = copy_checkpoint(saved.checkpoint)
self.checkpoint_metadata = saved.metadata
self.checkpoint_pending_writes = saved.pending_writes
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels = self.stack.enter_context(
ChannelsManager(self.graph.channels, self.checkpoint, self.config)
)
self.managed = self.stack.enter_context(
ManagedValuesManager(
self.graph.managed_values_dict, self.config, self.graph
)
)
self.status = "pending"
self.step = self.checkpoint_metadata["step"] + 1
return self
def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
del self.graph
return self.stack.__exit__(exc_type, exc_value, traceback)
class AsyncPregelLoop(PregelLoop, AsyncContextManager):
def __init__(
self,
input: Optional[Any],
*,
config: RunnableConfig,
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
) -> None:
self.stream = deque()
self.stack = AsyncExitStack()
self.input = input
self.config = config
self.checkpointer = checkpointer
self.checkpointer_get_next_version = (
checkpointer.get_next_version if checkpointer else increment
)
self.checkpointer_put_writes = (
checkpointer.aput_writes if checkpointer else None
)
self.checkpointer_put = checkpointer.aput if checkpointer else None
self.graph = graph
# TODO if managed values no longer needs graph we can replace with
# managed_specs, channel_specs
# context manager
async def __aenter__(self) -> Self:
saved = (
await self.checkpointer.aget_tuple(self.config)
if self.checkpointer
else None
) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
self.checkpoint_config = {
**self.config,
**saved.config,
"configurable": {
**self.config.get("configurable", {}),
**saved.config.get("configurable", {}),
},
}
self.checkpoint = copy_checkpoint(saved.checkpoint)
self.checkpoint_metadata = saved.metadata
self.checkpoint_pending_writes = saved.pending_writes
self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
self.channels = await self.stack.enter_async_context(
AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
)
self.managed = await self.stack.enter_async_context(
AsyncManagedValuesManager(
self.graph.managed_values_dict, self.config, self.graph
)
)
self.status = "pending"
self.step = self.checkpoint_metadata["step"] + 1
return self
async def __aexit__(
self,
exc_type: Optional[Type[BaseException]],
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
del self.graph
return await asyncio.shield(
self.stack.__aexit__(exc_type, exc_value, traceback)
)
libs/langgraph/langgraph/pregel/types.py
+9
View File
@@ -88,3 +88,12 @@ class StateSnapshot(NamedTuple):
All = Literal["*"]
StreamMode = Literal["values", "updates", "debug"]
"""How the stream method should emit outputs.
- 'values': Emit all values of the state for each step.
- 'updates': Emit only the node name(s) and updates
that were returned by the node(s) **after** each step.
- 'debug': Emit debug events for each step.
"""
libs/langgraph/poetry.lock
Generated
+39 -4
View File
@@ -747,6 +747,20 @@ files = [
[package.extras]
test = ["pytest (>=6)"]
[[package]]
name = "execnet"
version = "2.1.1"
description = "execnet: rapid multi-Python deployment"
optional = false
python-versions = ">=3.8"
files = [
{file = "execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc"},
{file = "execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3"},
]
[package.extras]
testing = ["hatch", "pre-commit", "pytest", "tox"]
[[package]]
name = "executing"
version = "2.0.1"
@@ -1746,13 +1760,13 @@ langchain-core = ">=0.2.2rc1,<0.3"
[[package]]
name = "langchain-core"
version = "0.2.19"
version = "0.2.22"
description = "Building applications with LLMs through composability"
optional = false
python-versions = "<4.0,>=3.8.1"
files = [
{file = "langchain_core-0.2.19-py3-none-any.whl", hash = "sha256:5b3cd34395be274c89e822c84f0e03c4da14168c177a83921c5b9414ac7a0651"},
{file = "langchain_core-0.2.19.tar.gz", hash = "sha256:13043a83e5c9ab58b9f5ce2a56896e7e88b752e8891b2958960a98e71801471e"},
{file = "langchain_core-0.2.22-py3-none-any.whl", hash = "sha256:7731a86440c0958b3186c003fb9b26b2d5a682a6344bda7bfb9174e2898f8b43"},
{file = "langchain_core-0.2.22.tar.gz", hash = "sha256:582d6f929a43b830139444e4124123cd415331ad62f25757b1406252958cdcac"},
]
[package.dependencies]
@@ -2784,6 +2798,27 @@ files = [
tomli = {version = ">=2.0.1,<3.0.0", markers = "python_version < \"3.11\""}
watchdog = ">=2.0.0"
[[package]]
name = "pytest-xdist"
version = "3.6.1"
description = "pytest xdist plugin for distributed testing, most importantly across multiple CPUs"
optional = false
python-versions = ">=3.8"
files = [
{file = "pytest_xdist-3.6.1-py3-none-any.whl", hash = "sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7"},
{file = "pytest_xdist-3.6.1.tar.gz", hash = "sha256:ead156a4db231eec769737f57668ef58a2084a34b2e55c4a8fa20d861107300d"},
]
[package.dependencies]
execnet = ">=2.1"
psutil = {version = ">=3.0", optional = true, markers = "extra == \"psutil\""}
pytest = ">=7.0.0"
[package.extras]
psutil = ["psutil (>=3.0)"]
setproctitle = ["setproctitle"]
testing = ["filelock"]
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"
@@ -4130,4 +4165,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "170eaa0e542a02d5f2fb0d42d1f04c5d010bd3735a44b927b28e6b742c689eb0"
content-hash = "0d877d3879473de43aca1e1d36a8f420ff3f4b140807cb5cc24935d9114947be"
libs/langgraph/pyproject.toml
+3 -2
View File
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
[tool.poetry.dependencies]
python = ">=3.9.0,<4.0"
langchain-core = ">=0.2.19,<0.3"
langchain-core = ">=0.2.22,<0.3"
[tool.poetry.group.dev.dependencies]
@@ -31,6 +31,7 @@ langchainhub = "^0.1.14"
langchain-openai = ">=0.1.2"
langchain-anthropic = ">=0.1.8"
dataclasses-json = "^0.6.7"
pytest-xdist = {extras = ["psutil"], version = "^3.6.1"}
[tool.poetry.group.dev]
optional = true
@@ -61,7 +62,7 @@ omit = ["tests/*"]
[tool.pytest-watcher]
now = true
delay = 0.1
runner_args = ["-x", "--ff", "-vv", "--snapshot-update"]
runner_args = ["-x", "--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
patterns = ["*.py"]
[build-system]
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+2 -2
View File
@@ -509,10 +509,10 @@
'''
# ---
# name: test_conditional_state_graph
'{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
'{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
# ---
# name: test_conditional_state_graph.1
'{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
'{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
# ---
# name: test_conditional_state_graph.2
'''
libs/langgraph/tests/test_prebuilt.py
+181 -16
View File
@@ -1,15 +1,11 @@
from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union
from typing import Annotated, Any, Callable, Dict, List, Optional, Sequence, Type, Union
import pytest
from langchain_core.callbacks import (
CallbackManagerForLLMRun,
)
from langchain_core.language_models import (
BaseChatModel,
LanguageModelInput,
)
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel, LanguageModelInput
from langchain_core.messages import (
AIMessage,
AnyMessage,
BaseMessage,
HumanMessage,
SystemMessage,
@@ -22,11 +18,11 @@ from langchain_core.tools import BaseTool
from langchain_core.tools import tool as dec_tool
from pydantic import BaseModel as BaseModelV2
from langgraph.prebuilt import (
ToolNode,
ValidationNode,
create_react_agent,
)
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
from langgraph.prebuilt.tool_node import InjectedState
from tests.any_str import AnyStr
from tests.memory_assert import MemorySaverAssertImmutable
class FakeToolCallingModel(BaseChatModel):
@@ -56,14 +52,117 @@ class FakeToolCallingModel(BaseChatModel):
return self
def test_no_modifier():
@pytest.mark.parametrize(
"checkpointer",
[
MemorySaverAssertImmutable(),
None,
],
ids=[
"memory",
"none",
],
)
def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]):
model = FakeToolCallingModel()
agent = create_react_agent(model, [])
agent = create_react_agent(model, [], checkpointer=checkpointer)
inputs = [HumanMessage("hi?")]
response = agent.invoke({"messages": inputs})
thread = {"configurable": {"thread_id": "123"}}
response = agent.invoke({"messages": inputs}, thread, debug=True)
expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
assert response == expected_response
if checkpointer:
saved = checkpointer.get_tuple(thread)
assert saved is not None
assert saved.checkpoint == {
"v": 1,
"ts": AnyStr(),
"id": AnyStr(),
"channel_values": {
"messages": [
HumanMessage(content="hi?", id=AnyStr()),
AIMessage(content="hi?", id="0"),
],
"agent": "agent",
},
"channel_versions": {
"__start__": 2,
"messages": 3,
"start:agent": 3,
"agent": 3,
},
"versions_seen": {
"__input__": {},
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
},
"pending_sends": [],
"current_tasks": {},
}
assert saved.metadata == {
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
}
assert saved.pending_writes == []
@pytest.mark.parametrize(
"checkpointer",
[
MemorySaverAssertImmutable(),
None,
],
ids=[
"memory",
"none",
],
)
async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]):
model = FakeToolCallingModel()
agent = create_react_agent(model, [], checkpointer=checkpointer)
inputs = [HumanMessage("hi?")]
thread = {"configurable": {"thread_id": "123"}}
response = await agent.ainvoke({"messages": inputs}, thread, debug=True)
expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
assert response == expected_response
if checkpointer:
saved = await checkpointer.aget_tuple(thread)
assert saved is not None
assert saved.checkpoint == {
"v": 1,
"ts": AnyStr(),
"id": AnyStr(),
"channel_values": {
"messages": [
HumanMessage(content="hi?", id=AnyStr()),
AIMessage(content="hi?", id="0"),
],
"agent": "agent",
},
"channel_versions": {
"__start__": 2,
"messages": 3,
"start:agent": 3,
"agent": 3,
},
"versions_seen": {
"__input__": {},
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
},
"pending_sends": [],
"current_tasks": {},
}
assert saved.metadata == {
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
}
assert saved.pending_writes == []
def test_passing_two_modifiers():
model = FakeToolCallingModel()
@@ -347,3 +446,69 @@ async def test_validation_node(tool_schema: Any, use_message_key: bool):
if use_message_key:
result_sync = result_sync["messages"]
check_results(result_sync)
def test_tool_node_inject_state() -> None:
def tool1(some_val: int, state: Annotated[dict, InjectedState]) -> str:
"""Tool 1 docstring."""
return state["foo"]
def tool2(some_val: int, state: Annotated[dict, InjectedState()]) -> str:
"""Tool 1 docstring."""
return state["foo"]
def tool3(
some_val: int,
foo: Annotated[str, InjectedState("foo")],
msgs: Annotated[List[AnyMessage], InjectedState("messages")],
) -> str:
"""Tool 1 docstring."""
return foo
def tool4(
some_val: int, msgs: Annotated[List[AnyMessage], InjectedState("messages")]
) -> str:
"""Tool 1 docstring."""
return msgs[0].content
node = ToolNode([tool1, tool2, tool3, tool4])
for tool_name in ("tool1", "tool2", "tool3"):
tool_call = {
"name": tool_name,
"args": {"some_val": 1},
"id": "some 0",
"type": "tool_call",
}
msg = AIMessage("hi?", tool_calls=[tool_call])
result = node.invoke({"messages": [msg], "foo": "bar"})
tool_message = result["messages"][-1]
assert tool_message.content == "bar"
if tool_name == "tool3":
with pytest.raises(KeyError):
node.invoke({"messages": [msg], "notfoo": "bar"})
with pytest.raises(ValueError):
node.invoke([msg])
else:
tool_message = node.invoke({"messages": [msg], "notfoo": "bar"})[
"messages"
][-1]
assert "KeyError" in tool_message.content
tool_message = node.invoke([msg])[-1]
assert "KeyError" in tool_message.content
tool_call = {
"name": "tool4",
"args": {"some_val": 1},
"id": "some 0",
"type": "tool_call",
}
msg = AIMessage("hi?", tool_calls=[tool_call])
result = node.invoke({"messages": [msg]})
tool_message = result["messages"][-1]
assert tool_message.content == "hi?"
result = node.invoke([msg])
tool_message = result[-1]
assert tool_message.content == "hi?"
libs/langgraph/tests/test_pregel.py
+15 -5
View File
@@ -520,10 +520,8 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert app.invoke(2) == 4
assert app.invoke(2, input_keys="inbox") == 3
with pytest.raises(GraphRecursionError):
app.invoke(2, {"recursion_limit": 1})
app.invoke(2, {"recursion_limit": 1}, debug=1)
graph = Graph()
graph.add_node("add_one", add_one)
@@ -535,7 +533,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert gapp.invoke(2) == 4
for step, values in enumerate(gapp.stream(2), start=1):
for step, values in enumerate(gapp.stream(2, debug=1), start=1):
if step == 1:
assert values == {
"add_one": 3,
@@ -6168,6 +6166,18 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
"my_key": "value ⛰️",
"market": "DE",
}
assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [
{
"source": "loop",
"step": 0,
"writes": None,
},
{
"source": "input",
"step": -1,
"writes": {"my_key": "value ⛰️", "market": "DE"},
},
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two_slow",),
@@ -6877,7 +6887,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
]
libs/langgraph/tests/test_pregel_async.py
+12 -2
View File
@@ -623,8 +623,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert await app.ainvoke(2) == 4
assert await app.ainvoke(2, input_keys="inbox") == 3
with pytest.raises(GraphRecursionError):
await app.ainvoke(2, {"recursion_limit": 1})
@@ -4737,6 +4735,18 @@ async def test_start_branch_then() -> None:
"my_key": "value",
"market": "DE",
}
assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
{
"source": "loop",
"step": 0,
"writes": None,
},
{
"source": "input",
"step": -1,
"writes": {"my_key": "value", "market": "DE"},
},
]
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value", "market": "DE"},
next=("tool_two_slow",),
libs/sdk-py/langgraph_sdk/client.py
+3
View File
@@ -257,6 +257,7 @@ class AssistantsClient:
*,
metadata: Metadata = None,
assistant_id: Optional[str] = None,
if_exists: Optional[OnConflictBehavior] = None,
) -> Assistant:
"""Create a new assistant."""
payload: Dict[str, Any] = {
@@ -268,6 +269,8 @@ class AssistantsClient:
payload["metadata"] = metadata
if assistant_id:
payload["assistant_id"] = assistant_id
if if_exists:
payload["if_exists"] = if_exists
return await self.http.post("/assistants", json=payload)
async def update(