feat(integrations): Add WandbTracer (langchain-ai#4521)

# WandbTracer This PR adds the `WandbTracer` and deprecates the existing `WandbCallbackHandler`. Added an example notebook under the docs section alongside the `LangchainTracer` Here's an example [colab](https://colab.research.google.com/drive/1pY13ym8ENEZ8Fh7nA99ILk2GcdUQu0jR?usp=sharing) with the same notebook and the [trace](https://wandb.ai/parambharat/langchain-tracing/runs/8i45cst6) generated from the colab run Co-authored-by: Bharat Ramanathan <ramanathan.parameshwaran@gohuddl.com> Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
Undertone0809 · Jun 19, 2023 · a900859 · a900859
1 parent 53adb8a
commit a900859
Show file tree

Hide file tree

Showing 10 changed files with 697 additions and 34 deletions.
diff --git a/docs/integrations/agent_with_wandb_tracing.ipynb b/docs/integrations/agent_with_wandb_tracing.ipynb
diff --git a/docs/integrations/wandb_tracking.ipynb b/docs/integrations/wandb_tracking.ipynb
@@ -10,7 +10,9 @@
     "\n",
     "Run in Colab: https://colab.research.google.com/drive/1DXH4beT4HFaRKy_Vm4PoxhXVDRf7Ym8L?usp=sharing\n",
     "\n",
-    "View Report: https://wandb.ai/a-sh0ts/langchain_callback_demo/reports/Prompt-Engineering-LLMs-with-LangChain-and-W-B--VmlldzozNjk1NTUw#👋-how-to-build-a-callback-in-langchain-for-better-prompt-engineering"
+    "View Report: https://wandb.ai/a-sh0ts/langchain_callback_demo/reports/Prompt-Engineering-LLMs-with-LangChain-and-W-B--VmlldzozNjk1NTUw#👋-how-to-build-a-callback-in-langchain-for-better-prompt-engineering\n",
+    "\n",
+    "**Note**: _the `WandbCallbackHandler` is being deprecated in favour of the `WandbTracer`_ . In future please use the `WandbTracer` as it is more flexible and allows for more granular logging. To know more about the `WandbTracer` refer to the agent_with_wandb_tracing.ipynb notebook in docs or use the following [colab](https://colab.research.google.com/drive/1pY13ym8ENEZ8Fh7nA99ILk2GcdUQu0jR?usp=sharing)."
    ]
   },
   {
@@ -107,7 +109,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mharrison-chase\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+      "\u001B[34m\u001B[1mwandb\u001B[0m: Currently logged in as: \u001B[33mharrison-chase\u001B[0m. Use \u001B[1m`wandb login --relogin`\u001B[0m to force relogin\n"
      ]
     },
     {
@@ -174,7 +176,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m The wandb callback is currently in beta and is subject to change based on updates to `langchain`. Please report any issues to https://github.com/wandb/wandb/issues with the tag `langchain`.\n"
+      "\u001B[34m\u001B[1mwandb\u001B[0m: \u001B[33mWARNING\u001B[0m The wandb callback is currently in beta and is subject to change based on updates to `langchain`. Please report any issues to https://github.com/wandb/wandb/issues with the tag `langchain`.\n"
      ]
     }
    ],
@@ -521,20 +523,20 @@
      "text": [
       "\n",
       "\n",
-      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
-      "\u001b[32;1m\u001b[1;3m I need to find out who Leo DiCaprio's girlfriend is and then calculate her age raised to the 0.43 power.\n",
+      "\u001B[1m> Entering new AgentExecutor chain...\u001B[0m\n",
+      "\u001B[32;1m\u001B[1;3m I need to find out who Leo DiCaprio's girlfriend is and then calculate her age raised to the 0.43 power.\n",
       "Action: Search\n",
-      "Action Input: \"Leo DiCaprio girlfriend\"\u001b[0m\n",
-      "Observation: \u001b[36;1m\u001b[1;3mDiCaprio had a steady girlfriend in Camila Morrone. He had been with the model turned actress for nearly five years, as they were first said to be dating at the end of 2017. And the now 26-year-old Morrone is no stranger to Hollywood.\u001b[0m\n",
-      "Thought:\u001b[32;1m\u001b[1;3m I need to calculate her age raised to the 0.43 power.\n",
+      "Action Input: \"Leo DiCaprio girlfriend\"\u001B[0m\n",
+      "Observation: \u001B[36;1m\u001B[1;3mDiCaprio had a steady girlfriend in Camila Morrone. He had been with the model turned actress for nearly five years, as they were first said to be dating at the end of 2017. And the now 26-year-old Morrone is no stranger to Hollywood.\u001B[0m\n",
+      "Thought:\u001B[32;1m\u001B[1;3m I need to calculate her age raised to the 0.43 power.\n",
       "Action: Calculator\n",
-      "Action Input: 26^0.43\u001b[0m\n",
-      "Observation: \u001b[33;1m\u001b[1;3mAnswer: 4.059182145592686\n",
-      "\u001b[0m\n",
-      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer.\n",
-      "Final Answer: Leo DiCaprio's girlfriend is Camila Morrone and her current age raised to the 0.43 power is 4.059182145592686.\u001b[0m\n",
+      "Action Input: 26^0.43\u001B[0m\n",
+      "Observation: \u001B[33;1m\u001B[1;3mAnswer: 4.059182145592686\n",
+      "\u001B[0m\n",
+      "Thought:\u001B[32;1m\u001B[1;3m I now know the final answer.\n",
+      "Final Answer: Leo DiCaprio's girlfriend is Camila Morrone and her current age raised to the 0.43 power is 4.059182145592686.\u001B[0m\n",
       "\n",
-      "\u001b[1m> Finished chain.\u001b[0m\n"
+      "\u001B[1m> Finished chain.\u001B[0m\n"
      ]
     },
     {

diff --git a/langchain/callbacks/__init__.py b/langchain/callbacks/__init__.py
@@ -6,6 +6,7 @@
 from langchain.callbacks.manager import (
     get_openai_callback,
     tracing_enabled,
+    wandb_tracing_enabled,
 )
 from langchain.callbacks.mlflow_callback import MlflowCallbackHandler
 from langchain.callbacks.openai_info import OpenAICallbackHandler
@@ -26,4 +27,5 @@
     "AsyncIteratorCallbackHandler",
     "get_openai_callback",
     "tracing_enabled",
+    "wandb_tracing_enabled",
 ]
diff --git a/langchain/callbacks/manager.py b/langchain/callbacks/manager.py
@@ -25,6 +25,7 @@
 from langchain.callbacks.tracers.langchain_v1 import LangChainTracerV1, TracerSessionV1
 from langchain.callbacks.tracers.schemas import TracerSession
 from langchain.callbacks.tracers.stdout import ConsoleCallbackHandler
+from langchain.callbacks.tracers.wandb import WandbTracer
 from langchain.schema import (
     AgentAction,
     AgentFinish,
@@ -44,6 +45,12 @@
 ] = ContextVar(  # noqa: E501
     "tracing_callback", default=None
 )
+wandb_tracing_callback_var: ContextVar[
+    Optional[WandbTracer]
+] = ContextVar(  # noqa: E501
+    "tracing_wandb_callback", default=None
+)
+
 tracing_v2_callback_var: ContextVar[
     Optional[LangChainTracer]
 ] = ContextVar(  # noqa: E501
@@ -76,6 +83,17 @@ def tracing_enabled(
     tracing_callback_var.set(None)
 
 
+@contextmanager
+def wandb_tracing_enabled(
+    session_name: str = "default",
+) -> Generator[None, None, None]:
+    """Get WandbTracer in a context manager."""
+    cb = WandbTracer()
+    wandb_tracing_callback_var.set(cb)
+    yield None
+    wandb_tracing_callback_var.set(None)
+
+
 @contextmanager
 def tracing_v2_enabled(
     session_name: Optional[str] = None,
@@ -831,12 +849,17 @@ def _configure(
             callback_manager.add_handler(handler, False)
 
     tracer = tracing_callback_var.get()
+    wandb_tracer = wandb_tracing_callback_var.get()
     open_ai = openai_callback_var.get()
     tracing_enabled_ = (
         os.environ.get("LANGCHAIN_TRACING") is not None
         or tracer is not None
         or os.environ.get("LANGCHAIN_HANDLER") is not None
     )
+    wandb_tracing_enabled_ = (
+        os.environ.get("LANGCHAIN_WANDB_TRACING") is not None
+        or wandb_tracer is not None
+    )
 
     tracer_v2 = tracing_v2_callback_var.get()
     tracing_v2_enabled_ = (
@@ -851,6 +874,7 @@ def _configure(
         or debug
         or tracing_enabled_
         or tracing_v2_enabled_
+        or wandb_tracing_enabled_
         or open_ai is not None
     ):
         if verbose and not any(
@@ -876,6 +900,14 @@ def _configure(
                 handler = LangChainTracerV1()
                 handler.load_session(tracer_session)
                 callback_manager.add_handler(handler, True)
+        if wandb_tracing_enabled_ and not any(
+            isinstance(handler, WandbTracer) for handler in callback_manager.handlers
+        ):
+            if wandb_tracer:
+                callback_manager.add_handler(wandb_tracer, True)
+            else:
+                handler = WandbTracer()
+                callback_manager.add_handler(handler, True)
         if tracing_v2_enabled_ and not any(
             isinstance(handler, LangChainTracer)
             for handler in callback_manager.handlers

diff --git a/langchain/callbacks/tracers/__init__.py b/langchain/callbacks/tracers/__init__.py
@@ -3,5 +3,11 @@
 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.wandb import WandbTracer
 
-__all__ = ["LangChainTracer", "LangChainTracerV1", "ConsoleCallbackHandler"]
+__all__ = [
+    "LangChainTracer",
+    "LangChainTracerV1",
+    "ConsoleCallbackHandler",
+    "WandbTracer",
+]