Optimize tracing output of Graph/StateGraph/MessageGraph

- control selection of relevant runs (needs langsmith release)
- see output of conditional edge function
- fix issue with conditional entry point not getting full state values as input
This commit is contained in:
Nuno Campos
2024-04-02 10:32:47 -07:00
parent 041ca2f784
commit 2a299d070e
15 changed files with 398 additions and 2571 deletions
+10 -10
View File
File diff suppressed because one or more lines are too long
+22 -30
View File
@@ -57,7 +57,7 @@
"metadata": {},
"outputs": [
{
"name": "stdin",
"name": "stdout",
"output_type": "stream",
"text": [
"OpenAI API Key: ········\n",
@@ -239,13 +239,6 @@
" return \"continue\"\n",
"\n",
"\n",
"# Define the function that calls the model\n",
"def call_model(messages):\n",
" response = model.invoke(messages)\n",
" # We return a list, because this will get added to the existing list\n",
" return response\n",
"\n",
"\n",
"# Define the function to execute tools\n",
"def call_tool(messages):\n",
" # Based on the continue condition\n",
@@ -289,7 +282,7 @@
"workflow = MessageGraph()\n",
"\n",
"# Define the two nodes we will cycle between\n",
"workflow.add_node(\"agent\", call_model)\n",
"workflow.add_node(\"agent\", model)\n",
"workflow.add_node(\"action\", call_tool)\n",
"\n",
"# Set the entrypoint as `agent`\n",
@@ -377,7 +370,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"content='Hello Bob! How can I assist you today?' id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0'\n"
"content='Hello Bob! How can I assist you today?' id='87e8b88a-b28e-4517-a62f-7b9baebd3329'\n"
]
}
],
@@ -408,7 +401,7 @@
{
"data": {
"text/plain": [
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='774c03a0-4f61-49f1-8447-6fa5cbc3e56e'), AIMessage(content='Hello Bob! How can I assist you today?', id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0')], next=(), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T00:09:43.801248+00:00'}}, parent_config=None)"
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='5a84a602-37d9-4903-9c3c-d0f892238580'), AIMessage(content='Hello Bob! How can I assist you today?', id='87e8b88a-b28e-4517-a62f-7b9baebd3329')], next=(), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T17:13:46.250299+00:00'}}, parent_config=None)"
]
},
"execution_count": 10,
@@ -438,7 +431,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"content='' additional_kwargs={'function_call': {'arguments': '{\"query\":\"current weather in San Francisco\"}', 'name': 'tavily_search_results_json'}} id='79bcb08c-0cb4-4b37-aa68-4df08b3a290f'\n"
"content='' additional_kwargs={'function_call': {'arguments': '{\"query\":\"current weather in San Francisco\"}', 'name': 'tavily_search_results_json'}} id='8e6ae8d5-a258-47aa-9e2c-1ca3e78add90'\n"
]
}
],
@@ -509,7 +502,7 @@
"data": {
"text/plain": [
"{'configurable': {'thread_id': '3',\n",
" 'thread_ts': '2024-04-02T00:09:44.323782+00:00'}}"
" 'thread_ts': '2024-04-02T16:58:19.692203+00:00'}}"
]
},
"execution_count": 14,
@@ -538,7 +531,7 @@
{
"data": {
"text/plain": [
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='774c03a0-4f61-49f1-8447-6fa5cbc3e56e'), AIMessage(content='Hello Bob! How can I assist you today?', id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0'), HumanMessage(content='what is the weather in sf currently', id='4138a1b0-3053-4b12-ad15-6198be969947'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"weather in San Francisco today\"}', 'name': 'tavily_search_results_json'}}, id='79bcb08c-0cb4-4b37-aa68-4df08b3a290f')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T00:09:44.323782+00:00'}}, parent_config=None)"
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='1493cdf4-b7b5-46c9-a3e7-ad3b661fcb92'), AIMessage(content='Hello Bob! How can I assist you today?', id='028e0185-2ca4-4c1d-8660-93c17b839275'), HumanMessage(content='what is the weather in sf currently', id='76d6a5ff-d400-4edb-a264-b584107e2231'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"weather in San Francisco today\"}', 'name': 'tavily_search_results_json'}}, id='b143d5f6-5216-4a48-b76a-68b2f3eb6ff2')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T16:58:19.692203+00:00'}}, parent_config=None)"
]
},
"execution_count": 15,
@@ -568,8 +561,8 @@
"name": "stdout",
"output_type": "stream",
"text": [
"content=\"[{'url': 'https://forecast.weather.gov/zipcity.php?inputstring=San francisco,CA', 'content': 'Detailed Forecast. Today. Mostly sunny, with a high near 62. Light and variable wind becoming west southwest 5 to 8 mph in the afternoon. Tonight. Mostly clear, with a low around 49. West wind 5 to 8 mph becoming north northwest after midnight. Monday. Sunny, with a high near 67.'}]\" name='tavily_search_results_json' id='388a34ca-6b79-4eb2-84e7-34849ce7c4cf'\n",
"content='The weather in San Francisco today is mostly sunny with a high near 62 degrees Fahrenheit. The wind is light and variable, becoming west-southwest at 5 to 8 mph in the afternoon. Tonight will be mostly clear with a low around 49 degrees Fahrenheit. Tomorrow is expected to be sunny with a high near 67 degrees Fahrenheit.' id='5d5ac23c-0e64-462d-b9ef-c499a417e906'\n"
"content=\"[{'url': 'https://forecast.weather.gov/zipcity.php?inputstring=San francisco,CA', 'content': 'Detailed Forecast. Today. Mostly sunny, with a high near 62. Light and variable wind becoming west southwest 5 to 8 mph in the afternoon. Tonight. Mostly clear, with a low around 49. West wind 5 to 8 mph becoming north northwest after midnight. Monday. Sunny, with a high near 67.'}]\" name='tavily_search_results_json' id='9b33eb74-b079-4615-9c8a-9f41b4e176ba'\n",
"content='The weather in San Francisco today is mostly sunny with a high near 62 degrees Fahrenheit. The wind is light and variable, becoming west-southwest at 5 to 8 mph in the afternoon. Tonight, it will be mostly clear with a low around 49 degrees Fahrenheit. Tomorrow is expected to be sunny with a high near 67 degrees Fahrenheit.' id='18ff870d-5d5d-40d9-9eca-72e345749b51'\n"
]
}
],
@@ -591,7 +584,7 @@
},
{
"cell_type": "code",
"execution_count": 19,
"execution_count": 17,
"id": "8578a66d-6489-4e03-8c23-fd0530278455",
"metadata": {},
"outputs": [
@@ -599,13 +592,13 @@
"name": "stdout",
"output_type": "stream",
"text": [
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='774c03a0-4f61-49f1-8447-6fa5cbc3e56e'), AIMessage(content='Hello Bob! How can I assist you today?', id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0'), HumanMessage(content='what is the weather in sf currently', id='4138a1b0-3053-4b12-ad15-6198be969947'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"weather in San Francisco today\"}', 'name': 'tavily_search_results_json'}}, id='79bcb08c-0cb4-4b37-aa68-4df08b3a290f'), FunctionMessage(content=\"[{'url': 'https://forecast.weather.gov/zipcity.php?inputstring=San francisco,CA', 'content': 'Detailed Forecast. Today. Mostly sunny, with a high near 62. Light and variable wind becoming west southwest 5 to 8 mph in the afternoon. Tonight. Mostly clear, with a low around 49. West wind 5 to 8 mph becoming north northwest after midnight. Monday. Sunny, with a high near 67.'}]\", name='tavily_search_results_json', id='388a34ca-6b79-4eb2-84e7-34849ce7c4cf'), AIMessage(content='The weather in San Francisco today is mostly sunny with a high near 62 degrees Fahrenheit. The wind is light and variable, becoming west-southwest at 5 to 8 mph in the afternoon. Tonight will be mostly clear with a low around 49 degrees Fahrenheit. Tomorrow is expected to be sunny with a high near 67 degrees Fahrenheit.', id='5d5ac23c-0e64-462d-b9ef-c499a417e906')], next=(), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T00:09:48.961894+00:00'}}, parent_config=None)\n",
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='1493cdf4-b7b5-46c9-a3e7-ad3b661fcb92'), AIMessage(content='Hello Bob! How can I assist you today?', id='028e0185-2ca4-4c1d-8660-93c17b839275'), HumanMessage(content='what is the weather in sf currently', id='76d6a5ff-d400-4edb-a264-b584107e2231'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"weather in San Francisco today\"}', 'name': 'tavily_search_results_json'}}, id='b143d5f6-5216-4a48-b76a-68b2f3eb6ff2'), FunctionMessage(content=\"[{'url': 'https://forecast.weather.gov/zipcity.php?inputstring=San francisco,CA', 'content': 'Detailed Forecast. Today. Mostly sunny, with a high near 62. Light and variable wind becoming west southwest 5 to 8 mph in the afternoon. Tonight. Mostly clear, with a low around 49. West wind 5 to 8 mph becoming north northwest after midnight. Monday. Sunny, with a high near 67.'}]\", name='tavily_search_results_json', id='9b33eb74-b079-4615-9c8a-9f41b4e176ba'), AIMessage(content='The weather in San Francisco today is mostly sunny with a high near 62 degrees Fahrenheit. The wind is light and variable, becoming west-southwest at 5 to 8 mph in the afternoon. Tonight, it will be mostly clear with a low around 49 degrees Fahrenheit. Tomorrow is expected to be sunny with a high near 67 degrees Fahrenheit.', id='18ff870d-5d5d-40d9-9eca-72e345749b51')], next=(), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T16:58:23.633015+00:00'}}, parent_config=None)\n",
"--\n",
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='774c03a0-4f61-49f1-8447-6fa5cbc3e56e'), AIMessage(content='Hello Bob! How can I assist you today?', id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0'), HumanMessage(content='what is the weather in sf currently', id='4138a1b0-3053-4b12-ad15-6198be969947'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"weather in San Francisco today\"}', 'name': 'tavily_search_results_json'}}, id='79bcb08c-0cb4-4b37-aa68-4df08b3a290f')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T00:09:44.323782+00:00'}}, parent_config=None)\n",
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='1493cdf4-b7b5-46c9-a3e7-ad3b661fcb92'), AIMessage(content='Hello Bob! How can I assist you today?', id='028e0185-2ca4-4c1d-8660-93c17b839275'), HumanMessage(content='what is the weather in sf currently', id='76d6a5ff-d400-4edb-a264-b584107e2231'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"weather in San Francisco today\"}', 'name': 'tavily_search_results_json'}}, id='b143d5f6-5216-4a48-b76a-68b2f3eb6ff2')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T16:58:19.692203+00:00'}}, parent_config=None)\n",
"--\n",
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='774c03a0-4f61-49f1-8447-6fa5cbc3e56e'), AIMessage(content='Hello Bob! How can I assist you today?', id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0'), HumanMessage(content='what is the weather in sf currently', id='4138a1b0-3053-4b12-ad15-6198be969947'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"current weather in San Francisco\"}', 'name': 'tavily_search_results_json'}}, id='79bcb08c-0cb4-4b37-aa68-4df08b3a290f')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T00:09:44.242352+00:00'}}, parent_config=None)\n",
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='1493cdf4-b7b5-46c9-a3e7-ad3b661fcb92'), AIMessage(content='Hello Bob! How can I assist you today?', id='028e0185-2ca4-4c1d-8660-93c17b839275'), HumanMessage(content='what is the weather in sf currently', id='76d6a5ff-d400-4edb-a264-b584107e2231'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"current weather in San Francisco\"}', 'name': 'tavily_search_results_json'}}, id='b143d5f6-5216-4a48-b76a-68b2f3eb6ff2')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T16:58:18.058645+00:00'}}, parent_config=None)\n",
"--\n",
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='774c03a0-4f61-49f1-8447-6fa5cbc3e56e'), AIMessage(content='Hello Bob! How can I assist you today?', id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0')], next=(), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T00:09:43.801248+00:00'}}, parent_config=None)\n",
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='1493cdf4-b7b5-46c9-a3e7-ad3b661fcb92'), AIMessage(content='Hello Bob! How can I assist you today?', id='028e0185-2ca4-4c1d-8660-93c17b839275')], next=(), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T16:58:16.834223+00:00'}}, parent_config=None)\n",
"--\n"
]
}
@@ -628,17 +621,17 @@
},
{
"cell_type": "code",
"execution_count": 20,
"execution_count": 18,
"id": "21e7fc18-6fd9-4e11-a84b-e0325c9640c8",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='774c03a0-4f61-49f1-8447-6fa5cbc3e56e'), AIMessage(content='Hello Bob! How can I assist you today?', id='019fe7a8-f43d-4f83-8bbd-b1c553f32aa0'), HumanMessage(content='what is the weather in sf currently', id='4138a1b0-3053-4b12-ad15-6198be969947'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"current weather in San Francisco\"}', 'name': 'tavily_search_results_json'}}, id='79bcb08c-0cb4-4b37-aa68-4df08b3a290f')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T00:09:44.242352+00:00'}}, parent_config=None)"
"StateSnapshot(values=[HumanMessage(content=\"hi! I'm bob\", id='1493cdf4-b7b5-46c9-a3e7-ad3b661fcb92'), AIMessage(content='Hello Bob! How can I assist you today?', id='028e0185-2ca4-4c1d-8660-93c17b839275'), HumanMessage(content='what is the weather in sf currently', id='76d6a5ff-d400-4edb-a264-b584107e2231'), AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"current weather in San Francisco\"}', 'name': 'tavily_search_results_json'}}, id='b143d5f6-5216-4a48-b76a-68b2f3eb6ff2')], next=('action',), config={'configurable': {'thread_id': '3', 'thread_ts': '2024-04-02T16:58:18.058645+00:00'}}, parent_config=None)"
]
},
"execution_count": 20,
"execution_count": 18,
"metadata": {},
"output_type": "execute_result"
}
@@ -657,7 +650,7 @@
},
{
"cell_type": "code",
"execution_count": 21,
"execution_count": 19,
"id": "e986f94f-706f-4b6f-b3c4-f95483b9e9b8",
"metadata": {},
"outputs": [
@@ -665,16 +658,15 @@
"name": "stdout",
"output_type": "stream",
"text": [
"content=\"[{'url': 'https://www.accuweather.com/en/us/san-francisco/94103/current-weather/347629', 'content': 'Current weather in San Francisco, CA. Check current conditions in San Francisco, CA with radar, hourly, and more.'}]\" name='tavily_search_results_json' id='6ce0ce23-4fb7-4d3a-acbf-f882f8d52054'\n",
"content='You can check the current weather in San Francisco by visiting [AccuWeather](https://www.accuweather.com/en/us/san-francisco/94103/current-weather/347629).' id='0f071adc-a949-4242-ac3e-7ccd2132653d'\n"
"content=\"[{'url': 'https://www.accuweather.com/en/us/san-francisco/94103/current-weather/347629', 'content': 'Get the latest weather conditions and forecast for San Francisco, CA. See the temperature, humidity, wind, pressure, cloud cover, and alerts for the current hour and the next few days.'}]\" name='tavily_search_results_json' id='73ab98d6-d7cf-4812-9572-f58698394c9f'\n",
"content='You can check the current weather conditions and forecast for San Francisco, CA on [AccuWeather](https://www.accuweather.com/en/us/san-francisco/94103/current-weather/347629). This will provide you with information on temperature, humidity, wind, pressure, cloud cover, and alerts for the current hour and the next few days.' id='a54c0531-1683-4d61-87b6-ee7f66971c25'\n"
]
}
],
"source": [
"for event in app.stream(None, to_replay.config):\n",
" for k, v in event.items():\n",
" if k != \"__end__\":\n",
" print(v)"
" for v in event.values():\n",
" print(v)"
]
},
{
+3
View File
@@ -1,3 +1,6 @@
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
INTERRUPT = "__interrupt__"
TAG_HIDDEN = "langsmith:hidden"
+28 -57
View File
@@ -1,16 +1,15 @@
import asyncio
import logging
from collections import defaultdict
from typing import (
Any,
Awaitable,
Callable,
Coroutine,
Dict,
NamedTuple,
Optional,
Sequence,
Union,
cast,
)
from langchain_core.runnables import Runnable
@@ -25,9 +24,11 @@ from langchain_core.runnables.graph import (
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.constants import TAG_HIDDEN
from langgraph.pregel import Channel, Pregel
from langgraph.pregel.read import PregelNode
from langgraph.pregel.write import ChannelWrite
from langgraph.utils import RunnableCallable
logger = logging.getLogger(__name__)
@@ -35,36 +36,8 @@ START = "__start__"
END = "__end__"
class RunnableCallable(Runnable):
def __init__(
self,
func: Callable[..., Optional[Runnable]],
afunc: Callable[..., Awaitable[Optional[Runnable]]],
name: str,
writer: Callable[[str], Optional[Runnable]],
) -> None:
self.name = name
self.func = func
self.afunc = afunc
self.writer = writer
def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
ret = self._call_with_config(self.func, input, config, writer=self.writer)
if isinstance(ret, Runnable):
return ret.invoke(input, config)
return ret
async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
ret = await self._acall_with_config(
self.afunc, input, config, writer=self.writer
)
if isinstance(ret, Runnable):
return await ret.ainvoke(input, config)
return ret
class Branch(NamedTuple):
condition: Union[Runnable[Any, str], Callable[..., str], Coroutine[Any, Any, str]]
condition: Runnable[Any, str]
ends: Optional[dict[str, str]]
def run(self, writer: Callable[[str], Optional[Runnable]]) -> None:
@@ -73,19 +46,19 @@ class Branch(NamedTuple):
func=self._route,
afunc=self._aroute,
writer=writer,
name=self.condition.name
if isinstance(self.condition, Runnable)
else self.condition.__name__,
name=None,
trace=False,
)
)
def _route(
self, input: Any, *, writer: Callable[[str], Optional[Runnable]]
self,
input: Any,
config: RunnableConfig,
*,
writer: Callable[[str], Optional[Runnable]],
) -> Runnable:
if isinstance(self.condition, Runnable):
result = self.condition.invoke(input, {"run_name": "condition"})
else:
result = self.condition(input)
result = self.condition.invoke(input, config)
if self.ends:
destination = self.ends[result]
else:
@@ -93,14 +66,13 @@ class Branch(NamedTuple):
return writer(destination)
async def _aroute(
self, input: Any, *, writer: Callable[[str], Optional[Runnable]]
self,
input: Any,
config: RunnableConfig,
*,
writer: Callable[[str], Optional[Runnable]],
) -> Runnable:
if isinstance(self.condition, Runnable):
result = await self.condition.ainvoke(input, {"run_name": "condition"})
elif asyncio.iscoroutinefunction(self.condition):
result = await self.condition(input)
else:
result = self.condition(input)
result = await self.condition.ainvoke(input, config)
if self.ends:
destination = self.ends[result]
else:
@@ -172,12 +144,8 @@ class Graph:
"not be reflected in the compiled graph."
)
# find a name for the condition
try:
name = (
condition.__name__ if condition.__name__ != "<lambda>" else "condition"
)
except AttributeError:
name = "condition"
condition = coerce_to_runnable(condition)
name = condition.name or "condition"
# validate the condition
if start_key not in self.nodes and start_key != START:
raise ValueError(f"Need to add_node `{start_key}` first")
@@ -293,14 +261,16 @@ class CompiledGraph(Pregel):
def attach_node(self, key: str, node: Runnable) -> None:
self.channels[key] = EphemeralValue(Any)
self.nodes[key] = (
PregelNode(channels=[], triggers=[]) | node | Channel.write_to(key)
PregelNode(channels=[], triggers=[])
| node
| Channel.write_to(key, tags=[TAG_HIDDEN])
)
self.stream_channels.append(key)
cast(list[str], self.stream_channels).append(key)
def attach_edge(self, start: str, end: str) -> None:
if end == END:
# publish to end channel
self.nodes[start].writers.append(Channel.write_to(END))
self.nodes[start].writers.append(Channel.write_to(END, tags=[TAG_HIDDEN]))
else:
# subscribe to start channel
self.nodes[end].triggers.append(start)
@@ -309,12 +279,13 @@ class CompiledGraph(Pregel):
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(end: str) -> Optional[ChannelWrite]:
return Channel.write_to(
f"branch:{start}:{name}:{end}" if end != END else END
f"branch:{start}:{name}:{end}" if end != END else END,
tags=[TAG_HIDDEN],
)
# add hidden start node
if start == START and start not in self.nodes:
self.nodes[start] = Channel.subscribe_to(START, tags=["langsmith:hidden"])
self.nodes[start] = Channel.subscribe_to(START, tags=[TAG_HIDDEN])
# attach branch writer
self.nodes[start] |= branch.run(branch_writer)
+34 -23
View File
@@ -12,8 +12,8 @@ from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.channels.named_barrier_value import NamedBarrierValue
from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.constants import TAG_HIDDEN
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph
from langgraph.pregel import Channel
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
@@ -139,12 +139,31 @@ class CompiledStateGraph(CompiledGraph):
else ChannelWriteEntry(key, RunnableLambda(partial(_get_state_key, key)))
for key in state_keys
]
# node that reads current state with (this node's) updates applied
state_reader = ChannelRead(
state_keys[0] if state_keys == ["__root__"] else state_keys,
tags=[TAG_HIDDEN],
fresh=True,
# coerce state dict to schema class (eg. pydantic model)
mapper=(
None
if state_keys == ["__root__"]
else partial(_coerce_state, self.graph.schema)
),
)
# add node and output channel
if key == START:
self.nodes[key] = Channel.subscribe_to(
START, tags=["langsmith:hidden"]
).pipe(ChannelWrite(state_write_entries))
self.nodes[key] = PregelNode(
tags=[TAG_HIDDEN],
triggers=[START],
channels=[START],
writers=[
ChannelWrite(state_write_entries, tags=[TAG_HIDDEN]),
# read back state with updates applied
state_reader,
],
)
else:
self.channels[key] = EphemeralValue(Any)
self.nodes[key] = PregelNode(
@@ -156,24 +175,15 @@ class CompiledStateGraph(CompiledGraph):
else {chan: chan for chan in state_keys}
),
# coerce state dict to schema class (eg. pydantic model)
mapper=(
None
if state_keys == ["__root__"]
else partial(_coerce_state, self.graph.schema)
),
# publish to this channel and state keys
mapper=state_reader.mapper,
writers=[
ChannelWrite([ChannelWriteEntry(key)] + state_write_entries),
# read back state with updates applied
ChannelRead(
state_keys[0] if state_keys == ["__root__"] else state_keys,
fresh=True,
mapper=(
None
if state_keys == ["__root__"]
else partial(_coerce_state, self.graph.schema)
),
# publish to this channel and state keys
ChannelWrite(
[ChannelWriteEntry(key)] + state_write_entries,
tags=[TAG_HIDDEN],
),
# read back state with updates applied
state_reader,
],
).pipe(node)
@@ -187,7 +197,7 @@ class CompiledStateGraph(CompiledGraph):
self.nodes[end].triggers.append(channel_name)
# publish to channel
self.nodes[START] |= ChannelWrite(
[ChannelWriteEntry(channel_name, START)]
[ChannelWriteEntry(channel_name, START)], tags=[TAG_HIDDEN]
)
elif end != END:
# subscribe to start channel
@@ -201,14 +211,15 @@ class CompiledStateGraph(CompiledGraph):
# publish to channel
for start in starts:
self.nodes[start] |= ChannelWrite(
[ChannelWriteEntry(channel_name, start)]
[ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN]
)
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(end: str) -> Optional[ChannelWrite]:
if end != END:
return ChannelWrite(
[ChannelWriteEntry(f"branch:{start}:{name}:{end}", start)]
[ChannelWriteEntry(f"branch:{start}:{name}:{end}", start)],
tags=[TAG_HIDDEN],
)
# attach branch publisher
+2
View File
@@ -430,6 +430,7 @@ class Pregel(
task.input,
patch_config(
config,
run_name=self.name + "UpdateState",
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
@@ -489,6 +490,7 @@ class Pregel(
task.input,
patch_config(
config,
run_name=self.name + "UpdateState",
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
+2 -3
View File
@@ -1,6 +1,7 @@
from typing import Any, Iterator, Mapping, Optional, Sequence, Union
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.constants import TAG_HIDDEN
from langgraph.pregel.log import logger
from langgraph.pregel.types import PregelExecutableTask
@@ -81,9 +82,7 @@ def map_output_updates(
) -> Optional[dict[str, Union[Any, dict[str, Any]]]]:
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
output_tasks = [
t
for t in tasks
if not t.config or "langsmith:hidden" not in t.config.get("tags")
t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags")
]
if isinstance(output_channels, str):
if updated := {
+22 -11
View File
@@ -6,7 +6,6 @@ from langchain_core.pydantic_v1 import Field
from langchain_core.runnables import (
Runnable,
RunnableConfig,
RunnableLambda,
RunnablePassthrough,
RunnableSequence,
RunnableSerializable,
@@ -17,11 +16,12 @@ from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_READ
from langgraph.pregel.write import ChannelWrite
from langgraph.utils import RunnableCallable
READ_TYPE = Callable[[str, bool], Union[Any, dict[str, Any]]]
class ChannelRead(RunnableLambda):
class ChannelRead(RunnableCallable):
channel: Union[str, list[str]]
fresh: bool = False
@@ -46,12 +46,23 @@ class ChannelRead(RunnableLambda):
*,
fresh: bool = False,
mapper: Optional[Callable[[Any], Any]] = None,
tags: Optional[list[str]] = None,
) -> None:
super().__init__(func=self._read, afunc=self._aread)
super().__init__(func=self._read, afunc=self._aread, tags=tags, name=None)
self.fresh = fresh
self.mapper = mapper
self.channel = channel
self.name = f"ChannelRead<{channel}>"
def get_name(
self, suffix: Optional[str] = None, *, name: Optional[str] = None
) -> str:
if name:
pass
elif isinstance(self.channel, str):
name = f"ChannelRead<{self.channel}>"
else:
name = f"ChannelRead<{','.join(self.channel)}>"
return super().get_name(suffix, name=name)
def _read(self, _: Any, config: RunnableConfig) -> Any:
try:
@@ -80,7 +91,7 @@ class ChannelRead(RunnableLambda):
return read(self.channel, self.fresh)
default_bound: RunnablePassthrough = RunnablePassthrough()
DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()
class PregelNode(RunnableBindingBase):
@@ -92,7 +103,7 @@ class PregelNode(RunnableBindingBase):
writers: list[Runnable] = Field(default_factory=list)
bound: Runnable[Any, Any] = Field(default=default_bound)
bound: Runnable[Any, Any] = Field(default=DEFAULT_BOUND)
kwargs: Mapping[str, Any] = Field(default_factory=dict)
@@ -125,11 +136,11 @@ class PregelNode(RunnableBindingBase):
def get_node(self) -> Optional[Runnable[Any, Any]]:
writers = self.get_writers()
if self.bound is default_bound and not writers:
if self.bound is DEFAULT_BOUND and not writers:
return None
elif self.bound is default_bound and len(writers) == 1:
elif self.bound is DEFAULT_BOUND and len(writers) == 1:
return writers[0]
elif self.bound is default_bound:
elif self.bound is DEFAULT_BOUND:
return RunnableSequence(*writers)
elif writers:
return RunnableSequence(self.bound, *writers)
@@ -154,7 +165,7 @@ class PregelNode(RunnableBindingBase):
triggers=triggers,
mapper=mapper,
writers=writers or [],
bound=bound or default_bound,
bound=bound or DEFAULT_BOUND,
kwargs=kwargs or {},
config=merge_configs(config, {"tags": tags or []}),
**other_kwargs,
@@ -201,7 +212,7 @@ class PregelNode(RunnableBindingBase):
kwargs=self.kwargs,
config=self.config,
)
elif self.bound is default_bound:
elif self.bound is DEFAULT_BOUND:
return PregelNode(
channels=self.channels,
triggers=self.triggers,
+15 -22
View File
@@ -3,14 +3,11 @@ from __future__ import annotations
import asyncio
from typing import Any, Callable, NamedTuple, Optional, Sequence, TypeVar, Union
from langchain_core.runnables import (
Runnable,
RunnableConfig,
RunnablePassthrough,
)
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_SEND
from langgraph.utils import RunnableCallable
TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
R = TypeVar("R", bound=Runnable)
@@ -25,7 +22,7 @@ class ChannelWriteEntry(NamedTuple):
skip_none: bool = False
class ChannelWrite(RunnablePassthrough):
class ChannelWrite(RunnableCallable):
writes: Sequence[ChannelWriteEntry]
"""
Sequence of write entries, each of which is a tuple of:
@@ -34,11 +31,11 @@ class ChannelWrite(RunnablePassthrough):
- whether to skip writing if the mapped value is None
"""
class Config:
arbitrary_types_allowed = True
def __init__(self, writes: Sequence[ChannelWriteEntry]):
super().__init__(func=self._write, afunc=self._awrite, writes=writes)
def __init__(
self, writes: Sequence[ChannelWriteEntry], *, tags: Optional[list[str]] = None
):
super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags)
self.writes = writes
def __repr_args__(self) -> Any:
return [("writes", self.writes)]
@@ -46,15 +43,9 @@ class ChannelWrite(RunnablePassthrough):
def get_name(
self, suffix: Optional[str] = None, *, name: Optional[str] = None
) -> str:
return super().get_name(
suffix,
name=name
or f"ChannelWrite<{','.join(chan for chan, _, _ in self.writes)}>",
)
@property
def is_channel_writer(self) -> bool:
return True
if not name:
name = f"ChannelWrite<{','.join(chan for chan, _, _ in self.writes)}>"
return super().get_name(suffix, name=name)
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
@@ -85,8 +76,8 @@ class ChannelWrite(RunnablePassthrough):
for write, (_, _, skip_none) in zip(values, self.writes)
if not skip_none or write[1] is not None
]
self.do_write(config, **dict(values))
return input
async def _awrite(self, input: Any, config: RunnableConfig) -> None:
values = await asyncio.gather(
@@ -104,8 +95,8 @@ class ChannelWrite(RunnablePassthrough):
for val, (chan, _, skip_none) in zip(values, self.writes)
if not skip_none or val is not None
]
self.do_write(config, **dict(values))
return input
@staticmethod
def do_write(config: RunnableConfig, **values: Any) -> None:
@@ -121,6 +112,8 @@ class ChannelWrite(RunnablePassthrough):
@staticmethod
def register_writer(runnable: R) -> R:
# using object.__setattr__ to work around objects that override __setattr__
# eg. pydantic models and dataclasses
object.__setattr__(runnable, "_is_channel_writer", True)
return runnable
+48
View File
@@ -1,4 +1,8 @@
import enum
from typing import Any, Awaitable, Callable, Optional
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.config import merge_configs
# Before Python 3.11 native StrEnum is not available
@@ -6,3 +10,47 @@ class StrEnum(str, enum.Enum):
"""A string enum."""
pass
class RunnableCallable(Runnable):
"""A much simpler version of RunnableLambda that requires sync and async functions."""
def __init__(
self,
func: Callable[..., Optional[Runnable]],
afunc: Callable[..., Awaitable[Optional[Runnable]]],
name: str,
tags: Optional[list[str]] = None,
trace: bool = True,
**kwargs: Any,
) -> None:
self.name = name
self.func = func
self.afunc = afunc
self.config = {"tags": tags} if tags else None
self.kwargs = kwargs
self.trace = trace
def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
if self.trace:
ret = self._call_with_config(
self.func, input, merge_configs(self.config, config), **self.kwargs
)
else:
ret = self.func(input, merge_configs(self.config, config), **self.kwargs)
if isinstance(ret, Runnable):
return ret.invoke(input, config)
return ret
async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
if self.trace:
ret = await self._acall_with_config(
self.afunc, input, merge_configs(self.config, config), **self.kwargs
)
else:
ret = await self.afunc(
input, merge_configs(self.config, config), **self.kwargs
)
if isinstance(ret, Runnable):
return await ret.ainvoke(input, config)
return ret
Generated
+5 -6
View File
@@ -1587,17 +1587,16 @@ extended-testing = ["aiosqlite (>=0.19.0,<0.20.0)", "aleph-alpha-client (>=2.15.
[[package]]
name = "langchain-core"
version = "0.1.30"
version = "0.1.38"
description = "Building applications with LLMs through composability"
optional = false
python-versions = ">=3.8.1,<4.0"
python-versions = "<4.0,>=3.8.1"
files = [
{file = "langchain_core-0.1.30-py3-none-any.whl", hash = "sha256:c9643505e41d25ba8f20a2e8bf083d0f0d50b9a098d901511fff8df79f831ada"},
{file = "langchain_core-0.1.30.tar.gz", hash = "sha256:e13a016e55e7f082ff3eeeda2d0cb505b89a8830e3a23c1d134d0a89d7871894"},
{file = "langchain_core-0.1.38-py3-none-any.whl", hash = "sha256:d881b2754254cb4bdb0d5bb56e5c138d032b6e75e5cb21f151b01224b322e02b"},
{file = "langchain_core-0.1.38.tar.gz", hash = "sha256:ee8da6d061c06cce7dc22fec224b6ecbc3a8de106d6dd9f409c7fe448ea41861"},
]
[package.dependencies]
anyio = ">=3,<5"
jsonpatch = ">=1.33,<2.0"
langsmith = ">=0.1.0,<0.2.0"
packaging = ">=23.2,<24.0"
@@ -3860,4 +3859,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "2d35e923bf3902e0e11a305f58d17b0efc3fbb444dff8d6cb92e070a993115c9"
content-hash = "3f31fdccb53a66dc294a53d63bcef0233f38c9909170113c44b8dc215fc77d7c"
+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.1.25"
langchain-core = "^0.1.38"
[tool.poetry.group.test.dependencies]
File diff suppressed because it is too large Load Diff
+14 -4
View File
@@ -1773,6 +1773,7 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None
class AgentState(TypedDict, total=False):
input: str
output: str
steps: Annotated[list[str], operator.add]
def left(data: AgentState) -> AgentState:
return {"output": data["input"] + "->left"}
@@ -1781,6 +1782,7 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None
return {"output": data["input"] + "->right"}
def should_start(data: AgentState) -> str:
assert data["steps"] == [], "Expected input to be read from the state"
# Logic to decide where to start
if len(data["input"]) > 10:
return "go-right"
@@ -1891,6 +1893,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
"messages": [
HumanMessage(content="what is weather in sf"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -1907,6 +1910,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -1931,7 +1935,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
ToolMessage(content="result for a third one", tool_call_id="tool_call567"),
AIMessage(content="answer"),
AIMessage(content="answer", id=AnyStr()),
]
}
@@ -1942,6 +1946,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -1970,6 +1975,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -2007,7 +2013,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
]
}
},
{"agent": {"messages": [AIMessage(content="answer")]}},
{"agent": {"messages": [AIMessage(content="answer", id=AnyStr())]}},
]
@@ -2065,6 +2071,7 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
"messages": [
HumanMessage(content="what is weather in sf"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"query"'}
@@ -2072,13 +2079,14 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
),
FunctionMessage(content="result for query", name="search_api"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"another"'}
},
),
FunctionMessage(content="result for another", name="search_api"),
AIMessage(content="answer"),
AIMessage(content="answer", id=AnyStr()),
]
}
@@ -2089,6 +2097,7 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {
@@ -2111,6 +2120,7 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {
@@ -2129,7 +2139,7 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
]
}
},
{"agent": {"messages": [AIMessage(content="answer")]}},
{"agent": {"messages": [AIMessage(content="answer", id=AnyStr())]}},
]
+14 -4
View File
@@ -1814,6 +1814,7 @@ async def test_conditional_entrypoint_graph_state() -> None:
class AgentState(TypedDict, total=False):
input: str
output: str
steps: Annotated[list[str], operator.add]
async def left(data: AgentState) -> AgentState:
return {"output": data["input"] + "->left"}
@@ -1822,6 +1823,7 @@ async def test_conditional_entrypoint_graph_state() -> None:
return {"output": data["input"] + "->right"}
def should_start(data: AgentState) -> str:
assert data["steps"] == [], "Expected input to be read from the state"
# Logic to decide where to start
if len(data["input"]) > 10:
return "go-right"
@@ -1922,6 +1924,7 @@ async def test_prebuilt_tool_chat() -> None:
"messages": [
HumanMessage(content="what is weather in sf"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -1938,6 +1941,7 @@ async def test_prebuilt_tool_chat() -> None:
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -1962,7 +1966,7 @@ async def test_prebuilt_tool_chat() -> None:
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
ToolMessage(content="result for a third one", tool_call_id="tool_call567"),
AIMessage(content="answer"),
AIMessage(content="answer", id=AnyStr()),
]
}
@@ -1976,6 +1980,7 @@ async def test_prebuilt_tool_chat() -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -2004,6 +2009,7 @@ async def test_prebuilt_tool_chat() -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"tool_calls": [
@@ -2041,7 +2047,7 @@ async def test_prebuilt_tool_chat() -> None:
]
}
},
{"agent": {"messages": [AIMessage(content="answer")]}},
{"agent": {"messages": [AIMessage(content="answer", id=AnyStr())]}},
]
@@ -2094,6 +2100,7 @@ async def test_prebuilt_chat() -> None:
"messages": [
HumanMessage(content="what is weather in sf"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"query"'}
@@ -2101,13 +2108,14 @@ async def test_prebuilt_chat() -> None:
),
FunctionMessage(content="result for query", name="search_api"),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"another"'}
},
),
FunctionMessage(content="result for another", name="search_api"),
AIMessage(content="answer"),
AIMessage(content="answer", id=AnyStr()),
]
}
@@ -2121,6 +2129,7 @@ async def test_prebuilt_chat() -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {
@@ -2143,6 +2152,7 @@ async def test_prebuilt_chat() -> None:
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {
@@ -2161,7 +2171,7 @@ async def test_prebuilt_chat() -> None:
]
}
},
{"agent": {"messages": [AIMessage(content="answer")]}},
{"agent": {"messages": [AIMessage(content="answer", id=AnyStr())]}},
]