langchain_community.callbacks.manager — 🦜🔗 LangChain 0.2.0

from __future__ import annotations

import logging
from contextlib import contextmanager
from contextvars import ContextVar
from typing import (
    Generator,
    Optional,
)

from langchain_core.tracers.context import register_configure_hook

from langchain_community.callbacks.bedrock_anthropic_callback import (
    BedrockAnthropicTokenUsageCallbackHandler,
)
from langchain_community.callbacks.openai_info import OpenAICallbackHandler
from langchain_community.callbacks.tracers.comet import CometTracer
from langchain_community.callbacks.tracers.wandb import WandbTracer

logger = logging.getLogger(__name__)

openai_callback_var: ContextVar[Optional[OpenAICallbackHandler]] = ContextVar(
    "openai_callback", default=None
)
bedrock_anthropic_callback_var: (ContextVar)[
    Optional[BedrockAnthropicTokenUsageCallbackHandler]
] = ContextVar("bedrock_anthropic_callback", default=None)
wandb_tracing_callback_var: ContextVar[Optional[WandbTracer]] = ContextVar(
    "tracing_wandb_callback", default=None
)
comet_tracing_callback_var: ContextVar[Optional[CometTracer]] = ContextVar(  # noqa: E501
    "tracing_comet_callback", default=None
)

register_configure_hook(openai_callback_var, True)
register_configure_hook(bedrock_anthropic_callback_var, True)
register_configure_hook(
    wandb_tracing_callback_var, True, WandbTracer, "LANGCHAIN_WANDB_TRACING"
)
register_configure_hook(
    comet_tracing_callback_var, True, CometTracer, "LANGCHAIN_COMET_TRACING"
)


[docs]@contextmanager
def get_openai_callback() -> Generator[OpenAICallbackHandler, None, None]:
    """在上下文管理器中获取OpenAI回调处理程序。
方便地公开令牌和成本信息。

返回：
    OpenAICallbackHandler：OpenAI回调处理程序。

示例：
    >>> with get_openai_callback() as cb:
    ...     # 使用OpenAI回调处理程序
"""
    cb = OpenAICallbackHandler()
    openai_callback_var.set(cb)
    yield cb
    openai_callback_var.set(None)


[docs]@contextmanager
def get_bedrock_anthropic_callback() -> (
    Generator[BedrockAnthropicTokenUsageCallbackHandler, None, None]
):
    """在上下文管理器中获取Bedrock人类回调处理程序。
方便地公开令牌和成本信息。

返回：
    BedrockAnthropicTokenUsageCallbackHandler:
        Bedrock人类回调处理程序。

示例：
    >>> with get_bedrock_anthropic_callback() as cb:
    ...     # 使用Bedrock人类回调处理程序
"""
    cb = BedrockAnthropicTokenUsageCallbackHandler()
    bedrock_anthropic_callback_var.set(cb)
    yield cb
    bedrock_anthropic_callback_var.set(None)


[docs]@contextmanager
def wandb_tracing_enabled(
    session_name: str = "default",
) -> Generator[None, None, None]:
    """在上下文管理器中获取WandbTracer。

参数:
    session_name (str, 可选): 会话的名称。
        默认为"default"。

返回值:
    None

示例:
    >>> with wandb_tracing_enabled() as session:
    ...     # 使用WandbTracer会话
"""
    cb = WandbTracer()
    wandb_tracing_callback_var.set(cb)
    yield None
    wandb_tracing_callback_var.set(None)
Source code for langchain_community.callbacks.manager
