Callback Handler that tracks LLMs info.
| 41 | |
| 42 | |
| 43 | class CustomCallbackHandler(BaseCallbackHandler): |
| 44 | """Callback Handler that tracks LLMs info.""" |
| 45 | |
| 46 | total_tokens: int = 0 |
| 47 | prompt_tokens: int = 0 |
| 48 | completion_tokens: int = 0 |
| 49 | successful_requests: int = 0 |
| 50 | total_cost: float = 0.0 |
| 51 | |
| 52 | def __init__(self, llm_model_name: str) -> None: |
| 53 | super().__init__() |
| 54 | self._lock = threading.Lock() |
| 55 | self.model_name = llm_model_name if llm_model_name else "unknown" |
| 56 | |
| 57 | def __repr__(self) -> str: |
| 58 | return ( |
| 59 | f"Tokens Used: {self.total_tokens}\n" |
| 60 | f"\tPrompt Tokens: {self.prompt_tokens}\n" |
| 61 | f"\tCompletion Tokens: {self.completion_tokens}\n" |
| 62 | f"Successful Requests: {self.successful_requests}\n" |
| 63 | f"Total Cost (USD): ${self.total_cost}" |
| 64 | ) |
| 65 | |
| 66 | @property |
| 67 | def always_verbose(self) -> bool: |
| 68 | """Whether to call verbose callbacks even if verbose is False.""" |
| 69 | return True |
| 70 | |
| 71 | def on_llm_start( |
| 72 | self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any |
| 73 | ) -> None: |
| 74 | """Print out the prompts.""" |
| 75 | pass |
| 76 | |
| 77 | def on_llm_new_token(self, token: str, **kwargs: Any) -> None: |
| 78 | """Print out the token.""" |
| 79 | pass |
| 80 | |
| 81 | def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: |
| 82 | """Collect token usage.""" |
| 83 | # Check for usage_metadata (langchain-core >= 0.2.2) |
| 84 | try: |
| 85 | generation = response.generations[0][0] |
| 86 | except IndexError: |
| 87 | generation = None |
| 88 | if isinstance(generation, ChatGeneration): |
| 89 | try: |
| 90 | message = generation.message |
| 91 | if isinstance(message, AIMessage): |
| 92 | usage_metadata = message.usage_metadata |
| 93 | else: |
| 94 | usage_metadata = None |
| 95 | except AttributeError: |
| 96 | usage_metadata = None |
| 97 | else: |
| 98 | usage_metadata = None |
| 99 | if usage_metadata: |
| 100 | token_usage = {"total_tokens": usage_metadata["total_tokens"]} |
no outgoing calls
no test coverage detected