community[minor]: add bedrock anthropic callback for token usage counting (#19864)

**Description:** add bedrock anthropic callback for token usage
counting, consulted openai callback.

---------

Co-authored-by: Massimiliano Pronesti <massimiliano.pronesti@gmail.com>
This commit is contained in:
Shuqian
2024-04-09 22:18:48 +08:00
committed by GitHub
parent 1f9f4d8742
commit ad9750403b
5 changed files with 202 additions and 2 deletions

View File

@@ -10,6 +10,9 @@ from typing import (
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
@@ -19,7 +22,10 @@ logger = logging.getLogger(__name__)
openai_callback_var: ContextVar[Optional[OpenAICallbackHandler]] = ContextVar(
"openai_callback", default=None
)
wandb_tracing_callback_var: ContextVar[Optional[WandbTracer]] = ContextVar( # noqa: E501
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
@@ -27,6 +33,7 @@ comet_tracing_callback_var: ContextVar[Optional[CometTracer]] = ContextVar( # n
)
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"
)
@@ -53,6 +60,27 @@ def get_openai_callback() -> Generator[OpenAICallbackHandler, None, None]:
openai_callback_var.set(None)
@contextmanager
def get_bedrock_anthropic_callback() -> (
Generator[BedrockAnthropicTokenUsageCallbackHandler, None, None]
):
"""Get the Bedrock anthropic callback handler in a context manager.
which conveniently exposes token and cost information.
Returns:
BedrockAnthropicTokenUsageCallbackHandler:
The Bedrock anthropic callback handler.
Example:
>>> with get_bedrock_anthropic_callback() as cb:
... # Use the Bedrock anthropic callback handler
"""
cb = BedrockAnthropicTokenUsageCallbackHandler()
bedrock_anthropic_callback_var.set(cb)
yield cb
bedrock_anthropic_callback_var.set(None)
@contextmanager
def wandb_tracing_enabled(
session_name: str = "default",