from__future__importannotationsfromcontextlibimportcontextmanagerfromcontextvarsimportContextVarfromtypingimport(TYPE_CHECKING,Any,Generator,List,Optional,Tuple,Type,Union,cast,)fromuuidimportUUIDfromlangsmithimportutilsasls_utilsfromlangsmith.run_helpersimportget_run_tree_contextfromlangchain_core.tracers.langchainimportLangChainTracerfromlangchain_core.tracers.run_collectorimportRunCollectorCallbackHandlerfromlangchain_core.tracers.schemasimportTracerSessionV1ifTYPE_CHECKING:fromlangsmithimportClientasLangSmithClientfromlangchain_core.callbacks.baseimportBaseCallbackHandler,Callbacksfromlangchain_core.callbacks.managerimportAsyncCallbackManager,CallbackManager# for backwards partial compatibility if this is imported by users but unusedtracing_callback_var:Any=Nonetracing_v2_callback_var:ContextVar[Optional[LangChainTracer]]=ContextVar("tracing_callback_v2",default=None)run_collector_var:ContextVar[Optional[RunCollectorCallbackHandler]]=ContextVar("run_collector",default=None)
[docs]@contextmanagerdeftracing_enabled(session_name:str="default",)->Generator[TracerSessionV1,None,None]:"""Throw an error because this has been replaced by tracing_v2_enabled."""raiseRuntimeError("tracing_enabled is no longer supported. Please use tracing_enabled_v2 instead.")
[docs]@contextmanagerdeftracing_v2_enabled(project_name:Optional[str]=None,*,example_id:Optional[Union[str,UUID]]=None,tags:Optional[List[str]]=None,client:Optional[LangSmithClient]=None,)->Generator[LangChainTracer,None,None]:"""Instruct LangChain to log all runs in context to LangSmith. Args: project_name (str, optional): The name of the project. Defaults to "default". example_id (str or UUID, optional): The ID of the example. Defaults to None. tags (List[str], optional): The tags to add to the run. Defaults to None. client (LangSmithClient, optional): The client of the langsmith. Defaults to None. Yields: LangChainTracer: The LangChain tracer. Example: >>> with tracing_v2_enabled(): ... # LangChain code will automatically be traced You can use this to fetch the LangSmith run URL: >>> with tracing_v2_enabled() as cb: ... chain.invoke("foo") ... run_url = cb.get_run_url() """ifisinstance(example_id,str):example_id=UUID(example_id)cb=LangChainTracer(example_id=example_id,project_name=project_name,tags=tags,client=client,)try:tracing_v2_callback_var.set(cb)yieldcbfinally:tracing_v2_callback_var.set(None)
[docs]@contextmanagerdefcollect_runs()->Generator[RunCollectorCallbackHandler,None,None]:"""Collect all run traces in context. Yields: run_collector.RunCollectorCallbackHandler: The run collector callback handler. Example: >>> with collect_runs() as runs_cb: chain.invoke("foo") run_id = runs_cb.traced_runs[0].id """cb=RunCollectorCallbackHandler()run_collector_var.set(cb)yieldcbrun_collector_var.set(None)
def_get_trace_callbacks(project_name:Optional[str]=None,example_id:Optional[Union[str,UUID]]=None,callback_manager:Optional[Union[CallbackManager,AsyncCallbackManager]]=None,)->Callbacks:if_tracing_v2_is_enabled():project_name_=project_nameor_get_tracer_project()tracer=tracing_v2_callback_var.get()orLangChainTracer(project_name=project_name_,example_id=example_id,)ifcallback_managerisNone:fromlangchain_core.callbacks.baseimportCallbackscb=cast(Callbacks,[tracer])else:ifnotany(isinstance(handler,LangChainTracer)forhandlerincallback_manager.handlers):callback_manager.add_handler(tracer,True)# If it already has a LangChainTracer, we don't need to add another one.# this would likely mess up the trace hierarchy.cb=callback_managerelse:cb=Nonereturncbdef_tracing_v2_is_enabled()->bool:iftracing_v2_callback_var.get()isnotNone:returnTruereturnls_utils.tracing_is_enabled()def_get_tracer_project()->str:run_tree=get_run_tree_context()returngetattr(run_tree,"session_name",getattr(# Note, if people are trying to nest @traceable functions and the# tracing_v2_enabled context manager, this will likely mess up the# tree structure.tracing_v2_callback_var.get(),"project",# Have to set this to a string even though it always will return# a string because `get_tracer_project` technically can return# None, but only when a specific argument is supplied.# Therefore, this just tricks the mypy type checkerstr(ls_utils.get_tracer_project()),),)_configure_hooks:List[Tuple[ContextVar[Optional[BaseCallbackHandler]],bool,Optional[Type[BaseCallbackHandler]],Optional[str],]]=[]
[docs]defregister_configure_hook(context_var:ContextVar[Optional[Any]],inheritable:bool,handle_class:Optional[Type[BaseCallbackHandler]]=None,env_var:Optional[str]=None,)->None:"""Register a configure hook. Args: context_var (ContextVar[Optional[Any]]): The context variable. inheritable (bool): Whether the context variable is inheritable. handle_class (Optional[Type[BaseCallbackHandler]], optional): The callback handler class. Defaults to None. env_var (Optional[str], optional): The environment variable. Defaults to None. Raises: ValueError: If env_var is set, handle_class must also be set to a non-None value. """ifenv_varisnotNoneandhandle_classisNone:raiseValueError("If env_var is set, handle_class must also be set to a non-None value.")fromlangchain_core.callbacks.baseimportBaseCallbackHandler_configure_hooks.append((# the typings of ContextVar do not have the generic arg set as covariant# so we have to cast itcast(ContextVar[Optional[BaseCallbackHandler]],context_var),inheritable,handle_class,env_var,))