diff --git a/permchain/pubsub.py b/permchain/pubsub.py index e066a86f2..63b96a653 100644 --- a/permchain/pubsub.py +++ b/permchain/pubsub.py @@ -20,7 +20,9 @@ from typing import ( from langchain.callbacks.manager import CallbackManagerForChainRun from langchain.load.serializable import Serializable from langchain.schema.runnable import Runnable, RunnableConfig, patch_config +from langchain.schema.runnable.base import Runnable from langchain.schema.runnable.config import get_executor_for_config +from tenacity import BaseRetrying from permchain.connection import PubSubConnection from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY @@ -62,6 +64,12 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): class Config: arbitrary_types_allowed = True + def with_retry(self, retry: BaseRetrying) -> Runnable[Any, Any]: + return self.__class__( + processes=[p.with_retry(retry) for p in self.processes], + connection=self.connection, + ) + def _transform( self, input: Iterator[Any],