Implement Function Callback tracer (langchain-ai#6835)

Description: I wanted to be able to redirect debug output to a function, but it wasn't very easy. I figured it would make sense to implement a `FunctionCallbackHandler`, and reimplement `ConsoleCallbackHandler` as a subclass that calls the `print` function. Now I can create a simple subclass in my project that calls `logging.info` or whatever I need. Tag maintainer: @agola11 Twitter handle: `@andandaraalex`
AlmyAI · Jul 12, 2023 · adc96d6 · adc96d6
1 parent 93a84f6
commit adc96d6
Showing 2 changed files with 31 additions and 14 deletions.
diff --git a/langchain/callbacks/tracers/__init__.py b/langchain/callbacks/tracers/__init__.py
@@ -2,12 +2,16 @@
 
 from langchain.callbacks.tracers.langchain import LangChainTracer
 from langchain.callbacks.tracers.langchain_v1 import LangChainTracerV1
-from langchain.callbacks.tracers.stdout import ConsoleCallbackHandler
+from langchain.callbacks.tracers.stdout import (
+    ConsoleCallbackHandler,
+    FunctionCallbackHandler,
+)
 from langchain.callbacks.tracers.wandb import WandbTracer
 
 __all__ = [
     "LangChainTracer",
     "LangChainTracerV1",
+    "FunctionCallbackHandler",
     "ConsoleCallbackHandler",
     "WandbTracer",
 ]
diff --git a/langchain/callbacks/tracers/stdout.py b/langchain/callbacks/tracers/stdout.py
@@ -1,5 +1,5 @@
 import json
-from typing import Any, List
+from typing import Any, Callable, List
 
 from langchain.callbacks.tracers.base import BaseTracer
 from langchain.callbacks.tracers.schemas import Run
@@ -41,10 +41,14 @@ def elapsed(run: Any) -> str:
     return f"{(milliseconds / 1000):.2f}s"
 
 
-class ConsoleCallbackHandler(BaseTracer):
-    """Tracer that prints to the console."""
+class FunctionCallbackHandler(BaseTracer):
+    """Tracer that calls a function with a single str parameter."""
 
-    name = "console_callback_handler"
+    name = "function_callback_handler"
+
+    def __init__(self, function: Callable[[str], None], **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.function_callback = function
 
     def _persist_run(self, run: Run) -> None:
         pass
@@ -74,15 +78,15 @@ def get_breadcrumbs(self, run: Run) -> str:
     # logging methods
     def _on_chain_start(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
-        print(
+        self.function_callback(
             f"{get_colored_text('[chain/start]', color='green')} "
             + get_bolded_text(f"[{crumbs}] Entering Chain run with input:\n")
             + f"{try_json_stringify(run.inputs, '[inputs]')}"
         )
 
     def _on_chain_end(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
-        print(
+        self.function_callback(
             f"{get_colored_text('[chain/end]', color='blue')} "
             + get_bolded_text(
                 f"[{crumbs}] [{elapsed(run)}] Exiting Chain run with output:\n"
@@ -92,7 +96,7 @@ def _on_chain_end(self, run: Run) -> None:
 
     def _on_chain_error(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
-        print(
+        self.function_callback(
             f"{get_colored_text('[chain/error]', color='red')} "
             + get_bolded_text(
                 f"[{crumbs}] [{elapsed(run)}] Chain run errored with error:\n"
@@ -107,15 +111,15 @@ def _on_llm_start(self, run: Run) -> None:
             if "prompts" in run.inputs
             else run.inputs
         )
-        print(
+        self.function_callback(
             f"{get_colored_text('[llm/start]', color='green')} "
             + get_bolded_text(f"[{crumbs}] Entering LLM run with input:\n")
             + f"{try_json_stringify(inputs, '[inputs]')}"
         )
 
     def _on_llm_end(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
-        print(
+        self.function_callback(
             f"{get_colored_text('[llm/end]', color='blue')} "
             + get_bolded_text(
                 f"[{crumbs}] [{elapsed(run)}] Exiting LLM run with output:\n"
@@ -125,7 +129,7 @@ def _on_llm_end(self, run: Run) -> None:
 
     def _on_llm_error(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
-        print(
+        self.function_callback(
             f"{get_colored_text('[llm/error]', color='red')} "
             + get_bolded_text(
                 f"[{crumbs}] [{elapsed(run)}] LLM run errored with error:\n"
@@ -135,7 +139,7 @@ def _on_llm_error(self, run: Run) -> None:
 
     def _on_tool_start(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
-        print(
+        self.function_callback(
             f'{get_colored_text("[tool/start]", color="green")} '
             + get_bolded_text(f"[{crumbs}] Entering Tool run with input:\n")
             + f'"{run.inputs["input"].strip()}"'
@@ -144,7 +148,7 @@ def _on_tool_start(self, run: Run) -> None:
     def _on_tool_end(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
         if run.outputs:
-            print(
+            self.function_callback(
                 f'{get_colored_text("[tool/end]", color="blue")} '
                 + get_bolded_text(
                     f"[{crumbs}] [{elapsed(run)}] Exiting Tool run with output:\n"
@@ -154,9 +158,18 @@ def _on_tool_end(self, run: Run) -> None:
 
     def _on_tool_error(self, run: Run) -> None:
         crumbs = self.get_breadcrumbs(run)
-        print(
+        self.function_callback(
             f"{get_colored_text('[tool/error]', color='red')} "
             + get_bolded_text(f"[{crumbs}] [{elapsed(run)}] ")
             + f"Tool run errored with error:\n"
             f"{run.error}"
         )
+
+
+class ConsoleCallbackHandler(FunctionCallbackHandler):
+    """Tracer that prints to the console."""
+
+    name = "console_callback_handler"
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(function=print, **kwargs)