diff --git a/README.md b/README.md index 464592a85..cc1c3bd09 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ Check `examples` for more examples. - [x] Test different input and output types (str, str sequence) - [x] Add tests for Stream, UniqueInbox - [ ] Add tests for subscribe_to_each().join() -- [ ] Add optional debug logging +- [x] Add optional debug logging - [ ] Implement checkpointing - [ ] Save checkpoints at end of each step - [ ] Load checkpoint at start of invocation diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index e60888166..48f8573da 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -110,8 +110,10 @@ draft_revise_loop = Pregel( }, # output will be a dict with keys "draft" and "notes" output=["draft", "notes"], - # input can be a dict with any of the channels as keys - input=None, + # input will be a dict with a single key, "question" + input=["question"], + # debug logging + debug=True, ) # run diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 02e04318e..ca95829f6 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -10,6 +10,8 @@ from langchain.callbacks.manager import ( AsyncCallbackManagerForChainRun, CallbackManagerForChainRun, ) +from langchain.globals import get_debug +from langchain.utils.input import get_bolded_text, get_colored_text from langchain.pydantic_v1 import BaseModel, create_model from langchain.schema.runnable import ( Runnable, @@ -48,6 +50,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): step_timeout: Optional[float] = None + debug: bool + class Config: arbitrary_types_allowed = True @@ -58,6 +62,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): output: str | Sequence[str], input: str | Sequence[str], step_timeout: Optional[float] = None, + debug: Optional[bool] = None, ) -> None: chains_flat: list[PregelInvoke | PregelBatch] = [] for chain in chains: @@ -74,6 +79,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): output=output, input=input, step_timeout=step_timeout, + debug=debug if debug is not None else get_debug(), ) @property @@ -181,6 +187,18 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps for step in range(config["recursion_limit"]): + if self.debug: + from pprint import pformat + + n_tasks = len(next_tasks) + print( + f"{get_colored_text('[pregel/step]', color='blue')} " + + get_bolded_text( + f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Current values:\n" + ) + + pformat({k: read(k) for k in channels}) + ) + # collect all writes to channels, without applying them yet pending_writes = deque[tuple[str, Any]]() @@ -281,6 +299,18 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # channels are guaranteed to be immutable for the duration of the step, # channel updates being applied only at the transition between steps for step in range(config["recursion_limit"]): + if self.debug: + from pprint import pformat + + n_tasks = len(next_tasks) + print( + f"{get_colored_text('[pregel/step]', color='blue')} " + + get_bolded_text( + f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Current values:\n" + ) + + pformat({k: read(k) for k in channels}) + ) + # collect all writes to channels, without applying them yet pending_writes = deque[tuple[str, Any]]()