langchain/callbacks/arize_callback.py/ArizeCallbackHandler/on_llm_end
class ArizeCallbackHandler:    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Log data to Arize when an LLM ends."""

        from arize.utils.types import Embedding, Environments, ModelTypes

        # Record token usage of the LLM
        if response.llm_output is not None:
            self.prompt_tokens = response.llm_output["token_usage"]["prompt_tokens"]
            self.total_tokens = response.llm_output["token_usage"]["total_tokens"]
            self.completion_tokens = response.llm_output["token_usage"][
                "completion_tokens"
            ]
        i = 0

        # Go through each prompt response pair and generate embeddings as
        # well as timestamp and prediction ids
        for generations in response.generations:
            for generation in generations:
                prompt = self.prompt_records[i]
                prompt_embedding = pd.Series(
                    self.generator.generate_embeddings(
                        text_col=pd.Series(prompt.replace("\n", " "))
                    ).reset_index(drop=True)
                )
                generated_text = generation.text.replace("\n", " ")
                response_embedding = pd.Series(
                    self.generator.generate_embeddings(
                        text_col=pd.Series(generation.text.replace("\n", " "))
                    ).reset_index(drop=True)
                )
                pred_id = str(uuid.uuid4())

                # Define embedding features for Arize ingestion
                embedding_features = {
                    "prompt_embedding": Embedding(
                        vector=pd.Series(prompt_embedding[0]), data=prompt
                    ),
                    "response_embedding": Embedding(
                        vector=pd.Series(response_embedding[0]), data=generated_text
                    ),
                }
                tags = {
                    "Prompt Tokens": self.prompt_tokens,
                    "Completion Tokens": self.completion_tokens,
                    "Total Tokens": self.total_tokens,
                }

                # Log each prompt response data into arize
                future = self.arize_client.log(
                    prediction_id=pred_id,
                    tags=tags,
                    prediction_label="1",
                    model_id=self.model_id,
                    model_type=ModelTypes.SCORE_CATEGORICAL,
                    model_version=self.model_version,
                    environment=Environments.PRODUCTION,
                    embedding_features=embedding_features,
                )

                result = future.result()
                if result.status_code == 200:
                    print("✅ Successfully logged data to Arize!")
                else:
                    print(
                        f"❌ Logging failed with status code {result.status_code}"
                        f' and message "{result.text}"'
                    )

                i = i + 1