def _handle_llm_payload(
        self, event_pair: List[CBEvent], span: "trace_tree.Span"
    ) -> Tuple[Dict[str, Any], Dict[str, Any], "trace_tree.Span"]:
        """Handle the payload of a LLM event."""
        inputs = event_pair[0].payload
        outputs = event_pair[-1].payload

        assert isinstance(inputs, dict) and isinstance(outputs, dict)

        # Get `original_template` from Prompt
        if EventPayload.PROMPT in inputs:
            inputs[EventPayload.PROMPT] = inputs[EventPayload.PROMPT]

        # Format messages
        if EventPayload.MESSAGES in inputs:
            inputs[EventPayload.MESSAGES] = "\n".join(
                [str(x) for x in inputs[EventPayload.MESSAGES]]
            )

        token_counts = get_llm_token_counts(self._token_counter, outputs)
        metadata = {
            "formatted_prompt_tokens_count": token_counts.prompt_token_count,
            "prediction_tokens_count": token_counts.completion_token_count,
            "total_tokens_used": token_counts.total_token_count,
        }
        span.attributes = metadata

        # Make `response` part of `outputs`
        outputs = {EventPayload.RESPONSE: str(outputs[EventPayload.RESPONSE])}

        return inputs, outputs, span
Handle the payload of a LLM event.
llama-index-integrations/callbacks/llama-index-callbacks-wandb/llama_index/callbacks/wandb/base.py