from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from typing import Any, Dict

class TokenUsageCallbackHandler(BaseCallbackHandler):
    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        token_usage = response.llm_output.get('token_usage', {})

        if token_usage:
            prompt_tokens = token_usage.get('prompt_tokens', 0)
            completion_tokens = token_usage.get('completion_tokens', 0)
            total_tokens = token_usage.get('total_tokens', 0)

            print("--- [Token Usage] ---")
            print(f"  Prompt Tokens: {prompt_tokens}")
            print(f"  Completion Tokens: {completion_tokens}")
            print(f"  Total Tokens: {total_tokens}")
            print("---------------------")