How can i find out that a pair of Callback events correspond to a same LLM call

[fernando-torres-blip-ai]

Im implementing a Custom Callbaback Handler and i want to group the events from on_llm_start and on_llm_end by LLM call. Is there any ID or implementation suggestion to do so ?

[dosubot[bot]]

Hey @fernando-torres-blip-ai! 👋 I'm here to help you with any bugs, questions, or contributions while you wait for a human maintainer. Let's tackle this together!

To group on_llm_start and on_llm_end events by LLM call in a Custom Callback Handler, use the run_id parameter. This UUID is passed to both methods, allowing you to associate the start and end of a particular LLM call.

Here's a basic implementation suggestion:

class CustomCallbackHandler(BaseCallbackHandler):
    def __init__(self):
        self.llm_calls = {}

    def on_llm_start(self, serialized, prompts, *, run_id, **kwargs):
        # Store the start event with the run_id
        self.llm_calls[run_id] = {'start': {'prompts': prompts}}

    def on_llm_end(self, response, *, run_id, **kwargs):
        # Retrieve the start event using the run_id and add the end event
        if run_id in self.llm_calls:
            self.llm_calls[run_id]['end'] = {'response': response}

This approach uses the run_id to track and group the start and end events of each LLM call ^[1][2].

To continue talking to Dosu, mention @dosu.