codekingpro/portable-devtools

1"""Tracers that print to the console."""2 3import json4from collections.abc import Callable5from typing import Any6 7from langchain_core.tracers.base import BaseTracer8from langchain_core.tracers.schemas import Run9from langchain_core.utils.input import get_bolded_text, get_colored_text10 11MILLISECONDS_IN_SECOND = 100012 13 14def try_json_stringify(obj: Any, fallback: str) -> str:15    """Try to stringify an object to JSON.16 17    Args:18        obj: Object to stringify.19        fallback: Fallback string to return if the object cannot be stringified.20 21    Returns:22        A JSON string if the object can be stringified, otherwise the fallback string.23    """24    try:25        return json.dumps(obj, indent=2, ensure_ascii=False)26    except Exception:27        return fallback28 29 30def elapsed(run: Any) -> str:31    """Get the elapsed time of a run.32 33    Args:34        run: any object with a `start_time` and `end_time` attribute.35 36    Returns:37        A string with the elapsed time in seconds or milliseconds if time is less than a38            second.39 40    """41    elapsed_time = run.end_time - run.start_time42    seconds = elapsed_time.total_seconds()43    if seconds < 1:44        return f"{seconds * MILLISECONDS_IN_SECOND:.0f}ms"45    return f"{seconds:.2f}s"46 47 48class FunctionCallbackHandler(BaseTracer):49    """Tracer that calls a function with a single str parameter."""50 51    name: str = "function_callback_handler"52    """The name of the tracer.53 54    This is used to identify the tracer in the logs.55    """56 57    def __init__(self, function: Callable[[str], None], **kwargs: Any) -> None:58        """Create a `FunctionCallbackHandler`.59 60        Args:61            function: The callback function to call.62        """63        super().__init__(**kwargs)64        self.function_callback = function65 66    def _persist_run(self, run: Run) -> None:67        pass68 69    def get_parents(self, run: Run) -> list[Run]:70        """Get the parents of a run.71 72        Args:73            run: The run to get the parents of.74 75        Returns:76            A list of parent runs.77        """78        parents = []79        current_run = run80        while current_run.parent_run_id:81            parent = self.run_map.get(str(current_run.parent_run_id))82            if parent:83                parents.append(parent)84                current_run = parent85            else:86                break87        return parents88 89    def get_breadcrumbs(self, run: Run) -> str:90        """Get the breadcrumbs of a run.91 92        Args:93            run: The run to get the breadcrumbs of.94 95        Returns:96            A string with the breadcrumbs of the run.97        """98        parents = self.get_parents(run)[::-1]99        return " > ".join(100            f"{parent.run_type}:{parent.name}"101            for i, parent in enumerate([*parents, run])102        )103 104    # logging methods105    def _on_chain_start(self, run: Run) -> None:106        crumbs = self.get_breadcrumbs(run)107        run_type = run.run_type.capitalize()108        self.function_callback(109            f"{get_colored_text('[chain/start]', color='green')} "110            + get_bolded_text(f"[{crumbs}] Entering {run_type} run with input:\n")111            + f"{try_json_stringify(run.inputs, '[inputs]')}"112        )113 114    def _on_chain_end(self, run: Run) -> None:115        crumbs = self.get_breadcrumbs(run)116        run_type = run.run_type.capitalize()117        self.function_callback(118            f"{get_colored_text('[chain/end]', color='blue')} "119            + get_bolded_text(120                f"[{crumbs}] [{elapsed(run)}] Exiting {run_type} run with output:\n"121            )122            + f"{try_json_stringify(run.outputs, '[outputs]')}"123        )124 125    def _on_chain_error(self, run: Run) -> None:126        crumbs = self.get_breadcrumbs(run)127        run_type = run.run_type.capitalize()128        self.function_callback(129            f"{get_colored_text('[chain/error]', color='red')} "130            + get_bolded_text(131                f"[{crumbs}] [{elapsed(run)}] {run_type} run errored with error:\n"132            )133            + f"{try_json_stringify(run.error, '[error]')}"134        )135 136    def _on_llm_start(self, run: Run) -> None:137        crumbs = self.get_breadcrumbs(run)138        inputs = (139            {"prompts": [p.strip() for p in run.inputs["prompts"]]}140            if "prompts" in run.inputs141            else run.inputs142        )143        self.function_callback(144            f"{get_colored_text('[llm/start]', color='green')} "145            + get_bolded_text(f"[{crumbs}] Entering LLM run with input:\n")146            + f"{try_json_stringify(inputs, '[inputs]')}"147        )148 149    def _on_llm_end(self, run: Run) -> None:150        crumbs = self.get_breadcrumbs(run)151        self.function_callback(152            f"{get_colored_text('[llm/end]', color='blue')} "153            + get_bolded_text(154                f"[{crumbs}] [{elapsed(run)}] Exiting LLM run with output:\n"155            )156            + f"{try_json_stringify(run.outputs, '[response]')}"157        )158 159    def _on_llm_error(self, run: Run) -> None:160        crumbs = self.get_breadcrumbs(run)161        self.function_callback(162            f"{get_colored_text('[llm/error]', color='red')} "163            + get_bolded_text(164                f"[{crumbs}] [{elapsed(run)}] LLM run errored with error:\n"165            )166            + f"{try_json_stringify(run.error, '[error]')}"167        )168 169    def _on_tool_start(self, run: Run) -> None:170        crumbs = self.get_breadcrumbs(run)171        self.function_callback(172            f"{get_colored_text('[tool/start]', color='green')} "173            + get_bolded_text(f"[{crumbs}] Entering Tool run with input:\n")174            + f'"{run.inputs["input"].strip()}"'175        )176 177    def _on_tool_end(self, run: Run) -> None:178        crumbs = self.get_breadcrumbs(run)179        if run.outputs:180            self.function_callback(181                f"{get_colored_text('[tool/end]', color='blue')} "182                + get_bolded_text(183                    f"[{crumbs}] [{elapsed(run)}] Exiting Tool run with output:\n"184                )185                + f'"{str(run.outputs["output"]).strip()}"'186            )187 188    def _on_tool_error(self, run: Run) -> None:189        crumbs = self.get_breadcrumbs(run)190        self.function_callback(191            f"{get_colored_text('[tool/error]', color='red')} "192            + get_bolded_text(f"[{crumbs}] [{elapsed(run)}] ")193            + f"Tool run errored with error:\n"194            f"{run.error}"195        )196 197 198class ConsoleCallbackHandler(FunctionCallbackHandler):199    """Tracer that prints to the console."""200 201    name: str = "console_callback_handler"202 203    def __init__(self, **kwargs: Any) -> None:204        """Create a ConsoleCallbackHandler."""205        super().__init__(function=print, **kwargs)206