self.arize_client = Client(space_key=SPACE_KEY, api_key=API_KEY)
        if SPACE_KEY == "SPACE_KEY" or API_KEY == "API_KEY":
            raise ValueError("❌ CHANGE SPACE AND API KEYS")
        else:
            print("✅ Arize client setup done! Now you can start using Arize!")
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        for prompt in prompts:
            self.prompt_records.append(prompt.replace("\n", ""))
[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Do nothing."""
        pass
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        pd = import_pandas()
        from arize.utils.types import (
            EmbeddingColumnNames,
            Environments,
            ModelTypes,
            Schema,
        )
        # Safe check if 'llm_output' and 'token_usage' exist
        if response.llm_output and "token_usage" in response.llm_output:
            self.prompt_tokens = response.llm_output["token_usage"].get(
                "prompt_tokens", 0
            )
            self.total_tokens = response.llm_output["token_usage"].get(
                "total_tokens", 0
            )
            self.completion_tokens = response.llm_output["token_usage"].get(
                "completion_tokens", 0
            )
        else:
            self.prompt_tokens = (
                self.total_tokens
            ) = self.completion_tokens = 0  # assign default value
        for generations in response.generations:
            for generation in generations: