Closed lionelchg closed 6 months ago
@lkuligin
Answer generated by a π€
I found related issues that might be helpful. I did my best to summarize the solutions, but I recommend looking at the related issues yourself.
No discussion on issue
This response is meant to be useful, save you time, and share context. It is not meant to be a precise solution, but rather a starting point for your own research.
Help me be more useful! Please leave a π if this is helpful and π if it is irrelevant.
It looks like a duplicate of #5904.
The only concern I have is that it would require hard-coding the pricing (plus the exact amount of tokens is unknown since Vertex SDK doens't return it) . So we can have only rough estimation. If the pricing changes or becomes more complicated, then estimations returned by the Langchain would be misleading. Maybe it's easier to add a callback that counts [approximate] amount of tokens instead.
@hwchase17 what are thoughts on this?
How is this progressing?
I found the following on about getting the token count and billable characters: https://cloud.google.com/vertex-ai/docs/generative-ai/get-token-count
Hi, @lionelchg,
I'm helping the LangChain team manage their backlog and am marking this issue as stale. From what I understand, you opened this issue to request a callback function for VertexAI to monitor cost and token consumption, similar to the existing function for OpenAI. There have been contributions and collaboration from the community, with lkuligin suggesting to add a callback that counts the approximate amount of tokens and Capsar sharing information about getting the token count and billable characters from the VertexAI documentation. The issue has been resolved with these contributions, and the feature has been finalized.
Could you please confirm if this issue is still relevant to the latest version of the LangChain repository? If it is, please let the LangChain team know by commenting on the issue. Otherwise, feel free to close the issue yourself, or it will be automatically closed in 7 days. Thank you!
Would be great to have this
I just stumbled over this thread, looking for a solution to token counting in Vertex.
After some digging I found out, there is progress in langchain, in that there is now a VertexAICallbackHandler
in langchain_google_vertexai.callbacks
. I figured, one can use it to count the tokens in the same manner as one can use the openai_callback
with the below approach.
However, a well designed solution integrated into langchain likely needs more thinking, for instance:
successful_requests
)Anyway, I hope this helps to reanimate this feature discussion, as it would be indeed very helpful, to have such a callback.
vertexai_callback_var: ContextVar[Optional[VertexAICallbackHandler]] = ContextVar("vertexai_callback", default=None)
register_configure_hook(vertexai_callback_var, True)
class HotfixedVertextAICallback(VertexAICallbackHandler):
def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
# TODO this is highly experimental!
# - as of the current version we do not get a llm_output['model_id'] with the LLMResult, thus we do not know for sure if we are called by a vertex model
# - we assume that if the model_id is not present, and we have a 'usage_metadata' key in the generation_info, we are called by a vertex model
if not response.llm_output and self.__check_for_metadata(response):
super().on_llm_end(response, **kwargs)
@property
def total_tokens(self) -> int:
# lock as we may have
# - read prompt tokens to sum
# - other thread writes to completion tokens
# - we read completion tokens to sum
# --> completion tokens and prompt tokens may be inconsistent
with self._lock:
return self.prompt_tokens + self.completion_tokens
def __check_for_metadata(self, response: LLMResult) -> bool:
for generations in response.generations:
for generation in generations:
if generation.generation_info:
return "usage_metadata" in generation.generation_info
return False
@contextmanager
def get_vertexai_callback():
cb = HotfixedVertextAICallback()
vertexai_callback_var.set(cb)
yield cb
vertexai_callback_var.set(None)
with (get_openai_callback() as openai_cb,
get_vertexai_callback() as vertexai_cb):
<call your chains here>
I have created one, in case someone fight it helpful
"""Callback Handler that prints to std out."""
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, Dict, Generator, List, Optional, Union
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import AgentAction, AgentFinish, LLMResult
MODEL_COST_PER_1K_TOKENS = {"chat-bison": 0.0010, "text-bison": 0.0005, ""}
def get_vertexai_token_cost_for_model(model_name: str, num_characters: int) -> float:
"""
Characters are counted by UTF-8 code points and white space is excluded from the count.
Source: https://cloud.google.com/vertex-ai/pricing
"""
if model_name not in MODEL_COST_PER_1K_TOKENS:
raise ValueError(
f"Unknown model: {model_name}. Please provide a valid VertexAI model name."
"Known models are: " + ", ".join(MODEL_COST_PER_1K_TOKENS.keys())
)
return MODEL_COST_PER_1K_TOKENS[model_name] * num_characters / 1000
def count_characters_without_space(text: str) -> int:
"""
Returns the number of characters in the given text string,
excluding spaces.
Parameters:
text (str): The input text string
Returns:
int: The number of characters in the input string, excluding spaces
"""
return len(text) - text.count(" ")
class VertexAICallbackHandler(BaseCallbackHandler):
"""Callback Handler that tracks VertexAI info."""
total_tokens: int = 0
prompt_tokens: int = 0
completion_tokens: int = 0
successful_requests: int = 0
total_cost: float = 0.0
current_prompt_token: int = 0
model_name: str = ""
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}"
)
@property
def always_verbose(self) -> bool:
"""Whether to call verbose callbacks even if verbose is False."""
return True
def on_llm_start(
self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
) -> None:
"""Print out the prompts."""
# Calculate for a single call to LLM
for prompt in prompts:
self.current_prompt_token += count_characters_without_space(prompt)
if "ChatVertexAI" in serialized["id"]:
self.model_name = "chat-bison"
else:
self.model_name = "text-bison"
def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
"""Collect token usage."""
completion_tokens = 0
self.successful_requests += 1
for generations in response.generations:
for generation in generations:
completion_tokens += count_characters_without_space(generation.text)
if self.model_name in MODEL_COST_PER_1K_TOKENS:
completion_cost = get_vertexai_token_cost_for_model(
self.model_name, completion_tokens
)
prompt_cost = get_vertexai_token_cost_for_model(
self.model_name, self.current_prompt_token
)
self.total_cost += prompt_cost + completion_cost
self.total_tokens += completion_tokens + self.current_prompt_token
self.prompt_tokens += self.current_prompt_token
self.completion_tokens += completion_tokens
Feature request
It would be nice to have a similar function as
get_openai_callback()
for VertexAI. This actually gives the input tokens, output tokens and cost of using OpenAI models:I would like to have:
Motivation
I would like to monitor my usage of VertexAI models
Your contribution
I already read quite a bit the openai version of the callback but if anyone has already thought about how to do it with vertex I would be really curious :). If someone else also planned to do it we could merge efforts!