from typing import Dict

from dashscope.api_entities.dashscope_response import GenerationUsage


class QWenCallbackHandler:

    total_tokens: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    total_cost: float = 0.0

    def on_llm_end(self, response: GenerationUsage) -> None:
        """Collect token usage."""


        self.input_tokens += response.input_tokens
        self.output_tokens += response.output_tokens
        self.total_tokens = self.input_tokens + self.output_tokens
        self.total_cost = 0.04 * self.input_tokens / 1000 + 0.12 * self.output_tokens / 1000

    def on_chain_end(self, outputs: Dict ) -> None:
        """Write token usage to redis."""
        outputs["total_tokens"] = self.total_tokens
        outputs["total_cost"] = self.total_cost
        outputs["prompt_tokens"] = self.input_tokens
        outputs["completion_tokens"] = self.output_tokens
        print("input_tokens : {} \noutput_tokens : {}".format(outputs["prompt_tokens"] , outputs["completion_tokens"]))