Merge pull request #1794 from langchain-ai/nc/21sep/remove-runtime-value-substitution

Remove runtime value substitution
This commit is contained in:
Nuno Campos
2024-09-23 07:03:44 -07:00
committed by GitHub
6 changed files with 18 additions and 111 deletions
-3
View File
@@ -68,8 +68,6 @@ PUSH = "__pregel_push"
# denotes push-style tasks, ie. those created by Send objects
PULL = "__pregel_pull"
# denotes pull-style tasks, ie. those triggered by edges
RUNTIME_PLACEHOLDER = "__pregel_runtime_placeholder__"
# placeholder for managed values replaced at runtime
NS_SEP = "|"
# for checkpoint_ns, separates each level (ie. graph|subgraph|subsubgraph)
NS_END = ":"
@@ -103,7 +101,6 @@ RESERVED = {
# other constants
PUSH,
PULL,
RUNTIME_PLACEHOLDER,
NS_SEP,
NS_END,
}
+1 -47
View File
@@ -16,16 +16,11 @@ from typing import (
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self, TypeGuard
from langgraph.constants import RUNTIME_PLACEHOLDER
V = TypeVar("V")
U = TypeVar("U")
class ManagedValue(ABC, Generic[V]):
runtime: bool = False
"""Whether the managed value is always created at runtime, ie. never stored."""
def __init__(self, config: RunnableConfig) -> None:
self.config = config
@@ -105,45 +100,4 @@ ChannelKeyPlaceholder = object()
ChannelTypePlaceholder = object()
class ManagedValueMapping(dict[str, ManagedValue]):
def replace_runtime_values(
self, step: int, values: Union[dict[str, Any], Any]
) -> None:
if not self or not values:
return
if all(not mv.runtime for mv in self.values()):
return
if isinstance(values, dict):
for key, value in values.items():
for chan, mv in self.items():
if mv.runtime and mv(step) is value:
values[key] = {RUNTIME_PLACEHOLDER: chan}
elif hasattr(values, "__dir__") and callable(values.__dir__):
for key in dir(values):
try:
value = getattr(values, key)
for chan, mv in self.items():
if mv.runtime and mv(step) is value:
setattr(values, key, {RUNTIME_PLACEHOLDER: chan})
except AttributeError:
pass
def replace_runtime_placeholders(
self, step: int, values: Union[dict[str, Any], Any]
) -> None:
if not self or not values:
return
if all(not mv.runtime for mv in self.values()):
return
if isinstance(values, dict):
for key, value in values.items():
if isinstance(value, dict) and RUNTIME_PLACEHOLDER in value:
values[key] = self[value[RUNTIME_PLACEHOLDER]](step)
elif hasattr(values, "__dir__") and callable(values.__dir__):
for key in dir(values):
try:
value = getattr(values, key)
if isinstance(value, dict) and RUNTIME_PLACEHOLDER in value:
setattr(values, key, self[value[RUNTIME_PLACEHOLDER]](step))
except AttributeError:
pass
ManagedValueMapping = dict[str, ManagedValue]
+2 -8
View File
@@ -852,11 +852,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write,
step + 1,
writes.extend,
self.nodes,
channels,
managed,
self.nodes.keys(),
),
CONFIG_KEY_READ: partial(
local_read,
@@ -1001,11 +998,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write,
step + 1,
writes.extend,
self.nodes,
channels,
managed,
self.nodes.keys(),
),
CONFIG_KEY_READ: partial(
local_read,
+5 -21
View File
@@ -155,11 +155,8 @@ def local_read(
def local_write(
step: int,
commit: Callable[[Sequence[tuple[str, Any]]], None],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
process_keys: Iterable[str],
writes: Sequence[tuple[str, Any]],
) -> None:
"""Function injected under CONFIG_KEY_SEND in task config, to write to channels.
@@ -167,15 +164,9 @@ def local_write(
for chan, value in writes:
if chan == TASKS:
if not isinstance(value, Send):
raise InvalidUpdateError(
f"Invalid packet type, expected Packet, got {value}"
)
if value.node not in processes:
raise InvalidUpdateError(f"Expected Send, got {value}")
if value.node not in process_keys:
raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
# replace any runtime values with placeholders
managed.replace_runtime_values(step, value.arg)
elif chan not in channels and chan not in managed:
logger.warning(f"Skipping write for channel '{chan}' which has no readers")
commit(writes)
@@ -411,7 +402,6 @@ def prepare_single_task(
if for_execution:
proc = processes[packet.node]
if node := proc.node:
managed.replace_runtime_placeholders(step, packet.arg)
if proc.metadata:
metadata.update(proc.metadata)
writes: deque[tuple[str, Any]] = deque()
@@ -433,11 +423,8 @@ def prepare_single_task(
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write,
step,
writes.extend,
processes,
channels,
managed,
processes.keys(),
),
CONFIG_KEY_READ: partial(
local_read,
@@ -543,11 +530,8 @@ def prepare_single_task(
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write,
step,
writes.extend,
processes,
channels,
managed,
processes.keys(),
),
CONFIG_KEY_READ: partial(
local_read,
+5 -14
View File
@@ -4440,25 +4440,16 @@ def test_state_graph_packets(
), "nodes can pass extra data to their cond edges, which isn't saved in state"
# Logic to decide whether to continue in the loop or exit
if tool_calls := data["messages"][-1].tool_calls:
return [
Send("tools", {"call": tool_call, "my_session": data["session"]})
for tool_call in tool_calls
]
return [Send("tools", tool_call) for tool_call in tool_calls]
else:
return END
class ToolInput(TypedDict):
call: ToolCall
my_session: httpx.Client
def tools_node(input: ToolInput, config: RunnableConfig) -> AgentState:
assert isinstance(input["my_session"], httpx.Client)
tool_call = input["call"]
time.sleep(tool_call["args"].get("idx", 0) / 10)
output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config)
def tools_node(input: ToolCall, config: RunnableConfig) -> AgentState:
time.sleep(input["args"].get("idx", 0) / 10)
output = tools_by_name[input["name"]].invoke(input["args"], config)
return {
"messages": ToolMessage(
content=output, name=tool_call["name"], tool_call_id=tool_call["id"]
content=output, name=input["name"], tool_call_id=input["id"]
)
}
+5 -18
View File
@@ -4204,12 +4204,6 @@ async def test_prebuilt_tool_chat() -> None:
]
# defined outside to allow deserializer to see it
class ToolInput(BaseModel, arbitrary_types_allowed=True):
call: ToolCall
my_session: httpx.AsyncClient
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_state_graph_packets(checkpointer_name: str) -> None:
from langchain_core.language_models.fake_chat_models import (
@@ -4273,23 +4267,16 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
assert isinstance(data["session"], httpx.AsyncClient)
# Logic to decide whether to continue in the loop or exit
if tool_calls := data["messages"][-1].tool_calls:
return [
Send("tools", ToolInput(call=tool_call, my_session=data["session"]))
for tool_call in tool_calls
]
return [Send("tools", tool_call) for tool_call in tool_calls]
else:
return END
async def tools_node(input: ToolInput, config: RunnableConfig) -> AgentState:
assert isinstance(input.my_session, httpx.AsyncClient)
tool_call = input.call
await asyncio.sleep(tool_call["args"].get("idx", 0) / 10)
output = await tools_by_name[tool_call["name"]].ainvoke(
tool_call["args"], config
)
async def tools_node(input: ToolCall, config: RunnableConfig) -> AgentState:
await asyncio.sleep(input["args"].get("idx", 0) / 10)
output = await tools_by_name[input["name"]].ainvoke(input["args"], config)
return {
"messages": ToolMessage(
content=output, name=tool_call["name"], tool_call_id=tool_call["id"]
content=output, name=input["name"], tool_call_id=input["id"]
)
}