From b2f920e89164db2df677e6b4d97f859835956380 Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Wed, 10 May 2023 11:08:29 -0700
Subject: [PATCH] add tracing v2 env var (#4465)

Co-authored-by: Ankush Gola <ankush.gola@gmail.com>
---
 langchain/callbacks/manager.py                | 20 ++++++++++++
 .../callbacks/test_langchain_tracer.py        | 31 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/langchain/callbacks/manager.py b/langchain/callbacks/manager.py
index cd03de43..745f6dd2 100644
--- a/langchain/callbacks/manager.py
+++ b/langchain/callbacks/manager.py
@@ -32,6 +32,11 @@ openai_callback_var: ContextVar[Optional[OpenAICallbackHandler]] = ContextVar(
 tracing_callback_var: ContextVar[Optional[LangChainTracer]] = ContextVar(  # noqa: E501
     "tracing_callback", default=None
 )
+tracing_v2_callback_var: ContextVar[
+    Optional[LangChainTracerV2]
+] = ContextVar(  # noqa: E501
+    "tracing_callback_v2", default=None
+)
 
 
 @contextmanager
@@ -727,6 +732,11 @@ def _configure(
         or tracer is not None
         or os.environ.get("LANGCHAIN_HANDLER") is not None
     )
+
+    tracer_v2 = tracing_v2_callback_var.get()
+    tracing_v2_enabled_ = (
+        os.environ.get("LANGCHAIN_TRACING_V2") is not None or tracer_v2 is not None
+    )
     tracer_session = os.environ.get("LANGCHAIN_SESSION")
     if tracer_session is None:
         tracer_session = "default"
@@ -746,6 +756,16 @@ def _configure(
                 handler = LangChainTracer()
                 handler.load_session(tracer_session)
                 callback_manager.add_handler(handler, True)
+        if tracing_v2_enabled_ and not any(
+            isinstance(handler, LangChainTracerV2)
+            for handler in callback_manager.handlers
+        ):
+            if tracer_v2:
+                callback_manager.add_handler(tracer_v2, True)
+            else:
+                handler = LangChainTracerV2()
+                handler.load_session(tracer_session)
+                callback_manager.add_handler(handler, True)
         if open_ai is not None and not any(
             isinstance(handler, OpenAICallbackHandler)
             for handler in callback_manager.handlers
diff --git a/tests/integration_tests/callbacks/test_langchain_tracer.py b/tests/integration_tests/callbacks/test_langchain_tracer.py
index cd830ea2..781da6f7 100644
--- a/tests/integration_tests/callbacks/test_langchain_tracer.py
+++ b/tests/integration_tests/callbacks/test_langchain_tracer.py
@@ -7,6 +7,7 @@ from aiohttp import ClientSession
 
 from langchain.agents import AgentType, initialize_agent, load_tools
 from langchain.callbacks import tracing_enabled
+from langchain.callbacks.manager import tracing_v2_enabled
 from langchain.llms import OpenAI
 
 questions = [
@@ -121,3 +122,33 @@ async def test_tracing_context_manager_async() -> None:
         await asyncio.gather(*tasks)
 
     await task
+
+
+@pytest.mark.asyncio
+async def test_tracing_v2_environment_variable() -> None:
+    os.environ["LANGCHAIN_TRACING_V2"] = "true"
+
+    aiosession = ClientSession()
+    llm = OpenAI(temperature=0)
+    async_tools = load_tools(["llm-math", "serpapi"], llm=llm, aiosession=aiosession)
+    agent = initialize_agent(
+        async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    tasks = [agent.arun(q) for q in questions[:3]]
+    await asyncio.gather(*tasks)
+    await aiosession.close()
+
+
+def test_tracing_v2_context_manager() -> None:
+    llm = OpenAI(temperature=0)
+    tools = load_tools(["llm-math", "serpapi"], llm=llm)
+    agent = initialize_agent(
+        tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    if "LANGCHAIN_TRACING_V2" in os.environ:
+        del os.environ["LANGCHAIN_TRACING_V2"]
+    with tracing_v2_enabled() as session:
+        assert session
+        agent.run(questions[0])  # this should be traced
+
+    agent.run(questions[0])  # this should not be traced