codekingpro/portable-devtools

1"""Callback handler that writes to a file."""2 3from __future__ import annotations4 5from pathlib import Path6from typing import TYPE_CHECKING, Any, TextIO, cast7 8from typing_extensions import Self, override9 10from langchain_core._api import warn_deprecated11from langchain_core.callbacks import BaseCallbackHandler12from langchain_core.utils.input import print_text13 14if TYPE_CHECKING:15    from langchain_core.agents import AgentAction, AgentFinish16 17 18_GLOBAL_DEPRECATION_WARNED = False19 20 21class FileCallbackHandler(BaseCallbackHandler):22    """Callback handler that writes to a file.23 24    This handler supports both context manager usage (recommended) and direct25    instantiation (deprecated) for backwards compatibility.26 27    Examples:28        Using as a context manager (recommended):29 30        ```python31        with FileCallbackHandler("output.txt") as handler:32            # Use handler with your chain/agent33            chain.invoke(inputs, config={"callbacks": [handler]})34        ```35 36        Direct instantiation (deprecated):37 38        ```python39        handler = FileCallbackHandler("output.txt")40        # File remains open until handler is garbage collected41        try:42            chain.invoke(inputs, config={"callbacks": [handler]})43        finally:44            handler.close()  # Explicit cleanup recommended45        ```46 47    Args:48        filename: The file path to write to.49        mode: The file open mode. Defaults to `'a'` (append).50        color: Default color for text output.51 52    !!! note53 54        When not used as a context manager, a deprecation warning will be issued on55        first use. The file will be opened immediately in `__init__` and closed in56        `__del__` or when `close()` is called explicitly.57 58    """59 60    def __init__(61        self, filename: str, mode: str = "a", color: str | None = None62    ) -> None:63        """Initialize the file callback handler.64 65        Args:66            filename: Path to the output file.67            mode: File open mode (e.g., `'w'`, `'a'`, `'x'`). Defaults to `'a'`.68            color: Default text color for output.69 70        """71        self.filename = filename72        self.mode = mode73        self.color = color74        self._file_opened_in_context = False75        self.file: TextIO = cast(76            "TextIO",77            # Open the file in the specified mode with UTF-8 encoding.78            Path(self.filename).open(self.mode, encoding="utf-8"),  # noqa: SIM11579        )80 81    def __enter__(self) -> Self:82        """Enter the context manager.83 84        Returns:85            The `FileCallbackHandler` instance.86 87        !!! note88 89            The file is already opened in `__init__`, so this just marks that the90            handler is being used as a context manager.91 92        """93        self._file_opened_in_context = True94        return self95 96    def __exit__(97        self,98        exc_type: type[BaseException] | None,99        exc_val: BaseException | None,100        exc_tb: object,101    ) -> None:102        """Exit the context manager and close the file.103 104        Args:105            exc_type: Exception type if an exception occurred.106            exc_val: Exception value if an exception occurred.107            exc_tb: Exception traceback if an exception occurred.108 109        """110        self.close()111 112    def __del__(self) -> None:113        """Destructor to cleanup when done."""114        self.close()115 116    def close(self) -> None:117        """Close the file if it's open.118 119        This method is safe to call multiple times and will only close120        the file if it's currently open.121 122        """123        if hasattr(self, "file") and self.file and not self.file.closed:124            self.file.close()125 126    def _write(127        self,128        text: str,129        color: str | None = None,130        end: str = "",131    ) -> None:132        """Write text to the file with deprecation warning if needed.133 134        Args:135            text: The text to write to the file.136            color: Optional color for the text. Defaults to `self.color`.137            end: String appended after the text.138            file: Optional file to write to. Defaults to `self.file`.139 140        Raises:141            RuntimeError: If the file is closed or not available.142 143        """144        global _GLOBAL_DEPRECATION_WARNED  # noqa: PLW0603145        if not self._file_opened_in_context and not _GLOBAL_DEPRECATION_WARNED:146            warn_deprecated(147                since="0.3.67",148                pending=True,149                message=(150                    "Using FileCallbackHandler without a context manager is "151                    "deprecated. Use 'with FileCallbackHandler(...) as "152                    "handler:' instead."153                ),154            )155            _GLOBAL_DEPRECATION_WARNED = True156 157        if not hasattr(self, "file") or self.file is None or self.file.closed:158            msg = "File is not open. Use FileCallbackHandler as a context manager."159            raise RuntimeError(msg)160 161        print_text(text, file=self.file, color=color, end=end)162 163    @override164    def on_chain_start(165        self, serialized: dict[str, Any], inputs: dict[str, Any], **kwargs: Any166    ) -> None:167        """Print that we are entering a chain.168 169        Args:170            serialized: The serialized chain information.171            inputs: The inputs to the chain.172            **kwargs: Additional keyword arguments that may contain `'name'`.173 174        """175        name = (176            kwargs.get("name")177            or serialized.get("name", serialized.get("id", ["<unknown>"])[-1])178            or "<unknown>"179        )180        self._write(f"\n\n> Entering new {name} chain...", end="\n")181 182    @override183    def on_chain_end(self, outputs: dict[str, Any], **kwargs: Any) -> None:184        """Print that we finished a chain.185 186        Args:187            outputs: The outputs of the chain.188            **kwargs: Additional keyword arguments.189 190        """191        self._write("\n> Finished chain.", end="\n")192 193    @override194    def on_agent_action(195        self, action: AgentAction, color: str | None = None, **kwargs: Any196    ) -> Any:197        """Handle agent action by writing the action log.198 199        Args:200            action: The agent action containing the log to write.201            color: Color override for this specific output.202 203                If `None`, uses `self.color`.204            **kwargs: Additional keyword arguments.205 206        """207        self._write(action.log, color=color or self.color)208 209    @override210    def on_tool_end(211        self,212        output: str,213        color: str | None = None,214        observation_prefix: str | None = None,215        llm_prefix: str | None = None,216        **kwargs: Any,217    ) -> None:218        """Handle tool end by writing the output with optional prefixes.219 220        Args:221            output: The tool output to write.222            color: Color override for this specific output.223 224                If `None`, uses `self.color`.225            observation_prefix: Optional prefix to write before the output.226            llm_prefix: Optional prefix to write after the output.227            **kwargs: Additional keyword arguments.228 229        """230        if observation_prefix is not None:231            self._write(f"\n{observation_prefix}")232        self._write(output)233        if llm_prefix is not None:234            self._write(f"\n{llm_prefix}")235 236    @override237    def on_text(238        self, text: str, color: str | None = None, end: str = "", **kwargs: Any239    ) -> None:240        """Handle text output.241 242        Args:243            text: The text to write.244            color: Color override for this specific output.245 246                If `None`, uses `self.color`.247            end: String appended after the text.248            **kwargs: Additional keyword arguments.249 250        """251        self._write(text, color=color or self.color, end=end)252 253    @override254    def on_agent_finish(255        self, finish: AgentFinish, color: str | None = None, **kwargs: Any256    ) -> None:257        """Handle agent finish by writing the finish log.258 259        Args:260            finish: The agent finish object containing the log to write.261            color: Color override for this specific output.262 263                If `None`, uses `self.color`.264            **kwargs: Additional keyword arguments.265 266        """267        self._write(finish.log, color=color or self.color, end="\n")268