From 6bca615e6a926ffc2cf97b9d4da497eb61ce08c9 Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Wed, 23 Apr 2025 14:04:54 -0400 Subject: [PATCH] langgraph: allow nested lists/dicts of primitives in RemoteGraph config (#4387) --- libs/langgraph/langgraph/pregel/remote.py | 36 ++++++++- libs/langgraph/tests/test_remote_graph.py | 98 +++++++++++++++++++++++ 2 files changed, 130 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index ba7cdb2c6..09d5f90a8 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -54,6 +54,28 @@ CONF_DROPLIST = frozenset( ) +def sanitize_config_value(v: Any) -> Any: + """Recursively sanitize a config value to ensure it contains only primitives.""" + if isinstance(v, (str, int, float, bool)): + return v + elif isinstance(v, dict): + sanitized_dict = {} + for k, val in v.items(): + if isinstance(k, str): + sanitized_value = sanitize_config_value(val) + if sanitized_value is not None: + sanitized_dict[k] = sanitized_value + return sanitized_dict + elif isinstance(v, (list, tuple)): + sanitized_list = [] + for item in v: + sanitized_item = sanitize_config_value(item) + if sanitized_item is not None: + sanitized_list.append(sanitized_item) + return sanitized_list + return None + + class RemoteException(Exception): """Exception raised when an error occurs in the remote graph.""" @@ -303,20 +325,26 @@ class RemoteGraph(PregelProtocol): sanitized["recursion_limit"] = config["recursion_limit"] if "tags" in config: sanitized["tags"] = [tag for tag in config["tags"] if isinstance(tag, str)] + if "metadata" in config: sanitized["metadata"] = {} for k, v in config["metadata"].items(): - if isinstance(k, str) and isinstance(v, (str, int, float, bool)): - sanitized["metadata"][k] = v + if ( + isinstance(k, str) + and (sanitized_value := sanitize_config_value(v)) is not None + ): + sanitized["metadata"][k] = sanitized_value + if "configurable" in config: sanitized["configurable"] = {} for k, v in config["configurable"].items(): if ( isinstance(k, str) and k not in CONF_DROPLIST - and isinstance(v, (str, int, float, bool)) + and (sanitized_value := sanitize_config_value(v)) is not None ): - sanitized["configurable"][k] = v + sanitized["configurable"][k] = sanitized_value + return sanitized def get_state( diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 878f5e6f1..b30ec3a05 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -1,6 +1,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from langchain_core.runnables import RunnableConfig from langchain_core.runnables.graph import ( Edge as DrawableEdge, ) @@ -861,3 +862,100 @@ async def test_langgraph_cloud_integration(): remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID graph = await remote_pregel.aget_graph(xray=True) print("graph:", graph) + + +def test_sanitize_config(): + # Create a test instance + remote = RemoteGraph("test-graph") + + # Test 1: Basic config with primitives + basic_config: RunnableConfig = { + "recursion_limit": 10, + "tags": ["tag1", "tag2"], + "metadata": {"str_key": "value", "int_key": 42, "bool_key": True}, + "configurable": {"param1": "value1", "param2": 123}, + } + sanitized = remote._sanitize_config(basic_config) + assert sanitized["recursion_limit"] == 10 + assert sanitized["tags"] == ["tag1", "tag2"] + assert sanitized["metadata"] == { + "str_key": "value", + "int_key": 42, + "bool_key": True, + } + assert sanitized["configurable"] == {"param1": "value1", "param2": 123} + + # Test 2: Config with non-string tags and complex metadata + complex_config: RunnableConfig = { + "tags": ["tag1", 123, {"obj": "tag"}, "tag2"], # Only string tags should remain + "metadata": { + "nested": { + "key": "value", + "num": 42, + "invalid": lambda x: x, + }, # Last item should be removed + "list": [1, 2, "three"], + "invalid": lambda x: x, # Should be removed + "tuple": (1, 2, 3), # Should be converted to list + }, + } + sanitized = remote._sanitize_config(complex_config) + assert sanitized["tags"] == ["tag1", "tag2"] + assert sanitized["metadata"] == { + "nested": {"key": "value", "num": 42}, + "list": [1, 2, "three"], + "tuple": [1, 2, 3], + } + assert "invalid" not in sanitized["metadata"] + + # Test 3: Config with configurable fields that should be dropped + config_with_drops: RunnableConfig = { + "configurable": { + "normal_param": "value", + "checkpoint_map": {"key": "value"}, # Should be dropped + "checkpoint_id": "123", # Should be dropped + "checkpoint_ns": "ns", # Should be dropped + } + } + sanitized = remote._sanitize_config(config_with_drops) + assert sanitized["configurable"] == {"normal_param": "value"} + assert "checkpoint_map" not in sanitized["configurable"] + assert "checkpoint_id" not in sanitized["configurable"] + assert "checkpoint_ns" not in sanitized["configurable"] + + # Test 4: Empty config + empty_config: RunnableConfig = {} + sanitized = remote._sanitize_config(empty_config) + assert sanitized == {} + + # Test 5: Config with non-string keys in configurable + invalid_keys_config: RunnableConfig = { + "configurable": { + "valid": "value", + 123: "invalid", # Should be dropped + ("tuple", "key"): "invalid", # Should be dropped + } + } + sanitized = remote._sanitize_config(invalid_keys_config) + assert sanitized["configurable"] == {"valid": "value"} + + # Test 6: Deeply nested structures + nested_config: RunnableConfig = { + "metadata": { + "level1": { + "level2": { + "level3": { + "str": "value", + "list": [1, [2, [3]]], + "dict": {"a": {"b": {"c": "d"}}}, + } + } + } + } + } + sanitized = remote._sanitize_config(nested_config) + assert sanitized["metadata"]["level1"]["level2"]["level3"]["str"] == "value" + assert sanitized["metadata"]["level1"]["level2"]["level3"]["list"] == [1, [2, [3]]] + assert sanitized["metadata"]["level1"]["level2"]["level3"]["dict"] == { + "a": {"b": {"c": "d"}} + }