mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-05 13:06:03 +00:00
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:
@@ -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",
|
||||
|
Reference in New Issue
Block a user