Merge pull request #1479 from langchain-ai/nc/26aug/with-config

Override with_config to store config in Pregel instance

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)

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,18 @@ 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) -> Self:
+        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 +305,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 +325,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 +352,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 +385,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 +434,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 +484,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 +526,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 +660,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 +909,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 +1149,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),