def get_callback_handler()

in 3_optimization-design-ptn/03_prompt-optimization/promptwizard/glue/common/llm/llm_helper.py [0:0]

• 6 lines of code
• 5 McCabe index (conditional complexity)

def get_callback_handler(llm_handle: LLM, class_name: str) -> BaseCallbackHandler:
    """
    Extract callback_manager from llm_handle, find out which call back manager is of class type `class_name`.
    Return that object.

    :param llm_handle: Object of class LLM, which is the handle to make all LLM related calls
    :param class_name: Name of class (without prefix file path) e.g. TokenCountingHandler
    :return: Object of BaseCallbackHandler, that's registered as callback_manager in LLM. If not found, return None
    """
    if llm_handle and llm_handle.callback_manager:
        for handler in llm_handle.callback_manager.handlers:
            if type(handler).__name__ == class_name:
                return handler

    return None