codekingpro/portable-devtools

1"""Context management for tracers."""2 3from __future__ import annotations4 5from contextlib import contextmanager6from contextvars import ContextVar7from typing import (8    TYPE_CHECKING,9    Any,10    Literal,11    cast,12)13from uuid import UUID14 15from langsmith import run_helpers as ls_rh16from langsmith import utils as ls_utils17 18from langchain_core.tracers.langchain import LangChainTracer19from langchain_core.tracers.run_collector import RunCollectorCallbackHandler20 21if TYPE_CHECKING:22    from collections.abc import Generator23 24    from langsmith import Client as LangSmithClient25 26    from langchain_core.callbacks.base import BaseCallbackHandler, Callbacks27    from langchain_core.callbacks.manager import AsyncCallbackManager, CallbackManager28 29# for backwards partial compatibility if this is imported by users but unused30tracing_callback_var: Any = None31tracing_v2_callback_var: ContextVar[LangChainTracer | None] = ContextVar(32    "tracing_callback_v2", default=None33)34run_collector_var: ContextVar[RunCollectorCallbackHandler | None] = ContextVar(35    "run_collector", default=None36)37 38 39@contextmanager40def tracing_v2_enabled(41    project_name: str | None = None,42    *,43    example_id: str | UUID | None = None,44    tags: list[str] | None = None,45    client: LangSmithClient | None = None,46) -> Generator[LangChainTracer, None, None]:47    """Instruct LangChain to log all runs in context to LangSmith.48 49    Args:50        project_name: The name of the project.51 52            Defaults to `'default'`.53        example_id: The ID of the example.54        tags: The tags to add to the run.55        client: The client of the langsmith.56 57    Yields:58        The LangChain tracer.59 60    Example:61        >>> with tracing_v2_enabled():62        ...     # LangChain code will automatically be traced63 64        You can use this to fetch the LangSmith run URL:65 66        >>> with tracing_v2_enabled() as cb:67        ...     chain.invoke("foo")68        ...     run_url = cb.get_run_url()69    """70    if isinstance(example_id, str):71        example_id = UUID(example_id)72    cb = LangChainTracer(73        example_id=example_id,74        project_name=project_name,75        tags=tags,76        client=client,77    )78    token = tracing_v2_callback_var.set(cb)79    try:80        yield cb81    finally:82        tracing_v2_callback_var.reset(token)83 84 85@contextmanager86def collect_runs() -> Generator[RunCollectorCallbackHandler, None, None]:87    """Collect all run traces in context.88 89    Yields:90        The run collector callback handler.91 92    Example:93        >>> with collect_runs() as runs_cb:94                chain.invoke("foo")95                run_id = runs_cb.traced_runs[0].id96    """97    cb = RunCollectorCallbackHandler()98    token = run_collector_var.set(cb)99    try:100        yield cb101    finally:102        run_collector_var.reset(token)103 104 105def _get_trace_callbacks(106    project_name: str | None = None,107    example_id: str | UUID | None = None,108    callback_manager: CallbackManager | AsyncCallbackManager | None = None,109) -> Callbacks:110    if _tracing_v2_is_enabled():111        project_name_ = project_name or _get_tracer_project()112        tracer = tracing_v2_callback_var.get() or LangChainTracer(113            project_name=project_name_,114            example_id=example_id,115        )116        if callback_manager is None:117            cb = cast("Callbacks", [tracer])118        else:119            if not any(120                isinstance(handler, LangChainTracer)121                for handler in callback_manager.handlers122            ):123                callback_manager.add_handler(tracer)124                # If it already has a LangChainTracer, we don't need to add another one.125                # this would likely mess up the trace hierarchy.126            cb = callback_manager127    else:128        cb = None129    return cb130 131 132def _tracing_v2_is_enabled() -> bool | Literal["local"]:133    if tracing_v2_callback_var.get() is not None:134        return True135    return ls_utils.tracing_is_enabled()136 137 138def _get_tracer_project() -> str:139    tracing_context = ls_rh.get_tracing_context()140    run_tree = tracing_context["parent"]141    if run_tree is None and tracing_context["project_name"] is not None:142        return cast("str", tracing_context["project_name"])143    return getattr(144        run_tree,145        "session_name",146        getattr(147            # Note, if people are trying to nest @traceable functions and the148            # tracing_v2_enabled context manager, this will likely mess up the149            # tree structure.150            tracing_v2_callback_var.get(),151            "project",152            # Have to set this to a string even though it always will return153            # a string because `get_tracer_project` technically can return154            # None, but only when a specific argument is supplied.155            # Therefore, this just tricks the mypy type checker156            str(ls_utils.get_tracer_project()),157        ),158    )159 160 161_configure_hooks: list[162    tuple[163        ContextVar[BaseCallbackHandler | None],164        bool,165        type[BaseCallbackHandler] | None,166        str | None,167    ]168] = []169 170 171def register_configure_hook(172    context_var: ContextVar[Any | None],173    inheritable: bool,  # noqa: FBT001174    handle_class: type[BaseCallbackHandler] | None = None,175    env_var: str | None = None,176) -> None:177    """Register a configure hook.178 179    Args:180        context_var: The context variable.181        inheritable: Whether the context variable is inheritable.182        handle_class: The callback handler class.183        env_var: The environment variable.184 185    Raises:186        ValueError: If `env_var` is set, `handle_class` must also be set to a non-`None`187            value.188    """189    if env_var is not None and handle_class is None:190        msg = "If env_var is set, handle_class must also be set to a non-None value."191        raise ValueError(msg)192 193    _configure_hooks.append(194        (195            # the typings of ContextVar do not have the generic arg set as covariant196            # so we have to cast it197            cast("ContextVar[BaseCallbackHandler | None]", context_var),198            inheritable,199            handle_class,200            env_var,201        )202    )203 204 205register_configure_hook(run_collector_var, inheritable=False)206