Merge branch 'main' into nc/19jul/nested-checkpoints

This commit is contained in:
Nuno Campos
2024-07-22 16:33:33 -07:00
39 changed files with 849 additions and 666 deletions
+2 -2
View File
@@ -73,8 +73,8 @@ def search(query: str):
"""Call to surf the web."""
# This is a placeholder, but don't tell the LLM that...
if "sf" in query.lower() or "san francisco" in query.lower():
return ["It's 60 degrees and foggy."]
return ["It's 90 degrees and sunny."]
return "It's 60 degrees and foggy."
return "It's 90 degrees and sunny."
tools = [search]
@@ -1,11 +1,11 @@
<a name="readmemd"></a>
**@langchain/langgraph-sdk** • **Docs**
**[@langchain/langgraph-sdk](https://github.com/langchain-ai/langgraph/tree/main/libs/sdk-js)** • **Docs**
***
## @langchain/langgraph-sdk
## [@langchain/langgraph-sdk](https://github.com/langchain-ai/langgraph/tree/main/libs/sdk-js)
### Modules
-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
+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
+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": {
@@ -9,7 +9,7 @@
"\n",
"This tutorial will show how to add a custom system prompt to the prebuilt ReAct agent. Please see [this tutorial](./create-react-agent.ipynb) for how to get started with the prebuilt ReAct agent\n",
"\n",
"You can add a custom system prompt by passing a string to the `messages_modifier` param."
"You can add a custom system prompt by passing a string to the `state_modifier` param."
]
},
{
@@ -112,7 +112,7 @@
"\n",
"from langgraph.prebuilt import create_react_agent\n",
"\n",
"graph = create_react_agent(model, tools=tools, messages_modifier=prompt)"
"graph = create_react_agent(model, tools=tools, state_modifier=prompt)"
]
},
{
+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",
@@ -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",
@@ -137,22 +137,27 @@
"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\", 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",
@@ -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\"])"
]
},
{
+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"
]
},
@@ -138,7 +138,7 @@
"\n",
"# Define the function that calls the model\n",
"def call_model(state: MessagesState):\n",
" response = model.invoke(state[\"messages\"])\n",
" response = bound_model.invoke(state[\"messages\"])\n",
" # We return a list, because this will get added to the existing list\n",
" return {\"messages\": response}\n",
"\n",
@@ -275,7 +275,7 @@
"# Define the function that calls the model\n",
"def call_model(state: MessagesState):\n",
" messages = filter_messages(state[\"messages\"])\n",
" response = model.invoke(messages)\n",
" response = bound_model.invoke(messages)\n",
" # We return a list, because this will get added to the existing list\n",
" return {\"messages\": response}\n",
"\n",
+10 -2
View File
@@ -65,23 +65,31 @@
"\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(\"query_database\",query_database, retry=RetryPolicy(retry_on=sqlite3.OperationalError))\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(\"model\", \"query_database\")\n",
"workflow.add_edge(\"query_database\", END)\n",
"\n",
"app = workflow.compile()"
+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",
+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"
+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",
+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",
+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",
@@ -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\"])"
@@ -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())"
+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()))"
]
},
{
+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",
+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()"
]
+26 -5
View File
@@ -131,21 +131,42 @@ async def monitor_stream(
if collect:
ba = bytearray()
def handle(line: bytes):
def handle(line: bytes, overrun: bool):
nonlocal on_line
nonlocal display
if display:
sys.stdout.buffer.write(line)
if overrun:
return
if collect:
ba.extend(line)
if display:
sys.stdout.write(line.decode())
if on_line:
if on_line(line.decode()):
on_line = None
display = True
async for line in stream:
await asyncio.to_thread(handle, line)
"""Adapted from asyncio.StreamReader.readline() to handle LimitOverrunError."""
sep = b"\n"
seplen = len(sep)
while True:
try:
line = await stream.readuntil(sep)
overrun = False
except asyncio.IncompleteReadError as e:
line = e.partial
overrun = False
except asyncio.LimitOverrunError as e:
if stream._buffer.startswith(sep, e.consumed):
line = stream._buffer[: e.consumed + seplen]
else:
line = stream._buffer.clear()
overrun = True
stream._maybe_resume_transport()
await asyncio.to_thread(handle, line, overrun)
if line == b"":
break
if collect:
return ba
else:
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-cli"
version = "0.1.49"
version = "0.1.50"
description = "CLI for interacting with LangGraph API"
authors = []
license = "MIT"
+1 -1
View File
@@ -2,7 +2,7 @@
"python_version": "3.12",
"pip_config_file": "pipconfig.txt",
"dockerfile_lines": [
"ARG meow"
"ARG meow=woof"
],
"dependencies": [
"langchain_openai",
+2 -2
View File
@@ -73,8 +73,8 @@ def search(query: str):
"""Call to surf the web."""
# This is a placeholder, but don't tell the LLM that...
if "sf" in query.lower() or "san francisco" in query.lower():
return ["It's 60 degrees and foggy."]
return ["It's 90 degrees and sunny."]
return "It's 60 degrees and foggy."
return "It's 90 degrees and sunny."
tools = [search]
+5 -3
View File
@@ -56,9 +56,11 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
pass
def __eq__(self, value: object) -> bool:
return (
isinstance(value, BinaryOperatorAggregate)
and value.operator == self.operator
return isinstance(value, BinaryOperatorAggregate) and (
value.operator is self.operator
if value.operator.__name__ != "<lambda>"
and self.operator.__name__ != "<lambda>"
else True
)
@property
@@ -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 (?, ?, ?, ?, ?, ?)",
+18 -3
View File
@@ -29,6 +29,8 @@ 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.
@@ -119,6 +121,8 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
class CheckpointTuple(NamedTuple):
"""A tuple containing a checkpoint and its associated data."""
config: RunnableConfig
checkpoint: Checkpoint
metadata: CheckpointMetadata
@@ -269,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
@@ -286,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
@@ -301,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
@@ -326,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
+28 -5
View File
@@ -123,15 +123,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:
@@ -184,6 +185,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.
@@ -218,6 +220,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.
@@ -281,6 +284,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
)
@@ -291,6 +304,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
)
@@ -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:
@@ -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",
]
+149 -10
View File
@@ -1,11 +1,24 @@
import asyncio
from typing import Any, Callable, Dict, Literal, Optional, Sequence, Tuple, Union, cast
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_config_list, get_executor_for_config
from langchain_core.tools import BaseTool
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,18 +73,18 @@ class ToolNode(RunnableCallable):
def _func(
self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
) -> Any:
message, output_type = self._parse_input(input)
config_list = get_config_list(config, len(message.tool_calls))
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(self._run_one, message.tool_calls, config_list)]
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:
message, output_type = self._parse_input(input)
tool_calls, output_type = self._parse_input(input)
outputs = await asyncio.gather(
*(self._arun_one(call, config) for call in message.tool_calls)
*(self._arun_one(call, config) for call in tool_calls)
)
return outputs if output_type == "list" else {"messages": outputs}
@@ -102,7 +115,7 @@ class ToolNode(RunnableCallable):
def _parse_input(
self, input: Union[list[AnyMessage], dict[str, Any]]
) -> Tuple[AIMessage, Literal["list", "dict"]]:
) -> Tuple[List[ToolCall], Literal["list", "dict"]]:
if isinstance(input, list):
output_type = "list"
message: AnyMessage = input[-1]
@@ -114,8 +127,12 @@ class ToolNode(RunnableCallable):
if not isinstance(message, AIMessage):
raise ValueError("Last message is not an AIMessage")
else:
return cast(AIMessage, message), output_type
tool_calls = [
self._inject_state(call, input)
for call in cast(AIMessage, message).tool_calls
]
return tool_calls, output_type
def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]:
if (requested_tool := call["name"]) not in self.tools_by_name:
@@ -127,6 +144,39 @@ class ToolNode(RunnableCallable):
else:
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(
state: Union[list[AnyMessage], dict[str, Any]],
@@ -183,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
+1
View File
@@ -257,6 +257,7 @@ class PregelLoop:
self.config,
self.step,
for_execution=True,
manager=None,
)
# apply input writes
apply_writes(
+1 -1
View File
@@ -4179,4 +4179,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "da02c84e7232f1aa2980537e71d88a757a62444a41b1873566ec81842fda8f82"
content-hash = "18b26895b05f2f7cdcd08d59ba164ac788e0e8005e3ec5d7089469b4f3a96aea"
+1 -1
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]
+72 -13
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,
@@ -23,11 +19,8 @@ from langchain_core.tools import tool as dec_tool
from pydantic import BaseModel as BaseModelV2
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.prebuilt import (
ToolNode,
ValidationNode,
create_react_agent,
)
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
@@ -453,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?"
+4
View File
@@ -392,6 +392,10 @@ class ThreadsClient:
json=payload,
)
async def copy(self, thread_id: str) -> None:
"""Copy a thread."""
return await self.http.post(f"/threads/{thread_id}/copy", json=None)
async def get_state(
self, thread_id: str, checkpoint_id: Optional[str] = None
) -> ThreadState: