From 25a72e77efb00c8dbed907e5e1baaad77295388f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 26 Aug 2024 11:57:57 -0700 Subject: [PATCH 1/3] Override with_config to store config in Pregel instance - This enables eg customizing callbacks/metadata in langgraph cloud deployments --- libs/langgraph/langgraph/pregel/__init__.py | 34 ++++++++++++++++++--- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 73db1bfb8..f127859ff 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -38,6 +38,7 @@ from langchain_core.runnables.config import ( ensure_config, get_async_callback_manager_for_config, get_callback_manager_for_config, + merge_configs, patch_config, ) from langchain_core.runnables.utils import ( @@ -219,11 +220,20 @@ class Pregel( config_type: Optional[Type[Any]] = None + config: Optional[RunnableConfig] = None + name: str = "LangGraph" class Config: arbitrary_types_allowed = True + def with_config( + self, config: RunnableConfig | None = None, **kwargs: Any + ) -> Runnable[Dict[All, Any] | Any, Dict[All, Any] | Any]: + return self.copy( + update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})} + ) + @classmethod def is_lc_serializable(cls) -> bool: """Return whether the graph can be serialized by Langchain.""" @@ -297,6 +307,7 @@ class Pregel( def get_input_schema( self, config: Optional[RunnableConfig] = None ) -> Type[BaseModel]: + config = merge_configs(self.config, config) if isinstance(self.input_channels, str): return super().get_input_schema(config) else: @@ -316,6 +327,7 @@ class Pregel( def get_output_schema( self, config: Optional[RunnableConfig] = None ) -> Type[BaseModel]: + config = merge_configs(self.config, config) if isinstance(self.output_channels, str): return super().get_output_schema(config) else: @@ -342,6 +354,7 @@ class Pregel( if not self.checkpointer: raise ValueError("No checkpointer set") + config = merge_configs(self.config, config) if self.config else config saved = self.checkpointer.get_tuple(config) checkpoint = saved.checkpoint if saved else empty_checkpoint() config = saved.config if saved else config @@ -374,6 +387,7 @@ class Pregel( if not self.checkpointer: raise ValueError("No checkpointer set") + config = merge_configs(self.config, config) if self.config else config saved = await self.checkpointer.aget_tuple(config) checkpoint = saved.checkpoint if saved else empty_checkpoint() @@ -422,7 +436,12 @@ class Pregel( metadata, parent_config, pending_writes, - ) in self.checkpointer.list(config, before=before, limit=limit, filter=filter): + ) in self.checkpointer.list( + merge_configs(self.config, config) if self.config else config, + before=before, + limit=limit, + filter=filter, + ): with ChannelsManager( self.channels, checkpoint, config, skip_context=True ) as (channels, managed): @@ -467,7 +486,12 @@ class Pregel( metadata, parent_config, pending_writes, - ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter): + ) in self.checkpointer.alist( + merge_configs(self.config, config) if self.config else config, + before=before, + limit=limit, + filter=filter, + ): async with AsyncChannelsManager( self.channels, checkpoint, config, skip_context=True ) as (channels, managed): @@ -504,6 +528,7 @@ class Pregel( raise ValueError("No checkpointer set") # get last checkpoint + config = merge_configs(self.config, config) if self.config else config saved = self.checkpointer.get_tuple(config) checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() checkpoint_previous_versions = ( @@ -637,6 +662,7 @@ class Pregel( raise ValueError("No checkpointer set") # get last checkpoint + config = merge_configs(self.config, config) if self.config else config saved = await self.checkpointer.aget_tuple(config) checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() checkpoint_previous_versions = ( @@ -885,7 +911,7 @@ class Pregel( {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}} ``` """ - config = ensure_config(config) + config = ensure_config(merge_configs(self.config, config)) callback_manager = get_callback_manager_for_config(config) run_manager = callback_manager.on_chain_start( dumpd(self), @@ -1125,7 +1151,7 @@ class Pregel( {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}} ``` """ - config = ensure_config(config) + config = ensure_config(merge_configs(self.config, config)) callback_manager = get_async_callback_manager_for_config(config) run_manager = await callback_manager.on_chain_start( dumpd(self), From 3979bdb792a2f1605b954cca93ea9ffd78e84edd Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 26 Aug 2024 13:14:11 -0700 Subject: [PATCH 2/3] Type as self --- libs/langgraph/langgraph/pregel/__init__.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f127859ff..e871815ea 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -227,9 +227,7 @@ class Pregel( class Config: arbitrary_types_allowed = True - def with_config( - self, config: RunnableConfig | None = None, **kwargs: Any - ) -> Runnable[Dict[All, Any] | Any, Dict[All, Any] | Any]: + def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: return self.copy( update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})} ) From 1d0f3577a7e9ee7e5086352560298791e55db9ff Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 26 Aug 2024 13:17:19 -0700 Subject: [PATCH 3/3] Catch any type error when reviving saved values --- libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index d52bfc8c1..e39a66e76 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -150,7 +150,7 @@ class JsonPlusSerializer(SerializerProtocol): return method(**value["kwargs"]) else: return method() - except (ImportError, AttributeError): + except (ImportError, AttributeError, TypeError): return None return LC_REVIVER(value)