add tracing v2 env var (#4465)

Co-authored-by: Ankush Gola <ankush.gola@gmail.com>
This commit is contained in:
Harrison Chase 2023-05-10 11:08:29 -07:00 committed by GitHub
parent 9231143f91
commit b2f920e891
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 51 additions and 0 deletions

View File

@ -32,6 +32,11 @@ openai_callback_var: ContextVar[Optional[OpenAICallbackHandler]] = ContextVar(
tracing_callback_var: ContextVar[Optional[LangChainTracer]] = ContextVar( # noqa: E501 tracing_callback_var: ContextVar[Optional[LangChainTracer]] = ContextVar( # noqa: E501
"tracing_callback", default=None "tracing_callback", default=None
) )
tracing_v2_callback_var: ContextVar[
Optional[LangChainTracerV2]
] = ContextVar( # noqa: E501
"tracing_callback_v2", default=None
)
@contextmanager @contextmanager
@ -727,6 +732,11 @@ def _configure(
or tracer is not None or tracer is not None
or os.environ.get("LANGCHAIN_HANDLER") 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") tracer_session = os.environ.get("LANGCHAIN_SESSION")
if tracer_session is None: if tracer_session is None:
tracer_session = "default" tracer_session = "default"
@ -746,6 +756,16 @@ def _configure(
handler = LangChainTracer() handler = LangChainTracer()
handler.load_session(tracer_session) handler.load_session(tracer_session)
callback_manager.add_handler(handler, True) 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( if open_ai is not None and not any(
isinstance(handler, OpenAICallbackHandler) isinstance(handler, OpenAICallbackHandler)
for handler in callback_manager.handlers for handler in callback_manager.handlers

View File

@ -7,6 +7,7 @@ from aiohttp import ClientSession
from langchain.agents import AgentType, initialize_agent, load_tools from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.callbacks import tracing_enabled from langchain.callbacks import tracing_enabled
from langchain.callbacks.manager import tracing_v2_enabled
from langchain.llms import OpenAI from langchain.llms import OpenAI
questions = [ questions = [
@ -121,3 +122,33 @@ async def test_tracing_context_manager_async() -> None:
await asyncio.gather(*tasks) await asyncio.gather(*tasks)
await task 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