codekingpro/portable-devtools

1"""Callback handler that prints to std out."""2 3from __future__ import annotations4 5from typing import TYPE_CHECKING, Any6 7from typing_extensions import override8 9from langchain_core.callbacks.base import BaseCallbackHandler10from langchain_core.utils import print_text11 12if TYPE_CHECKING:13    from langchain_core.agents import AgentAction, AgentFinish14 15 16class StdOutCallbackHandler(BaseCallbackHandler):17    """Callback handler that prints to std out."""18 19    def __init__(self, color: str | None = None) -> None:20        """Initialize callback handler.21 22        Args:23            color: The color to use for the text.24        """25        self.color = color26 27    @override28    def on_chain_start(29        self, serialized: dict[str, Any], inputs: dict[str, Any], **kwargs: Any30    ) -> None:31        """Print out that we are entering a chain.32 33        Args:34            serialized: The serialized chain.35            inputs: The inputs to the chain.36            **kwargs: Additional keyword arguments.37        """38        if "name" in kwargs:39            name = kwargs["name"]40        elif serialized:41            name = serialized.get("name", serialized.get("id", ["<unknown>"])[-1])42        else:43            name = "<unknown>"44        print(f"\n\n\033[1m> Entering new {name} chain...\033[0m")  # noqa: T20145 46    @override47    def on_chain_end(self, outputs: dict[str, Any], **kwargs: Any) -> None:48        """Print out that we finished a chain.49 50        Args:51            outputs: The outputs of the chain.52            **kwargs: Additional keyword arguments.53        """54        print("\n\033[1m> Finished chain.\033[0m")  # noqa: T20155 56    @override57    def on_agent_action(58        self, action: AgentAction, color: str | None = None, **kwargs: Any59    ) -> Any:60        """Run on agent action.61 62        Args:63            action: The agent action.64            color: The color to use for the text.65            **kwargs: Additional keyword arguments.66        """67        print_text(action.log, color=color or self.color)68 69    @override70    def on_tool_end(71        self,72        output: Any,73        color: str | None = None,74        observation_prefix: str | None = None,75        llm_prefix: str | None = None,76        **kwargs: Any,77    ) -> None:78        """If not the final action, print out observation.79 80        Args:81            output: The output to print.82            color: The color to use for the text.83            observation_prefix: The observation prefix.84            llm_prefix: The LLM prefix.85            **kwargs: Additional keyword arguments.86        """87        output = str(output)88        if observation_prefix is not None:89            print_text(f"\n{observation_prefix}")90        print_text(output, color=color or self.color)91        if llm_prefix is not None:92            print_text(f"\n{llm_prefix}")93 94    @override95    def on_text(96        self,97        text: str,98        color: str | None = None,99        end: str = "",100        **kwargs: Any,101    ) -> None:102        """Run when the agent ends.103 104        Args:105            text: The text to print.106            color: The color to use for the text.107            end: The end character to use.108            **kwargs: Additional keyword arguments.109        """110        print_text(text, color=color or self.color, end=end)111 112    @override113    def on_agent_finish(114        self, finish: AgentFinish, color: str | None = None, **kwargs: Any115    ) -> None:116        """Run on the agent end.117 118        Args:119            finish: The agent finish.120            color: The color to use for the text.121            **kwargs: Additional keyword arguments.122        """123        print_text(finish.log, color=color or self.color, end="\n")124