Source code for langchain_nvidia_ai_endpoints.callbacks

"""Callback Handler that prints to std out."""
from __future__ import annotations

import logging
import threading
from collections import defaultdict
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, Dict, Generator, List, Optional

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_core.tracers.context import register_configure_hook

logger = logging.getLogger(__name__)

## This module contains output parsers for OpenAI tools. Set here for version control

"""
### **Usage/Cost Tracking**

For tracking model usage and , you can use the `get_usage_callback` context manager to
track token information similar to `get_openai_callback`. Additionally, you can specify
custom price mappings as necessary (`price_map` argument), or provide a custom callback
manager for advanced use-cases (`callback` argument).

**NOTE:** This feature is currently not supported in streaming modes, but works fine
for non-streaming `invoke/ainvoke` queries.

```
from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings
from langchain_nvidia_ai_endpoints.callbacks import get_usage_callback

## Assume a price map per 1K tokens for a particular deployment plan
price_map = {
    "mixtral_8x7b": 0.00060,
    "gemma_7b": 0.0002,
    "nvolveqa_40k": 0.000016,
}

llm_large = ChatNVIDIA(model="mixtral_8x7b", temperature=0.1)
llm_small = ChatNVIDIA(model="gemma_2b", temperature=0.1)
embedding = NVIDIAEmbeddings(model="nvolveqa_40k")
models = [llm_large, llm_small, embedding]

with get_usage_callback(price_map=price_map) as cb:
    ## Reset either at beginning or end. Statistics will run until cleared
    cb.reset()

    llm_large.invoke("Tell me a joke")
    print(cb, end="\n\n")
    # llm_large.invoke("Tell me a short joke")
    # print(cb, end="\n\n")
    # ## Tracking through streaming coming soon
    # [_ for _ in llm_small.stream("Tell me a joke")]
    # print(cb, end="\n[Should not change yet]\n\n")
    ## Tracking for streaming supported
    embedding.embed_query("What a nice day :D")
    print(cb, end="\n\n")
    # ## Sanity check. Should still be tracked fine
    # llm_small.invoke("Tell me a long joke")
    # print(cb, end="\n\n")

## Out of scope. Will not be tracked
llm_small.invoke("Tell me a short joke")
print(cb, end="\n[Should not change ever]\n\n")
cb.model_usage
```
"""


DEFAULT_MODEL_COST_PER_1K_TOKENS: Dict[str, float] = {}



[docs]
def standardize_model_name(
    model_name: str,
    price_map: dict = {},
    is_completion: bool = False,
) -> str:
    """
    Standardize the model name to a format that can be used in the OpenAI API.

    Args:
        model_name: Model name to standardize.
        is_completion: Whether the model is used for completion or not.
            Defaults to False.

    Returns:
        Standardized model name.

    """
    model_name = model_name.lower()
    if ".ft-" in model_name:
        model_name = model_name.split(".ft-")[0] + "-azure-finetuned"
    if ":ft-" in model_name:
        model_name = model_name.split(":")[0] + "-finetuned-legacy"
    if "ft:" in model_name:
        model_name = model_name.split(":")[1] + "-finetuned"
    if model_name.startswith("playground_"):
        model_name = model_name.replace("playground_", "")
    if (
        is_completion
        and model_name + "-completion" in price_map
        and (
            model_name.startswith("gpt-4")
            or model_name.startswith("gpt-3.5")
            or model_name.startswith("gpt-35")
            or ("finetuned" in model_name and "legacy" not in model_name)
        )
    ):
        return model_name + "-completion"
    else:
        return model_name




[docs]
def get_token_cost_for_model(
    model_name: str, num_tokens: int, price_map: dict, is_completion: bool = False
) -> float:
    """
    Get the cost in USD for a given model and number of tokens.

    Args:
        model_name: Name of the model
        num_tokens: Number of tokens.
        price_map: Map of model names to cost per 1000 tokens.
            Defaults to AI Foundation Endpoint pricing per https://www.together.ai/pricing.
        is_completion: Whether the model is used for completion or not.
            Defaults to False.

    Returns:
        Cost in USD.
    """
    model_name = standardize_model_name(
        model_name,
        price_map,
        is_completion=is_completion,
    )
    if model_name not in price_map:
        raise ValueError(
            f"Unknown model: {model_name}. Please provide a valid model name."
            "Known models are: " + ", ".join(price_map.keys())
        )
    return price_map[model_name] * (num_tokens / 1000)




[docs]
class UsageCallbackHandler(BaseCallbackHandler):
    """Callback Handler that tracks OpenAI info."""

    ## Per-model statistics
    _model_usage: defaultdict = defaultdict(
        lambda: {
            "total_tokens": 0,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "successful_requests": 0,
            "total_cost": 0.0,
        }
    )

    llm_output: dict = {}
    price_map: dict = {k: v for k, v in DEFAULT_MODEL_COST_PER_1K_TOKENS.items()}

    ## Aggregate statistics, compatible with OpenAICallbackHandler
    @property
    def total_tokens(self) -> int:
        """Total tokens used."""
        return self._model_usage["total"]["total_tokens"]

    @property
    def prompt_tokens(self) -> int:
        """Prompt tokens used."""
        return self._model_usage["total"]["prompt_tokens"]

    @property
    def completion_tokens(self) -> int:
        """Completion tokens used."""
        return self._model_usage["total"]["completion_tokens"]

    @property
    def successful_requests(self) -> int:
        """Total successful requests."""
        return self._model_usage["total"]["successful_requests"]

    @property
    def total_cost(self) -> float:
        """Total cost in USD."""
        return self._model_usage["total"]["total_cost"]


[docs]
    def __init__(self) -> None:
        super().__init__()
        self._lock = threading.Lock()


    def __repr__(self) -> str:
        return (
            f"Tokens Used: {self.total_tokens}\n"
            f"\tPrompt Tokens: {self.prompt_tokens}\n"
            f"\tCompletion Tokens: {self.completion_tokens}\n"
            f"Successful Requests: {self.successful_requests}\n"
            f"Total Cost (USD): ${self.total_cost:.8g}"
        )

    @property
    def model_usage(self) -> dict:
        """Whether to call verbose callbacks even if verbose is False."""
        return dict(self._model_usage)


[docs]
    def reset(self) -> None:
        """Reset the model usage."""
        with self._lock:
            self._model_usage.clear()


    @property
    def always_verbose(self) -> bool:
        """Whether to call verbose callbacks even if verbose is False."""
        return True


[docs]
    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        pass



[docs]
    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        pass



[docs]
    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        if not response.llm_output:
            response.llm_output = {}
        if not self.llm_output:
            self.llm_output = {}

        response.llm_output = {**self.llm_output, **response.llm_output}
        self.llm_output = {}

        if not response.llm_output:
            return None

        # compute tokens and cost for this request
        token_usage = response.llm_output.get(
            "token_usage", response.llm_output.get("usage", {})
        )
        completion_tokens = token_usage.get("completion_tokens", 0)
        prompt_tokens = token_usage.get("prompt_tokens", 0)
        model_name = response.llm_output.get("model_name", "")
        if model_name in self.price_map:
            completion_cost = get_token_cost_for_model(
                model_name, completion_tokens, self.price_map, is_completion=True
            )
            prompt_cost = get_token_cost_for_model(
                model_name, prompt_tokens, self.price_map
            )
        else:
            completion_cost = 0
            prompt_cost = 0

        # update shared state behind lock
        with self._lock:
            for base in (self._model_usage["total"], self._model_usage[model_name]):
                base["total_tokens"] += token_usage.get("total_tokens", 0)
                base["prompt_tokens"] += prompt_tokens
                base["completion_tokens"] += completion_tokens
                base["total_cost"] += prompt_cost + completion_cost
                base["successful_requests"] += 1
                for key in base.keys():
                    base[key] = round(base[key], 10)


    def __copy__(self) -> "UsageCallbackHandler":
        """Return a copy of the callback handler."""
        return self

    def __deepcopy__(self, memo: Any) -> "UsageCallbackHandler":
        """Return a deep copy of the callback handler."""
        return self



## get_usage_callack variable construction, registration, management

usage_callback_var: ContextVar[Optional[UsageCallbackHandler]] = ContextVar(
    "usage_callback", default=None
)

register_configure_hook(usage_callback_var, True)



[docs]
@contextmanager
def get_usage_callback(
    price_map: dict = {},
    callback: Optional[UsageCallbackHandler] = None,
) -> Generator[UsageCallbackHandler, None, None]:
    """Get the OpenAI callback handler in a context manager.
    which conveniently exposes token and cost information.

    Returns:
        OpenAICallbackHandler: The OpenAI callback handler.

    Example:
        >>> with get_openai_callback() as cb:
        ...     # Use the OpenAI callback handler
    """
    if not callback:
        callback = UsageCallbackHandler()
    if hasattr(callback, "price_map"):
        if hasattr(callback, "_lock"):
            with callback._lock:
                callback.price_map.update(price_map)
        else:
            callback.price_map.update(price_map)
    usage_callback_var.set(callback)
    yield callback
    usage_callback_var.set(None)